distinct_tools/workspace.py

"""Tools that hold state for the model instead of asking it to.

**What the literature actually supports.** The oldest and best-established
result here is Nye et al., *Show Your Work: Scratchpads for Intermediate
Computation with Language Models* (ICLR 2022): a model allowed to emit
intermediate steps into a buffer before answering solves multi-step problems
that the same model cannot solve in one shot. The mechanism is not
mysterious. It moves working memory out of a single forward pass and into
text, and the smaller the model the more of its capacity a single pass has to
spend on the question itself.

The efficient-agent literature since then has converged on a second idea from
the same direction: an agent that keeps its own list of what remains to be
done, and reads it back, degrades far more slowly as a task lengthens than one
carrying that list in context. That is the case for a to-do list as a tool
rather than as a prompt convention.

**Where this stops, and why that matters here.** ``TOOL_CATEGORIES_REVIEW.md``
already records the one direct small-model measurement this project found, and
it is a negative result: Llama-1B performed *worse* when planning for itself.
So the claim these tools are built on is narrow and it is worth stating
exactly. Externalising state is well supported. Asking a small model to
*reason about strategy* is not. ``make_plan`` and the two tools here are
therefore all storage: they write down what the model tells them, hand it back
unchanged, and never summarise, rank, infer or advise. A tool that told a 0.6B
model what to do next would be adding an opinion nobody measured.

**Why per-run and not persistent.** Both of these live for the length of one
run and are erased with its sandbox. A scratchpad that survived would be a
second, unerasable copy of the user's prompt data sitting on a volunteer's
machine, which is the one thing this network promises does not happen. The
cost is that a model cannot carry notes between turns; the alternative is a
promise this product could not keep.

Like the rest of the library, nothing here opens a socket, spawns a process,
reads a file or consults a model. Every call answers from memory, which is the
same viability bar the local tool set applies below 7B.
"""

from __future__ import annotations

import threading
from collections.abc import Mapping

from .core import Registry, ToolContext, ToolInputError, ToolSpec

#: Bounds. A small model with a runaway loop must not be able to turn a
#: scratchpad into an unbounded allocation on a volunteer's machine.
MAX_NOTES = 40
MAX_NOTE_CHARACTERS = 2_000
MAX_SCRATCHPAD_CHARACTERS = 20_000
MAX_TASKS = 40
MAX_TASK_CHARACTERS = 400

_ACTIONS_NOTE = ("write", "read", "clear")
_ACTIONS_TASK = ("add", "complete", "list", "clear")


class _RunState:
    """Per-run storage, keyed by the run that owns it.

    A dictionary keyed by run id rather than a single buffer, because one
    worker runs several requests and two users' notes sharing a buffer would
    be both a correctness bug and a disclosure. Guarded by a lock for the same
    reason the queue is: a worker may execute concurrently.
    """

    def __init__(self) -> None:
        self._lock = threading.RLock()
        self._notes: dict[str, list[str]] = {}
        self._tasks: dict[str, list[dict[str, object]]] = {}

    def notes(self, run_id: str) -> list[str]:
        with self._lock:
            return list(self._notes.get(run_id, ()))

    def append_note(self, run_id: str, text: str) -> int:
        with self._lock:
            notes = self._notes.setdefault(run_id, [])
            if len(notes) >= MAX_NOTES:
                raise ToolInputError(
                    f"the scratchpad holds at most {MAX_NOTES} notes; clear it to continue"
                )
            if sum(len(item) for item in notes) + len(text) > MAX_SCRATCHPAD_CHARACTERS:
                raise ToolInputError(
                    f"the scratchpad holds at most {MAX_SCRATCHPAD_CHARACTERS} "
                    "characters; clear it to continue"
                )
            notes.append(text)
            return len(notes)

    def clear_notes(self, run_id: str) -> int:
        with self._lock:
            return len(self._notes.pop(run_id, ()))

    def tasks(self, run_id: str) -> list[dict[str, object]]:
        with self._lock:
            return [dict(task) for task in self._tasks.get(run_id, ())]

    def add_task(self, run_id: str, text: str) -> int:
        with self._lock:
            tasks = self._tasks.setdefault(run_id, [])
            if len(tasks) >= MAX_TASKS:
                raise ToolInputError(
                    f"the list holds at most {MAX_TASKS} tasks; complete or clear some"
                )
            tasks.append({"task": text, "done": False})
            return len(tasks)

    def complete_task(self, run_id: str, number: int) -> dict[str, object]:
        with self._lock:
            tasks = self._tasks.get(run_id, [])
            if not 1 <= number <= len(tasks):
                raise ToolInputError(
                    f"there is no task {number}; the list has {len(tasks)}"
                )
            tasks[number - 1]["done"] = True
            return dict(tasks[number - 1])

    def clear_tasks(self, run_id: str) -> int:
        with self._lock:
            return len(self._tasks.pop(run_id, ()))

    def forget(self, run_id: str) -> None:
        """Erase everything this run wrote. Called when its sandbox closes."""

        with self._lock:
            self._notes.pop(run_id, None)
            self._tasks.pop(run_id, None)


STATE = _RunState()


