assistant/chat.py

from __future__ import annotations

import re
from datetime import date
from typing import Iterator

from config import CONFIG
from assistant import llm, memory, tools

PACK = CONFIG.pack

# Speaker labels for the search-query transcript, per language.
_ROLE_LABELS = {
    "es": {"user": "Usuario", "assistant": "Amigo", "cue": "Consulta"},
    "en": {"user": "User", "assistant": "Amigo", "cue": "Query"},
}
_YEAR = re.compile(r"^(?:19|20)\d{2}$")


def build_messages(
    user_text: str, history: list[dict], web_context: str = "",
    profile: dict | None = None,
) -> list[dict]:
    """Assemble the message list for one turn.

    The system message stacks the persona, profile, recalled memories, and any
    web context; then come the prior history and the new user turn.
    """
    profile_block = memory.profile_to_prompt(profile or {})
    recall_block = memory.recall_block(user_text)

    system = PACK["persona"]
    if profile_block:
        system += "\n\n" + profile_block
    if recall_block:
        system += "\n\n" + recall_block
    if web_context:
        system += "\n\n" + PACK["web_header"] + "\n" + web_context

    messages = [{"role": "system", "content": system}]
    messages.extend(history)
    messages.append({"role": "user", "content": user_text})
    return messages


def _search_query(user_text: str, history: list[dict] | None = None) -> str:
    """Distill the latest message into search keywords.

    Passing the raw sentence to DuckDuckGo returns junk; a short keyword query
    returns the right article. The recent turns go in too, so a follow-up
    ("¿sabes quién es?") keeps the subject from earlier in the chat; the prompt
    tells the model to build the query for the LAST message only, and to drop
    the context when the subject changes.
    """
    labels = _ROLE_LABELS.get(CONFIG.lang, _ROLE_LABELS["es"])
    convo = ""
    if history:
        for m in [m for m in history if m.get("content")][-4:]:
            convo += f"{labels.get(m['role'], m['role'])}: {m['content']}\n"
    prompt = (f"{PACK['query_instruction']}{convo}"
              f"{labels['user']}: {user_text}\n{labels['cue']}:")
    raw = llm.complete(prompt)

    query = next((ln for ln in raw.splitlines() if ln.strip()), "").strip().strip('"')
    for cue in (labels["cue"].lower() + ":", "consulta:", "query:"):
        if query.lower().startswith(cue):
            query = query[len(cue):].strip()
            break
    # Drop filler, and any year the model guessed from stale memory (it tends to
    # emit its training-cutoff year); we anchor the real year ourselves below.
    words = [
        w for w in query.split()
        if w.lower() not in PACK["search_filler"] and not _YEAR.match(w.strip(".,"))
    ]
    query = " ".join(words)
    region = PACK["search_region"]
    if region and region.lower() not in query.lower():
        query += " " + region
    # Anchor to the current year so DuckDuckGo ranks current pages over evergreen
    # ones (old elections, old prices).
    query += " " + str(date.today().year)
    return query.strip()


def _as_text(content) -> str:
    """Flatten a chat message's content to plain text.

    Gradio's Chatbot round-trips a plain string into OpenAI-style content parts
    ([{"type": "text", "text": "..."}]), so the history we get back is not always
    a string. The gate, query distillation, and prompt all want text.
    """
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = [
            p if isinstance(p, str)
            else (p.get("text", "") if isinstance(p, dict) else "")
            for p in content
        ]
        return " ".join(p for p in parts if p)
    if isinstance(content, dict):
        return content.get("text", "")
    return ""


def _normalize_history(history: list[dict] | None) -> list[dict]:
    """Coerce Gradio chat history into clean [{'role', 'content': str}] dicts."""
    out = []
    for m in history or []:
        if isinstance(m, dict):
            out.append({"role": m.get("role"), "content": _as_text(m.get("content"))})
    return out


def respond(
    user_text: str, history: list[dict], profile_text: str = ""
) -> Iterator[str]:
    """Stream the assistant's text reply, then persist the exchange.

    `profile_text` is the YAML the user typed in the UI. It is parsed per turn,
    so edits take effect on the next message. When the turn needs current facts,
    search runs first and the results go into the prompt.
    """
    history = _normalize_history(history)
    profile = memory.parse_profile(profile_text)
    web_context = ""
    if tools.needs_search(user_text, history):
        web_context = tools.search_web(_search_query(user_text, history))
        if web_context:
            # Tell the model what "now" is, so it reads the results as current
            # and does not fall back to its stale training-time knowledge.
            today = date.today().isoformat()
            web_context = f"{PACK['today_header']}: {today}.\n{web_context}"
    messages = build_messages(user_text, history, web_context, profile)
    temperature = CONFIG.search_temperature if web_context else None
    full = ""
    for chunk in llm.stream_reply(messages, temperature=temperature):
        full += chunk
        yield chunk
    memory.remember(f"El usuario dijo: {user_text}", source="conversation")
    if full.strip():
        memory.remember(f"Yo respondi: {full.strip()}", source="conversation")