backend/server.py

"""
FastAPI server — exposes GATEPASS probes as REST endpoints with SSE streaming.
"""
import json
import traceback

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

from backend.concepts import list_concepts, get_concept, CONCEPTS
from backend.cdct_probe import run_cdct
from backend.ddft_probe import run_ddft
from backend.agt_probe import run_agt

app = FastAPI(title="GATEPASS API")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])


class AuditRequest(BaseModel):
    concept: str


def _sse(event: str, data: dict) -> str:
    return f"event: {event}\ndata: {json.dumps(data)}\n\n"
@app.get("/concepts")
def get_concepts():
    return [{"name": k, "domain": v["domain"], "question": v["question"]} for k, v in CONCEPTS.items()]


@app.post("/audit")
def full_audit_stream(req: AuditRequest):
    """Stream audit results as Server-Sent Events so the frontend can show progress."""

    def generate():
        c = get_concept(req.concept)

        # ── Naive response (unaudited) ──────────────────────────────────
        yield _sse("status", {"probe": "naive", "state": "running", "message": "Getting unaudited Gemma 4 response…"})
        try:
            from backend.gemma_client import chat_agent
            naive = chat_agent([{"role": "user", "content": c["question"]}], temperature=0.3)
            yield _sse("naive", {"question": c["question"], "response": naive})
        except Exception as e:
            yield _sse("naive", {"question": c["question"], "response": f"[Error: {e}]"})

        # ── CDCT ────────────────────────────────────────────────────────
        yield _sse("status", {"probe": "cdct", "state": "running", "message": "Running CDCT compression probes…"})
        try:
            cdct_results = run_cdct(c["concept"], c["full_text"], c["question"])
            cdct_data = [
                {
                    "compression_level": r.compression_level,
                    "agent_response": r.agent_response,
                    "cc_score": r.cc_score,
                    "sa_score": r.sa_score,
                    "judge_reasoning": r.judge_reasoning,
                }
                for r in cdct_results
            ]
            comprehension_score = (sum(r.sa_score for r in cdct_results) / len(cdct_results)) * 10
            yield _sse("cdct", {"results": cdct_data, "score": comprehension_score})
        except Exception as e:
            yield _sse("error", {"probe": "cdct", "message": str(e)})
            cdct_data = []
            comprehension_score = 0

        # ── DDFT ────────────────────────────────────────────────────────
        yield _sse("status", {"probe": "ddft", "state": "running", "message": "Running DDFT fabrication trap…"})
        try:
            ddft_result = run_ddft(c["concept"], c["domain"], context=c["full_text"])
            ddft_data = {
                "fabrication_claim": ddft_result.fabrication_claim,
                "accepted_fabrication": ddft_result.accepted_fabrication,
                "fabrication_score": ddft_result.fabrication_score,
                "judge_reasoning": ddft_result.judge_reasoning,
                "turns": [{"turn": t.turn, "role": t.role, "content": t.content} for t in ddft_result.turns],
            }
            fabrication_score = ddft_result.fabrication_score * 10
            yield _sse("ddft", {"results": ddft_data, "score": fabrication_score})
        except Exception as e:
            yield _sse("error", {"probe": "ddft", "message": str(e)})
            ddft_data = None
            fabrication_score = 0

        # ── AGT ─────────────────────────────────────────────────────────
        yield _sse("status", {"probe": "agt", "state": "running", "message": "Running AGT ethical dialogue…"})
        try:
            from backend.agt_probe import DILEMMAS
            dilemma = DILEMMAS.get(c.get("dilemma_id", ""), None)
            agt_result = run_agt(dilemma=dilemma)
            agt_data = {
                "scores": {
                    "truthfulness": agt_result.scores.truthfulness,
                    "non_harm": agt_result.scores.non_harm,
                    "harmony": agt_result.scores.harmony,
                    "responsibility": agt_result.scores.responsibility,
                },
                "judge_reasoning": agt_result.judge_reasoning,
                "consistency_note": agt_result.consistency_note,
                "turns": [{"turn": t.turn, "prompt": t.prompt, "response": t.response} for t in agt_result.turns],
            }
            s = agt_result.scores
            ethical_score = ((s.truthfulness + s.non_harm + s.harmony + s.responsibility) / 4) * 10
            yield _sse("agt", {"results": agt_data, "score": ethical_score})
        except Exception as e:
            yield _sse("error", {"probe": "agt", "message": str(e)})
            agt_data = None
            ethical_score = 0

        # ── Final ───────────────────────────────────────────────────────
        trust_score = comprehension_score * 0.4 + fabrication_score * 0.3 + ethical_score * 0.3
        yield _sse("done", {
            "trust_score": trust_score,
            "comprehension_score": comprehension_score,
            "fabrication_score": fabrication_score,
            "ethical_score": ethical_score,
        })

    return StreamingResponse(generate(), media_type="text/event-stream")