FastAPI Integration Examples

GUARDAI_SESSION_SECRET=$(openssl rand -base64 32) \
  docker compose -f ../../deploy/docker/docker-compose.yml up -d

docker run -p 3000:3000 -e GUARDAI_HOST=0.0.0.0 -e GUARDAI_ALLOW_DEV_NON_LOOPBACK=1 \
  ghcr.io/oronts/oguardai/server:0.1.0

pip install fastapi uvicorn httpx

export GUARDAI_BASE_URL="http://localhost:3000"
export GUARDAI_API_KEY="<key>"          # only if the server enforces auth
uvicorn main:app --port 9000

for tok in demo-customer-lena demo-agent-mara demo-auditor-per demo-analyst-nore; do
  echo "== $tok =="
  curl -s -X POST http://localhost:9000/support/messages/TW-58120/reply \
    -H "Authorization: Bearer $tok" | python3 -c 'import sys,json;print(json.load(sys.stdin)["reply"])'
done

# Ad-hoc message instead of an inbox item (any authenticated role)
curl -s -X POST http://localhost:9000/support/reply \
  -H "Authorization: Bearer demo-agent-mara" -H "Content-Type: application/json" \
  -d '{"message":"Hallo, hier ist Frau Erika Mustermann, Bestellung ORD-2026-5583, Erstattung auf DE97 9900 0000 0123 4567 89.","subject":"Erstattung"}'

# Detect-only triage (staff only, no restoration)
curl -s -X POST http://localhost:9000/support/triage \
  -H "Authorization: Bearer demo-agent-mara" -H "Content-Type: application/json" \
  -d '{"message":"Widerruf ORD-2026-5610, Herr Max Mustermann, max.mustermann@example.de","subject":"Widerruf"}'

# Health (open) and inbox (staff only)
curl -s http://localhost:9000/healthz
curl -s http://localhost:9000/support/inbox -H "Authorization: Bearer demo-agent-mara"

{"message_id": "TW-58120", "customer_subject": "cust_lena", "received_at": "2026-07-08T07:14:00+02:00", "subject": "Fehlbuchung bei Erstattung", "body": "Sehr geehrtes Team, mein Name ist Frau Erika Mustermann. Zu meiner Bestellung ORD-2026-5583 wurde mir versehentlich der Betrag 189,00 EUR doppelt abgebucht. Bitte erstatten Sie die Fehlbuchung auf mein Konto DE97 9900 0000 0123 4567 89. Sie erreichen mich unter erika.mustermann@example.de oder telefonisch unter 0176 44229180. Mit freundlichen Gruessen, Erika Mustermann"}
{"message_id": "TW-58121", "customer_subject": "cust_jonas", "received_at": "2026-07-08T07:41:00+02:00", "subject": "Lieferadresse aendern", "body": "Guten Tag, hier ist Herr Dr. Max Mustermann. Ich bin umgezogen und moechte die Lieferadresse fuer meine Bestellung ORD-2026-5591 auf Musterstrasse 8, 70173 Stuttgart aendern. Bitte bestaetigen Sie die Aenderung an max.mustermann@example.de. Vielen Dank, Max Mustermann"}
{"message_id": "TW-58122", "customer_subject": "cust_miriam", "received_at": "2026-07-08T08:03:00+02:00", "subject": "Wo bleibt meine Lieferung", "body": "Hallo, ich bin Frau Erika Mustermann und warte seit zehn Tagen auf meine Bestellung ORD-2026-5604. Koennen Sie mir bitte den Sendungsstatus mitteilen? Meine Kontaktdaten: erika.mustermann@example.de, 0151 33810274. Freundliche Gruesse"}
{"message_id": "TW-58123", "customer_subject": "cust_samuel", "received_at": "2026-07-08T08:26:00+02:00", "subject": "Widerruf und Rueckzahlung", "body": "Sehr geehrte Damen und Herren, ich moechte meine Bestellung ORD-2026-5610 widerrufen und bitte um Rueckzahlung des Kaufpreises von 74,50 EUR auf DE97 9900 0000 0123 4567 89. Mein Name ist Herr Max Mustermann, erreichbar unter max.mustermann@example.de. Mit freundlichen Gruessen"}

fastapi>=0.100.0
uvicorn>=0.20.0
httpx>=0.24.0

from __future__ import annotations

import json
import logging
import os
import re
from contextlib import asynccontextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import AsyncIterator

import httpx
from fastapi import Depends, FastAPI, HTTPException, Request, Security
from fastapi.responses import JSONResponse
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from pydantic import BaseModel

# Accept whichever base-URL var is set so the example drops into any of the repo's
# env conventions without edits.
GUARDAI_URL = (
    os.getenv("GUARDAI_BASE_URL")
    or os.getenv("GUARDAI_URL")
    or os.getenv("GUARDAI_BASE")
    or "http://localhost:3000"
).rstrip("/")
GUARDAI_API_KEY = os.getenv("GUARDAI_API_KEY", "")

# One vertical policy governs every request: German formal restore, gender/formality
# exposed, IBAN masked even under the widest restore. `de` pins honorifics to German.
POLICY = os.getenv("GUARDAI_POLICY", "german-support")
LANGUAGE = os.getenv("GUARDAI_LANGUAGE", "de")

# Presence of a provider key only flips the label; the model call is a deterministic
# offline stand-in so the example runs without network egress. See draft_reply().
LLM_LABEL = "OpenAI (token-only)" if os.getenv("OPENAI_API_KEY") else "offline template drafter"

HERE = Path(__file__).resolve().parent
TOKEN_RE = re.compile(r"\{\{([^}]+)\}\}")

log = logging.getLogger("musterfirma_support")

