wehub-resource-sync
131 lines
3.8 KiB
Python
131 lines
3.8 KiB
Python
import os
|
|
from contextlib import asynccontextmanager
|
|
|
|
from fastapi import FastAPI, HTTPException
|
|
from pydantic import BaseModel
|
|
|
|
from agents import Agent, Runner
|
|
from agents.mcp import MCPServer, MCPServerManager, MCPServerStreamableHttp
|
|
from agents.model_settings import ModelSettings
|
|
|
|
MCP_SERVER_URL = os.getenv("MCP_SERVER_URL", "http://localhost:8000/mcp")
|
|
INACTIVE_MCP_SERVER_URL = os.getenv("INACTIVE_MCP_SERVER_URL", "http://localhost:8001/mcp")
|
|
APP_HOST = "127.0.0.1"
|
|
APP_PORT = 9001
|
|
USE_MCP_MANAGER = os.getenv("USE_MCP_MANAGER", "1") != "0"
|
|
|
|
|
|
class AddRequest(BaseModel):
|
|
a: int
|
|
b: int
|
|
|
|
|
|
class RunRequest(BaseModel):
|
|
input: str
|
|
|
|
|
|
class ReconnectRequest(BaseModel):
|
|
failed_only: bool = True
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
server = MCPServerStreamableHttp({"url": MCP_SERVER_URL})
|
|
inactive_server = MCPServerStreamableHttp({"url": INACTIVE_MCP_SERVER_URL})
|
|
servers = [server, inactive_server]
|
|
if USE_MCP_MANAGER:
|
|
async with MCPServerManager(
|
|
servers=servers,
|
|
connect_in_parallel=True,
|
|
) as manager:
|
|
app.state.mcp_manager = manager
|
|
app.state.mcp_servers = servers
|
|
yield
|
|
return
|
|
|
|
await server.connect()
|
|
app.state.mcp_servers = servers
|
|
app.state.active_servers = [server]
|
|
try:
|
|
yield
|
|
finally:
|
|
await server.cleanup()
|
|
|
|
|
|
app = FastAPI(lifespan=lifespan)
|
|
|
|
|
|
@app.get("/health")
|
|
async def health() -> dict[str, object]:
|
|
if USE_MCP_MANAGER:
|
|
manager: MCPServerManager = app.state.mcp_manager
|
|
return {
|
|
"connected_servers": [server.name for server in manager.active_servers],
|
|
"failed_servers": [server.name for server in manager.failed_servers],
|
|
}
|
|
|
|
active_servers = _get_active_servers()
|
|
return {
|
|
"connected_servers": [server.name for server in active_servers],
|
|
"failed_servers": [],
|
|
}
|
|
|
|
|
|
@app.get("/tools")
|
|
async def list_tools() -> dict[str, object]:
|
|
active_servers = _get_active_servers()
|
|
if not active_servers:
|
|
return {"tools": []}
|
|
tools = await active_servers[0].list_tools()
|
|
return {"tools": [tool.name for tool in tools]}
|
|
|
|
|
|
@app.post("/add")
|
|
async def add(req: AddRequest) -> dict[str, object]:
|
|
active_servers = _get_active_servers()
|
|
if not active_servers:
|
|
raise HTTPException(status_code=503, detail="No MCP servers available")
|
|
result = await active_servers[0].call_tool("add", {"a": req.a, "b": req.b})
|
|
return {"result": result.model_dump(mode="json")}
|
|
|
|
|
|
@app.post("/run")
|
|
async def run_agent(req: RunRequest) -> dict[str, object]:
|
|
if not os.getenv("OPENAI_API_KEY"):
|
|
raise HTTPException(status_code=400, detail="OPENAI_API_KEY is required")
|
|
|
|
servers = _get_active_servers()
|
|
if not servers:
|
|
raise HTTPException(status_code=503, detail="No MCP servers available")
|
|
|
|
agent = Agent(
|
|
name="FastAPI Agent",
|
|
instructions="Use the MCP tools when needed.",
|
|
mcp_servers=servers,
|
|
model_settings=ModelSettings(tool_choice="auto"),
|
|
)
|
|
result = await Runner.run(starting_agent=agent, input=req.input)
|
|
return {"output": result.final_output}
|
|
|
|
|
|
@app.post("/reconnect")
|
|
async def reconnect(req: ReconnectRequest) -> dict[str, object]:
|
|
if not USE_MCP_MANAGER:
|
|
raise HTTPException(status_code=400, detail="MCPServerManager is disabled")
|
|
manager: MCPServerManager = app.state.mcp_manager
|
|
servers = await manager.reconnect(failed_only=req.failed_only)
|
|
return {"connected_servers": [server.name for server in servers]}
|
|
|
|
|
|
def _get_active_servers() -> list[MCPServer]:
|
|
if USE_MCP_MANAGER:
|
|
manager: MCPServerManager = app.state.mcp_manager
|
|
return list(manager.active_servers)
|
|
return list(app.state.active_servers)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
|
|
uvicorn.run(app, host=APP_HOST, port=APP_PORT)
|