tiny_gdn/cli_chat.py

"""Terminal chat: live reasoning stream plus SmolTalk tool-call rounds.

Tercet-R assistant turns may contain:

- a zero-loss `<|think|>` / `<|no_think|>` control prefix
- a `<think>…</think>` reasoning block
- one or more SmolTalk JSON `<tool_call>` blocks (NVIDIA XML is also parsed)

This module colours those regions as tokens arrive and, after a completed
turn, collects tool observations (prefixed with `<|tool_response|>`) so the
model can continue the same conversation.
"""

from __future__ import annotations

import json
import sys
from collections.abc import Callable, Sequence
from dataclasses import dataclass, field
from typing import Any, Literal, TextIO

from tiny_gdn.smoltalk_chat import (
    TOOL_RESPONSE_TOKEN,
    format_smoltalk_tool_call,
    wrap_smoltalk_tool_result,
)
from tiny_gdn.code_exec import (
    CALCULATOR_TOOL,
    NEMOTRON_PYTHON_EXEC_TOOL,
    execute_math_tool,
    is_auto_math_tool_name,
    reset_default_python_session,
)
from tiny_gdn.tools import ParsedToolCall, parse_tool_calls
from tiny_gdn.web_search import (
    NEMOTRON_WEB_SEARCH_TOOL,
    is_web_search_tool_name,
    query_from_arguments,
    search_web,
)


SegmentKind = Literal["answer", "think", "tool_call"]
MarkerKind = Literal[
    "think_open",
    "think_close",
    "think_control",
    "no_think",
    "tool_open",
    "tool_close",
]

THINK_OPEN = "<think>"
THINK_CLOSE = "</think>"
TOOL_CALL_OPEN = "<tool_call>"
TOOL_CALL_CLOSE = "</tool_call>"
THINK_CONTROL = "<|think|>"
NO_THINK_CONTROL = "<|no_think|>"

MARKERS: tuple[tuple[str, MarkerKind], ...] = (
    (THINK_CLOSE, "think_close"),
    (TOOL_CALL_CLOSE, "tool_close"),
    (THINK_OPEN, "think_open"),
    (TOOL_CALL_OPEN, "tool_open"),
    (NO_THINK_CONTROL, "no_think"),
    (THINK_CONTROL, "think_control"),
)

ANSI = {
    "think": "\033[2;33m",
    "tool_call": "\033[36m",
    "answer": "\033[0m",
    "reset": "\033[0m",
}

DEFAULT_MAX_TOOL_ROUNDS = 8


@dataclass(frozen=True)
class ContentSegment:
    kind: SegmentKind
    text: str
    open: bool = False


@dataclass(frozen=True)
class GeneratedTurn:
    text: str
    token_count: int
    stop_reason: str
    tool_calls: tuple[ParsedToolCall, ...]


def first_marker(text: str) -> tuple[int, str, MarkerKind] | None:
    best: tuple[int, str, MarkerKind] | None = None
    for marker, kind in MARKERS:
        at = text.find(marker)
        if at < 0:
            continue
        if (
            best is None
            or at < best[0]
            or (at == best[0] and len(marker) > len(best[1]))
        ):
            best = (at, marker, kind)
    return best


def holdback_prefix_length(text: str) -> int:
    if not text:
        return 0
    keep = 0
    for marker, _kind in MARKERS:
        limit = min(len(marker) - 1, len(text))
        for size in range(1, limit + 1):
            if marker.startswith(text[-size:]):
                keep = max(keep, size)
    return keep


def mode_after_marker(kind: MarkerKind, current: SegmentKind) -> SegmentKind:
    if kind in {"think_open", "think_control"}:
        return "think"
    if kind in {"think_close", "no_think", "tool_close"}:
        return "answer"
    if kind == "tool_open":
        return "tool_call"
    return current


def _segment_text(segment: ContentSegment) -> str:
    return segment.text.strip()


def format_assistant_markdown(source: str) -> str:
    """Render a streamed assistant turn as a single markdown block."""

    think_parts: list[str] = []
    answer_parts: list[str] = []
    tool_parts: list[str] = []
    for segment in split_assistant_segments(source):
        text = _segment_text(segment)
        if not text:
            continue
        if segment.kind == "think":
            think_parts.append(text)
        elif segment.kind == "tool_call":
            tool_parts.append(text)
        else:
            answer_parts.append(text)
    blocks: list[str] = []
    if think_parts:
        blocks.append("**Reasoning**\n\n" + "\n\n".join(think_parts))
    if answer_parts:
        blocks.append("\n\n".join(answer_parts))
    for tool in tool_parts:
        blocks.append(f"```tool_call\n{tool}\n```")
    return "\n\n".join(blocks).strip()


