nisse/nisse-dev/src/nisse_dev/loop/driver.py

"""A driver run (L2): one model with tools, looping until it calls a waiting
tool, runs out of turns, or is stopped.

Everything the run does is written to the item's transcript as the driver
(L7), and every model call sees the driver's view of it (L6). Execution tools
go to the build VM (L13); control tools are answered here, and a waiting one
ends the run, so the scheduler can put it to the operator.
"""

from collections.abc import Callable, Mapping
from dataclasses import dataclass
from typing import Any, Literal

from nisse_dev.cancel import Cancel, skip_calls
from nisse_dev.canonical import AssistantText, Fact, ToolCall, ToolResult
from nisse_dev.llm import LLMClient, LLMError
from nisse_dev.loop.tools import WAITING, names
from nisse_dev.state import Transcript
from nisse_dev.tools import Workspace

DEFAULT_MAX_TURNS = 30

# A reply in prose breaks the tool contract; the model is told once.
_PROSE_NUDGE = (
    "Your last reply was prose. Every reply must be a tool call: use one of "
    "your tools, or ask_human if you need the operator."
)


@dataclass(frozen=True)
class RunOutcome:
    # stuck: failed commands in a row hit the failure limit (L9's trigger)
    kind: Literal["wait", "capped", "stuck", "stopped", "failed"]
    tool: str = ""  # the waiting tool called
    arguments: dict[str, Any] | None = None
    note: str = ""


class Driver:
    def __init__(
        self,
        client: LLMClient,
        model: str,
        *,
        max_turns: int = DEFAULT_MAX_TURNS,
        cancel: Cancel | None = None,
        failure_limit: int = 0,  # 0: no limit
    ) -> None:
        self._client = client
        self._model = model
        self._max_turns = max_turns
        self._failure_limit = failure_limit
        self._cancel = cancel or Cancel()

    def run(
        self,
        transcript: Transcript,
        *,
        system: str,
        tools: list[dict[str, Any]],
        workspace: Workspace | None,
        handlers: Mapping[str, Callable[[dict[str, Any]], tuple[str, bool]]]
        | None = None,
        checks: Mapping[str, Callable[[dict[str, Any]], str | None]] | None = None,
        model: str | None = None,
        on_turn: Callable[[str], None] | None = None,
    ) -> RunOutcome:
        """``handlers`` answer control tools that don't end the run (e.g.
        start_milestone): each returns a result and whether it succeeded.
        ``checks`` vet a waiting tool's arguments: a problem is the call's
        failed result and the run goes on, so the model puts it right."""
        handlers = handlers or {}
        checks = checks or {}
        model = model or self._model  # the work item's pinned driver (L9)
        failures: list[str] = []  # the failed commands in a row, latest last
        offered = names(tools)
        nudged = False
        for turn in range(self._max_turns):
            # L2: stop takes effect at the next tool boundary.
            if self._cancel.requested:
                return RunOutcome("stopped", note=f"before turn {turn + 1}")
            try:
                produced = self._client.complete(
                    transcript.view(), tools=tools, model=model, system=system
                )
            except LLMError as exc:
                return RunOutcome("failed", note=f"the driver's call failed: {exc}")
            calls: list[ToolCall] = []
            for event in produced:
                transcript.append(event)
                if isinstance(event, ToolCall):
                    calls.append(event)
            # The model's own words for the turn: its progress, in prose (#42).
            said = " ".join(
                e.text.strip() for e in produced if isinstance(e, AssistantText)
            ).strip()
            if said and calls and on_turn is not None:
                on_turn(said)
            if not calls:
                if nudged:
                    return RunOutcome("failed", note="kept replying in prose")
                nudged = True
                # In the conversation, as the harness: a request must end on
                # the user's side (models refuse assistant prefill), and the
                # nudge is then in the log as what it was.
                transcript.as_actor("orchestrator").append(Fact(text=_PROSE_NUDGE))
                continue
            for index, call in enumerate(calls):
                if self._cancel.requested:
                    skip_calls(transcript, calls[index:])
                    return RunOutcome("stopped", note=f"during turn {turn + 1}")
                problem = (
                    checks[call.tool](dict(call.arguments))
                    if call.tool in checks and call.tool in offered
                    else None
                )
                if problem is not None:
                    transcript.append(
                        ToolResult(call_id=call.call_id, result=problem, success=False)
                    )
                    skip_calls(
                        transcript, calls[index + 1 :], "not run: fix the call above"
                    )
                    break
                if call.tool in WAITING and call.tool in offered:
                    transcript.append(
                        ToolResult(
                            call_id=call.call_id,
                            result="posted; your turn has ended",
                            success=True,
                        )
                    )
                    skip_calls(
                        transcript, calls[index + 1 :], "not run: the turn had ended"
                    )
                    return RunOutcome("wait", call.tool, dict(call.arguments))
                if call.tool in handlers and call.tool in offered:
                    result, ok = handlers[call.tool](dict(call.arguments))
                elif call.tool not in offered or workspace is None:
                    result, ok = (
                        f"no such tool: {call.tool}; the tools are "
                        + ", ".join(sorted(offered)),
                        False,
                    )
                else:
                    outcome = workspace.dispatch(call.tool, call.arguments)
                    result, ok = outcome.result, outcome.success
                transcript.append(
                    ToolResult(call_id=call.call_id, result=result, success=ok)
                )
                if call.tool == "run_command":
                    failures = [] if ok else [*failures, _brief(call, result)]
                    if self._failure_limit and len(failures) >= self._failure_limit:
                        skip_calls(
                            transcript, calls[index + 1 :], "not run: you're stuck"
                        )
                        return RunOutcome(
                            "stuck",
                            note=f"{len(failures)} commands failed in a row: "
                            + "; ".join(failures),
                        )
        return RunOutcome("capped", note=f"hit the {self._max_turns}-turn cap")


def _brief(call: ToolCall, result: str) -> str:
    """A failed command and the end of what it said, for a helper's brief."""
    command = str(call.arguments.get("command", ""))[:120]
    return f"`{command}`: {result.strip()[-300:]}"