219 lines
8.2 KiB
Python
219 lines
8.2 KiB
Python
"""The budgeted agent loop: model turn -> tool calls -> tool results -> repeat.
|
|
|
|
Determinism
|
|
-----------
|
|
The loop takes no clock and no randomness. A transcript is therefore a pure
|
|
function of ``(task, model, tools, budgets)``: run the same loop twice and the
|
|
two transcripts are byte-identical. Tests/test_agentloop.py asserts exactly
|
|
that, which is what makes the hash chain a usable integrity check rather than
|
|
decoration.
|
|
|
|
Stop reasons, and their precedence
|
|
-----------------------------------
|
|
step_budget checked BEFORE each model turn, so ``steps`` never exceeds
|
|
``max_steps`` and always equals the number of assistant records.
|
|
sentinel the assistant's own text contains ``stop_sentinel``.
|
|
no_tool_calls the assistant asked for no tools -- the ordinary end of a run.
|
|
token_budget cumulative estimated tokens exceeded ``max_tokens``; the
|
|
over-budget turn is recorded, then the loop stops WITHOUT
|
|
executing its tool calls.
|
|
model_error the model callable raised, or returned a non-ModelReply.
|
|
|
|
Within one turn the order is: record the turn, then sentinel, then
|
|
no_tool_calls, then token_budget. A model that says DONE and asks for a tool
|
|
in the same turn stops on the sentinel and never runs the tool; a model that
|
|
spends its last token on tool calls stops on the budget and never runs them.
|
|
|
|
Token accounting is a 4-chars-per-token heuristic, not a tokenizer. It bounds
|
|
the loop, which is the property that matters; it does not claim to measure a
|
|
real model's bill.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
from .tools import ToolCall, ToolError, ToolRegistry, ToolResult, ValidationError
|
|
from .transcript import KIND_ASSISTANT, KIND_FINAL, KIND_TOOL, KIND_USER, Transcript, canonical_json
|
|
|
|
__all__ = ["StopReason", "ModelReply", "LoopResult", "AgentLoop", "estimate_tokens"]
|
|
|
|
|
|
class StopReason(object):
|
|
NO_TOOL_CALLS = "no_tool_calls"
|
|
SENTINEL = "sentinel"
|
|
STEP_BUDGET = "step_budget"
|
|
TOKEN_BUDGET = "token_budget"
|
|
MODEL_ERROR = "model_error"
|
|
|
|
|
|
class ModelReply(object):
|
|
"""What a model returns for one turn."""
|
|
|
|
__slots__ = ("text", "tool_calls")
|
|
|
|
def __init__(self, text="", tool_calls=None):
|
|
self.text = text or ""
|
|
self.tool_calls = list(tool_calls) if tool_calls else []
|
|
|
|
def __repr__(self):
|
|
return "ModelReply(%r, %r)" % (self.text, self.tool_calls)
|
|
|
|
|
|
class LoopResult(object):
|
|
"""The outcome of a run: messages, the budget accounting, and the log."""
|
|
|
|
__slots__ = ("messages", "steps", "tokens", "stop_reason", "transcript", "model_error")
|
|
|
|
def __init__(self, messages, steps, tokens, stop_reason, transcript, model_error=None):
|
|
self.messages = messages
|
|
self.steps = steps
|
|
self.tokens = tokens
|
|
self.stop_reason = stop_reason
|
|
self.transcript = transcript
|
|
self.model_error = model_error
|
|
|
|
@property
|
|
def ok(self):
|
|
return self.stop_reason in (StopReason.NO_TOOL_CALLS, StopReason.SENTINEL)
|
|
|
|
def tool_results(self):
|
|
"""The ToolResult records, in order, as recorded in the transcript."""
|
|
out = []
|
|
for record in self.transcript.records:
|
|
if record["kind"] == KIND_TOOL:
|
|
p = record["payload"]
|
|
out.append(ToolResult(p["name"], p["ok"], p["value"], p["error"]))
|
|
return out
|
|
|
|
def __repr__(self):
|
|
return "LoopResult(stop=%r, steps=%d, tokens=%d)" % (
|
|
self.stop_reason,
|
|
self.steps,
|
|
self.tokens,
|
|
)
|
|
|
|
|
|
def estimate_tokens(text):
|
|
"""4 chars/token, floor 1 for non-empty text. A heuristic, documented."""
|
|
if not text:
|
|
return 0
|
|
return max(1, len(text) // 4)
|
|
|
|
|
|
class AgentLoop(object):
|
|
def __init__(self, registry, max_steps=8, max_tokens=4000, stop_sentinel=None):
|
|
if not isinstance(registry, ToolRegistry):
|
|
raise TypeError("registry must be a ToolRegistry, got %s" % type(registry).__name__)
|
|
if not isinstance(max_steps, int) or isinstance(max_steps, bool) or max_steps < 1:
|
|
raise ValueError("max_steps must be a positive integer")
|
|
if not isinstance(max_tokens, int) or isinstance(max_tokens, bool) or max_tokens < 1:
|
|
raise ValueError("max_tokens must be a positive integer")
|
|
self.registry = registry
|
|
self.max_steps = max_steps
|
|
self.max_tokens = max_tokens
|
|
self.stop_sentinel = stop_sentinel
|
|
|
|
# -- one tool call ---------------------------------------------------
|
|
|
|
def _execute(self, call):
|
|
try:
|
|
tool = self.registry.get(call.name)
|
|
except ToolError as exc:
|
|
return ToolResult(call.name, False, None, str(exc))
|
|
try:
|
|
value = tool(call.arguments)
|
|
except ValidationError as exc:
|
|
return ToolResult(call.name, False, None, "invalid arguments: %s" % exc)
|
|
except ToolError as exc:
|
|
return ToolResult(call.name, False, None, str(exc))
|
|
except Exception as exc: # defensive: a transcript must stay writable
|
|
return ToolResult(
|
|
call.name, False, None, "tool raised %s: %s" % (type(exc).__name__, exc)
|
|
)
|
|
# A value that cannot be serialised cannot enter the transcript, and a
|
|
# result that cannot be recorded is not a result.
|
|
try:
|
|
canonical_json(value)
|
|
except (TypeError, ValueError) as exc:
|
|
return ToolResult(
|
|
call.name, False, None, "tool returned a non-serialisable value: %s" % exc
|
|
)
|
|
return ToolResult(call.name, True, value)
|
|
|
|
# -- the loop --------------------------------------------------------
|
|
|
|
def run(self, task, model, transcript=None):
|
|
if not callable(model):
|
|
raise TypeError("model must be callable")
|
|
tr = transcript if transcript is not None else Transcript()
|
|
tr.append(KIND_USER, {"text": task})
|
|
|
|
messages = [{"role": "user", "content": task}]
|
|
steps = 0
|
|
tokens = 0
|
|
stop_reason = None
|
|
model_error = None
|
|
|
|
while True:
|
|
if steps >= self.max_steps:
|
|
stop_reason = StopReason.STEP_BUDGET
|
|
break
|
|
|
|
try:
|
|
reply = model(messages, self.registry.manifest())
|
|
except Exception as exc:
|
|
stop_reason = StopReason.MODEL_ERROR
|
|
model_error = "model raised %s: %s" % (type(exc).__name__, exc)
|
|
break
|
|
if not isinstance(reply, ModelReply):
|
|
stop_reason = StopReason.MODEL_ERROR
|
|
model_error = "model returned %s, expected ModelReply" % type(reply).__name__
|
|
break
|
|
|
|
steps += 1
|
|
tokens += estimate_tokens(reply.text)
|
|
|
|
calls = [
|
|
c if isinstance(c, ToolCall) else ToolCall(**c) for c in (reply.tool_calls or [])
|
|
]
|
|
recorded_calls = [c.as_dict() for c in calls]
|
|
tr.append(KIND_ASSISTANT, {"text": reply.text, "tool_calls": recorded_calls})
|
|
messages.append(
|
|
{"role": "assistant", "content": reply.text, "tool_calls": recorded_calls}
|
|
)
|
|
|
|
if self.stop_sentinel and self.stop_sentinel in reply.text:
|
|
stop_reason = StopReason.SENTINEL
|
|
break
|
|
if not calls:
|
|
stop_reason = StopReason.NO_TOOL_CALLS
|
|
break
|
|
if tokens > self.max_tokens:
|
|
stop_reason = StopReason.TOKEN_BUDGET
|
|
break
|
|
|
|
for call in calls:
|
|
result = self._execute(call)
|
|
tr.append(KIND_TOOL, result.as_dict())
|
|
content = json.dumps(result.value) if result.ok else result.error
|
|
messages.append(
|
|
{
|
|
"role": "tool",
|
|
"content": content,
|
|
"name": call.name,
|
|
"ok": result.ok,
|
|
}
|
|
)
|
|
|
|
tr.append(
|
|
KIND_FINAL,
|
|
{
|
|
"stop_reason": stop_reason,
|
|
"steps": steps,
|
|
"tokens": tokens,
|
|
"model_error": model_error,
|
|
},
|
|
)
|
|
return LoopResult(messages, steps, tokens, stop_reason, tr, model_error)
|