class GuardError(Exception):
    """A transform/rehydrate call that did not succeed. Carries the server's typed
    error code so the endpoint can fail closed instead of processing raw text."""

    def __init__(self, code: str, message: str) -> None:
        super().__init__(f"{code}: {message}")
        self.code = code
        self.message = message


@dataclass
class TransformResult:
    safe_text: str
    session_state: str
    entities: list[dict]
    entity_context: list[dict]


@dataclass
class Guard:
    """Request-scoped facade over the pooled OGuardAI client. Endpoints depend on
    this, never on httpx directly, so the whole app shares one connection pool and
    one fail-closed error surface."""

    client: httpx.AsyncClient
    policy: str = POLICY
    language: str = LANGUAGE

    async def _post(self, path: str, body: dict) -> dict:
        try:
            resp = await self.client.post(path, json=body)
        except httpx.HTTPError as exc:  # network / timeout: fail closed, never fall back to raw
            # The exception detail (internal host/URL) stays in the server log only.
            log.error("cannot reach OGuardAI: %s", exc)
            raise GuardError("GUARDAI_UNAVAILABLE", "cannot reach OGuardAI") from exc
        if resp.status_code >= 400:
            detail = resp.json() if resp.headers.get("content-type", "").startswith("application/json") else {}
            # Never fall back to the raw upstream body: it can echo request PII. Keep OGuardAI's own
            # structured error string when present, else a generic status-only message.
            error = detail.get("error")
            message = error if isinstance(error, str) else f"OGuardAI request failed (HTTP {resp.status_code})"
            raise GuardError(detail.get("code", f"HTTP_{resp.status_code}"), message)
        return resp.json()

    async def transform(self, text: str) -> TransformResult:
        data = await self._post(
            "/v1/transform",
            {"input": text, "input_type": "text", "language": self.language, "policy": self.policy},
        )
        return TransformResult(
            safe_text=data["safe_text"],
            session_state=data["session_state"],
            entities=data.get("entities", []),
            entity_context=data.get("entity_context", []),
        )

    async def detect(self, text: str) -> dict:
        # /v1/detect rejects unknown fields; it takes no input_type. Triage only needs entity
        # types and counts, so include_values=false keeps raw PII out of the detect response.
        return await self._post(
            "/v1/detect",
            {"input": text, "language": self.language, "policy": self.policy, "include_values": False},
        )

    async def rehydrate(
        self,
        output: str,
        session_state: str,
        *,
        channel: str | None = None,
        restore_mode: str | None = None,
    ) -> str:
        body: dict = {"output": output, "session_state": session_state}
        if channel:
            body["output_channel"] = channel
        if restore_mode:
            body["restore_mode"] = restore_mode
        return (await self._post("/v1/rehydrate", body))["restored_text"]

@dataclass(frozen=True)
class Audience:
    role: str
    channel: str | None
    restore_mode: str | None
    description: str


# One drafted reply, four sinks, from ONE sealed session. Restore is controlled by the
# output channel AND the restore mode, so the same tokens resolve differently per caller:
#   customer -> channel customer_email + full: their own name and email come back; the
#              IBAN stays masked because the policy caps it (a refund never echoes the
#              full account over email).
#   agent    -> channel internal_summary: real name for the console, contact masked.
#   auditor  -> channel log_safe: every value redacted for the immutable log.
#   analyst  -> restore_mode abstract (no channel): entity TYPES only, "(Kundin)".
AUDIENCES: dict[str, Audience] = {
    "customer": Audience("customer", "customer_email", "full", "reply sent to the customer (name + email restored, IBAN masked)"),
    "agent": Audience("agent", "internal_summary", None, "agent console (real name, contact masked)"),
    "auditor": Audience("auditor", "log_safe", None, "immutable audit log (all PII redacted)"),
    "analyst": Audience("analyst", None, "abstract", "analytics pipeline (entity types only, no raw values)"),
}


@dataclass(frozen=True)
class Principal:
    subject: str
    role: str


# Demo bearer tokens mapped to a principal. In production replace this table lookup with
# your IdP token verification: the role must come from the verified credential, never from
# caller-supplied input. The tokens are intentionally human-readable, not secrets.
PRINCIPALS: dict[str, Principal] = {
    "demo-customer-lena": Principal("cust_lena", "customer"),
    "demo-customer-jonas": Principal("cust_jonas", "customer"),
    "demo-agent-mara": Principal("agent_mara", "agent"),
    "demo-auditor-per": Principal("auditor_per", "auditor"),
    "demo-analyst-nore": Principal("analyst_nore", "analyst"),
}

STAFF_ROLES = frozenset({"agent", "auditor", "analyst"})

_bearer = HTTPBearer(auto_error=False)


def get_principal(cred: HTTPAuthorizationCredentials | None = Security(_bearer)) -> Principal:
    """Resolve the caller to a Principal or fail closed with 401. A missing credential and an
    unknown token take the same lookup path and detail, so neither response nor timing leaks it."""
    principal = PRINCIPALS.get(cred.credentials if cred else "")
    if principal is None:
        raise HTTPException(
            status_code=401,
            detail="authentication required",
            headers={"WWW-Authenticate": "Bearer"},
        )
    return principal


def resolve_audience(principal: Principal = Depends(get_principal)) -> Audience:
    """FastAPI dependency: the restore policy follows the authenticated role, nothing else."""
    return AUDIENCES[principal.role]


def require_staff(principal: Principal = Depends(get_principal)) -> Principal:
    """Gate a staff-only resource. Authenticated but non-staff roles get 403, not 401."""
    if principal.role not in STAFF_ROLES:
        raise HTTPException(status_code=403, detail="staff role required")
    return principal

