server.py•5.37 kB
#!/usr/bin/env python3
from __future__ import annotations
import os
from pathlib import Path
from typing import List
from contextlib import asynccontextmanager
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from mcp.server.fastmcp import FastMCP
from .datasets import DatasetLoadError, get_metadata, metadata_keys
from .glossary import get_glossary_term as _get_glossary_term
from .glossary import list_glossary_terms as _list_glossary_terms
from .inhibitors import find_inhibitor_plants as _find_inhibitor_plants
from .inhibitors import get_inhibitor_sources as _get_inhibitor_sources
from .inhibitors import list_inhibitors as _list_inhibitors_data
from .inhibitors import summarize_inhibitors as _summarize_inhibitors_data
from .session import build_session_snapshot as _build_session_snapshot
# ---------------------------
# FastMCP server setup
# ---------------------------
mcp = FastMCP("aurora-mcp")
mcp.settings.streamable_http_path = "/"
_mcp_http_app = mcp.streamable_http_app()
@asynccontextmanager
async def _lifespan(app: FastAPI):
"""Run the MCP session manager for the lifetime of the FastAPI app."""
async with mcp.session_manager.run():
yield
# ---------------------------
# FastAPI app + CORS (optional)
# ---------------------------
app = FastAPI(title="Aurora-MCP (HTTP)", version="1.0", lifespan=_lifespan)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"],
expose_headers=["Mcp-Session-Id"],
)
# ---------------------------
# Health endpoint - MOVE BEFORE MOUNTING MCP
# ---------------------------
@app.get("/healthz")
def healthz():
return {
"ok": True,
"mcp": "mounted at /",
"tools": sorted([t.__name__ for t in _EXPORTED_TOOL_FUNCS]),
}
# ---- Example tools
@mcp.tool()
def list_files(path: str = ".") -> str:
"""List files in a directory. Argument: path (default '.')"""
try:
entries = sorted(os.listdir(path))
return "\n".join(entries)
except Exception as exc:
return f"ERROR: {exc}"
@mcp.tool()
def read_text(path: str) -> str:
"""Read a small UTF-8 text file. Argument: path (required)"""
p = Path(path)
try:
return p.read_text(encoding="utf-8")
except Exception as exc:
return f"ERROR: {exc}"
@mcp.tool()
def list_dataset_metadata() -> List[str]:
"""List dataset metadata keys derived from data/*.meta.json."""
return metadata_keys()
@mcp.tool()
def get_dataset_metadata(name: str) -> dict:
"""Return dataset metadata for the given key (see list_dataset_metadata)."""
meta = get_metadata(name)
if meta is None:
return {"error": f"Unknown dataset '{name}'", "available": metadata_keys()}
return meta
@mcp.tool()
def summarize_inhibitors():
"""Summarize Complex I inhibitor counts by status/confidence."""
try:
return _summarize_inhibitors_data()
except DatasetLoadError as exc:
return {"error": str(exc)}
@mcp.tool()
def list_inhibitors(known_status: str | None = None, confidence: str | None = None, limit: int | None = 20):
"""List Complex I inhibitors with optional known_status/confidence filters."""
try:
return _list_inhibitors_data(known_status=known_status, confidence=confidence, limit=limit)
except DatasetLoadError as exc:
return {"error": str(exc)}
except ValueError as exc:
return {"error": str(exc)}
@mcp.tool()
def get_inhibitor_sources(compound: str):
"""Fetch PubMed sources (IDs/URLs) for a given inhibitor compound."""
try:
return _get_inhibitor_sources(compound)
except DatasetLoadError as exc:
return {"error": str(exc)}
except ValueError as exc:
return {"error": str(exc)}
@mcp.tool()
def find_compound_plants(compound: str, scope: str = "global"):
"""List plant organisms associated with a compound in the requested scope."""
try:
return _find_inhibitor_plants(compound, scope=scope)
except DatasetLoadError as exc:
return {"error": str(exc)}
except ValueError as exc:
return {"error": str(exc)}
@mcp.tool()
def list_glossary():
"""List available domain glossary keys."""
return _list_glossary_terms()
@mcp.tool()
def get_glossary(key: str):
"""Retrieve glossary information for a given abbreviation or alias."""
result = _get_glossary_term(key)
if result is None:
return {"error": f"No glossary entry found for '{key}'", "available": _list_glossary_terms()}
return result
@mcp.tool()
def init_session():
"""Bootstrap the session with glossary and dataset guidance."""
return _build_session_snapshot()
_EXPORTED_TOOL_FUNCS: List[callable] = [
list_files,
read_text,
list_dataset_metadata,
get_dataset_metadata,
summarize_inhibitors,
list_inhibitors,
get_inhibitor_sources,
find_compound_plants,
list_glossary,
get_glossary,
init_session,
]
# Mount the MCP streamable HTTP app at root - THIS MUST BE LAST
app.router.redirect_slashes = False
app.mount("/", _mcp_http_app)
# ---------------------------
# Entrypoint (local dev)
# ---------------------------
if __name__ == "__main__":
import uvicorn
port = int(os.getenv("PORT", "7860"))
uvicorn.run(app, host="0.0.0.0", port=port)