iac-code 0.1.0__py3-none-any.whl

iac_code/agent/agent_tool.py

@@ -0,0 +1,269 @@
+"""AgentTool — spawns sub-agents with tool filtering and progress tracking."""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+from typing import Any
+
+from iac_code.agent.agent_types import filter_tools, get_agent_definition, get_builtin_agents
+from iac_code.i18n import _
+from iac_code.tools.base import Tool, ToolContext, ToolResult
+
+
+@dataclass
+class AgentProgress:
+    """Tracks sub-agent execution progress."""
+
+    tool_use_count: int = 0
+    token_count: int = 0
+    last_activity: str = ""
+    summary: str = ""
+
+
+async def run_sub_agent(
+    *,
+    prompt: str,
+    agent_type: str = "general-purpose",
+    cwd: str | None = None,
+    parent_provider_manager: Any = None,
+    parent_tool_registry: Any = None,
+    parent_system_prompt: str = "",
+    event_queue: asyncio.Queue | None = None,
+) -> tuple[str, AgentProgress]:
+    """Run a sub-agent and return (final_text, progress)."""
+    from iac_code.agent.agent_loop import AgentLoop
+    from iac_code.agent.system_prompt import build_system_prompt
+    from iac_code.types.stream_events import TextDeltaEvent, ToolResultEvent, ToolUseEndEvent, ToolUseStartEvent
+
+    defn = get_agent_definition(agent_type)
+    if defn is None:
+        raise ValueError(f"Unknown agent type: {agent_type}")
+
+    sub_registry = filter_tools(parent_tool_registry, defn) if parent_tool_registry else parent_tool_registry
+    system_prompt = parent_system_prompt or build_system_prompt(cwd=cwd)
+
+    sub_loop = AgentLoop(
+        provider_manager=parent_provider_manager,
+        system_prompt=system_prompt,
+        tool_registry=sub_registry or parent_tool_registry,
+        max_turns=defn.max_turns,
+    )
+
+    progress = AgentProgress(summary=f"Running {agent_type} agent")
+    text_chunks: list[str] = []
+    # Track tool inputs: tool_use_id -> (name, input)
+    pending_tool_inputs: dict[str, tuple[str, dict]] = {}
+
+    async for event in sub_loop.run_streaming(prompt):
+        if isinstance(event, TextDeltaEvent):
+            text_chunks.append(event.text)
+        elif isinstance(event, ToolUseStartEvent):
+            pending_tool_inputs[event.tool_use_id] = (event.name, {})
+        elif isinstance(event, ToolUseEndEvent):
+            if event.tool_use_id in pending_tool_inputs:
+                name = pending_tool_inputs[event.tool_use_id][0]
+                pending_tool_inputs[event.tool_use_id] = (name, event.input)
+            if event_queue:
+                tool_input = event.input
+                tool_name = pending_tool_inputs.get(event.tool_use_id, ("", {}))[0]
+                await event_queue.put(
+                    {
+                        "child_tool_name": tool_name,
+                        "child_tool_input": tool_input,
+                        "is_done": False,
+                    }
+                )
+        elif isinstance(event, ToolResultEvent):
+            progress.tool_use_count += 1
+            progress.last_activity = event.tool_name
+            tool_input = pending_tool_inputs.pop(event.tool_use_id, ("", {}))[1]
+            if event_queue:
+                await event_queue.put(
+                    {
+                        "child_tool_name": event.tool_name,
+                        "child_tool_input": tool_input,
+                        "is_done": True,
+                        "is_error": event.is_error,
+                    }
+                )
+
+    progress.token_count = sub_loop.context_manager.get_total_tokens()
+    final_text = "".join(text_chunks)
+
+    words = final_text.split()
+    if len(words) > 500:
+        final_text = " ".join(words[:500]) + "\n\n[... truncated to 500 words]"
+
+    return final_text, progress
+
+
+class AgentTool(Tool):
+    """Tool that spawns sub-agents for complex tasks."""
+
+    def __init__(
+        self,
+        task_manager: Any = None,
+        notification_queue: Any = None,
+        provider_manager: Any = None,
+        tool_registry: Any = None,
+        system_prompt: str = "",
+    ):
+        self._task_manager = task_manager
+        self._notification_queue = notification_queue
+        self._provider_manager = provider_manager
+        self._tool_registry = tool_registry
+        self._system_prompt = system_prompt
+        self._event_queue: asyncio.Queue | None = None  # Set by ToolExecutor via ToolCallRequest
+
+    @property
+    def name(self) -> str:
+        return "agent"
+
+    @property
+    def description(self) -> str:
+        agents = get_builtin_agents()
+        agent_list = "\n".join(f"  - {a.agent_type}: {a.when_to_use}" for a in agents)
+        return f"Launch a sub-agent to handle complex tasks.\n\nAvailable agent types:\n{agent_list}"
+
+    @property
+    def input_schema(self) -> dict[str, Any]:
+        agent_types = [a.agent_type for a in get_builtin_agents()]
+        return {
+            "type": "object",
+            "properties": {
+                "prompt": {
+                    "type": "string",
+                    "description": "The task for the sub-agent to perform.",
+                },
+                "description": {
+                    "type": "string",
+                    "description": "Short (3-5 word) description of the task.",
+                },
+                "subagent_type": {
+                    "type": "string",
+                    "enum": agent_types,
+                    "description": "The type of specialized agent to use.",
+                },
+                "run_in_background": {
+                    "type": "boolean",
+                    "description": "Run agent in background, parent continues.",
+                },
+            },
+            "required": ["prompt", "description"],
+        }
+
+    async def execute(self, *, tool_input: dict[str, Any], context: ToolContext) -> ToolResult:
+        prompt = tool_input["prompt"]
+        agent_type = tool_input.get("subagent_type", tool_input.get("agent_type", "general-purpose"))
+        run_in_background = tool_input.get("run_in_background", False)
+
+        defn = get_agent_definition(agent_type)
+        if defn is None:
+            return ToolResult.error(f"Unknown agent type: '{agent_type}'")
+
+        if run_in_background and self._task_manager:
+            task_id = self._task_manager.register(
+                description=tool_input.get("description", "Sub-agent task"),
+                agent_type=agent_type,
+            )
+            asyncio.create_task(self._run_background(task_id, prompt, agent_type, context))
+            return ToolResult.success(f"Background agent launched (task_id: {task_id}, type: {agent_type})")
+
+        try:
+            result_text, progress = await run_sub_agent(
+                prompt=prompt,
+                agent_type=agent_type,
+                cwd=context.cwd,
+                parent_provider_manager=self._provider_manager,
+                parent_tool_registry=self._tool_registry,
+                parent_system_prompt=self._system_prompt,
+                event_queue=self._event_queue,
+            )
+            if self._event_queue:
+                await self._event_queue.put(None)
+            return ToolResult.success(
+                f"{result_text}\n\n[Agent stats: {progress.tool_use_count} tool calls, {progress.token_count} tokens]"
+            )
+        except Exception as e:
+            if self._event_queue:
+                await self._event_queue.put(None)
+            return ToolResult.error(f"Sub-agent failed: {e}")
+
+    async def _run_background(
+        self,
+        task_id: str,
+        prompt: str,
+        agent_type: str,
+        context: ToolContext,
+    ) -> None:
+        try:
+            result_text, progress = await run_sub_agent(
+                prompt=prompt,
+                agent_type=agent_type,
+                cwd=context.cwd,
+                parent_provider_manager=self._provider_manager,
+                parent_tool_registry=self._tool_registry,
+                parent_system_prompt=self._system_prompt,
+            )
+            self._task_manager.complete(task_id, result=result_text)
+            self._task_manager.update_progress(
+                task_id,
+                tool_use_count=progress.tool_use_count,
+                token_count=progress.token_count,
+            )
+            if self._notification_queue:
+                self._notification_queue.enqueue(
+                    task_id=task_id,
+                    message=f"Agent completed: {progress.tool_use_count} tool calls",
+                )
+        except Exception as e:
+            self._task_manager.fail(task_id, error=str(e))
+            if self._notification_queue:
+                self._notification_queue.enqueue(
+                    task_id=task_id,
+                    message=f"Agent failed: {e}",
+                )
+
+    def is_read_only(self, input: dict | None = None) -> bool:
+        return False
+
+    def is_concurrency_safe(self, tool_input: dict[str, Any]) -> bool:
+        return True
+
+    def render_tool_use_message(self, input: dict, *, verbose: bool = False) -> str | None:
+        return input.get("description", "")
+
+    def render_tool_result_message(self, output: str, *, is_error: bool = False, verbose: bool = False) -> str | None:
+        if is_error:
+            return f"Agent error: {output[:200]}"
+        if verbose:
+            return output
+        # Extract stats from the end of the output
+        import re
+
+        match = re.search(r"\[Agent stats: (\d+) tool calls, (\d+) tokens\]", output)
+        if match:
+            tool_count = match.group(1)
+            token_count = int(match.group(2))
+            token_display = f"{token_count / 1000:.1f}k" if token_count >= 1000 else str(token_count)
+            return _("Done ({tool_count} tool uses · {token_display} tokens)").format(
+                tool_count=tool_count,
+                token_display=token_display,
+            )
+        return None  # Let renderer handle as default
+
+    def user_facing_name(self, input: dict | None = None) -> str:
+        if input:
+            agent_type = input.get("subagent_type", input.get("agent_type", "general-purpose"))
+            return {
+                "explore": _("Explore"),
+                "plan": _("Plan"),
+                "general-purpose": _("Agent"),
+            }.get(agent_type, _("Agent"))
+        return _("Agent")
+
+    def get_activity_description(self, input: dict | None = None) -> str | None:
+        if input is None:
+            return None
+        return f"Running agent: {input.get('description', 'sub-agent')}"