INTENTS = [
    ("refund", ("fehlbuchung", "erstattung", "rueckzahlung", "widerruf", "gutschrift")),
    ("address_change", ("adresse", "anschrift", "lieferadresse", "umzug", "umgezogen")),
    ("order_status", ("sendungsstatus", "lieferung", "wo bleibt", "versand", "tracking")),
]

# The person token restores WITH the German honorific (Herr/Frau) from gender metadata.
# The drafter only has to pick the adjective ending that AGREES: "geehrter" for a man,
# "geehrte" for a woman. That agreement is exactly what a gender-blind drafter gets wrong.
SALUTATION = {"male": "Sehr geehrter", "female": "Sehr geehrte"}


def first_token(safe_text: str, prefix: str) -> str:
    """First `{{prefix...}}` token in a safe text, hash included. The token must be
    reused verbatim; short forms like `{{person:p_001}}` do not resolve."""
    for inner in TOKEN_RE.findall(safe_text):
        if inner.startswith(prefix):
            return "{{" + inner + "}}"
    return ""


def classify(safe_text: str, subject: str) -> str:
    hay = f"{subject} {safe_text}".lower()
    for intent, keys in INTENTS:
        if any(k in hay for k in keys):
            return intent
    return "general"


def gender_of(result: TransformResult, person_token: str) -> str:
    for ent in result.entity_context:
        if ent.get("token") == person_token:
            return ent.get("gender", "unknown")
    return "unknown"


def draft_reply(result: TransformResult, subject: str) -> str:
    """Build a token-only German reply. Every `{{...}}` is a placeholder the rehydrator
    fills per audience; the drafter never handles a raw value."""
    person = first_token(result.safe_text, "person:")
    order = first_token(result.safe_text, "order:")
    iban = first_token(result.safe_text, "iban:")
    email = first_token(result.safe_text, "email:")

    gender = gender_of(result, person)
    opener = SALUTATION.get(gender)
    salutation = f"{opener} {person}" if (opener and person) else (f"Guten Tag {person}" if person else "Sehr geehrte Damen und Herren")
    zu_ref = f" zu Bestellung {order}" if order else ""       # "... zu Bestellung ORD-..."
    fuer_ref = f" fuer Bestellung {order}" if order else ""   # "... fuer Bestellung ORD-..."
    order_ref = f" {order}" if order else ""                  # bare token after "Bestellung"

    intent = classify(result.safe_text, subject)
    if intent == "refund":
        pay = f" Die Rueckzahlung erfolgt auf {iban}." if iban else ""
        body = f"wir haben Ihre Anfrage{zu_ref} geprueft und veranlassen die Korrektur.{pay}"
    elif intent == "address_change":
        body = f"die Lieferadresse{fuer_ref} ist aktualisiert und ab sofort aktiv."
    elif intent == "order_status":
        body = f"Ihre Bestellung{order_ref} befindet sich im Versand. Sie erhalten die Sendungsverfolgung in Kuerze."
    else:
        body = f"vielen Dank fuer Ihre Nachricht{zu_ref}. Ein Mitarbeiter meldet sich zeitnah bei Ihnen."

    # A confirmation-to-email line. This is where customer vs agent restore diverges:
    # the customer channel returns the real address, internal_summary masks it.
    confirm = f" Eine Bestaetigung senden wir an {email}." if email else ""
    return f"{salutation}, {body}{confirm} Mit freundlichen Gruessen, Ihr Musterfirma Kundenservice"

@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
    headers = {"X-API-Key": GUARDAI_API_KEY} if GUARDAI_API_KEY else {}
    app.state.guardai = httpx.AsyncClient(base_url=GUARDAI_URL, timeout=30.0, headers=headers)
    app.state.inbox = _load_inbox()
    try:
        yield
    finally:
        await app.state.guardai.aclose()


def _load_inbox() -> dict[str, dict]:
    path = HERE / "support_inbox.jsonl"
    with path.open(encoding="utf-8") as fh:
        return {row["message_id"]: row for row in (json.loads(line) for line in fh if line.strip())}


app = FastAPI(title="Musterfirma Guarded Support Desk", lifespan=lifespan)


# The output guard keeps two outcomes apart and the handler must not blur them: a content
# verdict cannot be fixed by resending the same draft, while a scanner outage is retriable once
# the deployment recovers. Every other guard failure stays a generic 502 gateway error.
_OUTPUT_STATUS = {"GUARDAI_OUTPUT_BLOCKED": 422, "GUARDAI_OUTPUT_GUARD_UNAVAILABLE": 503}


@app.exception_handler(GuardError)
async def _guard_error_handler(_request: Request, exc: GuardError) -> JSONResponse:
    # Fail closed: the business logic upstream never saw raw text, and no partial reply is sent.
    status = _OUTPUT_STATUS.get(exc.code, 502)
    return JSONResponse(status_code=status, content={"error": "guard_failed", "code": exc.code, "detail": exc.message})


async def get_client(request: Request) -> httpx.AsyncClient:
    return request.app.state.guardai


async def get_guard(client: httpx.AsyncClient = Depends(get_client)) -> Guard:
    return Guard(client=client)

class ReplyRequest(BaseModel):
    message: str
    subject: str = ""


class ReplyResponse(BaseModel):
    audience: str
    channel: str | None
    restore_mode: str | None
    audience_note: str
    intent: str
    drafted_by: str
    entities_detected: int
    safe_text: str          # exactly what left the trust boundary (tokens + abstractions)
    token_only_draft: str   # the model's output, before rehydration
    reply: str              # rehydrated for THIS audience


