| """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 |
|
|
| |
|
|
| @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()] |
|
|
| |
|
|
| 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") |
| |
| 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] |
|
|
|
|
