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

"""Orchestration AI : route les requêtes des clients vers LiteLLM (s01).

LiteLLM (:4000) est OpenAI-compatible et route lui-même vers Qwen3 (g01) ou Claude
selon l'alias `hermes-default` / `hermes-switch`. L'intégration des outils Hermes
(« agir sur Funk » via le gateway :8080) est une étape ultérieure.
"""

from __future__ import annotations

import json
from collections.abc import Awaitable, Callable

import httpx

from stt_server.config import settings

# Client persistant (pooling + keep-alive) : évite un handshake TCP vers LiteLLM à chaque tour.
_client: httpx.AsyncClient | None = None


def _get_client() -> httpx.AsyncClient:
    global _client
    if _client is None or _client.is_closed:
        _client = httpx.AsyncClient(timeout=settings.request_timeout)
    return _client


async def aclose() -> None:
    global _client
    if _client is not None and not _client.is_closed:
        await _client.aclose()
    _client = None


async def _post(payload: dict) -> dict:
    """POST vers LiteLLM avec retry sur blip transitoire ; renvoie le message du 1er choix.

    Résilience : la chaîne LiteLLM→llama-server (ROCm) se fige par à-coups. Une 2ᵉ tentative
    récupère un blip transitoire (ex. connexion keep-alive périmée) au lieu de renvoyer 502.
    Timeout PAR tentative (`request_timeout`) → on ne dépasse pas ~2× ce budget au pire.
    """
    headers = {"Authorization": f"Bearer {settings.litellm_key}"}
    last_exc: Exception | None = None
    for _ in range(2):
        try:
            r = await _get_client().post(
                settings.litellm_url, json=payload, headers=headers,
                timeout=settings.request_timeout,
            )
            r.raise_for_status()
            return r.json()["choices"][0]["message"]
        except httpx.HTTPError as e:
            last_exc = e
            await aclose()  # repart sur une connexion neuve (pool potentiellement vicié)
    raise last_exc  # type: ignore[misc]


def _content_of(msg: dict) -> str:
    """content du message ; repli sur reasoning_content (modèle thinking → content vide)."""
    content = (msg.get("content") or "").strip()
    if not content:
        content = (msg.get("reasoning_content") or "").strip()
    return content


def _with_no_think(system: str) -> str:
    if settings.disable_thinking:
        # Qwen3 est un modèle « thinking » : sans ça il dépense tout le budget max_tokens
        # en raisonnement (`reasoning_content`) → `content` vide, ou part en réflexion longue
        # → timeout (502). Le token `/no_think` désactive le raisonnement (le tool-calling
        # survit — vérifié). Inoffensif pour les modèles non-Qwen (simple texte ignoré).
        return system + "\n/no_think"
    return system


async def ask(
    text: str,
    system: str,
    model: str | None = None,
    history: list[dict] | None = None,
) -> str:
    """Interroge le LLM avec un system prompt **déjà assemblé** (voir contexts.assemble).

    Le choix du contexte, l'injection des sources live et du RAG sont faits en amont
    (app.v1_ask + contexts) ; ici on ne fait que le réglage thinking + l'appel LLM.
    """
    messages = [{"role": "system", "content": _with_no_think(system)}]
    if history:
        messages += history
    messages.append({"role": "user", "content": text})
    msg = await _post({
        "model": model or settings.model,
        "messages": messages,
        "max_tokens": settings.max_tokens,
        "temperature": settings.temperature,
    })
    return _content_of(msg)


# dispatch(name, args) -> texte du résultat d'outil. Fourni par app.py (closure sur les deps).
Dispatch = Callable[[str, dict], Awaitable[str]]


async def ask_with_tools(
    text: str,
    system: str,
    *,
    schemas: list[dict],
    dispatch: Dispatch,
    model: str | None = None,
    history: list[dict] | None = None,
    confirm_tools: tuple[str, ...] = (),
) -> tuple[str, list[dict]]:
    """Boucle de function-calling : le modèle appelle des outils jusqu'à pouvoir répondre.

    Renvoie (réponse, trace) où `trace` = [{name, args, result}] des outils exécutés
    (alimente le visualiseur du HUD). Borné par `tool_max_iters` ; au-delà, un dernier
    appel SANS outils force une réponse texte. Si aucun schéma n'est fourni, repli sur `ask`.

    `confirm_tools` : outils « à confirmation » (ex. admin_action). Quand le modèle en appelle
    un, la boucle s'ARRÊTE et renvoie directement le texte de son dispatch (la question de
    confirmation) comme réponse — l'exécution réelle est différée au tour suivant (app.py).
    """
    if not schemas:
        return await ask(text, system, model, history), []

    messages: list[dict] = [{"role": "system", "content": _with_no_think(system)}]
    if history:
        messages += history
    messages.append({"role": "user", "content": text})
    mdl = model or settings.model
    trace: list[dict] = []

    for _ in range(settings.tool_max_iters):
        msg = await _post({
            "model": mdl,
            "messages": messages,
            "tools": schemas,
            "tool_choice": "auto",
            "max_tokens": settings.tool_max_tokens,
            "temperature": settings.tool_temperature,
        })
        calls = msg.get("tool_calls") or []
        if not calls:
            return _content_of(msg), trace
        # Réinjecte le message assistant (avec ses tool_calls) avant les résultats d'outils.
        messages.append({
            "role": "assistant",
            "content": msg.get("content") or "",
            "tool_calls": calls,
        })
        for tc in calls:
            fn = tc.get("function", {})
            name = fn.get("name", "")
            try:
                args = json.loads(fn.get("arguments") or "{}")
            except (ValueError, TypeError):
                args = {}
            result = await dispatch(name, args)
            trace.append({"name": name, "args": args, "result": result})
            if name in confirm_tools:
                # Outil à confirmation : on stoppe la boucle et on surface la question.
                return result, trace
            messages.append({
                "role": "tool",
                "tool_call_id": tc.get("id", ""),
                "content": result,
            })

    # Budget d'itérations épuisé : on force une réponse finale sans nouveaux appels d'outils.
    msg = await _post({
        "model": mdl,
        "messages": messages,
        "tool_choice": "none",
        "max_tokens": settings.max_tokens,
        "temperature": settings.temperature,
    })
    return _content_of(msg), trace