class TriageResponse(BaseModel):
    intent: str
    entity_types: dict[str, int]
    detectors_used: list[str]
    detection_time_ms: float


async def _guarded_reply(guard: Guard, audience: Audience, message: str, subject: str) -> ReplyResponse:
    # 1. Transform: PII -> tokens. Fail-closed via GuardError -> HTTP 502.
    result = await guard.transform(message)

    # 2. Draft on tokens only. In production this is the single LLM call.
    draft = draft_reply(result, subject)

    # 3. Rehydrate for the caller's audience from the same sealed session.
    reply = await guard.rehydrate(
        draft,
        result.session_state,
        channel=audience.channel,
        restore_mode=audience.restore_mode,
    )

    return ReplyResponse(
        audience=audience.role,
        channel=audience.channel,
        restore_mode=audience.restore_mode,
        audience_note=audience.description,
        intent=classify(result.safe_text, subject),
        drafted_by=LLM_LABEL,
        entities_detected=len(result.entities),
        safe_text=result.safe_text,
        token_only_draft=draft,
        reply=reply,
    )


@app.get("/healthz")
async def healthz(guard: Guard = Depends(get_guard)) -> dict:
    try:
        health = (await guard.client.get("/v1/health")).json()
    except httpx.HTTPError:
        return {"status": "unhealthy", "detail": f"cannot reach OGuardAI at {GUARDAI_URL}"}
    detector = health.get("components", {}).get("detector", {}).get("message", "")
    return {
        "status": health.get("status", "unknown"),
        "policy": POLICY,
        "ner_active": detector.startswith("builtin_and_ner"),
        "detector": detector,
    }


@app.get("/support/inbox")
async def inbox(request: Request, _principal: Principal = Depends(require_staff)) -> list[dict]:
    # List metadata only; the raw bodies stay put until a request tokenizes them.
    return [
        {"message_id": m["message_id"], "subject": m["subject"], "received_at": m["received_at"]}
        for m in request.app.state.inbox.values()
    ]


@app.post("/support/messages/{message_id}/reply", response_model=ReplyResponse)
async def reply_to_inbox_message(
    message_id: str,
    request: Request,
    guard: Guard = Depends(get_guard),
    principal: Principal = Depends(get_principal),
    audience: Audience = Depends(resolve_audience),
) -> ReplyResponse:
    row = request.app.state.inbox.get(message_id)
    # A message owned by another customer answers exactly like a missing one, so a customer
    # credential can never probe which message ids belong to someone else.
    if row is None or (principal.role == "customer" and row.get("customer_subject") != principal.subject):
        raise HTTPException(status_code=404, detail=f"no inbox message '{message_id}'")
    return await _guarded_reply(guard, audience, row["body"], row["subject"])


@app.post("/support/reply", response_model=ReplyResponse)
async def reply_to_message(
    req: ReplyRequest,
    guard: Guard = Depends(get_guard),
    audience: Audience = Depends(resolve_audience),
) -> ReplyResponse:
    return await _guarded_reply(guard, audience, req.message, req.subject)


@app.post("/support/triage", response_model=TriageResponse)
async def triage(
    req: ReplyRequest,
    guard: Guard = Depends(get_guard),
    _principal: Principal = Depends(require_staff),
) -> TriageResponse:
    # Detect-only: route/classify the message without restoring anything.
    detected = await guard.detect(req.message)
    counts: dict[str, int] = {}
    for ent in detected.get("entities", []):
        etype = ent["type"]
        etype = etype if isinstance(etype, str) else next(iter(etype.values()))
        counts[etype] = counts.get(etype, 0) + 1
    stats = detected.get("stats", {})
    return TriageResponse(
        intent=classify(req.message, req.subject),
        entity_types=counts,
        detectors_used=stats.get("detectors_used", []),
        detection_time_ms=stats.get("detection_time_ms", 0.0),
    )

== customer ==
Sehr geehrte Frau Erika Mustermann, wir haben Ihre Anfrage zu Bestellung ORD-2026-5583
geprueft und veranlassen die Korrektur. Die Rueckzahlung erfolgt auf
D*** **** **** **** **** *0. Eine Bestaetigung senden wir an erika.mustermann@example.de.
Mit freundlichen Gruessen, Ihr Musterfirma Kundenservice

== agent ==
Sehr geehrte Frau Erika Mustermann, ... Die Rueckzahlung erfolgt auf
D*** **** **** **** **** *0. Eine Bestaetigung senden wir an e********************e. ...

== auditor ==
Sehr geehrte [REDACTED], ... Die Rueckzahlung erfolgt auf [REDACTED].
Eine Bestaetigung senden wir an [REDACTED]. ...

== analyst ==
Sehr geehrte (Kundin), ... Die Rueckzahlung erfolgt auf (iban on file).
Eine Bestaetigung senden wir an (email on file). ...

{
  "intent": "refund",
  "entity_types": {"order": 1, "iban": 1, "person": 1, "email": 1},
  "detectors_used": ["builtin_regex", "python_ner"],
  "detection_time_ms": 97.9
}

pip install -r requirements.txt

export GUARDAI_BASE_URL="http://localhost:3000"
export GUARDAI_API_KEY="<key>"        # only if the server enforces auth
# export OPENAI_API_KEY="<key>"       # optional; only flips the drafter label

# 1. start the service
uvicorn app:app --port 9000

# 2. in another terminal, drive it over the fictional inbound cases
python client.py

fastapi>=0.110.0
uvicorn>=0.27.0
httpx>=0.27.0

