project-lyra/lyra/chat.py

"""The chat turn: assemble the prompt (lyra.mind) then speak + persist.

`mind.assemble()` runs the society of parts (perceive → route → compose →
deliberate) and hands back a ready message list + the active mode. Then:
  - the MIND (the chat backend/model) runs the tool/generation loop — decide,
    reason, run tools — and produces a draft.
  - the MOUTH (a separate character model, if configured) re-voices that draft in
    her own voice. Default: no mouth configured → the mind's draft IS the reply
    (bit-for-bit the old behavior). The mouth slot is where a fine-tuned voice lands.
"""
from __future__ import annotations

import threading
import time

from lyra import config, llm, logbus, memory, mind, modes, poker_prompts, summary
from lyra import tools as toolkit
from lyra.llm import Backend

MAX_TOOL_ROUNDS = 5  # cap tool-call iterations per turn

# --- turn de-duplication --------------------------------------------------
# The web UI hits TWO endpoints for one message: it POSTs the SSE stream, and if
# nothing streams to the browser (a dropped connection — most often because Brian
# locks his phone to go play the hand) it falls back to the blocking endpoint. But
# the server-side stream runs to completion regardless, so BOTH turns execute —
# double-persisting the message and double-logging the hand. This guard makes a turn
# idempotent: the first request owns it; a duplicate reuses the owner's result
# instead of running a second full turn.
#
# The UI stamps each send with a unique turn_id and passes the SAME id on the stream
# AND the fallback, so we dedupe on that — bulletproof no matter how long he's away
# (a genuine new message gets a fresh id, so nothing legit is ever swallowed). Requests
# with no id fall back to a short (session, message) window for near-simultaneous dupes.
_TURN_TTL_ID = 3600.0   # id-keyed: unique per send, so keep it long for fire-and-forget
_TURN_TTL_MSG = 20.0    # (session, msg) keyed: short — only near-simultaneous dupes
_turn_lock = threading.Lock()
_turns: dict[tuple, dict] = {}  # key -> {event, reply, ts, ttl}


def _turn_key(session_id: str, user_msg: str, turn_id: str | None):
    if turn_id:
        return ("tid", turn_id), _TURN_TTL_ID
    return (session_id, (user_msg or "").strip()), _TURN_TTL_MSG


def _claim_turn(session_id: str, user_msg: str, turn_id: str | None = None):
    """(is_owner, rec). Owner executes the turn then calls _finish_turn; a non-owner
    (a duplicate of the same send) waits on rec['event'] and reuses rec['reply']."""
    key, ttl = _turn_key(session_id, user_msg, turn_id)
    now = time.monotonic()
    with _turn_lock:
        for k in [k for k, r in _turns.items() if now - r["ts"] > r["ttl"]]:
            del _turns[k]
        rec = _turns.get(key)
        if rec is not None:
            return False, rec
        rec = {"event": threading.Event(), "reply": None, "ts": now, "ttl": ttl}
        _turns[key] = rec
        return True, rec


def _finish_turn(rec: dict, reply: str) -> None:
    rec["reply"] = reply
    rec["ts"] = time.monotonic()
    rec["event"].set()


_AWAIT_TIMEOUT = 120.0  # a duplicate waits at most this long for the owner to finish


def _await_duplicate(rec: dict) -> str:
    rec["event"].wait(timeout=_AWAIT_TIMEOUT)
    return rec["reply"] or _TANGLED
# Which backends get function-calling tools is config-driven (cfg.tool_backends,
# env TOOL_BACKENDS, default "cloud"). The MI50's llama.cpp server only does tools
# when launched with --jinja + a tool-capable model, else it 500s on the tools
# param — so enabling "mi50" is a config flip once that precondition holds (Phase C),
# not a code change. See docs/superpowers/specs/2026-07-01-poker-prompts-design.md.
_TANGLED = "(I got tangled using my tools there — say that again?)"


def _resolve_model(backend: Backend, model_override: str | None, cfg) -> str:
    """Live chat uses the stronger chat_model on cloud; local/mi50 use their own.
    The UI's cloud-model picker only applies on the cloud backend."""
    model = {"local": cfg.local_model, "cloud": cfg.chat_model, "mi50": cfg.mi50_model}.get(
        backend, backend
    )
    if model_override and backend == "cloud":
        model = model_override
    return model


def _mouth_target(cfg, mind_backend: Backend, mind_model: str | None):
    """The mouth (backend, model) if configured AND different from the mind; else None
    (mouth == mind → no separate voice pass)."""
    if not cfg.mouth_backend and not cfg.mouth_model:
        return None
    backend = cfg.mouth_backend or mind_backend
    model = cfg.mouth_model or None
    if backend == mind_backend and model == mind_model:
        return None
    return backend, model


