agentpool 2.1.9__py3-none-any.whl → 2.2.3__py3-none-any.whl

agentpool/agents/agent.py

@@ -19,37 +19,68 @@ from pydantic import ValidationError
 from pydantic._internal import _typing_extra
 from pydantic_ai import (
     Agent as PydanticAgent,
-    AgentRunResultEvent,
     BaseToolCallPart,
+    CallToolsNode,
     FunctionToolCallEvent,
     FunctionToolResultEvent,
-    PartDeltaEvent,
+    ModelRequestNode,
     PartStartEvent,
     RunContext,
-    TextPart,
-    TextPartDelta,
     ToolReturnPart,
 )
 
 from agentpool.agents.base_agent import BaseAgent
-from agentpool.agents.events import RunStartedEvent, StreamCompleteEvent, ToolCallCompleteEvent
+from agentpool.agents.events import (
+    RunStartedEvent,
+    StreamCompleteEvent,
+    ToolCallCompleteEvent,
+)
+from agentpool.agents.modes import ModeInfo
 from agentpool.log import get_logger
 from agentpool.messaging import ChatMessage, MessageHistory, MessageNode
 from agentpool.messaging.processing import prepare_prompts
 from agentpool.prompts.convert import convert_prompts
 from agentpool.storage import StorageManager
-from agentpool.talk.stats import MessageStats
 from agentpool.tools import Tool, ToolManager
 from agentpool.tools.exceptions import ToolError
 from agentpool.utils.inspection import call_with_context, get_argument_key
 from agentpool.utils.now import get_now
+from agentpool.utils.pydantic_ai_helpers import safe_args_as_dict
 from agentpool.utils.result_utils import to_type
-from agentpool.utils.streams import merge_queue_into_iterator
+from agentpool.utils.streams import FileTracker, merge_queue_into_iterator
 
 
 TResult = TypeVar("TResult")
 
 
+def _extract_text_from_messages(
+    messages: list[Any], include_interruption_note: bool = False
+) -> str:
+    """Extract text content from pydantic-ai messages.
+
+    Args:
+        messages: List of ModelRequest/ModelResponse messages
+        include_interruption_note: Whether to append interruption notice
+
+    Returns:
+        Concatenated text content from all ModelResponse TextParts
+    """
+    from pydantic_ai.messages import ModelResponse, TextPart as PydanticTextPart
+
+    content = "".join(
+        part.content
+        for msg in messages
+        if isinstance(msg, ModelResponse)
+        for part in msg.parts
+        if isinstance(part, PydanticTextPart)
+    )
+    if include_interruption_note:
+        if content:
+            content += "\n\n"
+        content += "[Request interrupted by user]"
+    return content
+
+
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator, Coroutine, Sequence
     from datetime import datetime
@@ -57,13 +88,18 @@ if TYPE_CHECKING:
 
     from exxec import ExecutionEnvironment
     from pydantic_ai import UsageLimits
+    from pydantic_ai.builtin_tools import AbstractBuiltinTool
     from pydantic_ai.output import OutputSpec
     from pydantic_ai.settings import ModelSettings
+    from slashed import BaseCommand
+    from tokonomics.model_discovery import ProviderType
+    from tokonomics.model_discovery.model_info import ModelInfo
     from toprompt import AnyPromptType
     from upathtools import JoinablePathLike
 
     from agentpool.agents import AgentContext
     from agentpool.agents.events import RichAgentStreamEvent
+    from agentpool.agents.modes import ModeCategory
     from agentpool.common_types import (
         AgentName,
         BuiltinEventHandlerType,
@@ -77,7 +113,7 @@ if TYPE_CHECKING:
     )
     from agentpool.delegation import AgentPool, Team, TeamRun
     from agentpool.hooks import AgentHooks
-    from agentpool.models.agents import AutoCache, NativeAgentConfig, ToolMode
+    from agentpool.models.agents import NativeAgentConfig, ToolMode
     from agentpool.prompts.prompts import PromptType
     from agentpool.resource_providers import ResourceProvider
     from agentpool.ui.base import InputProvider
@@ -111,9 +147,11 @@ class AgentKwargs(TypedDict, total=False):
     input_provider: InputProvider | None
     event_handlers: Sequence[IndividualEventHandler | BuiltinEventHandlerType] | None
     env: ExecutionEnvironment | None
