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 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:
    kind: Literal["wait", "capped", "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,
    ) -> None:
        self._client = client
        self._model = model
        self._max_turns = max_turns
        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,
    ) -> RunOutcome:
        """``handlers`` answer control tools that don't end the run (e.g.
        start_milestone): each returns a result and whether it succeeded."""
        handlers = handlers or {}
        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=self._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)
            if not calls:
                if nudged:
                    return RunOutcome("failed", note="kept replying in prose")
                nudged = True
                system = f"{system}\n\n{_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}")
                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)
                )
        return RunOutcome("capped", note=f"hit the {self._max_turns}-turn cap")