| """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 ( |
| SMOLTALK_WEB_SEARCH_TOOL, |
| TOOL_RESPONSE_TOKEN, |
| format_smoltalk_tool_call, |
| wrap_smoltalk_tool_result, |
| ) |
| from tiny_gdn.tools import ParsedToolCall, parse_tool_calls |
|
|
|
|
| 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 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, |
| ) -> list[str]: |
| out = writer if writer is not None else sys.stdout |
| results: list[str] = [] |
| out.write( |
| f"\n{len(calls)} 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)}] {format_smoltalk_tool_call(call.name, call.arguments)}\n") |
| out.flush() |
| raw = read_line(f"result[{call.name}]> ") |
| results.append(raw) |
| 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 key == "web_search": |
| tools.append(SMOLTALK_WEB_SEARCH_TOOL) |
| continue |
| raise ValueError( |
| f"Unknown built-in tool {name!r}. Use web_search 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 = [] |
| 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 |
|
|
