vector_search_example.py•5.44 kB
"""
Vector Search Integration Example
This example shows how to use the vector search functionality with the new
separated architecture (postgres_integration.py + vector_search.py).
"""
import asyncio
from sqlalchemy import create_engine
from server import MCPServer
async def example_usage():
"""Example of how to use the vector search functionality"""
# Initialize MCP Server with database integrations
mcp_server = MCPServer()
# Initialize both PostgreSQL and vector search integrations
# Replace with your actual database connection string
engine = create_engine("postgresql://user:password@localhost/dbname")
mcp_server.initialize_database(engine, "my_database")
# Vector search is now available as a tool through the MCP server
# You can call it via the MCP protocol or directly:
# Direct usage example:
if mcp_server.vector_search:
results = mcp_server.vector_search.semantic_search(
question="Find information about recent trips",
table_filter="descriptions", # Optional: filter by table
k=5 # Return top 5 results
)
print(f"Found {len(results)} results:")
for result in results:
print(f"- Table: {result['table_name']}")
print(f" Snippet: {result['snippet'][:100]}...")
print()
# SQL operations are now handled by PostgreSQL integration:
if mcp_server.postgres_integration:
sql_result = mcp_server.postgres_integration.safe_run_sql(
"SELECT * FROM trips ORDER BY start_date DESC",
limit_safe=True # Automatically adds LIMIT 100 if needed
)
if sql_result.get('success'):
print(f"SQL executed successfully:")
print(f"Rows returned: {sql_result['row_count']}")
print(f"Columns: {sql_result['columns']}")
if sql_result['data']:
print("First row:", sql_result['data'][0])
else:
print(f"SQL error: {sql_result['error']}")
# Available MCP tools:
# Tool name: "semantic_search"
# Parameters:
# - question (required): Search query
# - table_filter (optional): Filter by table name
# - fk_filter (optional): Filter by foreign key values
# - k (optional): Number of results (default: 10)
# Tool name: "safe_sql"
# Parameters:
# - sql (required): SQL query (SELECT only)
# - limit_safe (optional): Auto-inject LIMIT clause (default: True)
def populate_sample_data_example():
"""Example of how to populate sample document embeddings"""
from vector_search import VectorSearchEngine
engine = create_engine("postgresql://user:password@localhost/dbname")
vector_search = VectorSearchEngine(engine, "my_database")
# Populate with default configuration
result = vector_search.populate_sample_docs()
print(f"Populated {result['documents_created']} documents")
# Or with custom table configurations
custom_configs = [
{
'table': 'my_custom_table',
'pk_cols': ['id'],
'content_col': 'description',
'snippet_template': "Custom: {description}"
}
]
result = vector_search.populate_sample_docs(custom_configs)
print(f"Custom population: {result['documents_created']} documents")
if __name__ == "__main__":
print("Vector Search Integration Example")
print("=" * 40)
print()
print("NEW ARCHITECTURE:")
print("• postgres_integration.py - General PostgreSQL operations")
print("• vector_search.py - Semantic search functionality")
print("• server.py - MCP integration with both modules")
print()
print("Available MCP Tools:")
print("=" * 20)
print()
print("1. semantic_search - Vector similarity search")
print(' {"name": "semantic_search", "arguments": {"question": "recent trips", "k": 5}}')
print()
print("2. safe_sql - Safe SQL execution")
print(' {"name": "safe_sql", "arguments": {"sql": "SELECT * FROM trips LIMIT 10"}}')
print()
print("3. get_database_schema - Get complete database schema")
print(' {"name": "get_database_schema", "arguments": {}}')
print()
print("4. get_table_schema - Get specific table schema")
print(' {"name": "get_table_schema", "arguments": {"table_name": "users"}}')
print()
print("5. test_db_connection - Test database connectivity")
print(' {"name": "test_db_connection", "arguments": {}}')
print()
print("6. get_database_stats - Get database performance statistics")
print(' {"name": "get_database_stats", "arguments": {}}')
print()
print("7. get_table_size - Get table size information")
print(' {"name": "get_table_size", "arguments": {"table_name": "users"}}')
print()
print("8. explain_query - Analyze query performance")
print(' {"name": "explain_query", "arguments": {"sql": "SELECT * FROM users WHERE active = true"}}')
print()
print("9. generate_sql - Generate SQL from natural language")
print(' {"name": "generate_sql", "arguments": {"question": "Show me recent trips"}}')
print()
print("Initialization:")
print("=" * 15)
print("mcp_server = MCPServer()")
print("mcp_server.initialize_database(engine, 'db_key') # Initializes both integrations")
print()
print("For a comprehensive example, see: postgres_integration_example.py")