llm/engine.py

"""
Motor conversacional. El LLM NUNCA es la fuente de la verdad:
  - decide la intención (charlar / contar / formas / colores / animales / cuento)
  - si hay una actividad o cuento, los TEXTOS salen de `content/` (curados),
    el modelo solo los presenta con calidez y maneja el turno.

Runtime: transformers + GPU dinámica del Space (decorador @spaces.GPU, ZeroGPU)
por default. Para desarrollo local sin esa GPU, `LUMI_LLM_BACKEND=ollama` usa
un servidor Ollama local con el mismo Qwen2.5 (cuantizado, ver CLAUDE.md).
"""

import os
import time

LLM_BACKEND = os.environ.get("LUMI_LLM_BACKEND", "zerogpu")  # "zerogpu" | "ollama"
OLLAMA_MODEL = os.environ.get("LUMI_OLLAMA_MODEL", "qwen2.5:7b")

MODEL_ID = "build-small-hackathon/sofia-qwen2.5-7b"

# ZeroGPU puede devolver "No CUDA GPUs are available" cuando la flota está
# saturada (frecuente en hackatones con muchos Spaces concurrentes). Estos
# reintentos cubren picos cortos; si persiste, `respond()` cae a contenido curado.
# En backend "ollama" cubren además un servidor Ollama caído/lento.
GPU_RETRY_ATTEMPTS = 2
GPU_RETRY_DELAY_S = 2

FALLBACK_REPLY = "El modelo falló (probablemente límite de ZeroGPU). Probá de nuevo en algunos minutos."

SYSTEM_PROMPT = """Eres Sofía, una amiga de juego cálida y paciente para una niña de unos 3 años.
Reglas que SIEMPRE cumples:
- Hablas en español latinoamericano neutro, con frases MUY cortas y simples (máximo dos oraciones).
- Tono dulce, alegre y alentador. Celebras cada intento ("¡muy bien!", "¡qué lindo!").
- NUNCA inventas datos, cuentos ni números. Si hace falta contenido (un cuento, una
  cuenta, una figura), usas SOLO el texto que te pasa el sistema entre <contenido>...</contenido>.
- Si el sistema te pasa un bloque <contexto>...</contexto>, son datos reales sobre lo que
  ya vivieron juntos (para sonar cercana, ej. "¡como el cuento que te conté el otro día!").
  Úsalo solo para el tono cálido: NUNCA inventes detalles nuevos a partir de él.
- Si el sistema te pasa un bloque <nota>...</nota>, es una instrucción puntual para
  ESTE turno (ej. cómo responder a un saludo): seguila.
- Si no sabes algo o no tienes contenido, lo dices simple y propones un juego: no inventas.
- Nunca hablas de temas adultos, miedos, violencia ni nada inapropiado. Rediriges con cariño.
- Una pregunta por vez. Esperas la respuesta de la niña.
"""

if LLM_BACKEND == "ollama":
    import ollama

    def _generate(messages: list[dict], max_new_tokens: int = 120) -> str:
        response = ollama.chat(
            model=OLLAMA_MODEL,
            messages=messages,
            options={"temperature": 0.6, "num_predict": max_new_tokens},
        )
        return response["message"]["content"].strip()

else:
    # `spaces` debe importarse antes que `torch`: parchea la inicialización de
    # CUDA para el IPC de ZeroGPU (ver nota en app.py).
    import spaces
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    # Cargados a nivel módulo: con ZeroGPU, .to("cuda") en el arranque corre en
    # modo "emulación" (no hace falta GPU real para cargar pesos); la GPU real
    # solo se asigna dentro de funciones @spaces.GPU.
    _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    _model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16).to("cuda")

    def _gpu_duration(messages: list[dict], max_new_tokens: int = 120) -> int:
        """Duración declarada para la cola de ZeroGPU (no el límite real de generación).

        Las respuestas reales tardan ~6-12s; el default de 60s del decorador le
        resta prioridad de cola frente a otros Spaces. Pedimos algo más realista
        (con margen) según cuán grande es la generación."""
        return 15 if max_new_tokens <= 8 else 25

    @spaces.GPU(duration=_gpu_duration)
    def _generate(messages: list[dict], max_new_tokens: int = 120) -> str:
        text = _tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = _tokenizer(text, return_tensors="pt").to("cuda")
        with torch.no_grad():
            output = _model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=True,
                temperature=0.6,
                pad_token_id=_tokenizer.eos_token_id,
            )
        new_tokens = output[0][inputs["input_ids"].shape[1]:]
        return _tokenizer.decode(new_tokens, skip_special_tokens=True).strip()


