Investment Statement MCP Server

"""Rebuild database from JSON archives to ensure no duplicates."""

import asyncio
from src.database.sqlite_client import SQLiteClient
from src.database.lancedb_client import LanceDBClient
from src.embeddings.generator import EmbeddingGenerator
from src.tools.indexing import IndexingTools


async def main():
    """Rebuild database from JSON archives."""
    # Initialize clients (paths auto-resolved relative to module location)
    sqlite_client = SQLiteClient()
    lance_client = LanceDBClient()
    embedding_gen = EmbeddingGenerator(model_name="all-MiniLM-L6-v2")

    # Initialize indexing tools (paths auto-resolved relative to module location)
    indexing_tools = IndexingTools(
        sqlite_client=sqlite_client,
        lance_client=lance_client,
        embedding_gen=embedding_gen,
    )

    # Initialize database
    await sqlite_client.initialize()

    print("Rebuilding database from JSON archives...")
    print("This will clear existing data and rebuild from scratch.\n")

    result = await indexing_tools.rebuild_from_archives()

    if result["status"] == "success":
        print(f"✓ Successfully rebuilt database!")
        print(f"  Rebuilt: {result['rebuilt']}/{result['total_archives']} archives")

        if "errors" in result:
            print(f"\n⚠️  Errors: {len(result['errors'])}")
            for error in result["errors"]:
                print(f"    - {error['file']}: {error['error']}")

        print(f"\nFinal Statistics:")
        stats = result.get("stats", {})
        print(f"  Statements: {stats.get('total_statements', 0)}")
        print(f"  Holdings: {stats.get('total_holdings', 0)}")
        print(f"  Transactions: {stats.get('total_transactions', 0)}")
        print(f"  Vector chunks: {stats.get('total_chunks', 0)}")
    else:
        print(f"✗ Error: {result.get('message', 'Unknown error')}")
        return False

    return True


if __name__ == "__main__":
    success = asyncio.run(main())
    exit(0 if success else 1)