"""
Test script for Phase 2 implementation
Tests the smart search orchestrator and enhanced service layer
"""
import asyncio
import logging
import os
from sqlalchemy import create_engine
# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Import our new Phase 2 components
from services.smart_search_service import SmartSearchService
from presentation.mcp_server import MCPDatabaseServer
from presentation.tools.smart_search_tools import SmartSearchTools
# Import Phase 1 components for testing
from repositories.postgres_repository import PostgresRepository
from repositories.vector_repository import VectorRepository
from services.schema_service import SchemaService
from services.sql_service import SQLService
from services.semantic_service import SemanticService
from services.synthesis_service import SynthesisService
async def test_smart_search_service():
"""Test Smart Search Service"""
print("\n=== Testing SmartSearchService ===")
try:
# Create engine (you'll need to adjust connection string)
engine = create_engine("postgresql://postgres:password@localhost:5432/test_db")
# Create repositories
postgres_repo = PostgresRepository(engine)
vector_repo = VectorRepository(engine)
# Create basic services
schema_service = SchemaService(postgres_repo)
sql_service = SQLService(postgres_repo, schema_service)
semantic_service = SemanticService(vector_repo)
synthesis_service = SynthesisService(llm_config={})
# Create smart search service
smart_search = SmartSearchService(
schema_service, sql_service, semantic_service, synthesis_service
)
# Test question classification
print("Testing question classification...")
classification = smart_search._classify_question("How many users are in the database?")
print(f"Classification: {classification.strategy.value}, confidence: {classification.confidence}")
# Test search capabilities
print("\nTesting search capabilities...")
capabilities = await smart_search.get_search_capabilities()
print(f"SQL available: {capabilities['sql_search']['available']}")
print(f"Semantic available: {capabilities['semantic_search']['available']}")
# Test question suggestions
print("\nTesting question suggestions...")
suggestions = await smart_search.suggest_questions()
print(f"Generated {len(suggestions)} suggestions")
# Test full search (mock)
print("\nTesting full search...")
result = await smart_search.search("What tables are available?")
print(f"Search result: success={result['success']}")
return True
except Exception as e:
print(f"SmartSearchService test failed: {e}")
return False
async def test_mcp_server():
"""Test MCP Server Integration"""
print("\n=== Testing MCPDatabaseServer ===")
try:
# Example configuration
config = {
'database': {
'host': 'localhost',
'port': 5432,
'database': 'test_db',
'username': 'postgres',
'password': 'password'
},
'llm': {
'provider': 'openai',
'model': 'gpt-4',
'api_key': 'test_key'
}
}
# Create MCP server
mcp_server = MCPDatabaseServer(config)
# Test initialization
print("Testing MCP server initialization...")
await mcp_server.initialize()
print("MCP server initialized successfully")
# Test smart search tools
if mcp_server.smart_search_service:
print("\nTesting Smart Search via MCP server...")
result = await mcp_server.smart_search_service.search("Test question")
print(f"Smart search result: {result['success']}")
return True
except Exception as e:
print(f"MCPDatabaseServer test failed: {e}")
return False
async def test_smart_search_tools():
"""Test Smart Search MCP Tools"""
print("\n=== Testing SmartSearchTools ===")
try:
# Create engine (you'll need to adjust connection string)
engine = create_engine("postgresql://postgres:password@localhost:5432/test_db")
# Create repositories
postgres_repo = PostgresRepository(engine)
vector_repo = VectorRepository(engine)
# Create services
schema_service = SchemaService(postgres_repo)
sql_service = SQLService(postgres_repo, schema_service)
semantic_service = SemanticService(vector_repo)
synthesis_service = SynthesisService(llm_config={})
# Create smart search service and tools
smart_search_service = SmartSearchService(
schema_service, sql_service, semantic_service, synthesis_service
)
smart_search_tools = SmartSearchTools(smart_search_service)
# Test smart search tool
print("Testing smart search MCP tool...")
result = await smart_search_tools.smart_search(
question="What data is available?",
max_sql_queries=2
)
print(f"Smart search tool result: success={result['success']}")
# Test capabilities tool
print("\nTesting capabilities MCP tool...")
result = await smart_search_tools.get_search_capabilities()
print(f"Capabilities tool result: success={result['success']}")
# Test suggestions tool
print("\nTesting suggestions MCP tool...")
result = await smart_search_tools.suggest_questions()
print(f"Suggestions tool result: success={result['success']}, count={result['count']}")
return True
except Exception as e:
print(f"SmartSearchTools test failed: {e}")
return False
async def test_question_classification():
"""Test question classification logic"""
print("\n=== Testing Question Classification ===")
try:
# Create engine (you'll need to adjust connection string)
engine = create_engine("postgresql://postgres:password@localhost:5432/test_db")
# Create repositories
postgres_repo = PostgresRepository(engine)
vector_repo = VectorRepository(engine)
# Create services
schema_service = SchemaService(postgres_repo)
sql_service = SQLService(postgres_repo, schema_service)
semantic_service = SemanticService(vector_repo)
synthesis_service = SynthesisService(llm_config={})
# Create smart search service
smart_search = SmartSearchService(
schema_service, sql_service, semantic_service, synthesis_service
)
# Test different question types
test_questions = [
("How many users are in the database?", "SQL_ONLY"),
("What is a foreign key?", "SEMANTIC_ONLY"),
("Show me users and explain what they are", "HYBRID"),
("What tables contain user information?", "HYBRID"),
("Count the orders by status", "SQL_ONLY"),
("Explain database normalization", "SEMANTIC_ONLY")
]
correct_classifications = 0
for question, expected_strategy in test_questions:
classification = smart_search._classify_question(question)
actual_strategy = classification.strategy.value.upper()
print(f"Question: {question}")
print(f"Expected: {expected_strategy}, Got: {actual_strategy}")
print(f"Confidence: {classification.confidence}, Reasoning: {classification.reasoning}")
if actual_strategy == expected_strategy:
correct_classifications += 1
print()
accuracy = correct_classifications / len(test_questions)
print(f"Classification accuracy: {accuracy:.2%} ({correct_classifications}/{len(test_questions)})")
return accuracy > 0.5 # At least 50% accuracy
except Exception as e:
print(f"Question classification test failed: {e}")
return False
async def main():
"""Run all Phase 2 tests"""
print("Starting Phase 2 Implementation Tests")
print("=" * 50)
# Note: These tests require a database connection
# You may need to adjust connection strings or mock the database
test_results = []
# Test each Phase 2 component
test_results.append(await test_smart_search_service())
test_results.append(await test_mcp_server())
test_results.append(await test_smart_search_tools())
test_results.append(await test_question_classification())
# Summary
print("\n" + "=" * 50)
print("PHASE 2 TEST SUMMARY")
print("=" * 50)
passed = sum(test_results)
total = len(test_results)
print(f"Tests passed: {passed}/{total}")
if passed == total:
print("ā
All Phase 2 components working correctly!")
print("\nš PHASE 2 IMPLEMENTATION COMPLETE!")
print("\nKey features now available:")
print("- Smart Search orchestration with automatic strategy selection")
print("- Question classification and routing")
print("- Unified MCP server with all tools")
print("- HTTP endpoints for React integration")
print("- Clean architecture with proper separation of concerns")
else:
print("ā Some tests failed - check connection strings and database setup")
return passed == total
if __name__ == "__main__":
# Run the tests
success = asyncio.run(main())
exit(0 if success else 1)
