scrypt/warden/harness.py

"""The mi-style harness: a loop, some tools, and an LLM.

One decision per run() call. The model answers with either prose or a
single JSON tool call ({"tool": ..., "args": {...}}). Tool args are
schema-validated; anything malformed gets one retry, then the caller's
deterministic fallback. Tool handlers may only touch game state — they
receive whatever context object the game wires in, never the host.
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from typing import Callable

from scrypt.inference.backend import Backend, Message, complete

from .guardrails import GuardrailViolation, validate_args


@dataclass
class Tool:
    name: str
    description: str
    schema: dict
    handler: Callable[[dict], str]  # returns a result line fed back to the model


@dataclass
class HarnessResult:
    text: str
    tool_calls: list[tuple[str, dict]] = field(default_factory=list)
    fell_back: bool = False


def extract_json(text: str) -> dict | None:
    """First balanced {...} object in text, or None."""
    start = text.find("{")
    while start != -1:
        depth = 0
        for i in range(start, len(text)):
            if text[i] == "{":
                depth += 1
            elif text[i] == "}":
                depth -= 1
                if depth == 0:
                    try:
                        return json.loads(text[start : i + 1])
                    except json.JSONDecodeError:
                        break
        start = text.find("{", start + 1)
    return None


def render_tools(tools: list[Tool]) -> str:
    lines = ["You may answer in prose, OR call exactly one tool by answering"
             ' only JSON: {"tool": "<name>", "args": {...}}', "Tools:"]
    for t in tools:
        props = ", ".join(
            f"{k}: {v.get('type', v.get('enum'))}" for k, v in t.schema.get("properties", {}).items()
        )
        lines.append(f"- {t.name}({props}) — {t.description}")
    return "\n".join(lines)


class Harness:
    def __init__(
        self,
        backend: Backend,
        tools: list[Tool] | None = None,
        *,
        max_steps: int = 3,
        max_tokens: int = 200,
    ):
        self.backend = backend
        self.tools = {t.name: t for t in (tools or [])}
        self.max_steps = max_steps
        self.max_tokens = max_tokens

    async def run(self, messages: list[Message]) -> HarnessResult:
        """messages: fully constructed prompt (persona + digest + frame)."""
        history = list(messages)
        if self.tools:
            history[0] = {
                "role": "system",
                "content": history[0]["content"] + "\n\n" + render_tools(list(self.tools.values())),
            }
        result = HarnessResult(text="")
        retried = False
        for _ in range(self.max_steps):
            reply = await complete(self.backend, history, max_tokens=self.max_tokens)
            call = extract_json(reply) if self.tools else None
            if not call or "tool" not in call:
                result.text = reply.strip()
                return result
            name = str(call.get("tool", ""))
            tool = self.tools.get(name)
            try:
                if tool is None:
                    raise GuardrailViolation(f"unknown tool {name!r}")
                args = validate_args(tool.schema, call.get("args", {}))
            except GuardrailViolation as err:
                if retried:
                    result.fell_back = True
                    return result
                retried = True
                history.append({"role": "assistant", "content": reply})
                history.append({
                    "role": "user",
                    "content": f"That call was invalid ({err}). Answer in prose, or retry with valid JSON.",
                })
                continue
            outcome = tool.handler(args)
            result.tool_calls.append((name, args))
            history.append({"role": "assistant", "content": reply})
            history.append({"role": "user", "content": f"TOOL RESULT: {outcome}\nNow answer in prose."})
        result.fell_back = not result.text
        return result