Funk-lab/stt/server/stt_server/app.py

"""STT-server — API FastAPI pour les clients STT.

Endpoints :
  GET  /healthz       → état du service
  POST /v1/ask {text} → {reply}   (requête AI, orchestrée vers LiteLLM)
"""

from __future__ import annotations

import logging
import time
from contextlib import asynccontextmanager

import httpx
from fastapi import BackgroundTasks, FastAPI, HTTPException
from pydantic import BaseModel

from stt_server import __version__
from stt_server import brain
from stt_server.brain import ask as brain_ask
from stt_server.config import settings
from stt_server.knowledge import Knowledge
from stt_server.longterm import LongTermMemory
from stt_server.memory import SessionStore

log = logging.getLogger("stt_server")
# uvicorn ne configure que ses propres loggers : on attache notre handler en INFO
# pour que les lignes de timing (`ask … recall/gen/total`) sortent dans les logs du pod.
if not log.handlers:
    log.setLevel(logging.INFO)
    _h = logging.StreamHandler()
    _h.setFormatter(logging.Formatter("%(levelname)s:     %(name)s: %(message)s"))
    log.addHandler(_h)
    log.propagate = False

sessions = SessionStore()
longterm = LongTermMemory() if settings.memory_longterm else None
knowledge = Knowledge() if settings.docs_rag else None


@asynccontextmanager
async def lifespan(app: FastAPI):
    yield
    # Fermeture propre des clients HTTP persistants (pooling).
    await brain.aclose()
    if longterm:
        await longterm.aclose()
    if knowledge:
        await knowledge.aclose()


app = FastAPI(title="STT-server", version=__version__, lifespan=lifespan)


class AskRequest(BaseModel):
    text: str
    model: str | None = None      # alias LiteLLM ; défaut serveur si absent
    session_id: str | None = None  # mémoire court-terme : fil de conversation


class AskReply(BaseModel):
    reply: str
    model: str


@app.get("/healthz")
async def healthz() -> dict:
    return {"status": "ok", "version": __version__}


@app.get("/v1/models")
async def v1_models() -> dict:
    return {"default": settings.model, "available": settings.allowed_models}


@app.get("/v1/memory/health")
async def v1_memory_health() -> dict:
    """État de la mémoire long-terme (embeddings + Qdrant + collection), erreurs exposées."""
    if not longterm:
        return {"enabled": False}
    return await longterm.health()


@app.post("/v1/reset")
async def v1_reset(req: AskRequest) -> dict:
    if req.session_id:
        sessions.reset(req.session_id)
    return {"status": "reset"}


@app.post("/v1/ask", response_model=AskReply)
async def v1_ask(req: AskRequest, background: BackgroundTasks) -> AskReply:
    text = req.text.strip()
    if not text:
        raise HTTPException(status_code=400, detail="text vide")
    model = req.model or settings.model
    if model not in settings.allowed_models:
        raise HTTPException(
            status_code=400,
            detail=f"modèle '{model}' non autorisé ; dispo : {settings.allowed_models}",
        )
    t0 = time.perf_counter()
    history = sessions.history(req.session_id) if req.session_id else None
    # recall : timeout serré, dégrade vite ; renvoie aussi le vecteur (réutilisé ci-dessous)
    memories, qvec = await longterm.recall(text) if longterm else ([], None)
    # RAG doc : on réutilise qvec (même embedder nomic) → ancre la réponse dans funk-docs
    docs = await knowledge.search(text, qvec) if knowledge else []
    t_recall = time.perf_counter()
    try:
        reply = await brain_ask(text, model, history, memories, docs)
    except httpx.HTTPError as e:
        raise HTTPException(status_code=502, detail=f"upstream LiteLLM : {e}") from e
    t_gen = time.perf_counter()
    if req.session_id:
        sessions.add(req.session_id, "user", text)
        sessions.add(req.session_id, "assistant", reply)
    # store : APRÈS la réponse (BackgroundTasks) → hors latence perçue, et on réutilise qvec
    if longterm:
        background.add_task(longterm.store, req.session_id or "anon", text, qvec)
    log.info(
        "ask model=%s recall=%.0fms gen=%.0fms total=%.0fms mem=%d docs=%d",
        model,
        (t_recall - t0) * 1000,
        (t_gen - t_recall) * 1000,
        (t_gen - t0) * 1000,
        len(memories),
        len(docs),
    )
    return AskReply(reply=reply, model=model)


def run() -> None:
    """Entrypoint `stt-server` (dev local). En prod : uvicorn via le conteneur."""
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)  # noqa: S104 — service interne au cluster