-    auto_cache: AutoCache
+
     hooks: AgentHooks | None
     model_settings: ModelSettings | None
+    usage_limits: UsageLimits | None
+    providers: Sequence[ProviderType] | None
 
 
 class Agent[TDeps = None, OutputDataT = str](BaseAgent[TDeps, OutputDataT]):
@@ -164,9 +202,12 @@ class Agent[TDeps = None, OutputDataT = str](BaseAgent[TDeps, OutputDataT]):
         knowledge: Knowledge | None = None,
         agent_config: NativeAgentConfig | None = None,
         env: ExecutionEnvironment | None = None,
-        auto_cache: AutoCache = "off",
         hooks: AgentHooks | None = None,
         tool_confirmation_mode: ToolConfirmationMode = "per_tool",
+        builtin_tools: Sequence[AbstractBuiltinTool] | None = None,
+        usage_limits: UsageLimits | None = None,
+        providers: Sequence[ProviderType] | None = None,
+        commands: Sequence[BaseCommand] | None = None,
     ) -> None:
         """Initialize agent.
 
@@ -206,9 +247,13 @@ class Agent[TDeps = None, OutputDataT = str](BaseAgent[TDeps, OutputDataT]):
             knowledge: Knowledge sources for this agent
             agent_config: Agent configuration
             env: Execution environment for code/command execution and filesystem access
-            auto_cache: Automatic caching configuration ("off", "5m", or "1h")
             hooks: AgentHooks instance for intercepting agent behavior at run and tool events
             tool_confirmation_mode: Tool confirmation mode
+            builtin_tools: PydanticAI builtin tools (WebSearchTool, CodeExecutionTool, etc.)
+            usage_limits: Usage limits for the agent
+            providers: Model providers for model discovery (e.g., ["openai", "anthropic"]).
+                Defaults to ["models.dev"] if not specified.
+            commands: Slash commands
         """
         from agentpool.agents.interactions import Interactions
         from agentpool.agents.sys_prompts import SystemPrompts
@@ -241,11 +286,15 @@ class Agent[TDeps = None, OutputDataT = str](BaseAgent[TDeps, OutputDataT]):
             output_type=to_type(output_type),  # type: ignore[arg-type]
             tool_confirmation_mode=tool_confirmation_mode,
             event_handlers=event_handlers,
+            commands=commands,
         )
 
         # Store config for context creation
         self._agent_config = agent_config or NativeAgentConfig(name=name)
 
+        # Store builtin tools for pydantic-ai
+        self._builtin_tools = list(builtin_tools) if builtin_tools else []
+
         # Override tools with Agent-specific ToolManager (with tools and tool_mode)
         all_tools = list(tools or [])
         self.tools = ToolManager(all_tools, tool_mode=tool_mode)
@@ -284,8 +333,11 @@ class Agent[TDeps = None, OutputDataT = str](BaseAgent[TDeps, OutputDataT]):
         # Store hooks
         self.hooks = hooks
 
-        # Store auto_cache setting
-        self._auto_cache: AutoCache = auto_cache
+        # Store default usage limits
+        self._default_usage_limits = usage_limits
+
+        # Store providers for model discovery
+        self._providers = list(providers) if providers else None
 
     def __repr__(self) -> str:
         desc = f", {self.description!r}" if self.description else ""
@@ -574,6 +626,7 @@ class Agent[TDeps = None, OutputDataT = str](BaseAgent[TDeps, OutputDataT]):
             output_retries=self._output_retries,
             deps_type=self.deps_type or NoneType,
             output_type=final_type,
+            builtin_tools=self._builtin_tools,
         )
 
         base_context = self.get_context()
@@ -715,6 +768,7 @@ class Agent[TDeps = None, OutputDataT = str](BaseAgent[TDeps, OutputDataT]):
         wait_for_connections: bool | None = None,
         deps: TDeps | None = None,
         instructions: str | None = None,