def _generate_with_retry(messages: list[dict], max_new_tokens: int = 120) -> str | None:
    """Llama a `_generate` con reintentos. Devuelve None si todos fallan,
    para que el caller use un fallback curado en vez de romper el turno."""
    for attempt in range(GPU_RETRY_ATTEMPTS + 1):
        try:
            return _generate(messages, max_new_tokens=max_new_tokens)
        except Exception as e:
            # `spaces` puede re-lanzar la falla de asignación de GPU con un
            # tipo propio (no necesariamente RuntimeError) al cruzar el IPC
            # del worker de ZeroGPU; cualquier excepción acá debe degradar
            # a contenido curado, no romper el turno de la nena.
            # Se imprime (va a "Logs" del Space) para poder distinguir cuota
            # de ZeroGPU agotada ("No CUDA GPUs are available") de otros
            # errores sin gastar cuota extra para diagnosticar.
            print(f"[lumi] _generate falló (intento {attempt + 1}/{GPU_RETRY_ATTEMPTS + 1}): {e!r}")
            if attempt < GPU_RETRY_ATTEMPTS:
                time.sleep(GPU_RETRY_DELAY_S)
    return None


class LumiEngine:
    def __init__(self, content, store=None):
        self.content = content
        self.store   = store

    # Verbos/frases que indican pedido de cambio (no solo charla sobre colores).
    _COLOR_CHANGE_WORDS = (
        "cambi", "pon", "volv", "vuelv", "transform", "vist", "vest",
        "pint", "hace", "hazt", "convert", "sea", "seas", "ser",
    )

    def _intent(self, message: str) -> str:
        m = message.lower()
        if any(w in m for w in ("cuento", "cuentito", "historia")):
            return "story"
        if any(w in m for w in ("contar", "número", "numero", "cuántos", "cuantos")):
            return "counting"
        if any(w in m for w in self._COLOR_CHANGE_WORDS) and self.content.match_color(m):
            return "sofia_color"
        if any(w in m for w in ("color", "colores")):
            return "colors"
        if any(w in m for w in ("animal", "animales", "perro", "gato", "vaca")):
            return "animals"
        if any(w in m for w in ("canción", "cancion", "canta", "cantar", "rima")):
            return "song"
        if any(w in m for w in ("emoción", "emocion", "siento", "sentir", "contenta", "contento", "triste", "enojada", "enojado")):
            return "emotion"
        if any(w in m for w in ("hola", "holis", "buenas", "buen día", "buen dia",
                                 "buenos días", "buenos dias", "buenas tardes",
                                 "buenas noches", "hey", "ey")):
            return "greeting"
        if any(w in m for w in ("jugar", "juego", "jugamos", "aprender", "adivina", "adivinar")):
            return "learn"
        return "chat"

    def _context_block(self, child_age: int) -> str:
        if not self.store:
            return ""
        return self.store.memory_block(child_age, limit=3)

    def _history_messages(self, child_age: int, limit: int = 3) -> list[dict]:
        """Últimos turnos reales (user/assistant) de esta conversación, para que
        el modelo mantenga el hilo (referencias, tono) además del resumen que
        da `_context_block`."""
        if not self.store:
            return []
        turns = self.store.recent_turns(child_age, limit=limit)
        messages: list[dict] = []
        for turn in reversed(turns):
            if turn["blocked"] or not turn["message"] or not turn["reply"]:
                continue
            messages.append({"role": "user", "content": turn["message"]})
            messages.append({"role": "assistant", "content": turn["reply"]})
        return messages

    def respond(self, message: str, child_age: int = 3):
        intent   = self._intent(message)
        activity = None
        content_block = ""

        if intent == "sofia_color":
            # Paleta curada (content/colors.json): el LLM solo confirma el
            # cambio con calidez, nunca elige ni inventa el color/paleta.
            activity = self.content.match_color(message.lower())
            if activity:
                content_block = f"\n<contenido>\n{activity['script']}\n</contenido>\n"
        elif intent not in ("chat", "greeting"):
            seen = self.store.seen_ids(child_age) if self.store else None
            activity = self.content.pick(intent, child_age, exclude_ids=seen)
            if activity:
                # Cuentos y canciones NO se pasan como <contenido>: el cuento es
                # muy largo para los 120 tokens de respuesta y la canción es un
                # mp3 curado, no texto. "Aprender" tampoco: la pregunta (emoji +
                # prompt) se muestra y dice aparte, curada. En ambos casos el
                # modelo solo anuncia. Para el resto de actividades sí es el
                # texto curado que el modelo presenta.
                if intent not in ("story", "song", "learn"):
                    content_block = f"\n<contenido>\n{activity['script']}\n</contenido>\n"
                if self.store:
                    self.store.mark_seen(child_age, activity["id"], intent, activity.get("title"))

        if intent == "greeting":
            # Saludo: arrancamos "en limpio", sin <contexto> ni historial, para que
            # la respuesta sea siempre una pregunta abierta y no retome actividades
            # pasadas (ver CLAUDE.md / diseño del flujo conversacional).
            context_block = ""
            note_block = (
                "\n<nota>\nEste turno es un saludo. Saludá con calidez, en pocas "
                "palabras, y preguntale qué quiere hacer hoy. No propongas una "
                "actividad puntual ni menciones juegos anteriores.\n</nota>\n"
            )
            history = []
        elif intent == "story" and activity:
            context_block = self._context_block(child_age)
            note_block = (
                f"\n<nota>\nLe vas a presentar el cuento \"{activity['title']}\". "
                "NO lo cuentes ni inventes su contenido: se reproduce aparte, "
                "grabado, automáticamente apenas termines de hablar. Anuncialo "
                "con cariño en una frase muy corta, sin pedirle que toque nada.\n</nota>\n"
            )
            history = self._history_messages(child_age)
        elif intent == "song" and activity:
            context_block = self._context_block(child_age)
            note_block = (
                f"\n<nota>\nLe vas a proponer la canción \"{activity['title']}\". "
                "NO la cantes ni inventes la letra: se reproduce aparte, grabada, "
                "automáticamente apenas termines de hablar. Anunciala con cariño "
                "en una frase muy corta, sin pedirle que toque nada.\n</nota>\n"
            )
            history = self._history_messages(child_age)
        elif intent == "learn" and activity:
            context_block = self._context_block(child_age)
            note_block = (
                "\n<nota>\nLe vas a proponer un juego de adivinar: un emoji y "
                "una pregunta muy simple. NO hagas la pregunta ni digas la "
                "respuesta: eso se muestra y se dice aparte, automáticamente "
                "apenas termines de hablar. Anuncialo con alegría en una frase "
                "muy corta, por ejemplo \"¡Juguemos a adivinar!\".\n</nota>\n"
            )
            history = self._history_messages(child_age)
        else:
            context_block = self._context_block(child_age)
            note_block = ""
            history = self._history_messages(child_age)

        user_turn = message + content_block + context_block + note_block

        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            *history,
            {"role": "user", "content": user_turn},
        ]
        reply = _generate_with_retry(messages, max_new_tokens=120)
        if reply is None:
            # ZeroGPU no respondió tras los reintentos: si hay contenido
            # curado para esta actividad, lo devolvemos directo (sigue
            # siendo la fuente de la verdad); si no, un mensaje cálido.
            if activity and intent == "story":
                reply = f'¡Tengo un cuento para vos! Se llama "{activity["title"]}". ¡Ya empieza!'
            elif activity and intent == "song":
                reply = f'¡Tengo una canción para vos! Se llama "{activity["title"]}". ¡Ya empieza!'
            elif activity and intent == "learn":
                reply = "¡Juguemos a adivinar!"
            elif activity:
                reply = activity["script"]
            else:
                reply = FALLBACK_REPLY
        return reply, activity

    def warmup(self):
        """Corre una generación mínima al arrancar la app.

        La primera llamada a `_generate` paga el costo de JIT/compilación de
        kernels CUDA; sin esto, esa demora caería sobre el primer turno real
        de la nena. Con reintentos: si ZeroGPU está saturado al arrancar,
        no queremos que un RuntimeError tire abajo el proceso."""
        _generate_with_retry(
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": "Hola"},
            ],
            max_new_tokens=1,
        )