Spaces:
Sleeping
Sleeping
| 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") | |
| 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) ────────────────────────── | |
| 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."}] | |
| } | |
| 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 | |
| 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) ────────────────── | |
| 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"] | |
| } | |
| 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 ───────────────────────────────────────────────────────────── | |
| def get_patients(): return PATIENTS | |
| async def upload_patient(file: UploadFile = File(...)): | |
| return {"id": "UP-" + str(int(time.time()) % 1000), "raw_text": (await file.read()).decode()} | |
| 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") | |