#!/usr/bin/env python3 """ LanceDB Initialization Script for ATOM Platform This script initializes the LanceDB database with sample data to enable search functionality in the ATOM platform. Usage: python initialize_lancedb.py """ from datetime import datetime, timezone import json import logging import os import sys from typing import Any, Dict, List import uuid # Configure logging logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" ) logger = logging.getLogger(__name__) def initialize_lancedb(): """Initialize LanceDB with sample data for search functionality""" try: # Import LanceDB import lancedb import numpy as np import pandas as pd import pyarrow as pa except ImportError as e: logger.error(f"LanceDB dependencies not available: {e}") logger.info( "Please install required packages: pip install lancedb pyarrow pandas numpy" ) return False # Get database path from environment or use default db_path = os.environ.get("LANCEDB_URI", "/tmp/atom_lancedb") logger.info(f"Initializing LanceDB at: {db_path}") try: # Connect to LanceDB db = lancedb.connect(db_path) logger.info("Connected to LanceDB") # Define schema for document chunks chunk_schema = pa.schema( [ pa.field("chunk_id", pa.string()), pa.field("doc_id", pa.string()), pa.field("user_id", pa.string()), pa.field("chunk_index", pa.int32()), pa.field("chunk_text", pa.string()), pa.field("metadata", pa.string()), pa.field( "vector_embedding", pa.list_(pa.float32(), 1536) ), # OpenAI embedding dimension pa.field("created_at", pa.string()), ] ) # Create or open the document_chunks table table_name = "document_chunks" if table_name in db.table_names(): logger.info(f"Table {table_name} already exists, opening...") table = db.open_table(table_name) else: logger.info(f"Creating new table: {table_name}") table = db.create_table(table_name, schema=chunk_schema) # Generate sample data sample_documents = [ { "title": "Project Planning Meeting", "content": "Discussed project timelines, resource allocation, and milestones for Q4 2024. Team agreed on aggressive but achievable deadlines.", "tags": ["meeting", "planning", "project"], }, { "title": "Technical Architecture Review", "content": "Reviewed the microservices architecture and API design patterns. Decided to use GraphQL for frontend communication.", "tags": ["technical", "architecture", "api"], }, { "title": "Customer Feedback Analysis", "content": "Analyzed customer feedback from Q3. Key themes: improved UI/UX, faster response times, better documentation.", "tags": ["customer", "feedback", "analysis"], }, { "title": "Security Audit Report", "content": "Completed security audit with penetration testing. Identified vulnerabilities in authentication system.", "tags": ["security", "audit", "vulnerabilities"], }, { "title": "Team Standup Notes", "content": "Daily standup: backend team working on authentication, frontend team implementing search UI, QA team testing workflows.", "tags": ["standup", "team", "progress"], }, ] # Generate embeddings for sample data logger.info("Generating embeddings for sample data...") try: from note_utils import get_text_embedding_openai embedding_function_available = True except ImportError: logger.warning("note_utils not available, using mock embeddings") embedding_function_available = False # Prepare data for insertion data_to_insert = [] user_id = "default_user" for doc_idx, doc in enumerate(sample_documents): doc_id = str(uuid.uuid4()) # Split content into chunks (simplified - just use the whole content) chunks = [ { "text": doc["content"], "metadata": { "title": doc["title"], "tags": doc["tags"], "source": "sample_data", }, } ] for chunk_idx, chunk in enumerate(chunks): # Generate embedding if embedding_function_available: embedding_result = get_text_embedding_openai(chunk["text"]) if embedding_result["status"] == "success": embedding = embedding_result["data"] else: logger.warning( f"Failed to generate embedding: {embedding_result.get('message')}" ) # Use mock embedding embedding = [0.01] * 1536 else: # Use mock embedding embedding = [0.01] * 1536 # Create chunk record chunk_record = { "chunk_id": str(uuid.uuid4()), "doc_id": doc_id, "user_id": user_id, "chunk_index": chunk_idx, "chunk_text": chunk["text"], "metadata": json.dumps(chunk["metadata"]), "vector_embedding": embedding, "created_at": datetime.now(timezone.utc).isoformat(), } data_to_insert.append(chunk_record) # Insert data into LanceDB if data_to_insert: logger.info(f"Inserting {len(data_to_insert)} document chunks...") table.add(data_to_insert) logger.info("Sample data inserted successfully!") else: logger.warning("No data to insert") # Verify the data was inserted count = table.count_rows() logger.info(f"Table now contains {count} rows") # Test search functionality logger.info("Testing search functionality...") # Generate embedding for test query test_query = "project planning" if embedding_function_available: embedding_result = get_text_embedding_openai(test_query) if embedding_result["status"] == "success": query_embedding = embedding_result["data"] # Perform search results = table.search(query_embedding).limit(3).to_list() logger.info(f"Search test returned {len(results)} results") # Display results for i, result in enumerate(results): logger.info( f"Result {i + 1}: {result.get('chunk_text', '')[:100]}..." ) else: logger.warning("Could not test search - embedding generation failed") else: logger.info("Search test skipped - mock embeddings in use") logger.info("LanceDB initialization completed successfully!") return True except Exception as e: logger.error(f"Failed to initialize LanceDB: {e}") return False def create_meeting_transcripts_table(): """Create meeting_transcripts_embeddings table if needed""" try: import lancedb import pyarrow as pa db_path = os.environ.get("LANCEDB_URI", "/tmp/atom_lancedb") db = lancedb.connect(db_path) table_name = "meeting_transcripts_embeddings" if table_name not in db.table_names(): logger.info(f"Creating table: {table_name}") # Define schema for meeting transcripts schema = pa.schema( [ pa.field("transcript_id", pa.string()), pa.field("user_id", pa.string()), pa.field("meeting_title", pa.string()), pa.field("content", pa.string()), pa.field("timestamp", pa.string()), pa.field("vector_embedding", pa.list_(pa.float32(), 1536)), pa.field("metadata", pa.string()), ] ) db.create_table(table_name, schema=schema) logger.info(f"Table {table_name} created successfully") else: logger.info(f"Table {table_name} already exists") return True except Exception as e: logger.error(f"Failed to create meeting transcripts table: {e}") return False def main(): """Main execution function""" logger.info("Starting LanceDB initialization...") # Initialize main document chunks table if not initialize_lancedb(): logger.error("LanceDB initialization failed") sys.exit(1) # Create meeting transcripts table (for compatibility) create_meeting_transcripts_table() logger.info("LanceDB setup completed!") logger.info( "You can now test search functionality at: http://localhost:5058/semantic_search_meetings" ) if __name__ == "__main__": main()