def _maybe_switch_mode(session_id: str, tool_name: str) -> None:
    """Opening a poker session auto-flips this chat into Poker mode. Manual UI switching
    still overrides anytime."""
    if tool_name == "start_session":
        memory.set_session_mode(session_id, modes.CASH.key)
        logbus.log("info", "mode auto-switch", session=session_id, mode=modes.CASH.key)


def _mind_loop(messages, backend: Backend, model: str | None, tool_specs,
               ctx: dict, session_id: str) -> tuple[str, list[str]]:
    """Run the tool/generation loop on the MIND model (non-streaming). Mutates
    `messages` with tool calls/results. Returns (draft_reply, tool_names_run)."""
    tools_run: list[str] = []
    reply = ""
    for _ in range(MAX_TOOL_ROUNDS):
        assistant_msg, tool_calls = llm.chat_call(
            messages, backend=backend, model=model, tools=tool_specs
        )
        if not tool_calls:
            reply = assistant_msg.get("content") or ""
            break
        messages.append(assistant_msg)
        for tc in tool_calls:
            result = toolkit.dispatch(tc["name"], tc["arguments"], ctx)
            memory.add_tool_event(session_id, tc["name"], tc["arguments"], result)
            logbus.log("info", "tool call", session=session_id, tool=tc["name"], result=result[:80])
            messages.append({"role": "tool", "tool_call_id": tc["id"], "content": result})
            _maybe_switch_mode(session_id, tc["name"])
            tools_run.append(tc["name"])
    return reply, tools_run


_FORCE_LOG = (
    "You have not logged Brian's hand yet — and a hand must ALWAYS be recorded, no exceptions. "
    "Call record_hand now: pass his ENTIRE hand description as one `shorthand` string."
)


def _ensure_hand_logged(messages, user_msg: str, msg_type: str | None, tools_run: list,
                        backend: Backend, model: str | None, ctx: dict, session_id: str) -> list:
    """Guarantee the ledger. If this turn was Brian's OWN hand and the model didn't log it,
    force the record_hand call — the log can't be left to the model's discretion, because
    mid-session the history few-shot-conditions it to skip logging (see mind._history_with_tools;
    even a maximal 'LOG FIRST' prompt scored 0/5 under a polluted history). Guarded to hero
    hands so an observed hand is never force-logged as his. Returns forced tool names."""
    if msg_type != "HAND" or backend not in config.load().tool_backends:
        return []
    if any(t in ("record_hand", "log_hand") for t in tools_run):
        return []
    if not poker_prompts.looks_like_hero_hand(user_msg):
        return []
    try:
        _, tcs = llm.chat_call(
            messages + [{"role": "system", "content": _FORCE_LOG}],
            backend=backend, model=model, tools=toolkit.specs(["record_hand"]),
            tool_choice={"type": "function", "function": {"name": "record_hand"}},
        )
    except Exception as exc:
        logbus.log("error", "forced hand-log failed", session=session_id, error=str(exc)[:160])
        return []
    forced = []
    for tc in (tcs or []):
        result = toolkit.dispatch(tc["name"], tc["arguments"], ctx)
        memory.add_tool_event(session_id, tc["name"], tc["arguments"], result)
        logbus.log("info", "forced hand log", session=session_id, tool=tc["name"], result=result[:80])
        forced.append(tc["name"])
    return forced


def _voice_pass(messages, draft: str, backend: Backend, model: str | None) -> str:
    """Mouth: re-render the mind's draft in her voice. Falls back to the draft on failure."""
    try:
        out = llm.complete(mind.voice_messages(messages, draft), backend=backend, model=model)
        return (out or "").strip() or draft
    except Exception as exc:
        logbus.log("error", "voice pass failed", error=str(exc)[:160])
        return draft


def respond(session_id: str, user_msg: str, backend: Backend = "cloud",
            model_override: str | None = None, turn_id: str | None = None) -> str:
    """Produce Lyra's reply to a single user message and persist the exchange."""
    cfg = config.load()
    model = _resolve_model(backend, model_override, cfg)
    logbus.log("info", "chat request", session=session_id, backend=backend,
               model=model, embed=cfg.embed_backend)

    # A duplicate of the same send (the UI's stream + blocking fallback) reuses the
    # owner's result instead of running a second full turn.
    is_owner, rec = _claim_turn(session_id, user_msg, turn_id)
    if not is_owner:
        logbus.log("info", "duplicate turn deduped", session=session_id, path="respond")
        return _await_duplicate(rec)

    reply = _TANGLED
    try:
        turn = mind.assemble(session_id, user_msg, backend, model)
        messages = turn.messages
        tool_specs = toolkit.specs(turn.mode.tools) if backend in cfg.tool_backends else None
        ctx = {"session_id": session_id, "backend": backend}

        # Persist the user turn before the tool loop so its timestamp precedes any
        # tool events fired mid-turn (keeps the transcript export in true order).
        memory.remember(session_id, "user", user_msg)
        reply, tools_run = _mind_loop(messages, backend, model, tool_specs, ctx, session_id)
        _ensure_hand_logged(messages, user_msg, turn.msg_type, tools_run, backend, model, ctx, session_id)
        mouth = _mouth_target(cfg, backend, model)
        if mouth and reply:
            reply = _voice_pass(messages, reply, *mouth)
        if not reply:
            reply = _TANGLED
        logbus.log("info", "reply", session=session_id, chars=len(reply), voiced=bool(mouth))

        memory.remember(session_id, "assistant", reply)
        summary.maybe_summarize_async(session_id)  # compact once enough new turns pile up
        return reply
    finally:
        _finish_turn(rec, reply)


