joelgilbert/NL2SQL

1"""2Initialize the vector store with database schemas.3 4This script should be run once after setting up your database to populate5the Upstash Vector store with schema embeddings for semantic search.6 7Usage:8    python scripts/init_vector_store.py9"""10 11import sys12import os13 14# Add parent directory to path15sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))16 17from database.connection import db18from database.schema_manager import schema_manager19from vector_store.upstash_client import upstash_client20from vector_store.embeddings import embedding_helper21import logging22 23logging.basicConfig(level=logging.INFO)24logger = logging.getLogger(__name__)25 26 27def initialize_vector_store():28    """Initialize Upstash Vector store with database schemas."""29    30    logger.info("=== Initializing Vector Store ===")31    32    # Step 1: Test database connection33    logger.info("Step 1: Testing database connection...")34    if not db.test_connection():35        logger.error("❌ Database connection failed! Check your connection string.")36        return False37    logger.info("✅ Database connection successful")38    39    # Step 2: Initialize Upstash Vector40    logger.info("\nStep 2: Initializing Upstash Vector store...")41    try:42        upstash_client.initialize_vector_store()43        logger.info("✅ Upstash Vector initialized")44    except Exception as e:45        logger.error(f"❌ Upstash initialization failed: {e}")46        return False47    48    # Step 3: Fetch database schema49    logger.info("\nStep 3: Fetching database schema...")50    try:51        schema_info = schema_manager.fetch_schema()52        logger.info(f"✅ Found {len(schema_info)} tables")53        54        for table in schema_info:55            logger.info(f"  - {table['table_name']} ({len(table['columns'])} columns)")56    except Exception as e:57        logger.error(f"❌ Schema fetch failed: {e}")58        return False59    60    # Step 4: Store schema embeddings61    logger.info("\nStep 4: Storing schema embeddings in vector database...")62    stored_count = 063    64    for table_info in schema_info:65        try:66            # Store in Upstash Vector (it will create embeddings automatically)67            upstash_client.store_schema_embeddings(table_info)68            69            stored_count += 170            logger.info(f"  ✅ Stored: {table_info['table_name']}")71            72        except Exception as e:73            logger.error(f"  ❌ Failed to store {table_info['table_name']}: {e}")74    75    logger.info(f"\n✅ Successfully stored {stored_count}/{len(schema_info)} schemas")76    77    # Step 5: Test semantic search78    logger.info("\nStep 5: Testing semantic search...")79    try:80        test_query = "find user information"81        results = upstash_client.search_similar_schemas(test_query, top_k=3)82        83        logger.info(f"✅ Search test successful! Found {len(results)} results for '{test_query}':")84        for i, result in enumerate(results, 1):85            logger.info(f"  {i}. {result['table_name']} (score: {result.get('score', 'N/A')})")86    87    except Exception as e:88        logger.error(f"❌ Search test failed: {e}")89    90    logger.info("\n=== Initialization Complete ===")91    logger.info("\n🎉 Your vector store is ready! You can now run the main application.")92    93    return True94 95 96if __name__ == "__main__":97    success = initialize_vector_store()98    sys.exit(0 if success else 1)99