pv-agent/agent/api.py

"""FastAPI-Oberfläche des PV RAG Agent.

Endpunkte:
  POST /ask      — Frage -> belegte Antwort (oder Verweigerung)
  GET  /health   — Index- und Ollama-Status
  POST /reindex  — Index-Neuaufbau (nach neuem Wissensbasis-Batch)
"""
from __future__ import annotations

import threading
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field

from .config import Config
from .generate import answer_question
from .ingest import build_index
from .ollama_client import OllamaClient
from .retrieve import Retriever


class AskRequest(BaseModel):
    question: str = Field(min_length=3, max_length=2000)
    top_k: int | None = Field(default=None, ge=1, le=20)


class SourceOut(BaseModel):
    id: str
    title: str
    section: str | None = None
    stand: str | None = None
    work: str | None = None


class AskResponse(BaseModel):
    question: str
    answer: str
    refused: bool
    verified: bool
    citations: list[str]
    sources: list[SourceOut]
    n_context: int
    model: str
    latency_ms: int
    regenerations: int = 0


class AppState:
    def __init__(self) -> None:
        self.cfg: Config | None = None
        self.client: OllamaClient | None = None
        self.retriever: Retriever | None = None
        self.lock = threading.Lock()

    def ensure(self) -> Config:
        if self.cfg is None:
            self.cfg = Config.from_env()
        return self.cfg

    def get_client(self) -> OllamaClient:
        if self.client is None:
            cfg = self.ensure()
            self.client = OllamaClient(
                cfg.ollama_url,
                embed_timeout_s=cfg.embed_timeout_s,
                chat_timeout_s=cfg.chat_timeout_s,
            )
        return self.client

    def get_retriever(self) -> Retriever:
        if self.retriever is None:
            cfg = self.ensure()
            self.retriever = Retriever(cfg)
        return self.retriever

    def reset_retriever(self) -> None:
        if self.retriever is not None:
            self.retriever.close()
        self.retriever = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    app.state.rag = AppState()
    yield
    rag: AppState = app.state.rag
    if rag.retriever:
        rag.retriever.close()
    if rag.client:
        rag.client.close()


app = FastAPI(title="PV RAG Agent", version="0.1.0", lifespan=lifespan)


@app.post("/ask", response_model=AskResponse)
def ask(req: AskRequest) -> AskResponse:
    rag: AppState = app.state.rag
    cfg = rag.ensure()
    try:
        result = answer_question(
            req.question, cfg,
            client=rag.get_client(),
            retriever=rag.get_retriever(),
            top_k=req.top_k,
        )
    except RuntimeError as e:  # Index fehlt
        raise HTTPException(status_code=503, detail=str(e)) from e
    except Exception as e:  # Ollama nicht erreichbar o. Ä.
        raise HTTPException(
            status_code=503,
            detail=f"Antwortgenerierung fehlgeschlagen: {type(e).__name__}: {e}",
        ) from e
    result.pop("draft", None)
    return AskResponse(**result)


@app.get("/health")
def health() -> dict:
    rag: AppState = app.state.rag
    cfg = rag.ensure()
    out: dict = {"service": "pv-rag-agent", "config": {
        "ollama_url": cfg.ollama_url,
        "answer_model": cfg.answer_model,
        "embed_model": cfg.embed_model,
    }}
    try:
        retriever = rag.get_retriever()
        out["index"] = retriever.stats()
    except RuntimeError as e:
        out["index"] = {"error": str(e)}
    client = rag.get_client()
    out["ollama_up"] = client.is_up()
    if out["ollama_up"]:
        try:
            out["ollama_models"] = client.list_models()
        except Exception:
            out["ollama_models"] = None
    return out


@app.post("/reindex")
def reindex() -> dict:
    rag: AppState = app.state.rag
    cfg = rag.ensure()
    with rag.lock:
        stats = build_index(cfg, client=rag.get_client())
        rag.reset_retriever()
    result = stats.as_dict()
    result["warning"] = (
        "Index ohne Dense-Vektoren aufgebaut (Ollama-Embedding nicht verfügbar) — "
        "BM25-only. 'ollama pull " + cfg.embed_model + "' prüfen und erneut reindexen."
        if stats.embed_error else None
    )
    return result