Spaces:
Sleeping
Sleeping
File size: 4,042 Bytes
f84a02d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | #!/usr/bin/env python3
"""
Populate LanceDB for Hybrid Search Validation
Ensures rich data exists for Docs, Meetings, and Tasks with proper embeddings.
"""
import asyncio
import logging
import os
from pathlib import Path
import sys
# Add project root to path
sys.path.append(str(Path(__file__).parent.parent.parent))
from backend.core.lancedb_handler import LanceDBHandler
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
def populate_db():
logger.info("Initializing LanceDB Handler...")
# Force local embeddings to ensure consistency without API keys if possible
handler = LanceDBHandler(embedding_provider="local")
if not handler.db:
logger.error("Failed to connect to LanceDB")
return False
# Test Data Categories
documents = [
# DOCUMENTS
{
"text": "The Q4 Marketing Strategy focuses on organic growth through content marketing and SEO optimization. Key targets include a 20% increase in inbound leads.",
"source": "doc",
"metadata": {"title": "Q4 Marketing Plan", "type": "document", "author": "Sarah J."}
},
{
"text": "API Documentation v2.0: All endpoints now require Bearer token authentication. Rate limits are set to 100 requests per minute.",
"source": "doc",
"metadata": {"title": "API Docs v2.0", "type": "document", "author": "Dev Team"}
},
# MEETINGS
{
"text": "Meeting Transcript: Team discussed the new frontend architecture. decided to migrate to Next.js 14 for better server-side rendering performance. Action items: JIRA-123, JIRA-124.",
"source": "meeting",
"metadata": {"title": "Frontend Architecture Review", "type": "meeting", "attendees": ["Alice", "Bob"]}
},
{
"text": "Client Call Notes: Client requested a new feature for exporting reports to PDF. Timeline agreed for delivery is next Friday.",
"source": "meeting",
"metadata": {"title": "Weekly Client Sync", "type": "meeting", "attendees": ["Client", "PM"]}
},
# TASKS
{
"text": "Task: Fix the login page layout issue on mobile devices. The submit button is overlapping with the footer.",
"source": "task",
"metadata": {"title": "Fix Mobile Login", "type": "task", "priority": "high"}
},
{
"text": "Task: Update the database schema to support multi-tenant architecture. Migration script needed.",
"source": "task",
"metadata": {"title": "DB Schema Migration", "type": "task", "priority": "critical"}
}
]
logger.info(f"Seeding {len(documents)} items into 'document_chunks' table...")
# Use the batch add method
count = handler.add_documents_batch("document_chunks", documents)
if count > 0:
logger.info(f"✅ Successfully added {count} documents.")
return True
else:
logger.error("❌ Failed to add documents.")
return False
def verify_search():
logger.info("Verifying Search Functionality...")
handler = LanceDBHandler(embedding_provider="local")
# Test Query 1: "marketing plan" (Should find Doc)
results = handler.search("document_chunks", "marketing strategy", limit=1)
if results and "Marketing" in results[0]['text']:
logger.info("✅ Search Verification 1 (Doc): PASS")
else:
logger.warning(f"❌ Search Verification 1 (Doc): FAIL - {results}")
# Test Query 2: "next.js" (Should find Meeting)
results = handler.search("document_chunks", "frontend framework", limit=1)
if results and "Next.js" in results[0]['text']:
logger.info("✅ Search Verification 2 (Meeting): PASS")
else:
logger.warning(f"❌ Search Verification 2 (Meeting): FAIL - {results}")
if __name__ == "__main__":
if populate_db():
verify_search()
else:
sys.exit(1)
|