def respond_stream(session_id: str, user_msg: str, backend: Backend = "cloud",
                   model_override: str | None = None, turn_id: str | None = None):
    """Streaming generator version of `respond`. Yields ("delta", text), ("tool", name),
    and a final ("done", reply). Same side effects as `respond`."""
    cfg = config.load()
    model = _resolve_model(backend, model_override, cfg)
    logbus.log("info", "chat request (stream)", session=session_id, backend=backend,
               model=model, embed=cfg.embed_backend)

    # A duplicate of the same send (this stream + the UI's blocking fallback) reuses
    # the owner's result instead of running a second full turn.
    is_owner, rec = _claim_turn(session_id, user_msg, turn_id)
    if not is_owner:
        logbus.log("info", "duplicate turn deduped", session=session_id, path="stream")
        reply = _await_duplicate(rec)
        yield ("delta", reply)
        yield ("done", reply)
        return

    reply = _TANGLED
    try:
        turn = mind.assemble(session_id, user_msg, backend, model)
        messages = turn.messages
        tool_specs = toolkit.specs(turn.mode.tools) if backend in cfg.tool_backends else None
        ctx = {"session_id": session_id, "backend": backend}
        mouth = _mouth_target(cfg, backend, model)

        # Persist the user turn up front (see respond): keeps tool events, which fire
        # mid-turn, chronologically after the user message in the exported transcript.
        memory.remember(session_id, "user", user_msg)

        if mouth is None:
            # No separate voice: stream the mind directly (the original path, unchanged).
            parts: list[str] = []
            tools_run: list[str] = []
            for _ in range(MAX_TOOL_ROUNDS):
                assistant_msg = None
                tool_calls = None
                for ev, payload in llm.chat_call_stream(
                    messages, backend=backend, model=model, tools=tool_specs
                ):
                    if ev == "delta":
                        parts.append(payload)
                        yield ("delta", payload)
                    elif ev == "message":
                        assistant_msg = payload
                    elif ev == "tool_calls":
                        tool_calls = payload
                if not tool_calls:
                    break
                messages.append(assistant_msg)
                for tc in tool_calls:
                    result = toolkit.dispatch(tc["name"], tc["arguments"], ctx)
                    memory.add_tool_event(session_id, tc["name"], tc["arguments"], result)
                    logbus.log("info", "tool call", session=session_id, tool=tc["name"], result=result[:80])
                    messages.append({"role": "tool", "tool_call_id": tc["id"], "content": result})
                    _maybe_switch_mode(session_id, tc["name"])
                    tools_run.append(tc["name"])
                    yield ("tool", tc["name"])
            for name in _ensure_hand_logged(messages, user_msg, turn.msg_type, tools_run,
                                            backend, model, ctx, session_id):
                yield ("tool", name)
            reply = "".join(parts)
            if not reply:
                reply = _TANGLED
                yield ("delta", reply)
        else:
            # Mind decides + runs tools (non-streamed); mouth re-voices, streamed.
            draft, tools_run = _mind_loop(messages, backend, model, tool_specs, ctx, session_id)
            tools_run += _ensure_hand_logged(messages, user_msg, turn.msg_type, tools_run,
                                             backend, model, ctx, session_id)
            for name in tools_run:
                yield ("tool", name)
            parts = []
            try:
                for ev, payload in llm.chat_call_stream(
                    mind.voice_messages(messages, draft), backend=mouth[0], model=mouth[1], tools=None
                ):
                    if ev == "delta":
                        parts.append(payload)
                        yield ("delta", payload)
            except Exception as exc:
                logbus.log("error", "voice stream failed", error=str(exc)[:160])
            reply = "".join(parts).strip() or draft or _TANGLED
            if not parts:
                yield ("delta", reply)

        logbus.log("info", "reply", session=session_id, chars=len(reply), voiced=bool(mouth))
        memory.remember(session_id, "assistant", reply)
        summary.maybe_summarize_async(session_id)
        yield ("done", reply)
    finally:
        _finish_turn(rec, reply)