[
  {
    "case_id": "CR-2026-0442",
    "channel_in": "email",
    "subject": "Doppelte Abbuchung meiner Rate",
    "message": "Sehr geehrtes Team, mein Name ist Frau Erika Mustermann. Meine Bestellung ORD-2026-4471 wurde diesen Monat zweimal von meinem Konto abgebucht. Bitte erstatten Sie den doppelten Betrag auf IBAN DE97 9900 0000 0123 4567 89. Sie erreichen mich unter erika.mustermann@beispiel.de oder telefonisch unter +49 151 23456789. Meine Anschrift lautet Musterstrasse 12, 80331 Muenchen. Mit freundlichen Gruessen, Erika Mustermann"
  },
  {
    "case_id": "CR-2026-0443",
    "channel_in": "email",
    "subject": "Adressaenderung nach Umzug",
    "message": "Guten Tag, hier ist Herr Max Mustermann. Ich bin umgezogen und moechte meine hinterlegte Anschrift aktualisieren. Zu meiner letzten Bestellung ORD-2026-8890 hatte ich noch eine offene Rueckfrage. Meine neue E-Mail ist max.mustermann@beispiel.de, Telefon +49 160 55512345. Meine hinterlegte IBAN DE97 9900 0000 0123 4567 89 soll unveraendert bleiben. Vielen Dank, Max Mustermann"
  },
  {
    "case_id": "CR-2026-0444",
    "channel_in": "email",
    "subject": "Identitaetspruefung fuer Kontoeroeffnung",
    "message": "Sehr geehrte Damen und Herren, ich bin Frau Erika Mustermann und moechte die Verifizierung abschliessen. Zur Sicherheit sende ich Ihnen direkt meinen Reisepass C12345678 und meine SSN 123-45-6789 mit. Meine Bestellung ORD-2026-1207 wartet noch auf Freigabe. Antwort bitte an erika.mustermann@beispiel.de. Freundliche Gruesse, Erika Mustermann"
  }
]

from __future__ import annotations

import logging
import os
import re
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from typing import AsyncIterator

import httpx
from fastapi import FastAPI, Request
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse

# Accept whichever base-URL var is set so the example drops into any of the repo's env
# conventions without edits.
GUARDAI_URL = (
    os.getenv("GUARDAI_BASE_URL")
    or os.getenv("GUARDAI_URL")
    or os.getenv("GUARDAI_BASE")
    or "http://localhost:3000"
).rstrip("/")
GUARDAI_API_KEY = os.getenv("GUARDAI_API_KEY", "")

# One vertical policy governs the whole service: German formal restore (Herr/Frau),
# IBAN capped at masked, passports and national IDs hard-redacted. `financial` is a
# drop-in alternative when the salutation should stay language-neutral.
GUARDAI_POLICY = os.getenv("GUARDAI_POLICY", "german-support")
GUARDAI_LANG = os.getenv("GUARDAI_LANG", "de")

# Routes the middleware guards. Anything else (health) passes through untouched.
GUARDED_PATHS = frozenset({"/support/reply"})

TOKEN_RE = re.compile(r"\{\{([^}]+)\}\}")
REDACTED_RE = re.compile(r"\[REDACTED:([^\]]+)\]")

log = logging.getLogger("musterbank_support")

@dataclass(frozen=True)
class Principal:
    subject: str
    role: str


@dataclass(frozen=True)
class AudienceView:
    channel: str | None
    restore_mode: str | None


# Demo bearer tokens mapped to a principal. In production replace this table lookup with
# your IdP token verification: the role must come from the verified credential, never from
# caller-supplied input. The tokens are intentionally human-readable, not secrets.
PRINCIPALS: dict[str, Principal] = {
    "demo-customer-token": Principal("cust_lena", "customer"),
    "demo-agent-token": Principal("agent_mara", "agent"),
    "demo-supervisor-token": Principal("sup_per", "supervisor"),
    "demo-auditor-token": Principal("audit_nore", "auditor"),
}

# One authenticated role resolves to exactly one restore sink.
AUDIENCES: dict[str, AudienceView] = {
    "customer": AudienceView("customer_email", None),
    "agent": AudienceView("internal_summary", None),
    "supervisor": AudienceView(None, "partial"),
    "auditor": AudienceView("log_safe", None),
}


def _authenticate(header: str | None) -> Principal | None:
    """Resolve `Authorization: Bearer <token>` (case-insensitive scheme) to a Principal,
    or None when the header is missing, malformed, or the token is unknown. A missing header, a
    wrong scheme, and an unknown token all take the same parse-and-lookup path so neither the
    response nor the timing confirms whether a token exists."""
    scheme, _, token = (header or "").partition(" ")
    candidate = token.strip() if scheme.lower() == "bearer" else ""
    return PRINCIPALS.get(candidate)

@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
    """Own the OGuardAI HTTP client for the app's lifetime.

    One pooled client, created at startup and closed at shutdown, so every request
    reuses the same connections instead of paying a fresh TCP/TLS handshake."""
    headers = {"X-API-Key": GUARDAI_API_KEY} if GUARDAI_API_KEY else {}
    app.state.guardai = httpx.AsyncClient(
        base_url=GUARDAI_URL, headers=headers, timeout=30.0
    )
    yield
    await app.state.guardai.aclose()


app = FastAPI(title="Musterbank support-desk (OGuardAI middleware)", lifespan=lifespan)

# Gender metadata only picks the German adjective ending that AGREES with the honorific
# the person restore inserts: "geehrter" for Herr, "geehrte" for Frau. That agreement is
# exactly what a gender-blind drafter gets wrong. Unknown gender uses the neutral opener.
_SALUTATION = {"male": "Sehr geehrter", "female": "Sehr geehrte"}