def format_assistant_chat_messages(source: str) -> list[dict[str, Any]]:
    """Split a turn into Gradio thoughts plus a normal assistant chat message.

    Messages with ``metadata.title`` render as collapsible thoughts. The reply
    has no metadata so Gradio shows it as the chat bubble.
    """

    think_parts: list[str] = []
    answer_parts: list[str] = []
    tool_parts: list[str] = []
    think_pending = False
    for segment in split_assistant_segments(source):
        text = _segment_text(segment)
        if segment.kind == "think":
            think_pending = segment.open
            if text:
                think_parts.append(text)
            continue
        think_pending = False
        if not text:
            continue
        if segment.kind == "tool_call":
            tool_parts.append(text)
        else:
            answer_parts.append(text)

    messages: list[dict[str, Any]] = []
    if think_parts:
        messages.append(
            {
                "role": "assistant",
                "content": "\n\n".join(think_parts),
                "metadata": {
                    "title": "Reasoning",
                    "status": "pending" if think_pending else "done",
                },
            }
        )
    for tool in tool_parts:
        messages.append(
            {
                "role": "assistant",
                "content": f"```json\n{tool}\n```",
                "metadata": {"title": "Tool call", "status": "done"},
            }
        )
    answer = "\n\n".join(answer_parts)
    if answer or not messages:
        messages.append({"role": "assistant", "content": answer})
    return messages


def split_assistant_segments(source: str) -> list[ContentSegment]:
    """Split a completed (or in-progress) assistant turn for tests / replay."""

    segments: list[ContentSegment] = []
    mode: SegmentKind = "answer"
    cursor = 0
    while cursor < len(source):
        found = first_marker(source[cursor:])
        if found is None:
            tail = source[cursor:]
            if tail:
                segments.append(ContentSegment(kind=mode, text=tail, open=True))
            break
        at, marker, kind = found
        at += cursor
        if at > cursor:
            segments.append(
                ContentSegment(kind=mode, text=source[cursor:at], open=False)
            )
        mode = mode_after_marker(kind, mode)
        cursor = at + len(marker)
    return [segment for segment in segments if segment.text]


class LiveReasoningStreamer:
    """Colour reasoning and tool-call regions as decoded text grows."""

    def __init__(
        self,
        writer: TextIO | None = None,
        *,
        color: bool | None = None,
    ) -> None:
        self.writer = writer if writer is not None else sys.stdout
        if color is None:
            color = bool(getattr(self.writer, "isatty", lambda: False)())
        self.color = color
        self._seen = ""
        self._hold = ""
        self._mode: SegmentKind = "answer"
        self._style: SegmentKind | None = None
        self._emitted_think_label = False
        self._emitted_tool_label = False

    def update(self, decoded: str) -> None:
        if decoded.startswith(self._seen):
            delta = decoded[len(self._seen) :]
        else:
            delta = decoded
        self._seen = decoded
        if delta:
            self._consume(delta, final=False)

    def finish(self) -> str:
        if self._hold:
            self._emit(self._hold)
            self._hold = ""
        self._set_style(None)
        self.writer.write("\n")
        self.writer.flush()
        return self._seen

    def _consume(self, delta: str, *, final: bool) -> None:
        buffer = self._hold + delta
        self._hold = ""
        while buffer:
            found = first_marker(buffer)
            if found is None:
                keep = 0 if final else holdback_prefix_length(buffer)
                if keep:
                    self._emit(buffer[:-keep])
                    self._hold = buffer[-keep:]
                else:
                    self._emit(buffer)
                return
            at, marker, kind = found
            if at:
                self._emit(buffer[:at])
            self._switch(kind)
            buffer = buffer[at + len(marker) :]
        if final:
            return

    def _switch(self, kind: MarkerKind) -> None:
        nxt = mode_after_marker(kind, self._mode)
        if nxt != self._mode and nxt == "answer":
            self._emit_plain("\n")
        self._mode = nxt
        if nxt == "think" and not self._emitted_think_label:
            self._emit_plain("\n")
            self._set_style("think")
            self._emit_plain("reasoning  ")
            self._emitted_think_label = True
        elif nxt == "tool_call" and not self._emitted_tool_label:
            self._emit_plain("\n")
            self._set_style("tool_call")
            self._emit_plain("tool_call  ")
            self._emitted_tool_label = True
        elif nxt == "answer":
            self._set_style("answer")

    def _emit(self, text: str) -> None:
        if not text:
            return
        self._set_style(self._mode)
        self.writer.write(text)
        self.writer.flush()

    def _emit_plain(self, text: str) -> None:
        if not text:
            return
        self._set_style(None)
        self.writer.write(text)
        self.writer.flush()

    def _set_style(self, kind: SegmentKind | None) -> None:
        if not self.color:
            self._style = kind
            return
        if kind == self._style:
            return
        self.writer.write(ANSI["reset"])
        if kind in {"think", "tool_call"}:
            self.writer.write(ANSI[kind])
        self._style = kind


