Investment Statement MCP Server

"""Index all IBKR CSV files into the database."""

import asyncio
import sys
from pathlib import Path
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

# Calculate CSV directory relative to this file
_SCRIPT_DIR = Path(__file__).parent.resolve()
_CSV_DIR = _SCRIPT_DIR.parent / "data" / "csvs"


async def main():
    """Index all CSV files from data/csvs directory."""
    # 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()

    # Get all CSV files
    csv_files = sorted(_CSV_DIR.glob("*.csv"))

    if not csv_files:
        print(f"No CSV files found in {_CSV_DIR}")
        return

    print(f"Found {len(csv_files)} CSV files to index\n")

    # Track results
    indexed = 0
    errors = []

    # Index each CSV file
    for csv_file in csv_files:
        print(f"Indexing {csv_file.name}...", end=" ")
        try:
            result = await indexing_tools.index_statement(str(csv_file))

            if result["status"] == "success":
                print(f"✓ {result['statement_id']}")
                print(f"  Account: {result['account_number']}")
                print(f"  Date: {result['statement_date']}")
                indexed += 1
            else:
                print(f"✗ {result.get('message', 'Unknown error')}")
                errors.append({
                    "file": csv_file.name,
                    "error": result.get("message"),
                })

        except Exception as e:
            print(f"✗ Error: {e}")
            errors.append({
                "file": csv_file.name,
                "error": str(e),
            })

        print()

    # Show summary
    print("=" * 60)
    print(f"Indexing complete: {indexed}/{len(csv_files)} files indexed")

    if errors:
        print(f"\nErrors ({len(errors)}):")
        for error in errors:
            print(f"  - {error['file']}: {error['error']}")

    # Get final stats
    stats = await indexing_tools.get_indexing_stats()
    print(f"\nDatabase Statistics:")
    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)}")


if __name__ == "__main__":
    asyncio.run(main())