Update start_mcp.py

start_mcp.py

@@ -1,166 +1,188 @@
-#!/usr/bin/env python3
-"""
-MaTableGPT MCP Server Launcher
-==============================
-
-This script starts the MaTableGPT MCP service for extracting 
-table data from materials science literature.
-
-Usage:
-    python start_mcp.py [--host HOST] [--port PORT] [--mode MODE]
-    
-Arguments:
-    --host      Host address (default: 0.0.0.0)
-    --port      Port number (default: 7860)
-    --mode      Run mode: 'stdio' or 'sse' (default: sse for HuggingFace Space)
-
-Environment Variables:
-    LLM_API_KEY / OPENAI_API_KEY     - API key for LLM service
-    LLM_API_BASE / OPENAI_API_BASE   - Custom API base URL (for third-party services)
-    LLM_MODEL / OPENAI_MODEL         - Model name (default: gpt-4-turbo-preview)
-    MCP_HOST                          - Server host (default: 0.0.0.0)
-    MCP_PORT                          - Server port (default: 7860)
-"""
-
-import os
-import sys
-import argparse
-import logging
-
-# Add parent directory to path for imports
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
-# Configure logging
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
-)
-logger = logging.getLogger("matablgpt-mcp-launcher")
-
-
-def check_environment():
-    """Check if required environment variables are set."""
-    warnings = []
-    
-    # Check for API key (support both naming conventions)
-    api_key = os.environ.get('LLM_API_KEY') or os.environ.get('OPENAI_API_KEY')
-    if not api_key:
-        warnings.append(
-            "LLM_API_KEY/OPENAI_API_KEY not set. GPT extraction features will not work. "
-            "Set it in HuggingFace Space secrets or environment variables."
-        )
-    
-    # Check for API base (for third-party services)
-    api_base = os.environ.get('LLM_API_BASE') or os.environ.get('OPENAI_API_BASE')
-    if api_base:
-        logger.info(f"Using custom API base: {api_base}")
-    
-    return warnings
-
-
-def check_dependencies():
-    """Check if required packages are installed."""
-    missing = []
-    
-    required = [
-        ('mcp', 'mcp[cli]'),
-        ('openai', 'openai'),
-        ('bs4', 'beautifulsoup4'),
-        ('pandas', 'pandas'),
-        ('lxml', 'lxml')
-    ]
-    
-    for module, package in required:
-        try:
-            __import__(module)
-        except ImportError:
-            missing.append(package)
-    
-    return missing
-
-
-def main():
-    """Main entry point."""
-    # Get default values from environment variables
-    default_host = os.environ.get('MCP_HOST', '0.0.0.0')
-    default_port = int(os.environ.get('MCP_PORT', '7860'))
-    
-    parser = argparse.ArgumentParser(
-        description="MaTableGPT MCP Server - Table Data Extraction from Materials Science Literature"
-    )
-    parser.add_argument(
-        '--host', 
-        default=default_host, 
-        help=f'Host address (default: {default_host})'
-    )
-    parser.add_argument(
-        '--port', 
-        type=int, 
-        default=default_port, 
-        help=f'Port number (default: {default_port})'
-    )
-    parser.add_argument(
-        '--mode',
-        choices=['stdio', 'sse'],
-        default='sse',
-        help='Run mode: stdio for standard I/O, sse for Server-Sent Events (default: sse)'
-    )
-    parser.add_argument(
-        '--debug',
-        action='store_true',
-        help='Enable debug logging'
-    )
-    
-    args = parser.parse_args()
-    
-    if args.debug:
-        logging.getLogger().setLevel(logging.DEBUG)
-    
-    # Check dependencies
-    missing = check_dependencies()
-    if missing:
-        logger.error(f"Missing required packages: {', '.join(missing)}")
-        logger.error(f"Install with: pip install {' '.join(missing)}")
-        sys.exit(1)
-    
-    # Check environment
-    warnings = check_environment()
-    for warning in warnings:
-        logger.warning(warning)
-    
-    # Display startup info
-    logger.info("=" * 60)
-    logger.info("MaTableGPT MCP Server")
-    logger.info("=" * 60)
-    logger.info(f"Mode: {args.mode}")
-    if args.mode == 'sse':
-        logger.info(f"Host: {args.host}")
-        logger.info(f"Port: {args.port}")
-        logger.info(f"SSE Endpoint: http://{args.host}:{args.port}/sse")
-    logger.info("=" * 60)
-    
-    # Import and run MCP service
-    try:
-        from mcp_service import mcp
-        
-        if args.mode == 'stdio':
-            logger.info("Starting MCP server in stdio mode...")
-            mcp.run()
-        else:
-            logger.info(f"Starting MCP server in SSE mode on {args.host}:{args.port}...")
-            logger.info("MCP SSE service is ready to accept connections!")
-            mcp.run(transport='sse', host=args.host, port=args.port)
-            
-    except ImportError as e:
-        logger.error(f"Failed to import MCP service: {e}")
-        logger.error("Make sure mcp_service.py is in the same directory")
-        sys.exit(1)
-    except Exception as e:
-        logger.error(f"Error starting MCP server: {e}")
-        import traceback
-        traceback.print_exc()
-        sys.exit(1)
-
-
-if __name__ == "__main__":
-    main()
+#!/usr/bin/env python3
+"""
+MaTableGPT MCP Server Launcher
+==============================
+
+This script starts the MaTableGPT MCP service for extracting 
+table data from materials science literature.
+
+Usage:
+    python start_mcp.py [--host HOST] [--port PORT] [--mode MODE]
+    
+Arguments:
+    --host      Host address (default: 0.0.0.0)
+    --port      Port number (default: 7860)
+    --mode      Run mode: 'stdio' or 'sse' (default: sse for HuggingFace Space)
+
+Environment Variables:
+    LLM_API_KEY / OPENAI_API_KEY     - API key for LLM service
+    LLM_API_BASE / OPENAI_API_BASE   - Custom API base URL (for third-party services)
+    LLM_MODEL / OPENAI_MODEL         - Model name (default: gpt-4-turbo-preview)
+    MCP_HOST                          - Server host (default: 0.0.0.0)
+    MCP_PORT                          - Server port (default: 7860)
+"""
+
+import os
+import sys
+import argparse
+import logging
+
+# Add current directory to path for imports
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger("matablgpt-mcp-launcher")
+
+
+def check_environment():
+    """Check if required environment variables are set."""
+    warnings = []
+    
+    # Check for API key (support both naming conventions)
+    api_key = os.environ.get('LLM_API_KEY') or os.environ.get('OPENAI_API_KEY')
+    if not api_key:
+        warnings.append(
+            "LLM_API_KEY/OPENAI_API_KEY not set. GPT extraction features will not work. "
+            "Set it in HuggingFace Space secrets or environment variables."
+        )
+    
+    # Check for API base (for third-party services)
+    api_base = os.environ.get('LLM_API_BASE') or os.environ.get('OPENAI_API_BASE')
+    if api_base:
+        logger.info(f"Using custom API base: {api_base}")
+    
+    return warnings
+
+
+def check_dependencies():
+    """Check if required packages are installed."""
+    missing = []
+    
+    required = [
+        ('mcp', 'mcp[cli]'),
+        ('openai', 'openai'),
+        ('bs4', 'beautifulsoup4'),
+        ('pandas', 'pandas'),
+        ('lxml', 'lxml')
+    ]
+    
+    for module, package in required:
+        try:
+            __import__(module)
+        except ImportError:
+            missing.append(package)
+    
+    return missing
+
+
+def run_sse_server(host: str, port: int):
+    """Run MCP server in SSE mode using uvicorn."""
+    import uvicorn
+    from mcp_service import mcp
+    
+    # Get the SSE app from FastMCP
+    sse_app = mcp.sse_app()
+    
+    logger.info(f"Starting MCP SSE server on {host}:{port}")
+    logger.info(f"SSE Endpoint: http://{host}:{port}/sse")
+    
+    # Run with uvicorn
+    uvicorn.run(
+        sse_app,
+        host=host,
+        port=port,
+        log_level="info"
+    )
+
+
+def run_stdio_server():
+    """Run MCP server in stdio mode."""
+    from mcp_service import mcp
+    logger.info("Starting MCP server in stdio mode...")
+    mcp.run()
+
+
+def main():
+    """Main entry point."""
+    # Get default values from environment variables
+    default_host = os.environ.get('MCP_HOST', '0.0.0.0')
+    default_port = int(os.environ.get('MCP_PORT', '7860'))
+    
+    parser = argparse.ArgumentParser(
+        description="MaTableGPT MCP Server - Table Data Extraction from Materials Science Literature"
+    )
+    parser.add_argument(
+        '--host', 
+        default=default_host, 
+        help=f'Host address (default: {default_host})'
+    )
+    parser.add_argument(
+        '--port', 
+        type=int, 
+        default=default_port, 
+        help=f'Port number (default: {default_port})'
+    )
+    parser.add_argument(
+        '--mode',
+        choices=['stdio', 'sse'],
+        default='sse',
+        help='Run mode: stdio for standard I/O, sse for Server-Sent Events (default: sse)'
+    )
+    parser.add_argument(
+        '--debug',
+        action='store_true',
+        help='Enable debug logging'
+    )
+    
+    args = parser.parse_args()
+    
+    if args.debug:
+        logging.getLogger().setLevel(logging.DEBUG)
+    
+    # Check dependencies
+    missing = check_dependencies()
+    if missing:
+        logger.error(f"Missing required packages: {', '.join(missing)}")
+        logger.error(f"Install with: pip install {' '.join(missing)}")
+        sys.exit(1)
+    
+    # Check environment
+    warnings = check_environment()
+    for warning in warnings:
+        logger.warning(warning)
+    
+    # Display startup info
+    logger.info("=" * 60)
+    logger.info("MaTableGPT MCP Server")
+    logger.info("=" * 60)
+    logger.info(f"Mode: {args.mode}")
+    if args.mode == 'sse':
+        logger.info(f"Host: {args.host}")
+        logger.info(f"Port: {args.port}")
+        logger.info(f"SSE Endpoint: http://{args.host}:{args.port}/sse")
+    logger.info("=" * 60)
+    
+    # Run the server
+    try:
+        if args.mode == 'stdio':
+            run_stdio_server()
+        else:
+            run_sse_server(args.host, args.port)
+            
+    except ImportError as e:
+        logger.error(f"Failed to import required module: {e}")
+        logger.error("Make sure all dependencies are installed: pip install -r requirements.txt")
+        sys.exit(1)
+    except Exception as e:
+        logger.error(f"Error starting MCP server: {e}")
+        import traceback
+        traceback.print_exc()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()