def prompt_tool_results(
    calls: Sequence[ParsedToolCall],
    *,
    read_line: Callable[[str], str],
    writer: TextIO | None = None,
    execute_web_search: bool = True,
) -> list[str]:
    out = writer if writer is not None else sys.stdout
    results: list[str] = []
    auto_search = sum(
        1
        for call in calls
        if execute_web_search and is_web_search_tool_name(call.name)
    )
    auto_math = sum(1 for call in calls if is_auto_math_tool_name(call.name))
    if auto_search:
        out.write(
            f"\n{auto_search} web-search call(s) will run automatically "
            f"(Tavily JSON, sent as {TOOL_RESPONSE_TOKEN}).\n"
        )
        out.flush()
    if auto_math:
        out.write(
            f"\n{auto_math} python/calculator call(s) will run automatically "
            f"(sent as {TOOL_RESPONSE_TOKEN}).\n"
        )
        out.flush()
    manual = len(calls) - auto_search - auto_math
    if manual:
        out.write(
            f"\n{manual} tool call(s). Paste each observation; "
            f"it is sent as {TOOL_RESPONSE_TOKEN}.\n"
        )
        out.flush()
    for index, call in enumerate(calls, start=1):
        out.write(
            f"\n[{index}/{len(calls)}] "
            f"{format_smoltalk_tool_call(call.name, call.arguments)}\n"
        )
        out.flush()
        if execute_web_search and is_web_search_tool_name(call.name):
            query = query_from_arguments(call.arguments)
            out.write(f"searching {query!r}…\n")
            out.flush()
            results.append(search_web(query))
            continue
        if is_auto_math_tool_name(call.name):
            out.write("running python…\n")
            out.flush()
            try:
                results.append(execute_math_tool(call.name, call.arguments))
            except Exception as error:
                results.append(f"Error: {error}")
            continue
        results.append(read_line(f"result[{call.name}]> "))
    return results


def append_tool_round(
    messages: list[dict[str, Any]],
    assistant_text: str,
    raw_results: Sequence[str],
) -> None:
    messages.append({"role": "assistant", "content": assistant_text})
    for raw in raw_results:
        wrapped = wrap_smoltalk_tool_result(raw)
        if not wrapped:
            raise ValueError("Tool result cannot be empty")
        messages.append({"role": "tool", "content": raw})


def resolve_cli_tools(spec: str | None, tools_json: str | None) -> list[dict[str, Any]] | None:
    tools: list[dict[str, Any]] = []
    if spec:
        for name in spec.split(","):
            key = name.strip().lower()
            if not key or key in {"none", "off"}:
                continue
            if is_web_search_tool_name(key):
                tools.append(NEMOTRON_WEB_SEARCH_TOOL)
                continue
            if key in {"python", "python-exec", "code-interpreter"}:
                tools.append(NEMOTRON_PYTHON_EXEC_TOOL)
                continue
            if key in {"calculator", "calc"}:
                tools.append(CALCULATOR_TOOL)
                continue
            raise ValueError(
                f"Unknown built-in tool {name!r}. "
                "Use web-search, python, calculator, or --tools-json."
            )
    if tools_json:
        payload = json.loads(tools_json)
        if isinstance(payload, dict):
            tools.append(payload)
        elif isinstance(payload, list):
            tools.extend(payload)
        else:
            raise ValueError("tools JSON must be an object or array")
    return tools or None


@dataclass
class ChatLoopState:
    messages: list[dict[str, Any]] = field(default_factory=list)
    system: str = ""
    enable_thinking: bool = True
    tools: list[dict[str, Any]] | None = None

    def reset(self) -> None:
        self.messages = []
        reset_default_python_session()
        if self.system.strip():
            self.messages.append({"role": "system", "content": self.system.strip()})

    def add_user(self, text: str) -> None:
        self.messages.append({"role": "user", "content": text})


def apply_slash_command(state: ChatLoopState, text: str) -> str | None:
    """Return a status string if `text` is a slash command, else None."""

    command = text.strip()
    lowered = command.lower()
    if lowered in {"/exit", "/quit"}:
        return "exit"
    if lowered == "/reset":
        state.reset()
        return "history cleared"
    if lowered == "/think":
        if state.messages:
            return "thinking can only be changed on a fresh conversation (/reset first)"
        state.enable_thinking = True
        return "thinking on — next assistant turn is prefixed with <|think|>"
    if lowered in {"/no_think", "/nothink"}:
        if state.messages:
            return "thinking can only be changed on a fresh conversation (/reset first)"
        state.enable_thinking = False
        return "thinking off — next assistant turn is prefixed with <|no_think|>"
    if lowered.startswith("/system"):
        rest = command[len("/system") :].strip()
        state.system = rest
        state.reset()
        return "system prompt updated" if rest else "system prompt cleared"
    return None