+        event_handlers: Sequence[IndividualEventHandler | BuiltinEventHandlerType] | None = None,
     ) -> AsyncIterator[RichAgentStreamEvent[OutputDataT]]:
         """Run agent with prompt and get a streaming response.
 
@@ -734,6 +788,8 @@ class Agent[TDeps = None, OutputDataT = str](BaseAgent[TDeps, OutputDataT]):
             wait_for_connections: Whether to wait for connected agents to complete
             deps: Optional dependencies for the agent
             instructions: Optional instructions to override the agent's system prompt
+            event_handlers: Optional event handlers for this run (overrides agent's handlers)
+
         Returns:
             An async iterator yielding streaming events with final message embedded.
 
@@ -741,9 +797,22 @@ class Agent[TDeps = None, OutputDataT = str](BaseAgent[TDeps, OutputDataT]):
             UnexpectedModelBehavior: If the model fails or behaves unexpectedly
         """
         conversation = message_history if message_history is not None else self.conversation
+        # Use provided event handlers or fall back to agent's handlers
+        if event_handlers is not None:
+            from anyenv import MultiEventHandler
+
+            from agentpool.agents.events import resolve_event_handlers
+
+            handler: MultiEventHandler[IndividualEventHandler] = MultiEventHandler(
+                resolve_event_handlers(event_handlers)
+            )
+        else:
+            handler = self.event_handler
         message_id = message_id or str(uuid4())
         run_id = str(uuid4())
-        user_msg, prompts, original_message = await prepare_prompts(*prompt)
+        # Get parent_id from last message in history for tree structure
+        last_msg_id = conversation.get_last_message_id()
+        user_msg, prompts, original_message = await prepare_prompts(*prompt, parent_id=last_msg_id)
         self.message_received.emit(user_msg)
         start_time = time.perf_counter()
         history_list = conversation.get_history()
@@ -753,9 +822,6 @@ class Agent[TDeps = None, OutputDataT = str](BaseAgent[TDeps, OutputDataT]):
         self._cancelled = False
         self._current_stream_task = asyncio.current_task()
 