def _run_id(context: ToolContext) -> str:
    value = getattr(context, "job_id", "") or getattr(context, "run_id", "")
    if not isinstance(value, str) or not value:
        raise ToolInputError("this tool needs a run to belong to")
    return value


def _text_argument(arguments: Mapping[str, object], key: str, limit: int) -> str:
    value = arguments.get(key)
    if not isinstance(value, str) or not value.strip():
        raise ToolInputError(f"{key} must be a non-empty string")
    text = " ".join(value.split())
    if len(text) > limit:
        raise ToolInputError(f"{key} must be at most {limit} characters")
    return text


def _scratchpad(arguments: Mapping[str, object], context: ToolContext) -> Mapping[str, object]:
    """Write a note, read the notes back, or throw them away.

    Read returns the notes verbatim and in order. It does not summarise them,
    because a summary is a judgement and the model asking for its own notes
    back is the one making the judgement.
    """

    run_id = _run_id(context)
    action = str(arguments.get("action") or "write")
    if action not in _ACTIONS_NOTE:
        raise ToolInputError(f"action must be one of {', '.join(_ACTIONS_NOTE)}")
    if action == "write":
        text = _text_argument(arguments, "note", MAX_NOTE_CHARACTERS)
        count = STATE.append_note(run_id, text)
        return {"ok": True, "result": {"written": text, "note_count": count}}
    if action == "read":
        notes = STATE.notes(run_id)
        return {
            "ok": True,
            "result": {
                "notes": notes,
                "note_count": len(notes),
                "empty": not notes,
            },
        }
    removed = STATE.clear_notes(run_id)
    return {"ok": True, "result": {"cleared": removed}}


def _todo(arguments: Mapping[str, object], context: ToolContext) -> Mapping[str, object]:
    """Keep the list of what is left to do outside the model's context."""

    run_id = _run_id(context)
    action = str(arguments.get("action") or "list")
    if action not in _ACTIONS_TASK:
        raise ToolInputError(f"action must be one of {', '.join(_ACTIONS_TASK)}")
    if action == "add":
        text = _text_argument(arguments, "task", MAX_TASK_CHARACTERS)
        number = STATE.add_task(run_id, text)
        return {"ok": True, "result": {"added": text, "number": number}}
    if action == "complete":
        number = arguments.get("number")
        if isinstance(number, bool) or not isinstance(number, int):
            raise ToolInputError("number must be an integer")
        task = STATE.complete_task(run_id, number)
        remaining = [item for item in STATE.tasks(run_id) if not item["done"]]
        return {
            "ok": True,
            "result": {"completed": task["task"], "remaining": len(remaining)},
        }
    if action == "list":
        tasks = STATE.tasks(run_id)
        return {
            "ok": True,
            "result": {
                "tasks": [
                    {"number": index, "task": item["task"], "done": bool(item["done"])}
                    for index, item in enumerate(tasks, start=1)
                ],
                "remaining": sum(1 for item in tasks if not item["done"]),
                "empty": not tasks,
            },
        }
    removed = STATE.clear_tasks(run_id)
    return {"ok": True, "result": {"cleared": removed}}


WORKSPACE_SPECS: tuple[ToolSpec, ...] = (
    ToolSpec(
        tool_id="scratchpad",
        version="1",
        kind="tool",
        description=(
            "Write a short note to yourself and read your notes back later in this "
            "same request. Use it to hold a partial result, a number you worked out, "
            "or something you found, so you do not have to carry it in your head "
            "while you do the next step. Arguments: action ('write', 'read' or "
            "'clear') and, for 'write', note. Notes are erased when this request "
            "ends and are never shared with another request."
        ),
        input_schema={
            "type": "object",
            "additionalProperties": False,
            "properties": {
                "action": {
                    "type": "string",
                    "enum": list(_ACTIONS_NOTE),
                    "description": "write a note, read them all back, or clear them",
                },
                "note": {"type": "string", "description": "the note, when writing"},
            },
            "required": ["action"],
        },
    ),
    ToolSpec(
        tool_id="todo",
        version="1",
        kind="tool",
        description=(
            "Keep a numbered list of what is left to do in this request, and tick "
            "items off as you finish them. Use it when the request has several parts "
            "so you do not lose one. Arguments: action ('add', 'complete', 'list' or "
            "'clear'), task when adding, and number when completing. The list is "
            "erased when this request ends."
        ),
        input_schema={
            "type": "object",
            "additionalProperties": False,
            "properties": {
                "action": {
                    "type": "string",
                    "enum": list(_ACTIONS_TASK),
                    "description": "add a task, complete one by number, list, or clear",
                },
                "task": {"type": "string", "description": "the task, when adding"},
                "number": {
                    "type": "integer",
                    "description": "which task to complete, counting from 1",
                },
            },
            "required": ["action"],
        },
    ),
)


def register_workspace_tools(registry: Registry, *, only: frozenset[str] | None = None) -> None:
    handlers = {"scratchpad": _scratchpad, "todo": _todo}
    for spec in WORKSPACE_SPECS:
        if only is not None and spec.ref not in only:
            continue
        registry.register(spec, handlers[spec.tool_id])


__all__ = ["WORKSPACE_SPECS", "STATE", "register_workspace_tools"]