narrator/client.py

"""Minimal llama.cpp llama-server client — LOCALHOST ONLY.

The narrator is the only component permitted to call a model, and only for prose
(Non-negotiable #1). This client speaks the OpenAI-compatible /v1/chat/completions
API exposed by `llama-server`. It refuses any non-localhost base URL so trace
content can never leave the machine (Non-negotiable #2: no network egress, ever).

It never reads raw JSONL and never produces or mutates a number — it only carries
prompts to the local model and returns the text the model wrote.
"""

from __future__ import annotations

import json
import time
import urllib.error
import urllib.parse
import urllib.request
from typing import Any, Callable

DEFAULT_BASE = "http://127.0.0.1:12781"

# Only these hosts are ever allowed. localhost model calls are fine; anything
# else is a hard refusal so trace content cannot be exfiltrated.
_ALLOWED_HOSTS = {"127.0.0.1", "localhost", "::1", "[::1]"}


class NarratorClient:
    def __init__(self, base_url: str = DEFAULT_BASE, *, timeout: float = 120.0):
        host = urllib.parse.urlparse(base_url).hostname or ""
        if host not in _ALLOWED_HOSTS:
            raise ValueError(
                f"NarratorClient refuses non-localhost host {host!r}. "
                "Trace content must never leave the machine."
            )
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout

    # --- liveness -----------------------------------------------------------
    def wait_until_ready(self, max_wait: float = 90.0, interval: float = 2.0) -> bool:
        """Poll /health then /v1/models until the server can serve, or time out."""
        deadline = time.time() + max_wait
        while time.time() < deadline:
            if self._health_ok() and self._models_ok():
                return True
            time.sleep(interval)
        return False

    def _health_ok(self) -> bool:
        try:
            with urllib.request.urlopen(
                self.base_url + "/health", timeout=5
            ) as resp:
                if resp.status != 200:
                    return False
                body = json.loads(resp.read().decode("utf-8") or "{}")
                # llama-server reports {"status":"ok"} once the model is loaded;
                # while loading it returns 503, so a 200 here is enough.
                return body.get("status", "ok") in ("ok", "loading", "no slot available") or True
        except (urllib.error.URLError, ValueError, ConnectionError, OSError):
            return False

    def _models_ok(self) -> bool:
        try:
            with urllib.request.urlopen(
                self.base_url + "/v1/models", timeout=5
            ) as resp:
                if resp.status != 200:
                    return False
                body = json.loads(resp.read().decode("utf-8") or "{}")
                return bool(body.get("data"))
        except (urllib.error.URLError, ValueError, ConnectionError, OSError):
            return False

    def model_id(self) -> str:
        try:
            with urllib.request.urlopen(
                self.base_url + "/v1/models", timeout=5
            ) as resp:
                body = json.loads(resp.read().decode("utf-8") or "{}")
                data = body.get("data") or []
                if data:
                    return str(data[0].get("id", "local-model"))
        except (urllib.error.URLError, ValueError, ConnectionError, OSError):
            pass
        return "local-model"

    # --- generation ---------------------------------------------------------
    def chat(
        self,
        system: str,
        user: str,
        *,
        temperature: float = 0.1,
        max_tokens: int = 320,
        seed: int = 7,
    ) -> str:
        """One deterministic-ish completion. Low temperature + fixed seed so the
        prose is stable across runs; the model writes English only."""
        payload = {
            "messages": [
                {"role": "system", "content": system},
                {"role": "user", "content": user},
            ],
            "temperature": temperature,
            "top_p": 0.9,
            "max_tokens": max_tokens,
            "seed": seed,
            "stream": False,
            # Qwen3 thinking mode off — we want prose, not chain-of-thought.
            "chat_template_kwargs": {"enable_thinking": False},
        }
        data = json.dumps(payload).encode("utf-8")
        req = urllib.request.Request(
            self.base_url + "/v1/chat/completions",
            data=data,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=self.timeout) as resp:
            body = json.loads(resp.read().decode("utf-8"))
        choices = body.get("choices") or []
        if not choices:
            return ""
        text = (choices[0].get("message") or {}).get("content", "") or ""
        return _strip_think(text).strip()

    # --- tool-calling generation -------------------------------------------- #
    def chat_with_tools(
        self,
        system: str,
        user: str,
        *,
        tools: list[dict[str, Any]],
        exec_tool: Callable[[str, dict[str, Any]], str],
        temperature: float = 0.2,
        max_tokens: int = 400,
        max_rounds: int = 4,
        seed: int = 7,
    ) -> str:
        """An OpenAI-style tool-calling loop, driven entirely by the LOCAL model.

        The model decides which of `tools` to call (the schemas are passed with
        tool_choice 'auto'); for each tool_call we run `exec_tool(name, args)` — a
        callable supplied by the caller that returns a SHORT string — and feed the
        result back as a 'tool' message, looping until the model stops calling tools
        or `max_rounds` is reached. Returns the final assistant prose (think-stripped).

        Still localhost-only and still prose-only at the seam: this client carries
        tool schemas + results but never decides what a tool does. If the model has
        no tool support (it never emits tool_calls), this degrades to a plain
        completion — we just return whatever content it produced (graceful, #7).
        """
        messages: list[dict[str, Any]] = [
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ]
        last_content = ""
        for _ in range(max(1, max_rounds)):
            payload = {
                "messages": messages,
                "tools": tools,
                "tool_choice": "auto",
                "temperature": temperature,
                "top_p": 0.9,
                "max_tokens": max_tokens,
                "seed": seed,
                "stream": False,
                # thinking off — we want decisions + prose, not chain-of-thought.
                "chat_template_kwargs": {"enable_thinking": False},
            }
            data = json.dumps(payload).encode("utf-8")
            req = urllib.request.Request(
                self.base_url + "/v1/chat/completions",
                data=data,
                headers={"Content-Type": "application/json"},
                method="POST",
            )
            with urllib.request.urlopen(req, timeout=self.timeout) as resp:
                body = json.loads(resp.read().decode("utf-8"))
            choices = body.get("choices") or []
            if not choices:
                break
            msg = (choices[0].get("message") or {})
            last_content = msg.get("content", "") or last_content
            tool_calls = msg.get("tool_calls") or []
            if not tool_calls:
                # model is done (or never supported tools) — return its prose.
                break
            # echo the assistant turn back verbatim, then answer each tool_call.
            messages.append(msg)
            for call in tool_calls:
                fn = (call.get("function") or {})
                name = fn.get("name") or ""
                try:
                    args = json.loads(fn.get("arguments") or "{}")
                    if not isinstance(args, dict):
                        args = {}
                except (ValueError, TypeError):
                    args = {}
                try:
                    result = exec_tool(name, args)
                except Exception as exc:  # a tool must never crash the loop
                    result = f"tool error: {exc}"
                messages.append({
                    "role": "tool",
                    "tool_call_id": call.get("id", ""),
                    "content": str(result),
                })
        return _strip_think(last_content).strip()


def _strip_think(text: str) -> str:
    """Drop any <think>...</think> block a reasoning model may emit, keeping prose."""
    if "</think>" in text:
        text = text.rsplit("</think>", 1)[-1]
    return text.replace("<think>", "").strip()