iac_code/agent/agent_types.py

@@ -0,0 +1,87 @@
+"""Agent definition model, built-in agent types, and tool filtering."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from iac_code.tools.base import ToolRegistry
+
+
+@dataclass
+class AgentDefinition:
+    """Structured definition of an agent type."""
+
+    agent_type: str
+    when_to_use: str
+    tools: list[str] | None = field(default_factory=lambda: ["*"])
+    disallowed_tools: list[str] | None = field(default_factory=list)
+    max_turns: int = 50
+    model: str = "inherit"
+
+    @property
+    def allows_all_tools(self) -> bool:
+        return self.tools is not None and "*" in self.tools
+
+    def is_tool_allowed(self, tool_name: str) -> bool:
+        if self.disallowed_tools and tool_name in self.disallowed_tools:
+            return False
+        if self.allows_all_tools:
+            return True
+        if self.tools:
+            return tool_name in self.tools
+        return False
+
+
+def filter_tools(registry: "ToolRegistry", agent_def: AgentDefinition) -> "ToolRegistry":
+    """Create a new ToolRegistry containing only tools allowed by the agent definition."""
+    from iac_code.tools.base import ToolRegistry
+
+    filtered = ToolRegistry()
+    for tool in registry.list_tools():
+        if agent_def.is_tool_allowed(tool.name):
+            filtered.register(tool)
+    return filtered
+
+
+def get_builtin_agents() -> list[AgentDefinition]:
+    return [
+        AgentDefinition(
+            agent_type="general-purpose",
+            when_to_use=(
+                "Use for complex, multi-step tasks that require research, "
+                "code changes, or coordinating multiple operations."
+            ),
+            tools=["*"],
+            disallowed_tools=["agent"],
+            max_turns=100,
+        ),
+        AgentDefinition(
+            agent_type="explore",
+            when_to_use=(
+                "Use to quickly find files, search code, or answer questions "
+                "about the codebase. Read-only — cannot modify files."
+            ),
+            tools=["read_file", "glob", "grep", "list_files", "bash"],
+            disallowed_tools=["write_file", "edit_file", "agent"],
+            max_turns=30,
+        ),
+        AgentDefinition(
+            agent_type="plan",
+            when_to_use=(
+                "Use to plan implementation strategy, review architecture, "
+                "or design solutions. Read-only, no execution."
+            ),
+            tools=["read_file", "glob", "grep", "list_files"],
+            disallowed_tools=["bash", "write_file", "edit_file", "agent"],
+            max_turns=20,
+        ),
+    ]
+
+
+def get_agent_definition(agent_type: str) -> AgentDefinition | None:
+    for defn in get_builtin_agents():
+        if defn.agent_type == agent_type:
+            return defn
+    return None

iac_code/agent/message.py

@@ -0,0 +1,153 @@
+"""Core message types for the agent system."""
+
+from __future__ import annotations
+
+import uuid
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field
+
+
+class TextBlock(BaseModel):
+    """A text content block."""
+
+    type: Literal["text"] = "text"
+    text: str
+
+
+class ToolUseBlock(BaseModel):
+    """A tool use request block from the assistant."""
+
+    type: Literal["tool_use"] = "tool_use"
+    id: str = Field(default_factory=lambda: f"toolu_{uuid.uuid4().hex[:24]}")
+    name: str
+    input: dict[str, Any] = Field(default_factory=dict)
+
+
+class ToolResultBlock(BaseModel):
+    """A tool result block sent back to the assistant."""
+
+    type: Literal["tool_result"] = "tool_result"
+    tool_use_id: str
+    content: str
+    is_error: bool = False
+
+
+class ThinkingBlock(BaseModel):
+    """An assistant reasoning/thinking block.
+
+    Used to round-trip reasoning content for models that require it
+    (DeepSeek V4 thinking mode, Qwen thinking mode, Anthropic extended thinking).
+    """
+
+    type: Literal["thinking"] = "thinking"
+    thinking: str
+
+
+# Union type for all content blocks
+ContentBlock = TextBlock | ToolUseBlock | ToolResultBlock | ThinkingBlock
+
+
+class Message(BaseModel):
+    """A single message in the conversation."""
+
+    role: Literal["user", "assistant"]
+    content: str | list[ContentBlock]
+    token_count: int = 0
+    elapsed_seconds: float = 0.0
+
+    def get_text(self) -> str:
+        """Extract text content from the message."""
+        if isinstance(self.content, str):
+            return self.content
+        return "\n".join(block.text for block in self.content if isinstance(block, TextBlock))
+
+    def get_tool_use_blocks(self) -> list[ToolUseBlock]:
+        """Extract tool use blocks from the message."""
+        if isinstance(self.content, str):
+            return []
+        return [block for block in self.content if isinstance(block, ToolUseBlock)]
+
+    def has_tool_use(self) -> bool:
+        """Check if this message contains tool use blocks."""
+        return len(self.get_tool_use_blocks()) > 0
+
+    def to_dict(self) -> dict:
+        """Serialize to a JSON-compatible dict for JSONL persistence."""
+        return self.model_dump()
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "Message":
+        """Deserialize from a dict."""
+        return cls.model_validate(data)
+
+    def to_api_format(self) -> dict:
+        """Convert to API-compatible format for litellm."""
+        if isinstance(self.content, str):
+            return {"role": self.role, "content": self.content}
+
+        content_list = []
+        for block in self.content:
+            if isinstance(block, TextBlock):
+                content_list.append({"type": "text", "text": block.text})
+            elif isinstance(block, ToolUseBlock):
+                content_list.append(
+                    {
+                        "type": "tool_use",
+                        "id": block.id,
+                        "name": block.name,
+                        "input": block.input,
+                    }
+                )
+            elif isinstance(block, ToolResultBlock):
+                content_list.append(
+                    {
+                        "type": "tool_result",
+                        "tool_use_id": block.tool_use_id,
+                        "content": block.content,
+                        "is_error": block.is_error,
+                    }
+                )
+            elif isinstance(block, ThinkingBlock):
+                content_list.append({"type": "thinking", "thinking": block.thinking})
+        return {"role": self.role, "content": content_list}
+
+
+class Conversation(BaseModel):
+    """Manages the conversation message history."""
+
+    messages: list[Message] = Field(default_factory=list)
+
+    def add_user_message(self, content: str) -> Message:
+        """Add a user message to the conversation."""
+        msg = Message(role="user", content=content)
+        self.messages.append(msg)
+        return msg
+
+    def add_assistant_message(self, content: str | list[ContentBlock]) -> Message:
+        """Add an assistant message to the conversation."""
+        msg = Message(role="assistant", content=content)
+        self.messages.append(msg)
+        return msg
+
+    def add_tool_results(self, tool_results: list[ToolResultBlock]) -> Message:
+        """Add tool results as a user message."""
+        msg = Message(role="user", content=list(tool_results))
+        self.messages.append(msg)
+        return msg
+
+    def to_api_format(self) -> list[dict]:
+        """Convert the entire conversation to API format."""
+        return [msg.to_api_format() for msg in self.messages]
+
+    def to_api_messages(self) -> list[dict]:
+        """Alias for to_api_format() for clarity in context management."""
+        return self.to_api_format()
+
+    def get_total_tokens(self) -> int:
+        """Sum token counts across all messages."""
+        return sum(msg.token_count for msg in self.messages)
+
+    def replace_messages(self, messages: list[Message]) -> None:
+        """Replace all messages (used after compaction)."""
+        self.messages = messages