-        # Track accumulated content for partial message on cancellation
-        accumulated_text: list[str] = []
-
         # Execute pre-run hooks
         if self.hooks:
             pre_run_result = await self.hooks.run_pre_run_hooks(
@@ -770,112 +836,120 @@ class Agent[TDeps = None, OutputDataT = str](BaseAgent[TDeps, OutputDataT]):
                 msg = f"Run blocked: {reason}"
                 raise RuntimeError(msg)
 
-        yield RunStartedEvent(thread_id=self.conversation_id, run_id=run_id, agent_name=self.name)
+        run_started = RunStartedEvent(
+            thread_id=self.conversation_id, run_id=run_id, agent_name=self.name
+        )
+        await handler(None, run_started)
+        yield run_started
         try:
+            from pydantic_graph import End
+
             agentlet = await self.get_agentlet(tool_choice, model, output_type, input_provider)
             content = await convert_prompts(prompts)
             response_msg: ChatMessage[Any] | None = None
             # Prepend pending context parts (content is already pydantic-ai format)
             converted = [*pending_parts, *content]
 
-            # Add CachePoint if auto_cache is enabled
-            if self._auto_cache != "off":
-                from pydantic_ai.messages import CachePoint
+            # Use provided usage_limits or fall back to default
+            effective_limits = usage_limits or self._default_usage_limits
+            history = [m for run in history_list for m in run.to_pydantic_ai()]
 
-                cache_point = CachePoint(ttl=self._auto_cache)
-                converted.append(cache_point)
-            stream_events = agentlet.run_stream_events(
+            # Track tool call starts to combine with results later
+            pending_tcs: dict[str, BaseToolCallPart] = {}
+            file_tracker = FileTracker()
+
+            async with agentlet.iter(
                 converted,
                 deps=deps,  # type: ignore[arg-type]
-                message_history=[m for run in history_list for m in run.to_pydantic_ai()],
-                usage_limits=usage_limits,
+                message_history=history,
+                usage_limits=effective_limits,
                 instructions=instructions,
-            )
-
-            # Stream events through merge_queue for progress events
-            async with merge_queue_into_iterator(stream_events, self._event_queue) as events:
-                # Track tool call starts to combine with results later
-                pending_tcs: dict[str, BaseToolCallPart] = {}
+            ) as agent_run:
                 try:
-                    async for event in events:
-                        # Check for cancellation
+                    async for node in agent_run:
                         if self._cancelled:
                             self.log.info("Stream cancelled by user")
                             break
+                        if isinstance(node, End):
+                            break
 
-                        # Call event handlers for all events
-                        for handler in self.event_handler._wrapped_handlers:
-                            await handler(None, event)
-
-                        yield event  # type: ignore[misc]
-
-                        # Accumulate text content for partial message
-                        match event:
-                            case PartDeltaEvent(delta=TextPartDelta(content_delta=delta)):
-                                accumulated_text.append(delta)
-                            case PartStartEvent(part=TextPart(content=text)) if text:
-                                accumulated_text.append(text)
-
-                        match event:
-                            case (
-                                PartStartEvent(part=BaseToolCallPart() as tool_part)
-                                | FunctionToolCallEvent(part=tool_part)
+                        # Stream events from model request node
+                        if isinstance(node, ModelRequestNode):
+                            async with (
+                                node.stream(agent_run.ctx) as agent_stream,
+                                merge_queue_into_iterator(
+                                    agent_stream,  # type: ignore[arg-type]
+                                    self._event_queue,
+                                ) as merged,
+                            ):
+                                async for event in file_tracker.track(merged):
+                                    if self._cancelled:
+                                        break
+                                    await handler(None, event)
+                                    yield event
+                                    combined = self._process_tool_event(
+                                        event, pending_tcs, message_id
+                                    )
+                                    if combined:
+                                        await handler(None, combined)
+                                        yield combined
+
+                        # Stream events from tool call node
+                        elif isinstance(node, CallToolsNode):
+                            async with (
+                                node.stream(agent_run.ctx) as tool_stream,
+                                merge_queue_into_iterator(tool_stream, self._event_queue) as merged,
                             ):
-                                # Store tool call start info for later combination with result
-                                pending_tcs[tool_part.tool_call_id] = tool_part
-                            case FunctionToolResultEvent(tool_call_id=call_id) as result_event:
-                                # Check if we have a pending tool call to combine with
-                                if call_info := pending_tcs.pop(call_id, None):
-                                    # Create and yield combined event
-                                    combined_event = ToolCallCompleteEvent(
-                                        tool_name=call_info.tool_name,
-                                        tool_call_id=call_id,
-                                        tool_input=call_info.args_as_dict(),
-                                        tool_result=result_event.result.content
-                                        if isinstance(result_event.result, ToolReturnPart)
-                                        else result_event.result,
-                                        agent_name=self.name,
-                                        message_id=message_id,
+                                async for event in file_tracker.track(merged):
+                                    if self._cancelled:
+                                        break
+                                    await handler(None, event)
+                                    yield event
+                                    combined = self._process_tool_event(
+                                        event, pending_tcs, message_id
                                     )
-                                    yield combined_event
-                            case AgentRunResultEvent():
-                                # Capture final result data, Build final response message
-                                response_time = time.perf_counter() - start_time
-                                response_msg = await ChatMessage.from_run_result(
-                                    event.result,
-                                    agent_name=self.name,
-                                    message_id=message_id,
-                                    conversation_id=conversation_id or user_msg.conversation_id,
-                                    response_time=response_time,
-                                )
+                                    if combined:
+                                        await handler(None, combined)
+                                        yield combined
                 except asyncio.CancelledError:
                     self.log.info("Stream cancelled via task cancellation")
                     self._cancelled = True
 
-            # Handle cancellation - emit partial message
-            if self._cancelled:
+                # Build response message
                 response_time = time.perf_counter() - start_time
-                partial_content = "".join(accumulated_text)
-                if partial_content:
-                    partial_content += "\n\n"
-                partial_content += "[Request interrupted by user]"
-                response_msg = ChatMessage(
-                    content=partial_content,
-                    role="assistant",
-                    name=self.name,
-                    message_id=message_id,
-                    conversation_id=conversation_id or user_msg.conversation_id,
-                    response_time=response_time,
-                    finish_reason="stop",
-                )
-                yield StreamCompleteEvent(message=response_msg)
-                self._current_stream_task = None
-                return
-
-            # Only finalize if we got a result (stream may exit early on error)
-            if response_msg is None:
-                msg = "Stream completed without producing a result"
-                raise RuntimeError(msg)  # noqa: TRY301
+                if self._cancelled:
+                    partial_content = _extract_text_from_messages(
+                        agent_run.all_messages(), include_interruption_note=True
+                    )
+                    response_msg = ChatMessage(
+                        content=partial_content,
+                        role="assistant",
+                        name=self.name,
+                        message_id=message_id,
+                        conversation_id=conversation_id or user_msg.conversation_id,
+                        parent_id=user_msg.message_id,
+                        response_time=response_time,
+                        finish_reason="stop",
+                    )
+                    complete_event = StreamCompleteEvent(message=response_msg)
+                    await handler(None, complete_event)
+                    yield complete_event
+                    self._current_stream_task = None
+                    return
+
+                if agent_run.result:
+                    response_msg = await ChatMessage.from_run_result(
+                        agent_run.result,
+                        agent_name=self.name,
+                        message_id=message_id,
+                        conversation_id=conversation_id or user_msg.conversation_id,
+                        parent_id=user_msg.message_id,
+                        response_time=response_time,
+                        metadata=file_tracker.get_metadata(),
+                    )
+                else:
+                    msg = "Stream completed without producing a result"
+                    raise RuntimeError(msg)  # noqa: TRY301
 
             # Execute post-run hooks
             if self.hooks:
@@ -893,7 +967,9 @@ class Agent[TDeps = None, OutputDataT = str](BaseAgent[TDeps, OutputDataT]):
             if original_message:
                 response_msg = response_msg.forwarded(original_message)
             # Send additional enriched completion event
-            yield StreamCompleteEvent(message=response_msg)
+            complete_event = StreamCompleteEvent(message=response_msg)
+            await handler(None, complete_event)
+            yield complete_event
             self.message_sent.emit(response_msg)
             await self.log_message(response_msg)
             if store_history:
@@ -907,6 +983,41 @@ class Agent[TDeps = None, OutputDataT = str](BaseAgent[TDeps, OutputDataT]):
         finally:
             self._current_stream_task = None
 
+    def _process_tool_event(
+        self,
+        event: RichAgentStreamEvent[Any],
+        pending_tool_calls: dict[str, BaseToolCallPart],
+        message_id: str,
+    ) -> ToolCallCompleteEvent | None:
+        """Process tool-related events and return combined event when complete.
+
+        Args:
+            event: The streaming event to process
+            pending_tool_calls: Dict tracking in-progress tool calls by ID
+            message_id: Message ID for the combined event
+
+        Returns:
+            ToolCallCompleteEvent if a tool call completed, None otherwise
+        """
+        match event:
+            case PartStartEvent(part=BaseToolCallPart() as tool_part):
+                pending_tool_calls[tool_part.tool_call_id] = tool_part
+            case FunctionToolCallEvent(part=tool_part):
+                pending_tool_calls[tool_part.tool_call_id] = tool_part
+            case FunctionToolResultEvent(tool_call_id=call_id) as result_event:
+                if call_info := pending_tool_calls.pop(call_id, None):
+                    return ToolCallCompleteEvent(
+                        tool_name=call_info.tool_name,
+                        tool_call_id=call_id,
+                        tool_input=safe_args_as_dict(call_info),
+                        tool_result=result_event.result.content
+                        if isinstance(result_event.result, ToolReturnPart)
+                        else result_event.result,
+                        agent_name=self.name,
+                        message_id=message_id,
+                    )
+        return None
+
     async def run_iter(
         self,
         *prompt_groups: Sequence[PromptCompatible],
@@ -1131,7 +1242,7 @@ class Agent[TDeps = None, OutputDataT = str](BaseAgent[TDeps, OutputDataT]):
             parent=self if pass_message_history else None,
         )
 
-    def set_model(self, model: ModelType) -> None:
+    async def set_model(self, model: ModelType) -> None:
         """Set the model for this agent.
 
         Args:
@@ -1166,11 +1277,6 @@ class Agent[TDeps = None, OutputDataT = str](BaseAgent[TDeps, OutputDataT]):
         )
         self.agent_reset.emit(event)
 
-    async def get_stats(self) -> MessageStats:
-        """Get message statistics (async version)."""
-        messages = await self.get_message_history()
-        return MessageStats(messages=messages)
-
     @asynccontextmanager
     async def temporary_state[T](
         self,
@@ -1247,6 +1353,107 @@ class Agent[TDeps = None, OutputDataT = str](BaseAgent[TDeps, OutputDataT]):
         else:
             return True
 
+    async def get_available_models(self) -> list[ModelInfo] | None:
+        """Get available models for this agent.
+
+        Uses tokonomics model discovery to fetch models from configured providers.
+        Defaults to openai, anthropic, and gemini if no providers specified.
+
+        Returns:
+            List of tokonomics ModelInfo, or None if discovery fails
+        """
+        from datetime import timedelta
+
+        from tokonomics.model_discovery import get_all_models
+
+        try:
+            max_age = timedelta(days=200)
+            return await get_all_models(
+                providers=self._providers or ["models.dev"], max_age=max_age
+            )
+        except Exception:
+            self.log.exception("Failed to discover models")
+            return None
+
+    def get_modes(self) -> list[ModeCategory]:
+        """Get available mode categories for this agent.
+
+        Native agents expose tool confirmation modes.
+
+        Returns:
+            List with single ModeCategory for tool confirmation
+        """
+        from agentpool.agents.modes import ModeCategory, ModeInfo
+
+        # Map current confirmation mode to mode ID
+        mode_id_map = {
+            "per_tool": "default",
+            "always": "default",
+            "never": "acceptEdits",
+        }
+        current_id = mode_id_map.get(self.tool_confirmation_mode, "default")
+
+        category_id = "permissions"
+        return [
+            ModeCategory(
+                id=category_id,
+                name="Permissions",
+                available_modes=[
+                    ModeInfo(
+                        id="default",
+                        name="Default",
+                        description="Require confirmation for tools marked as needing it",
+                        category_id=category_id,
+                    ),
+                    ModeInfo(
+                        id="acceptEdits",
+                        name="Accept Edits",
+                        description="Auto-approve all tool calls without confirmation",
+                        category_id=category_id,
+                    ),
+                ],
+                current_mode_id=current_id,
+            )
+        ]
+
+    async def set_mode(self, mode: ModeInfo | str, category_id: str | None = None) -> None:
+        """Set a mode for this agent.
+
+        Native agents support the "permissions" category with modes:
+        - "default": per_tool confirmation
+        - "acceptEdits": never confirm (auto-approve)
+
+        Args:
+            mode: Mode to activate - ModeInfo object or mode ID string
+            category_id: Optional category (only "permissions" supported)
+
+        Raises:
+            ValueError: If mode_id is invalid
+        """
+        # Extract mode_id and category from ModeInfo if provided
+        if isinstance(mode, ModeInfo):
+            mode_id = mode.id
+            category_id = category_id or mode.category_id or None
+        else:
+            mode_id = mode
+
+        # Validate category if provided
+        if category_id is not None and category_id != "permissions":
+            msg = f"Unknown category: {category_id}. Only 'permissions' is supported."
+            raise ValueError(msg)
+
+        # Map mode_id to confirmation mode
+        mode_map: dict[str, ToolConfirmationMode] = {
+            "default": "per_tool",
+            "acceptEdits": "never",
+        }
+
+        if mode_id not in mode_map:
+            msg = f"Unknown mode: {mode_id}. Available: {list(mode_map.keys())}"
+            raise ValueError(msg)
+
+        await self.set_tool_confirmation_mode(mode_map[mode_id])
+
 
 if __name__ == "__main__":
     import logging

agentpool/agents/agui_agent/__init__.py

@@ -2,7 +2,6 @@
 
 from agentpool.agents.agui_agent.agui_agent import AGUIAgent
 from agentpool.agents.agui_agent.agui_converters import (
-    ToolCallAccumulator,
     agui_to_native_event,
     to_agui_input_content,
     to_agui_tool,
@@ -12,7 +11,6 @@ from agentpool.agents.agui_agent.chunk_transformer import ChunkTransformer
 __all__ = [
     "AGUIAgent",
     "ChunkTransformer",
-    "ToolCallAccumulator",
     "agui_to_native_event",
     "to_agui_input_content",
     "to_agui_tool",