yau-plant-assistant/api/main.py

"""FastAPI entrypoint for ai-api.

Reached at https://api.yokogawa.tech, behind Caddy and Authelia. There is no
authentication in this application because Authelia does it at the edge - which
also means this app must never be given a published host port, and never a
Caddyfile block without `import authelia`.

Every request is traced to Langfuse with its class, confidence, tool calls,
retrieved chunks, tokens, latency and contract result. Tracing failures never
break the request path.
"""

from __future__ import annotations

import logging
import time
import uuid

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

import agent
from config import settings
from contracts import ContractViolation
from guardrails import GuardrailViolation

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(name)s %(message)s",
)
log = logging.getLogger("api")

app = FastAPI(
    title="WRPS Plant Operations Assistant",
    version="0.1.0",
    description=(
        "Information retrieval and analysis for the Waterloo Road Pump Station. "
        "Not a control system, not an advisory controller, not a substitute for "
        "a competent person."
    ),
)

# The browser reaches the API by its own public hostname, so a same-site
# origin is all that is ever needed.
app.add_middleware(
    CORSMiddleware,
    allow_origins=["https://ai.yokogawa.tech"],
    allow_methods=["GET", "POST"],
    allow_headers=["Content-Type"],
)


def _langfuse():
    """Langfuse client, or None. Never let observability break the answer path."""
    cfg = settings()
    if not (cfg.langfuse_public_key and cfg.langfuse_secret_key):
        return None
    try:
        from langfuse import Langfuse

        return Langfuse(
            public_key=cfg.langfuse_public_key,
            secret_key=cfg.langfuse_secret_key,
            host=cfg.langfuse_host,
        )
    except Exception:
        log.exception("Langfuse unavailable - continuing untraced")
        return None


class AskRequest(BaseModel):
    question: str = Field(min_length=3, max_length=1000)


class AskResponse(BaseModel):
    request_id: str
    question_class: str
    confidence: float
    downgraded_reason: str | None = None
    answer: dict
    latency_ms: int


@app.get("/healthz")
def healthz() -> dict:
    """Liveness plus configuration shape. No secrets, ever - values are
    reported as set/unset."""
    return {"status": "ok", "config": settings().redacted()}


@app.post("/ask", response_model=AskResponse)
def ask(request: AskRequest) -> AskResponse:
    request_id = str(uuid.uuid4())
    started = time.perf_counter()
    client = _langfuse()
    trace = None
    if client is not None:
        try:
            trace = client.trace(
                id=request_id, name="ask", input={"question": request.question}
            )
        except Exception:
            log.exception("could not open a Langfuse trace")

    try:
        answer, classification = agent.answer(request.question, trace=trace)
    except ContractViolation as violation:
        # The offending output has already gone to Langfuse with the whole
        # generated text. What comes back here says nothing about it: an error
        # message is not a side channel for content that failed a safety check.
        log.warning("contract failure request_id=%s rule=%s", request_id, violation.rule)
        raise HTTPException(
            status_code=422,
            detail={
                "request_id": request_id,
                "error": "contract_not_met",
                "message": (
                    "I could not produce an answer that meets the safety contract "
                    "for this question. The attempt has been logged for review."
                ),
            },
        ) from violation
    except GuardrailViolation as violation:
        log.warning("guardrail refusal request_id=%s rule=%s", request_id, violation.rule)
        raise HTTPException(
            status_code=400,
            detail={
                "request_id": request_id,
                "error": "query_refused",
                "message": f"The query was refused: {violation.rule}.",
            },
        ) from violation
    except Exception as exc:
        log.exception("unhandled error request_id=%s", request_id)
        raise HTTPException(
            status_code=500,
            detail={"request_id": request_id, "error": "internal_error"},
        ) from exc

    latency_ms = int((time.perf_counter() - started) * 1000)
    payload = answer.model_dump(mode="json")

    if trace is not None:
        try:
            trace.update(
                output=payload,
                metadata={
                    "question_class": classification.question_class.value,
                    "confidence": classification.confidence,
                    "downgraded_reason": classification.downgraded_reason,
                    "contract_result": "passed",
                    "used_fixture_data": payload.get("used_fixture_data", False),
                    "latency_ms": latency_ms,
                },
            )
        except Exception:
            log.exception("could not finalise the Langfuse trace")

    return AskResponse(
        request_id=request_id,
        question_class=classification.question_class.value,
        confidence=classification.confidence,
        downgraded_reason=classification.downgraded_reason,
        answer=payload,
        latency_ms=latency_ms,
    )