oj_queue.py

"""Cola por lotes para Gemini (moderacion por API, generateContent).

Un mensaje = una llamada Gemini secuencial. La cola acumula mensajes (ventana)
y los modera en lote; cada fallo individual se marca fail-closed
({"_error"} -> uncertain + escala) sin tumbar el resto. Solo si TODO el
lote falla (red/API caida) se reencola hasta 3 veces.

Sin GPU local: Space y local usan el mismo camino (solo cambia la key).
El registro al recibir garantiza que nada se pierde aunque la API caiga.
"""

from __future__ import annotations

import asyncio
import logging
import time
from dataclasses import dataclass, field
from typing import Any, Callable

import gemini_client
import store
from space_client import Item

log = logging.getLogger("kmbot.ojqueue")


@dataclass
class _Pending:
    item: Item
    retries: int = 0
    first_seen: float = field(default_factory=time.time)


def moderate_batch(items: list[Item]) -> dict[str, dict[str, Any]]:
    """{item.key: veredicto parseado}. Llamadas secuenciales (volumen bajo)."""
    out: dict[str, dict[str, Any]] = {}
    for it in items:
        try:
            hist = []
            for h in (it.history or []):
                if isinstance(h, dict):
                    a = str(h.get("author") or "?")
                    t = str(h.get("text") or "")
                    if t:
                        hist.append(f"{a}: {t}")
                elif h:
                    hist.append(str(h))
            out[it.key] = gemini_client.moderate(
                it.text, flags="", reply_to=it.reply_to or "", history=hist,
            )
        except Exception as exc:  # noqa: BLE001
            out[it.key] = {"unsafe": True, "categories": [], "raw": "",
                           "_error": f"{type(exc).__name__}: {exc}"}
    return out


class OJBatchQueue:
    def __init__(
        self,
        client: Any | None,
        on_results: Callable[[list[Item], dict[str, dict[str, Any]]], None],
        *,
        max_messages: int = 16,
        window: float = 8.0,
        max_retries: int = 3,
        loop: asyncio.AbstractEventLoop | None = None,
        predict_fn: Callable[[list[Item]], dict[str, dict]] | None = None,
    ) -> None:
        """`predict_fn` (default: moderate_batch) recibe items y devuelve
        {key: veredicto}. `client` se ignora (compat): ya no hay Space GPU."""
        self.client = client
        self.predict_fn = predict_fn or moderate_batch
        # Diagnostico visible en /status (sin esto vamos ciegos).
        self.last_error: str | None = None
        self.processed: int = 0
        self.dropped: int = 0
        self.on_results = on_results
        self.max_messages = max_messages
        self.window = window
        self.max_retries = max_retries
        self._pending: list[_Pending] = []
        self._task: asyncio.Task | None = None
        self._wake = asyncio.Event()
        self._loop = loop

    def submit(self, item: Item) -> None:
        self._pending.append(_Pending(item=item))
        self._kick()

    def _kick(self) -> None:
        if self._loop is None:
            try:
                self._loop = asyncio.get_running_loop()
            except RuntimeError:
                return
        if self._task is None or self._task.done():
            self._task = self._loop.create_task(self._run())

    def start(self) -> None:
        self._kick()
        self._anchor = self._task

    async def _run(self) -> None:
        while True:
            if not self._pending:
                self._wake.clear()
                try:
                    await asyncio.wait_for(self._wake.wait(), timeout=self.window)
                except asyncio.TimeoutError:
                    return
                continue
            oldest = self._pending[0].first_seen
            wait = self.window - (time.time() - oldest)
            if len(self._pending) < self.max_messages and wait > 0:
                try:
                    await asyncio.wait_for(self._wake.wait(), timeout=wait)
                except asyncio.TimeoutError:
                    pass
                continue
            await self._flush()

    async def _flush(self) -> None:
        batch = self._pending[: self.max_messages]
        self._pending = self._pending[self.max_messages :]
        items = [p.item for p in batch]

        try:
            loop = asyncio.get_running_loop()
            merged = await loop.run_in_executor(None, self.predict_fn, items)
            self.last_error = None
            self.processed += len(items)
        except Exception as exc:  # noqa: BLE001
            # UN solo camino de reintento para cualquier fallo del lote.
            self.last_error = f"{type(exc).__name__}: {exc}"
            log.warning(
                "Fallo el lote (%s: %s). Reintentare hasta %d veces; %d mensajes en cola.",
                type(exc).__name__, exc, self.max_retries, len(items),
            )
            requeue: list[_Pending] = []
            exhausted: list[Item] = []
            for p in batch:
                if p.retries < self.max_retries:
                    p.retries += 1
                    p.first_seen = time.time()  # reintentar en el proximo ciclo
                    requeue.append(p)
                else:
                    self.dropped += 1
                    exhausted.append(p.item)
                    log.error("Descartando %s tras %d intentos.", p.item.key, p.retries)
            self._pending = requeue + self._pending
            if exhausted:
                # Marca la fila para que no se confunda con "aun sin puntuar".
                for it in exhausted:
                    try:
                        store.apply_decision(
                            chat_id=it.chat_id,
                            message_id=it.message_id,
                            flags=[],
                            probabilities={},
                            summary=f"modelo no respondio: {self.last_error}",
                            max_severity=0,
                            escalate=False,
                            is_violation=False,
                            # scored=False a proposito: no hay veredicto. Si
                            # se marcara como puntuada aqui, la fila dejaria de
                            # contar en `sin puntuar` y la caida del modelo
                            # quedaria escondida en el contador de /status.
                            scored=False,
                        )
                    except Exception:  # noqa: BLE001
                        log.exception("No se pudo marcar el descarte de %s.", it.key)
            if requeue:
                self._kick()
            return

        by_msg: dict[str, Any] = {}
        for it in items:
            got = merged.get(it.key)
            # OJO: `got` NO tiene por que ser un dict. El backend Decision-2.0
            # devuelve un `rules.ModerationDecision` ya evaluado, y coercionarlo
            # a dict tiraba el veredicto (medido 2026-10-08 en produccion: el
            # Space juzgaba mensajes reales del grupo y la bitacora se quedaba
            # en "sin respuesta", con dec2_rows = 0). Solo se rellena el
            # marcador cuando no hay NADA para esa clave.
            by_msg[it.key] = got if got is not None else {"_error": "sin respuesta"}

        try:
            self.on_results(items, by_msg)
        except Exception:  # noqa: BLE001
            log.exception("Fallo el manejo de resultados.")