@dataclass
class GuardSession:
    """Everything the handler is allowed to see about one customer message.

    The raw text is deliberately absent: only `safe_text` (tokens), the sealed
    `session_state`, and the safe per-entity metadata cross into the handler."""

    client: httpx.AsyncClient
    safe_text: str
    session_state: str
    entity_context: list[dict]
    redacted_types: list[str] = field(default_factory=list)

    def token_of(self, entity_type: str) -> str | None:
        """The exact token string (hash included) for a type, e.g. the `{{person:...}}`
        the LLM must echo verbatim. Short forms without the hash do not resolve."""
        for inner in TOKEN_RE.findall(self.safe_text):
            if _type_key(inner) == entity_type:
                return "{{" + inner + "}}"
        return None

    def salutation(self) -> str:
        """A grammatically correct formal opener built from safe metadata.

        The person restore already inserts "Herr/Frau {name}", so the draft only supplies
        the agreeing adjective. Prepending "Herr/Frau" here as well would double it."""
        person_token = self.token_of("person")
        if not person_token:
            return "Sehr geehrte Damen und Herren"
        gender = next(
            (e.get("gender") for e in self.entity_context if e.get("token") == person_token),
            "unknown",
        )
        opener = _SALUTATION.get(gender)
        return f"{opener} {person_token}" if opener else f"Guten Tag {person_token}"

    async def rehydrate(
        self,
        draft: str,
        *,
        channel: str | None = None,
        restore_mode: str | None = None,
    ) -> str:
        """Restore one draft for one audience. Fail-closed: on any rehydrate error we
        raise, never return the unresolved draft. The policy is the ceiling, so a channel
        or mode can only tighten what the base rule already allows, never widen it."""
        body: dict = {"output": draft, "session_state": self.session_state}
        if channel:
            body["output_channel"] = channel
        if restore_mode:
            body["restore_mode"] = restore_mode
        resp = await self.client.post("/v1/rehydrate", json=body)
        resp.raise_for_status()
        return resp.json()["restored_text"]


def _type_key(inner: str) -> str:
    """Type from a token body: 'person:p_001:hash' -> 'person'."""
    return inner.rsplit(":", 2)[0] if inner.count(":") >= 2 else inner.split(":")[0]

class PIIGuardMiddleware(BaseHTTPMiddleware):
    """Transforms the inbound `message` for guarded routes before the handler runs.

    Design notes:
      * BaseHTTPMiddleware buffers the request body, so we parse it once here and forward
        the non-PII fields on `request.state.payload`. The handler reads from there and
        never re-reads the stream, and the raw `message` is dropped at this boundary.
      * Fail-closed: a malformed body is a 400 and a transform failure is a 502; in both
        cases the handler never runs, so the model can never be reached with raw text."""

    async def dispatch(self, request: Request, call_next):
        if request.url.path not in GUARDED_PATHS:
            return await call_next(request)

        principal = _authenticate(request.headers.get("authorization"))
        if principal is None:
            return _error(401, "GUARDAI_UNAUTHENTICATED", "Bearer credential required.",
                          headers={"WWW-Authenticate": "Bearer"})
        request.state.principal = principal

        try:
            payload = await request.json()
        except Exception:
            return _error(400, "GUARDAI_INVALID_INPUT", "Request body must be JSON.")
        message = payload.get("message")
        if not isinstance(message, str) or not message.strip():
            return _error(400, "GUARDAI_INVALID_INPUT", "Field 'message' is required.")

        client: httpx.AsyncClient = request.app.state.guardai
        try:
            resp = await client.post(
                "/v1/transform",
                json={"input": message, "input_type": "text",
                      "language": GUARDAI_LANG, "policy": GUARDAI_POLICY},
            )
            resp.raise_for_status()
        except httpx.HTTPError as exc:
            # No fallback to raw text: reject the request so no downstream logic runs.
            # The exception detail (internal host/URL) stays in the server log only.
            log.error("transform call failed: %s", exc)
            return _error(502, "GUARDAI_TRANSFORM_FAILED",
                          "OGuardAI transform unavailable.")

        result = resp.json()
        safe_text = result["safe_text"]
        request.state.guard = GuardSession(
            client=client,
            safe_text=safe_text,
            session_state=result["session_state"],
            entity_context=result.get("entity_context", []),
            redacted_types=sorted(set(REDACTED_RE.findall(safe_text))),
        )
        # Forward only the non-PII envelope; the raw message never crosses into the route.
        request.state.payload = {k: v for k, v in payload.items() if k != "message"}
        return await call_next(request)


app.add_middleware(PIIGuardMiddleware)


def _error(status: int, code: str, message: str, headers: dict[str, str] | None = None) -> JSONResponse:
    return JSONResponse(status_code=status, content={"error": message, "code": code}, headers=headers)

def _rehydrate_error(exc: httpx.HTTPError) -> JSONResponse:
    """Map a rehydrate failure to the right status. The output guard keeps a content verdict
    apart from a scanner outage, and the handler must preserve that: 422 GUARDAI_OUTPUT_BLOCKED
    was judged and refused (resending the same draft will not help), 503 means the scan never
    ran and the caller may retry. Anything else stays a generic 502 gateway failure."""
    resp = getattr(exc, "response", None)
    if resp is not None:
        code = ""
        try:
            body = resp.json()
            code = body.get("code", "") if isinstance(body, dict) else ""
        except ValueError:
            code = ""
        if code == "GUARDAI_OUTPUT_BLOCKED" or resp.status_code == 422:
            return _error(422, code or "GUARDAI_OUTPUT_BLOCKED", "Output blocked by data policy.")
        if code == "GUARDAI_OUTPUT_GUARD_UNAVAILABLE" or resp.status_code == 503:
            return _error(503, code or "GUARDAI_OUTPUT_GUARD_UNAVAILABLE", "Output guard unavailable.")
    return _error(502, "GUARDAI_REHYDRATE_FAILED", "Rehydrate failed.")


