Spaces:
Paused
Paused
File size: 3,562 Bytes
f647629 ceeb737 f647629 ceeb737 f647629 ceeb737 f647629 ceeb737 f647629 ceeb737 f647629 2325f96 ceeb737 f647629 ceeb737 f647629 ceeb737 |
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 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
#!/usr/bin/env python3
"""
HuggingFace Spaces entry point for the Weights & Biases MCP Server.
This script creates a FastAPI application with a landing page and mounts
the MCP server from the existing server module.
"""
import os
import sys
import logging
from pathlib import Path
from contextlib import asynccontextmanager
# Add the src directory to Python path
sys.path.insert(0, str(Path(__file__).parent / "src"))
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi.middleware.cors import CORSMiddleware
import uvicorn
# Import the existing MCP server instance and utilities
from wandb_mcp_server.server import mcp
from wandb_mcp_server.utils import get_rich_logger
# Configure logging for the app
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = get_rich_logger("huggingface-spaces-app")
# Lifecycle manager for the app
@asynccontextmanager
async def lifespan(app: FastAPI):
"""Manage application lifecycle."""
# Startup
logger.info("Starting Weights & Biases MCP Server on HuggingFace Spaces")
# Check for required environment variables
wandb_api_key = os.environ.get("WANDB_API_KEY")
if not wandb_api_key:
logger.warning("WANDB_API_KEY environment variable is not set!")
logger.warning("Please set your Weights & Biases API key in the Space's environment variables.")
logger.warning("You can get your API key from: https://wandb.ai/authorize")
else:
logger.info("WANDB_API_KEY configured: Yes")
logger.info(f"Starting HTTP server on port {os.environ.get('PORT', '7860')}")
logger.info("Landing page available at: /")
logger.info("MCP endpoint available at: /mcp/sse")
yield
# Shutdown
logger.info("Shutting down Weights & Biases MCP Server")
# Create FastAPI app with lifecycle management
app = FastAPI(
title="Weights & Biases MCP Server",
description="Model Context Protocol server for querying W&B data",
lifespan=lifespan
)
# Add CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Mount static files (create directory if it doesn't exist)
static_dir = Path(__file__).parent / "static"
static_dir.mkdir(exist_ok=True)
app.mount("/static", StaticFiles(directory=str(static_dir)), name="static")
# Set up templates
templates_dir = Path(__file__).parent / "templates"
templates = Jinja2Templates(directory=str(templates_dir))
@app.get("/", response_class=HTMLResponse)
async def index(request: Request):
"""Serve the landing page."""
return templates.TemplateResponse("index.html", {"request": request})
@app.get("/health")
async def health():
"""Health check endpoint."""
return {"status": "healthy", "service": "wandb-mcp-server"}
# Mount the existing MCP server at /mcp path
# The MCP server from server.py already has all tools registered
mcp.run(
transport="streamable-http",
http_app=app,
http_path="/mcp"
)
def main():
"""Main entry point for HuggingFace Spaces."""
# Run the FastAPI app with Uvicorn
port = int(os.environ.get("PORT", "7860"))
host = os.environ.get("HOST", "0.0.0.0")
uvicorn.run(
"app:app",
host=host,
port=port,
log_level="info",
reload=False
)
if __name__ == "__main__":
main() |