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.
#
# NO_LLM_STUB additionally allows a tunnelled origin, because the public
# hostnames do not resolve yet and the whole chain has to be exercisable
# without them. It is deliberately tied to the stub flag so it disappears when
# the flag does.
#
# The alternative - proxying /api through ai-web's nginx to make the page
# same-origin - was rejected: it creates a second route to the API that does
# not pass through the api.yokogawa.tech Caddy block, and that block is where
# the Phase 9 publisher rule for ^/docs/.* lives. A convenience path around an
# authorisation rule is how the rule stops meaning anything.
_ORIGINS = ["https://ai.yokogawa.tech"]
if settings().no_llm_stub:
    _ORIGINS.append("http://localhost:8080")

app.add_middleware(
    CORSMiddleware,
    allow_origins=_ORIGINS,
    # The UI sends credentials: "include" so the browser attaches the Authelia
    # session cookie. Without this the preflight fails and NO cross-origin call
    # succeeds - the browser refuses a credentialed request unless the response
    # says Access-Control-Allow-Credentials: true. Found driving the real UI in
    # a browser; it would have failed identically at ai.yokogawa.tech, and it
    # cannot be found by calling the API directly with curl.
    allow_credentials=True,
    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)
    question_class: str | None = Field(
        default=None,
        description=(
            "Force the class instead of classifying. NO-LLM STUB MODE ONLY - "
            "ignored unless NO_LLM_STUB is on, because letting a caller choose "
            "its own contract is exactly what the classifier exists to prevent."
        ),
    )


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, forced_class=request.question_class
        )
    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,
    )