# In production this is the LLM call. It reads OPENAI_API_KEY exactly like a real client
# would and would send `guard.safe_text` as the user turn, seeing only tokens, never raw
# PII. To stay deterministic and offline the example classifies on safe keywords and fills
# a template with the message's own tokens; the real provider is never contacted here.
_LLM_LABEL = "openai (token-only)" if os.getenv("OPENAI_API_KEY") else "offline template drafter"

_INTENTS = [
    ("refund", ("doppelt", "doppelte abbuchung", "erstattung", "rueckerstattung", "gutschrift")),
    ("address_change", ("adressaenderung", "anschrift", "umzug", "umgezogen", "adresse")),
    ("verification", ("verifizierung", "identitaetspruefung", "reisepass", "kontoeroeffnung")),
]


def _classify(safe_text: str, subject: str) -> str:
    hay = f"{subject} {safe_text}".lower()
    for intent, keys in _INTENTS:
        if any(k in hay for k in keys):
            return intent
    return "general"


def _draft_reply(intent: str, guard: GuardSession) -> str:
    """A token-only German reply. Every {{...}} is a placeholder the rehydrator fills per
    audience; the drafter never handles a raw value."""
    opener = guard.salutation()
    order = guard.token_of("order")
    iban = guard.token_of("iban")
    email = guard.token_of("email")
    ref = f" zu Bestellung {order}" if order else ""
    # A confirmation line that carries the email token makes the channel difference
    # visible: the customer sees their full address, the agent console sees it tightened.
    conf = f" Eine Bestaetigung senden wir an {email}." if email else ""
    if intent == "refund":
        pay = f" Die Erstattung erfolgt auf {iban}." if iban else ""
        return (f"{opener}, wir haben die doppelte Abbuchung{ref} geprueft und "
                f"veranlassen die Korrektur.{pay}{conf}")
    if intent == "address_change":
        return (f"{opener}, Ihre neue Anschrift ist hinterlegt und die Aenderung ist ab "
                f"sofort aktiv. Ihre Rueckfrage{ref} pruefen wir und melden uns zeitnah.{conf}")
    if intent == "verification":
        return (f"{opener}, aus Sicherheitsgruenden pruefen wir Ihre Verifizierung{ref} "
                f"ueber den gesicherten Upload. Bitte senden Sie Ausweisdaten nie per E-Mail.")
    return f"{opener}, vielen Dank fuer Ihre Nachricht. Ein Mitarbeiter meldet sich in Kuerze."

@app.post("/support/reply")
async def support_reply(request: Request) -> JSONResponse:
    """Draft one reply and restore it into the single view the caller's role entitles.

    The handler never calls transform (the middleware did) and never sees the raw message.
    It drafts on tokens, then asks the GuardSession for this caller's one audience view."""
    guard: GuardSession = request.state.guard
    principal: Principal = request.state.principal
    envelope: dict = request.state.payload
    subject = str(envelope.get("subject", ""))

    view_spec = AUDIENCES[principal.role]
    intent = _classify(guard.safe_text, subject)
    draft = _draft_reply(intent, guard)

    try:
        restored = await guard.rehydrate(
            draft, channel=view_spec.channel, restore_mode=view_spec.restore_mode
        )
    except httpx.HTTPError as exc:
        log.error("rehydrate call failed: %s", exc)
        return _rehydrate_error(exc)

    return JSONResponse({
        "case_id": envelope.get("case_id"),
        "intent": intent,
        "drafted_by": _LLM_LABEL,
        "entities_detected": len(TOKEN_RE.findall(guard.safe_text)),
        "hard_redacted": guard.redacted_types,
        "safe_draft": draft,
        "audience": principal.role,
        "view": restored,
    })


@app.get("/health")
async def health(request: Request) -> dict:
    """Proxy OGuardAI's health so a probe on this service reflects the runtime too."""
    client: httpx.AsyncClient = request.app.state.guardai
    try:
        resp = await client.get("/v1/health")
        return {"service": "ok", "guardai": resp.json()}
    except httpx.HTTPError:
        return {"service": "ok", "guardai": {"status": "unreachable"}}

"""Drive the running support-desk service over the fictional inbound cases.

Start the service first (`uvicorn app:app --port 9000`), then run this. For each case in
fixtures/inbound_cases.json it POSTs to /support/reply once per demo credential and prints
the four audience views side by side. The service only ever returns the ONE view the
caller's authenticated role entitles; the four-way comparison is assembled here, from four
credentials, not from one privileged response. Every case shows the trust boundary: the
`safe_draft` the model produced carries only tokens, never a raw value.

The demo tokens below mirror the service's PRINCIPALS table. In production these are real
IdP-issued bearer tokens and the role always comes from the verified credential.
"""

from __future__ import annotations

import json
import os
import re
from pathlib import Path

import httpx

SERVICE_URL = os.getenv("SERVICE_URL", "http://localhost:9000").rstrip("/")
HERE = Path(__file__).resolve().parent

EMAIL_RE = re.compile(r"[\w.+-]+@[\w.-]+\.[A-Za-z]{2,}")

# (label, demo bearer token) per audience. One credential yields exactly one view.
AUDIENCES = [
    ("customer email       ", "demo-customer-token"),
    ("agent console        ", "demo-agent-token"),
    ("supervisor escalation", "demo-supervisor-token"),
    ("audit log            ", "demo-auditor-token"),
]


