| 1 | from __future__ import annotations |
| 2 | |
| 3 | import json |
| 4 | import asyncio |
| 5 | from datetime import datetime |
| 6 | from pathlib import Path |
| 7 | from typing import Any, AsyncIterator |
| 8 | |
| 9 | from .context_compactor import ContextCompactor, CompactionResult |
| 10 | from .llm import OpenAICompatibleLLM |
| 11 | from .models import ToolCall, ToolResult, TurnControl |
| 12 | from .prompts import PromptBuilder |
| 13 | from .session_index import SessionIndex |
| 14 | from .tool_executor import ToolExecutor |
| 15 | from .tools import ToolRegistry, build_builtin_registry |
| 16 | |
| 17 | MAX_TOOL_PASSES = 50 |
| 18 | |
| 19 | |
| 20 | class AgentLoop: |
| 21 | def __init__(self, session_index: SessionIndex, prompt_builder: PromptBuilder, tool_registry: ToolRegistry, tool_executor: ToolExecutor, llm: Any, context_compactor: ContextCompactor | None = None) -> None: |
| 22 | self.session_index = session_index |
| 23 | self.prompt_builder = prompt_builder |
| 24 | self.tool_registry = tool_registry |
| 25 | self.tool_executor = tool_executor |
| 26 | self.llm = llm |
| 27 | self.context_compactor = context_compactor or ContextCompactor(llm) |
| 28 | self.history: list[dict[str, Any]] = [] |
| 29 | |
| 30 | async def compact_history(self, *, reason: str = "manual") -> str: |
| 31 | if not self.history: |
| 32 | return "No conversation history to compact." |
| 33 | session = self.session_index.active() or self.session_index.create() |
| 34 | result = await self.context_compactor.compact( |
| 35 | self.history, |
| 36 | previous_summary=str(session.get("compacted_summary", "") or ""), |
| 37 | reason=reason, |
| 38 | ) |
| 39 | self.history = [self.context_compactor.synthetic_summary_message(result.summary), *result.preserved_messages] |
| 40 | self.session_index.update_compaction(session["session_id"], _compaction_record(result)) |
| 41 | return f"Compacted context {result.estimated_tokens_before} -> {result.estimated_tokens_after} ({result.mode})." |
| 42 | |
| 43 | async def stream_events(self, user_input: str) -> AsyncIterator[dict[str, Any]]: |
| 44 | control = TurnControl() |
| 45 | yield {"type": "turn", "turn_id": control.turn_id, "turn": {"id": control.turn_id}} |
| 46 | tool_schemas = self.tool_registry.list_function_tools() |
| 47 | parts = self.prompt_builder.build_parts(user_input) |
| 48 | system = "\n\n".join(f"## {part.title}\n{part.body}" for part in parts if part.id != "request.user") |
| 49 | if self.context_compactor.should_preflight_compact( |
| 50 | [*self.history, {"role": "user", "content": user_input}], |
| 51 | system_tokens=_prompt_tokens(parts), |
| 52 | tools_tokens=_tool_schema_tokens(tool_schemas), |
| 53 | ): |
| 54 | yield {"type": "status", "turn_id": control.turn_id, "phase": "compact", "message": "Compacting context before sampling"} |
| 55 | await self.compact_history(reason="token-pressure") |
| 56 | parts = self.prompt_builder.build_parts(user_input) |
| 57 | system = "\n\n".join(f"## {part.title}\n{part.body}" for part in parts if part.id != "request.user") |
| 58 | yield {"type": "prompt_trace", "turn_id": control.turn_id, "prompt_trace": self.prompt_builder.trace(parts)} |
| 59 | runtime_messages: list[dict[str, Any]] = [{"role": "system", "content": system}, *self.history, {"role": "user", "content": user_input}] |
| 60 | assistant_turns: list[dict[str, Any]] = [] |
| 61 | tool_rounds: list[dict[str, Any]] = [] |
| 62 | transitions: list[dict[str, str]] = [] |
| 63 | all_tool_results: list[ToolResult] = [] |
| 64 | final_text = "" |
| 65 | status = "completed" |
| 66 | tool_round = 0 |
| 67 | |
| 68 | while True: |
| 69 | yield {"type": "status", "turn_id": control.turn_id, "phase": "sampling_assistant", "message": "Sampling assistant"} |
| 70 | try: |
| 71 | assistant = await self.llm.complete(runtime_messages, tools=tool_schemas) |
| 72 | except Exception as exc: |
| 73 | status = "failed" |
| 74 | final_text = f"Agent LLM request failed: {exc}" |
| 75 | transitions.append(_transition("sampling_assistant", "finalizing_answer", "llm_sampling_failed")) |
| 76 | yield {"type": "error", "turn_id": control.turn_id, "message": final_text, "metadata": {"error_type": "llm_sampling_failed"}} |
| 77 | break |
| 78 | assistant_turns.append({"phase": "initial" if tool_round == 0 else f"followup_{tool_round}", "text": assistant.text, "tool_calls": [call.as_dict() for call in assistant.tool_calls]}) |
| 79 | if not assistant.tool_calls: |
| 80 | transitions.append(_transition("sampling_assistant", "finalizing_answer", "assistant_finished_without_tools")) |
| 81 | final_text = assistant.text |
| 82 | if final_text: |
| 83 | yield {"type": "token", "turn_id": control.turn_id, "delta": final_text} |
| 84 | break |
| 85 | transitions.append(_transition("sampling_assistant", "executing_tools", "assistant_requested_tools")) |
| 86 | if tool_round >= MAX_TOOL_PASSES: |
| 87 | status = "halted" |
| 88 | final_text = "Tool loop halted after max tool passes." |
| 89 | transitions.append(_transition("executing_tools", "finalizing_answer", "max_tool_passes_reached")) |
| 90 | yield {"type": "error", "turn_id": control.turn_id, "message": final_text, "metadata": {"max_tool_passes": MAX_TOOL_PASSES}} |
| 91 | break |
| 92 | tool_round += 1 |
| 93 | yield {"type": "status", "turn_id": control.turn_id, "phase": "executing_tools", "message": f"Running tools (round {tool_round})"} |
| 94 | runtime_messages.append({"role": "assistant", "content": assistant.text or "", "tool_calls": [_openai_tool_call(call) for call in assistant.tool_calls]}) |
| 95 | round_results: list[ToolResult] = [] |
| 96 | round_model_content: list[dict[str, Any]] = [] |
| 97 | |
| 98 | for call in assistant.tool_calls: |
| 99 | yield {"type": "tool_start", "turn_id": control.turn_id, "tool": call.as_dict()} |
| 100 | progress_queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue() |
| 101 | |
| 102 | def on_progress(event: dict[str, Any]) -> None: |
| 103 | progress_queue.put_nowait(event) |
| 104 | |
| 105 | task = asyncio.create_task(self.tool_executor.execute(call, control, progress_callback=on_progress)) |
| 106 | while not task.done(): |
| 107 | try: |
| 108 | yield await asyncio.wait_for(progress_queue.get(), timeout=0.1) |
| 109 | except asyncio.TimeoutError: |
| 110 | continue |
| 111 | while not progress_queue.empty(): |
| 112 | yield progress_queue.get_nowait() |
| 113 | record = await task |
| 114 | result = record.result |
| 115 | round_results.append(result) |
| 116 | all_tool_results.append(result) |
| 117 | yield {"type": "tool_result", "turn_id": control.turn_id, "tool_result": result.as_dict()} |
| 118 | runtime_messages.append({"role": "tool", "tool_call_id": call.id, "name": result.name, "content": json.dumps(result.as_dict(), ensure_ascii=False)}) |
| 119 | if result.model_content: |
| 120 | round_model_content.extend(result.model_content) |
| 121 | if round_model_content: |
| 122 | runtime_messages.append( |
| 123 | { |
| 124 | "role": "user", |
| 125 | "content": [ |
| 126 | { |
| 127 | "type": "text", |
| 128 | "text": "Tool-provided image observation(s). Inspect these pixels as evidence for the active task; this is not a new user request.", |
| 129 | }, |
| 130 | *round_model_content, |
| 131 | ], |
| 132 | } |
| 133 | ) |
| 134 | tool_rounds.append({"tool_round": tool_round, "requested_tools": [call.as_dict() for call in assistant.tool_calls], "tool_results": [result.as_dict() for result in round_results]}) |
| 135 | transitions.append(_transition("executing_tools", "post_tool_decision", "tool_round_completed")) |
| 136 | transitions.append(_transition("post_tool_decision", "sampling_assistant", "runtime_continuation_after_tools")) |
| 137 | |
| 138 | self.history.extend([{"role": "user", "content": user_input}, {"role": "assistant", "content": final_text}]) |
| 139 | turn_record = {"turn_id": control.turn_id, "status": status, "raw_user_input": user_input, "assistant_turns": assistant_turns, "tool_rounds": tool_rounds, "transitions": transitions, "final_assistant_text": final_text, "created_at": datetime.now().isoformat(timespec="seconds")} |
| 140 | final_session = self.session_index.active() or self.session_index.create() |
| 141 | self.session_index.append_turn_record(final_session["session_id"], turn_record) |
| 142 | yield {"type": "done", "turn_id": control.turn_id, "assistant": final_text, "tool_results": [result.as_dict() for result in all_tool_results]} |
| 143 | yield {"type": "session", "turn_id": control.turn_id, "session": self.session_index.snapshot()} |
| 144 | |
| 145 | |
| 146 | def _compaction_record(result: CompactionResult) -> dict[str, Any]: |
| 147 | return { |
| 148 | "summary": result.summary, |
| 149 | "preserved_message_count": len(result.preserved_messages), |
| 150 | "compacted_message_count": result.compacted_message_count, |
| 151 | "estimated_tokens_before": result.estimated_tokens_before, |
| 152 | "estimated_tokens_after": result.estimated_tokens_after, |
| 153 | "reason": result.reason, |
| 154 | "mode": result.mode, |
| 155 | "created_at": result.created_at, |
| 156 | } |
| 157 | |
| 158 | |
| 159 | def _prompt_tokens(parts: list[Any]) -> int: |
| 160 | return sum(max(1, len(str(getattr(part, "body", ""))) // 4) for part in parts) |
| 161 | |
| 162 | |
| 163 | def _tool_schema_tokens(tool_schemas: list[dict[str, Any]]) -> int: |
| 164 | try: |
| 165 | return max(0, len(json.dumps(tool_schemas, ensure_ascii=False, default=str)) // 4) |
| 166 | except TypeError: |
| 167 | return max(0, len(str(tool_schemas)) // 4) |
| 168 | |
| 169 | |
| 170 | def _transition(src: str, dst: str, reason: str) -> dict[str, str]: |
| 171 | return {"from": src, "to": dst, "reason": reason} |
| 172 | |
| 173 | |
| 174 | def _openai_tool_call(call: ToolCall) -> dict[str, Any]: |
| 175 | return {"id": call.id, "type": "function", "function": {"name": call.name, "arguments": json.dumps(call.arguments, ensure_ascii=False)}} |
| 176 | |
| 177 | |
| 178 | def build_runtime(workspace_root: str | Path = ".", llm: Any | None = None, adapter_specs: list[Any] | None = None) -> AgentLoop: |
| 179 | from .vimax_adapters import build_vimax_adapter_specs |
| 180 | root = Path(workspace_root).resolve() |
| 181 | session_index = SessionIndex(root) |
| 182 | specs = adapter_specs if adapter_specs is not None else build_vimax_adapter_specs(root, session_index) |
| 183 | registry = build_builtin_registry(root, session_index, specs) |
| 184 | executor = ToolExecutor(registry, session_index) |
| 185 | prompt_builder = PromptBuilder(root / "prompts", session_index, registry) |
| 186 | resolved_llm = llm or OpenAICompatibleLLM() |
| 187 | return AgentLoop(session_index, prompt_builder, registry, executor, resolved_llm, ContextCompactor(resolved_llm)) |
| 188 |