from __future__ import annotations import json, os, re, time, asyncio from pathlib import Path from typing import Any, Dict, List, Optional from datetime import datetime # Load .env _env_file = Path(__file__).resolve().parent.parent / ".env" if _env_file.exists(): for _line in _env_file.read_text().splitlines(): _line = _line.strip() if _line and not _line.startswith("#") and "=" in _line: _k, _, _v = _line.partition("=") os.environ.setdefault(_k.strip(), _v.strip()) from fastapi import FastAPI, HTTPException, UploadFile, File, Body, Header, Request, Response, Form from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse, FileResponse, StreamingResponse from fastapi.staticfiles import StaticFiles import requests as _requests _KEY = os.getenv("gemini_api_key", "") or os.getenv("GEMINI_API_KEY", "") BASE_DIR = Path(__file__).resolve().parent def _load_patients() -> List[Dict[str, Any]]: p = BASE_DIR / "sample_patients.json" if p.exists(): try: data = json.loads(p.read_text()) return data["patients"] if isinstance(data, dict) else data except: return [] return [] PATIENTS: List[Dict[str, Any]] = _load_patients() app = FastAPI(title="CodeBlue Context API", version="3.1.2") @app.exception_handler(Exception) async def global_exception_handler(request: Request, exc: Exception): import traceback return JSONResponse(status_code=500, content={"detail": str(exc), "traceback": traceback.format_exc()}) app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"]) # ── Gemini helper ───────────────────────────────────────────────────────────── def _gemini(prompt: str, retries: int = 4, override_key: str = None) -> Optional[Dict]: key = (override_key or _KEY).strip() if not key: return None payload = {"contents": [{"parts": [{"text": prompt}]}], "generationConfig": {"temperature": 0.1, "maxOutputTokens": 2048}} for model in ["gemini-2.5-flash-lite"]: for attempt in range(retries): try: url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={key}" resp = _requests.post(url, json=payload, timeout=25) if resp.status_code == 200: text = resp.json()["candidates"][0]["content"]["parts"][0]["text"].strip() text = re.sub(r"^```(?:json)?\s*", "", text); text = re.sub(r"\s*```$", "", text) try: return json.loads(text) except: m = re.search(r"\{.*\}", text, re.DOTALL) if m: return json.loads(m.group(0)) time.sleep(1) except: time.sleep(1) return None # ── Agents ──────────────────────────────────────────────────────────────────── def agent_triage(r, k=None): p = f"Triage PATIENT: {json.dumps(r)}\nReturn JSON: {{'severity_label':'string','critical_risks':[]}}" return _gemini(p, override_key=k) or {"severity_label":"URGENT","critical_risks":[]} def agent_action(r, t, k=None): p = f"Attending. PATIENT: {json.dumps(r)}\nTRIAGE: {json.dumps(t)}\nReturn JSON: {{'primary_directive':'string','clinical_summary':'string'}}" return _gemini(p, override_key=k) or {"primary_directive":"Stabilize.","clinical_summary":"Offline."} async def process_synthesis(text: str, key: str): record = {"raw_text": text} t = agent_triage(record, k=key) a = agent_action(record, t, k=key) return {"triage": t, "action": a} # ── MCP Standard (Optimized Discovery & Handshake) ────────────────────────── @app.api_route("/mcp", methods=["GET", "POST", "OPTIONS"]) async def mcp_discovery(req: Request): if req.method == "OPTIONS": return Response(status_code=204) # Discovery should point to a relative path for the SSE stream return { "transport": "sse", "endpoint": "/mcp/sse", "tools": [{"name": "analyze_clinical_record", "description": "Synthesize clinical records."}] } @app.get("/mcp/sse") async def mcp_sse_stream(req: Request): """Handshake with manual CORS and Proxy Fixes""" async def event_generator(): # Immediate endpoint event yield f"event: endpoint\ndata: /mcp/invoke\n\n" while True: await asyncio.sleep(20) yield ": keep-alive\n\n" response = StreamingResponse(event_generator(), media_type="text/event-stream") response.headers["Access-Control-Allow-Origin"] = "*" response.headers["Cache-Control"] = "no-cache" response.headers["Connection"] = "keep-alive" response.headers["X-Accel-Buffering"] = "no" # Fix for Hugging Face/Nginx buffering return response @app.post("/mcp/invoke") async def mcp_invoke(req: Request, x_gemini_key: str = Header(None)): try: body = await req.json() except: return {"error": "Invalid JSON"} method = body.get("method") params = body.get("params", {}) msg_id = body.get("id") if method == "initialize" or method == "list_tools" or method == "notifications/initialized": return { "jsonrpc": "2.0", "id": msg_id, "result": { "protocolVersion": "2024-11-05", "capabilities": { "tools": {"listChanged": False}, "resources": {"subscribe": False, "listChanged": False}, "prompts": {"listChanged": False} }, "serverInfo": { "name": "CodeBlue-Context", "version": "1.0.0" }, "tools": [ { "name": "analyze_clinical_record", "description": "Performs a 5-agent clinical synthesis of fragmented patient records to identify life-threatening risks.", "input_schema": { "type": "object", "properties": { "text": {"type": "string", "description": "Raw clinical data or patient history."} }, "required": ["text"] } } ] } } if method == "call_tool": args = params.get("arguments", {}) res = await process_synthesis(args.get("text", ""), x_gemini_key) return { "jsonrpc": "2.0", "id": msg_id, "result": {"content": [{"type": "text", "text": f"DIRECTIVE: {res['action']['primary_directive']}"}]} } return {"jsonrpc": "2.0", "result": {}, "id": msg_id} # ── A2A Standard (Strict Prompt Opinion Compliance) ────────────────── @app.get("/.well-known/agent-card.json") def get_agent_card(): """Official Prompt Opinion Discovery Card""" return { "name": "CodeBlue Context", "version": "1.0.0", "description": "ER Clinical Intelligence Synthesis Agent.", "capabilities": ["clinical_synthesis"], "endpoints": { "invoke": "https://anil988-codeblue-context.hf.space/a2a/invoke" }, "protocols": ["A2A", "MCP"], "data_standards": ["FHIR-R4"] } @app.post("/a2a/invoke") async def a2a_invoke(request: Dict[str, Any], x_gemini_key: str = Header(None)): p = request.get("params", {}) res = await process_synthesis(p.get("text", ""), x_gemini_key) return {"jsonrpc": "2.0", "result": {"directive": res["action"]["primary_directive"]}, "id": request.get("id")} # ── App Endpoints ───────────────────────────────────────────────────────────── @app.get("/patients") def get_patients(): return PATIENTS @app.post("/patients/upload") async def upload_patient(file: UploadFile = File(...)): return {"id": "UP-" + str(int(time.time()) % 1000), "raw_text": (await file.read()).decode()} @app.post("/analyze") async def analyze_patient(data: Dict[str, Any], x_gemini_key: str = Header(None)): res = await process_synthesis(data.get("raw_text", ""), x_gemini_key) return {"triage": res["triage"], "action": res["action"], "contraindication": {"treatment_verdict":"SAFE"}, "labs": {"critical_labs":[]}, "differentials": {"differentials":[]}} if os.path.exists("frontend/dist"): app.mount("/", StaticFiles(directory="frontend/dist", html=True), name="frontend")