def main() -> None:
    cases = json.loads((HERE / "fixtures" / "inbound_cases.json").read_text("utf-8"))

    try:
        health = httpx.get(f"{SERVICE_URL}/health", timeout=10).json()
    except httpx.HTTPError as exc:
        raise SystemExit(f"Cannot reach the service at {SERVICE_URL}: {exc}")
    detector = health.get("guardai", {}).get("components", {}).get("detector", {}).get("message", "")
    ner_active = detector.startswith("builtin_and_ner")
    if not ner_active:
        print("NOTE: NER detector not active; the German honorific/name restore will be "
              "missing. Structured PII (email/phone/IBAN/order/passport/SSN) still works.\n")

    leaked_emails: set[str] = set()

    with httpx.Client(base_url=SERVICE_URL, timeout=30.0) as client:
        for case in cases:
            print("=" * 80)
            print(f"CASE {case['case_id']}  |  {case['subject']}")
            print("=" * 80)

            declared_emails = set(EMAIL_RE.findall(case["message"]))
            summary_shown = False
            for label, token in AUDIENCES:
                resp = client.post(
                    "/support/reply", json=case,
                    headers={"Authorization": f"Bearer {token}"},
                )
                if resp.status_code != 200:
                    # Fail-closed surfaces here as a typed error, never a raw-text fallback.
                    print(f"    {label}: ERROR {resp.status_code} {resp.json()}")
                    continue
                r = resp.json()
                if not summary_shown:
                    print(f"  intent={r['intent']}  entities={r['entities_detected']}  "
                          f"drafted_by={r['drafted_by']}")
                    if r["hard_redacted"]:
                        print(f"  hard-redacted (never reached the model): {', '.join(r['hard_redacted'])}")
                    print(f"  safe draft (tokens only, what the model produced):\n    {r['safe_draft']}")
                    # An email is a builtin-detectable identifier, so a raw one in the draft the
                    # model produced is a detection failure, not a caveat.
                    leaked_emails.update(e for e in declared_emails if e in r["safe_draft"])
                    print("  restored views (one credential each, one session per call):")
                    summary_shown = True
                print(f"    {label}: {r['view']}")
            print()

    if leaked_emails:
        raise SystemExit("FAIL: a raw email address reached the token-only draft the model "
                         "produced: " + ", ".join(sorted(leaked_emails)))
    # The client sees the model's draft, not the tokenized prompt it read; without NER a name
    # stays raw in that prompt and reaches the model, so the NER signal is what reports it.
    if not ner_active:
        print("Done. Structured PII stayed inside the runtime; customer names reached the model "
              "because this runtime runs without NER. Start the full stack with the NER detector "
              "to tokenize names too.")
        return
    print("Done. Raw PII stayed inside the trusted zone; the model and every downstream sink "
          "saw tokens only.")


if __name__ == "__main__":
    main()

================================================================================
CASE CR-2026-0442  |  Doppelte Abbuchung meiner Rate
================================================================================
  intent=refund  entities=8  drafted_by=offline template drafter
  safe draft (tokens only, what the model produced):
    Sehr geehrte {{person:p_001:...}}, wir haben die doppelte Abbuchung zu Bestellung {{order:o_001:...}} geprueft und veranlassen die Korrektur. Die Erstattung erfolgt auf {{iban:ib_001:...}}. Eine Bestaetigung senden wir an {{email:e_001:...}}.
  restored views (one credential each, one session per call):
    customer email        : Sehr geehrte Frau Erika Mustermann, ... Die Erstattung erfolgt auf D*** **** **** **** **** *0. Eine Bestaetigung senden wir an erika.mustermann@beispiel.de.
    agent console         : Sehr geehrte Frau Erika Mustermann, ... Eine Bestaetigung senden wir an e************************e.
    supervisor escalation : Sehr geehrte E. Mustermann, ... zu Bestellung ************* ...
    audit log             : Sehr geehrte [REDACTED], ... zu Bestellung [REDACTED] ... auf [REDACTED]. ... an [REDACTED].

================================================================================
CASE CR-2026-0443  |  Adressaenderung nach Umzug
================================================================================
  intent=address_change  entities=7  drafted_by=offline template drafter
    customer email        : Sehr geehrter Herr Max Mustermann, Ihre neue Anschrift ist hinterlegt ...
    agent console         : Sehr geehrter Herr Max Mustermann, ... an m*************************e.

================================================================================
CASE CR-2026-0444  |  Identitaetspruefung fuer Kontoeroeffnung
================================================================================
  intent=verification  entities=4  drafted_by=offline template drafter
  hard-redacted (never reached the model): passport, ssn
    customer email        : Sehr geehrte Frau Erika Mustermann, aus Sicherheitsgruenden pruefen wir Ihre Verifizierung zu Bestellung ORD-2026-1207 ...

Done. Raw PII stayed inside the trusted zone; the model and every downstream sink saw tokens only.

# no bearer credential -> 401, decided before any transform
curl -s -X POST http://localhost:9000/support/reply \
  -H "Content-Type: application/json" -d '{"case_id":"X","message":"Hallo"}'
# {"error":"Bearer credential required.","code":"GUARDAI_UNAUTHENTICATED"}

# missing 'message' -> 400, handler never runs
curl -s -X POST http://localhost:9000/support/reply \
  -H "Authorization: Bearer demo-customer-token" \
  -H "Content-Type: application/json" -d '{"case_id":"X"}'
# {"error":"Field 'message' is required.","code":"GUARDAI_INVALID_INPUT"}

# runtime unreachable -> 502, no model call on raw text
# (point GUARDAI_BASE_URL at a dead port and POST a case with a valid token)
# {"error":"OGuardAI transform unavailable.","code":"GUARDAI_TRANSFORM_FAILED"}