File size: 8,839 Bytes
5ef3875 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 | """The agent loop: model -> tool calls -> results -> model ... -> final answer."""
from __future__ import annotations
import enum
from dataclasses import dataclass, field
from ycode.agent.planner import Plan, UpdatePlanTool
from ycode.errors import YCodeError
from ycode.llm.base import (
AssistantMessage,
LLMProvider,
Message,
StopReason,
StreamCallbacks,
ToolCall,
ToolResultBlock,
ToolResultsMessage,
ToolSpec,
Usage,
UserMessage,
)
from ycode.tools.base import ToolContext, ToolRegistry, ToolResult
class TaskStatus(enum.Enum):
COMPLETED = "completed"
STEP_LIMIT = "step_limit"
INTERRUPTED = "interrupted"
REFUSED = "refused"
ERROR = "error"
@dataclass
class TaskResult:
status: TaskStatus
final_text: str = ""
steps: int = 0
changed_files: list[str] = field(default_factory=list)
commands_run: list[tuple[str, int | None]] = field(default_factory=list)
error: YCodeError | None = None
class AgentEvents:
"""Callbacks the UI implements. All methods are optional no-ops."""
def on_step_start(self, step: int, max_steps: int) -> None: ...
def on_text(self, chunk: str) -> None: ...
def on_tool_call_streaming(self, name: str) -> None: ...
def on_response_done(self, message: AssistantMessage) -> None: ...
def on_tool_start(self, call: ToolCall, description: str) -> None: ...
def on_tool_end(self, call: ToolCall, result: ToolResult) -> None: ...
def on_plan(self, plan: Plan) -> None: ...
def on_notice(self, message: str, level: str = "info") -> None: ...
_TRUNCATED_TOOL_INPUT = (
"Your response hit the output token limit while writing this tool call, so its input was "
"truncated and it was NOT executed. Retry with smaller pieces (e.g. several edit_file calls "
"instead of one huge write_file)."
)
_CONTINUE_AFTER_MAX_TOKENS = (
"Your previous response was cut off by the output token limit. Continue exactly where you left off."
)
_INTERRUPTED_NOTE = "[Note: the user interrupted the previous request before it finished.]"
class Agent:
def __init__(
self,
*,
provider: LLMProvider,
registry: ToolRegistry,
tool_context: ToolContext,
system_prompt: str,
max_steps: int = 50,
events: AgentEvents | None = None,
) -> None:
self.provider = provider
self.registry = registry
self.ctx = tool_context
self.system_prompt = system_prompt
self.max_steps = max_steps
self.events = events or AgentEvents()
self.messages: list[Message] = []
self.usage = Usage()
self.last_steps = 0
self.total_steps = 0
self._pending_note: str | None = None
self.plan_tool = UpdatePlanTool(on_update=self.events.on_plan)
if registry.get(self.plan_tool.name) is None:
registry.register(self.plan_tool)
else:
existing = registry.get(self.plan_tool.name)
if isinstance(existing, UpdatePlanTool):
self.plan_tool = existing
existing.on_update = self.events.on_plan
# ------------------------------------------------------------- state
@property
def plan(self) -> Plan:
return self.plan_tool.plan
def reset(self) -> None:
"""Forget the conversation (the /clear command)."""
self.messages.clear()
self.plan_tool.plan = Plan()
self.last_steps = 0
self._pending_note = None
def tool_specs(self) -> list[ToolSpec]:
return [ToolSpec(s["name"], s["description"], s["parameters"]) for s in self.registry.specs()]
# -------------------------------------------------------------- loop
def run(self, task: str) -> TaskResult:
text = task.strip()
if self._pending_note:
text = f"{self._pending_note}\n\n{text}"
self._pending_note = None
self.messages.append(UserMessage(text))
self.ctx.commands_run = []
changed_before = set(self.ctx.changed_files)
self.plan_tool.plan = Plan()
result = self._loop()
result.changed_files = sorted(self.ctx.changed_files - changed_before)
result.commands_run = list(self.ctx.commands_run)
self.last_steps = result.steps
self.total_steps += result.steps
return result
def _loop(self) -> TaskResult:
tools = self.tool_specs() if self.provider.supports_tools else []
callbacks = StreamCallbacks(on_text=self.events.on_text, on_tool_call_start=self.events.on_tool_call_streaming)
steps = 0
final_text = ""
while True:
if steps >= self.max_steps:
self.events.on_notice(
f"Reached the step limit ({self.max_steps}). Say 'continue' to keep going.", "warning"
)
return TaskResult(TaskStatus.STEP_LIMIT, final_text, steps)
steps += 1
self.events.on_step_start(steps, self.max_steps)
try:
response = self.provider.complete(
system=self.system_prompt, messages=self.messages, tools=tools, callbacks=callbacks
)
except KeyboardInterrupt:
self._pending_note = _INTERRUPTED_NOTE
return TaskResult(TaskStatus.INTERRUPTED, final_text, steps)
except YCodeError as exc:
return TaskResult(TaskStatus.ERROR, final_text, steps, error=exc)
self.usage.add(response.usage)
message = response.message
self.messages.append(message)
self.events.on_response_done(message)
if message.text:
final_text = message.text
if response.stop_reason is StopReason.REFUSAL:
detail = f" ({response.stop_detail})" if response.stop_detail else ""
self.events.on_notice(f"The model declined this request{detail}.", "error")
# Drop the refused turn so the conversation stays usable.
self.messages.pop()
return TaskResult(TaskStatus.REFUSED, final_text, steps)
if not message.tool_calls:
if response.stop_reason is StopReason.MAX_TOKENS:
self.messages.append(UserMessage(_CONTINUE_AFTER_MAX_TOKENS))
continue
if response.stop_reason is StopReason.PAUSE:
continue
return TaskResult(TaskStatus.COMPLETED, final_text, steps)
truncated = response.stop_reason is StopReason.MAX_TOKENS
try:
results = self._execute(message.tool_calls, truncated=truncated)
except _Interrupted as interrupted:
self.messages.append(ToolResultsMessage(interrupted.results))
self._pending_note = _INTERRUPTED_NOTE
return TaskResult(TaskStatus.INTERRUPTED, final_text, steps)
self.messages.append(ToolResultsMessage(results))
def _execute(self, calls: list[ToolCall], *, truncated: bool) -> list[ToolResultBlock]:
results: list[ToolResultBlock] = []
for index, call in enumerate(calls):
if truncated:
results.append(ToolResultBlock(call.id, _TRUNCATED_TOOL_INPUT, is_error=True))
continue
self.events.on_tool_start(call, self.registry.describe(call.name, call.arguments))
try:
outcome = self.registry.dispatch(call.name, call.arguments, self.ctx)
except KeyboardInterrupt:
done = results + [ToolResultBlock(call.id, "Interrupted by the user.", is_error=True)]
done += [
ToolResultBlock(c.id, "Not run: the user interrupted the task.", is_error=True)
for c in calls[index + 1:]
]
interrupted_result = ToolResult.error("interrupted by user", summary="Interrupted")
self.events.on_tool_end(call, interrupted_result)
raise _Interrupted(done) from None
self.events.on_tool_end(call, outcome)
results.append(ToolResultBlock(call.id, outcome.content, outcome.is_error))
return results
class _Interrupted(Exception):
def __init__(self, results: list[ToolResultBlock]) -> None:
super().__init__("interrupted")
self.results = results
def describe_status(result: TaskResult) -> str:
return {
TaskStatus.COMPLETED: "YCode completed the task.",
TaskStatus.STEP_LIMIT: "Stopped: step limit reached.",
TaskStatus.INTERRUPTED: "Interrupted.",
TaskStatus.REFUSED: "The model declined the request.",
TaskStatus.ERROR: "Stopped because of an error.",
}[result.status]
|