start_mcp_server.py

#!/usr/bin/env python3
"""
PyVector MCP Server Startup Script

This script starts PyVector as an MCP-compatible server that can be used with
local LLMs and other applications that support the Model Context Protocol.

Copyright (C) 2025 PyVector

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""

import sys
import json
import asyncio
import argparse
from typing import Dict, Any, Optional
import logging

# Try to import MCP if available, otherwise use fallback
try:
    from pyvector.server import PyVectorMCPServer, main as mcp_main
    MCP_AVAILABLE = True
except ImportError:
    MCP_AVAILABLE = False

# Always import HTTP server for fallback mode
from pyvector.simple_server import PyVectorHTTPServer


def print_banner():
    """Print startup banner with connection information."""
    print("=" * 60)
    print("🚀 PyVector MCP Server")
    print("=" * 60)
    print("A lightweight vector database for local LLMs")
    print()


def print_mcp_connection_info():
    """Print MCP server connection details."""
    print("📡 MCP Server Configuration:")
    print()
    print("Add this to your MCP client configuration:")
    print()
    print("```json")
    print("{")
    print('  "mcpServers": {')
    print('    "pyvector": {')
    print('      "command": "python",')
    print(f'      "args": ["{sys.argv[0]}"],')
    print(f'      "cwd": "{sys.path[0]}"')
    print('    }')
    print('  }')
    print("}")
    print("```")
    print()
    print("🔧 Available MCP Tools:")
    tools = [
        "create_database - Initialize a new vector database",
        "add_text - Add text documents with automatic embedding",
        "add_vector - Add vectors directly to the database",
        "search_text - Semantic search using text queries",
        "search_vector - Search using vector queries",
        "delete_vector - Remove vectors from database",
        "get_database_info - Get database statistics",
        "save_database - Persist database to disk",
        "load_database - Load saved database"
    ]
    
    for tool in tools:
        print(f"  • {tool}")
    print()


def print_http_connection_info(host: str, port: int):
    """Print HTTP server connection details."""
    print("🌐 HTTP Server Configuration:")
    print()
    print(f"Server running at: http://{host}:{port}")
    print()
    print("📋 Available Endpoints:")
    endpoints = [
        ("GET", "/health", "Health check"),
        ("GET", "/info", "Database information"),
        ("POST", "/create_database", "Create new database"),
        ("POST", "/add_text", "Add text to database"),
        ("POST", "/search_text", "Search for similar texts"),
        ("POST", "/save_database", "Save database to disk"),
        ("POST", "/load_database", "Load database from disk")
    ]
    
    for method, endpoint, description in endpoints:
        print(f"  {method:4} {endpoint:20} - {description}")
    print()
    
    print("💡 Example Usage:")
    print()
    print("# Create database")
    print(f'curl -X POST http://{host}:{port}/create_database \\')
    print('  -H "Content-Type: application/json" \\')
    print('  -d \'{"dimension": 384, "index_type": "flat"}\'')
    print()
    print("# Add text")
    print(f'curl -X POST http://{host}:{port}/add_text \\')
    print('  -H "Content-Type: application/json" \\')
    print('  -d \'{"text": "Hello world", "metadata": {"source": "example"}}\'')
    print()
    print("# Search")
    print(f'curl -X POST http://{host}:{port}/search_text \\')
    print('  -H "Content-Type: application/json" \\')
    print('  -d \'{"query": "greeting", "k": 5}\'')
    print()


def print_fallback_info():
    """Print information about fallback implementations."""
    print("⚠️  Fallback Mode Information:")
    print()
    
    # Check for optional dependencies
    missing_deps = []
    
    try:
        import faiss
    except ImportError:
        missing_deps.append("faiss-cpu (for high-performance indexing)")
    
    try:
        import sentence_transformers
    except ImportError:
        missing_deps.append("sentence-transformers (for quality embeddings)")
    
    if not MCP_AVAILABLE:
        missing_deps.append("mcp (for full MCP protocol support)")
    
    if missing_deps:
        print("Missing optional dependencies:")
        for dep in missing_deps:
            print(f"  • {dep}")
        print()
        print("PyVector will use fallback implementations.")
        print("For better performance, install with: pip install -e .[full]")
        print()


async def start_mcp_server():
    """Start the MCP server."""
    if MCP_AVAILABLE:
        print("✅ Starting MCP server with full protocol support...")
        print()
        print_mcp_connection_info()
        await mcp_main()
    else:
        print("⚠️  MCP package not available, using HTTP server fallback...")
        print("   (Requires Python 3.10+ for full MCP support)")
        print()
        return False
    return True


def start_http_server(host: str, port: int):
    """Start the HTTP server."""
    print(f"✅ Starting HTTP server on {host}:{port}...")
    print()
    print_http_connection_info(host, port)
    
    server = PyVectorHTTPServer(host, port)
    
    try:
        server.start()
        print("🎯 Server is ready! Press Ctrl+C to stop.")
        print()
        
        # Keep the server running
        while True:
            import time
            time.sleep(1)
            
    except KeyboardInterrupt:
        print("\n🛑 Shutting down server...")
        server.stop()
        print("✅ Server stopped successfully")


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Start PyVector as an MCP-compatible server",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python start_mcp_server.py                    # Start MCP server (if available)
  python start_mcp_server.py --http             # Force HTTP server mode
  python start_mcp_server.py --http --port 9000 # HTTP server on port 9000
  python start_mcp_server.py --verbose          # Enable verbose logging
        """
    )
    
    parser.add_argument(
        "--http", 
        action="store_true",
        help="Force HTTP server mode instead of MCP"
    )
    
    parser.add_argument(
        "--host",
        default="localhost",
        help="Host to bind HTTP server to (default: localhost)"
    )
    
    parser.add_argument(
        "--port",
        type=int,
        default=8080,
        help="Port for HTTP server (default: 8080)"
    )
    
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Enable verbose logging"
    )
    
    args = parser.parse_args()
    
    # Configure logging
    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)
    else:
        logging.basicConfig(level=logging.INFO)
    
    print_banner()
    print_fallback_info()
    
    # Determine server mode
    if args.http or not MCP_AVAILABLE:
        # HTTP server mode
        start_http_server(args.host, args.port)
    else:
        # MCP server mode
        try:
            asyncio.run(start_mcp_server())
        except Exception as e:
            print(f"❌ Failed to start MCP server: {e}")
            print("🔄 Falling back to HTTP server...")
            print()
            start_http_server(args.host, args.port)


if __name__ == "__main__":
    main()