Blame for backend/main.py

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")