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

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",

agentpool/agents/agui_agent/agui_agent.py

@@ -31,11 +31,11 @@ from agentpool.agents.agui_agent.chunk_transformer import ChunkTransformer
 from agentpool.agents.agui_agent.helpers import execute_tool_calls, parse_sse_stream
 from agentpool.agents.base_agent import BaseAgent
 from agentpool.agents.events import RunStartedEvent, StreamCompleteEvent
+from agentpool.agents.events.processors import FileTracker
 from agentpool.log import get_logger
 from agentpool.messaging import ChatMessage
-from agentpool.messaging.processing import prepare_prompts
-from agentpool.talk.stats import MessageStats
 from agentpool.tools import ToolManager
+from agentpool.utils.token_breakdown import calculate_usage_from_parts
 
 
 if TYPE_CHECKING:
@@ -44,22 +44,26 @@ if TYPE_CHECKING:
     from types import TracebackType
 
     from ag_ui.core import Message, ToolMessage
-    from evented.configs import EventConfig
+    from evented_config import EventConfig
+    from pydantic_ai import UserContent
+    from slashed import BaseCommand
+    from tokonomics.model_discovery.model_info import ModelInfo
 
-    from agentpool.agents.base_agent import ToolConfirmationMode
     from agentpool.agents.context import AgentContext
     from agentpool.agents.events import RichAgentStreamEvent
+    from agentpool.agents.modes import ModeCategory, ModeInfo
     from agentpool.common_types import (
         BuiltinEventHandlerType,
         IndividualEventHandler,
-        PromptCompatible,
         ToolType,
     )
     from agentpool.delegation import AgentPool
     from agentpool.messaging import MessageHistory
+    from agentpool.models.agui_agents import AGUIAgentConfig
     from agentpool.tools import Tool
     from agentpool.ui.base import InputProvider
     from agentpool_config.mcp_server import MCPServerConfig
+    from agentpool_config.nodes import ToolConfirmationMode
 
 
 logger = get_logger(__name__)
@@ -132,6 +136,7 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
         event_configs: Sequence[EventConfig] | None = None,
         event_handlers: Sequence[IndividualEventHandler | BuiltinEventHandlerType] | None = None,
         tool_confirmation_mode: ToolConfirmationMode = "per_tool",
+        commands: Sequence[BaseCommand] | None = None,
     ) -> None:
         """Initialize AG-UI agent client.
 
@@ -153,6 +158,7 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
             event_configs: Event trigger configurations
             event_handlers: Sequence of event handlers to register
             tool_confirmation_mode: Tool confirmation mode
+            commands: Slash commands
         """
         super().__init__(
             name=name,
@@ -164,6 +170,7 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
             event_configs=event_configs,
             tool_confirmation_mode=tool_confirmation_mode,
             event_handlers=event_handlers,
+            commands=commands,
         )
 
         # AG-UI specific configuration
@@ -187,6 +194,50 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
         # Chunk transformer for normalizing CHUNK events
         self._chunk_transformer = ChunkTransformer()
 
+    @classmethod
+    def from_config(
+        cls,
+        config: AGUIAgentConfig,
+        *,
+        event_handlers: Sequence[IndividualEventHandler | BuiltinEventHandlerType] | None = None,
+        input_provider: InputProvider | None = None,
+        agent_pool: AgentPool[Any] | None = None,
+    ) -> Self:
+        """Create an AGUIAgent from a config object.
+
+        This is the preferred way to instantiate an AGUIAgent from configuration.
+
+        Args:
+            config: AG-UI agent configuration
+            event_handlers: Optional event handlers (merged with config handlers)
+            input_provider: Optional input provider for user interactions
+            agent_pool: Optional agent pool for coordination
+
+        Returns:
+            Configured AGUIAgent instance
+        """
+        # Merge config-level handlers with provided handlers
+        config_handlers = config.get_event_handlers()
+        merged_handlers: list[IndividualEventHandler | BuiltinEventHandlerType] = [
+            *config_handlers,
+            *(event_handlers or []),
+        ]
+        return cls(
+            endpoint=config.endpoint,
+            name=config.name or "agui-agent",
+            description=config.description,
+            display_name=config.display_name,
+            event_handlers=merged_handlers or None,
+            timeout=config.timeout,
+            headers=config.headers,
+            startup_command=config.startup_command,
+            startup_delay=config.startup_delay,
+            tools=[tool_config.get_tool() for tool_config in config.tools],
+            mcp_servers=config.mcp_servers,
+            tool_confirmation_mode=config.requires_tool_confirmation,
+            agent_pool=agent_pool,
+        )
+
     def get_context(self, data: Any = None) -> AgentContext:
         """Create a new context for this agent.
 
@@ -200,7 +251,7 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
         from agentpool.models.agui_agents import AGUIAgentConfig
         from agentpool.models.manifest import AgentsManifest
 
-        cfg = AGUIAgentConfig(
+        cfg = AGUIAgentConfig(  # type: ignore[call-arg]
             name=self.name,
             description=self.description,
             display_name=self.display_name,
@@ -212,7 +263,14 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
             startup_delay=self._startup_delay,
         )
         defn = self.agent_pool.manifest if self.agent_pool else AgentsManifest()
-        return AgentContext(node=self, pool=self.agent_pool, config=cfg, definition=defn, data=data)
+        return AgentContext(
+            node=self,
+            pool=self.agent_pool,
+            config=cfg,
+            definition=defn,
+            input_provider=self._input_provider,
+            data=data,
+        )
 
     async def __aenter__(self) -> Self:
         """Enter async context - initialize client and base resources."""
@@ -302,71 +360,38 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
             self._startup_process = None
             self.log.info("AG-UI server stopped")
 
-    async def run(
-        self,
-        *prompts: PromptCompatible,
-        message_id: str | None = None,
-        input_provider: InputProvider | None = None,
-        message_history: MessageHistory | None = None,
-        **kwargs: Any,
-    ) -> ChatMessage[str]:
-        """Execute prompt against AG-UI agent.
-
-        Sends the prompt to the AG-UI server and waits for completion.
-        Events are collected via run_stream and event handlers are called.
-        The final text content is returned as a ChatMessage.
-
-        Args:
-            prompts: Prompts to send (will be joined with spaces)
-            message_id: Optional message id for the returned message
-            input_provider: Optional input provider for tool confirmation requests
-            message_history: Optional MessageHistory to use instead of agent's own
-            **kwargs: Additional arguments (ignored for compatibility)
-
-        Returns:
-            ChatMessage containing the agent's aggregated text response
-        """
-        final_message: ChatMessage[str] | None = None
-        async for event in self.run_stream(
-            *prompts,
-            message_id=message_id,
-            input_provider=input_provider,
-            message_history=message_history,
-        ):
-            if isinstance(event, StreamCompleteEvent):
-                final_message = event.message
-
-        if final_message is None:
-            raise RuntimeError("No final message received from stream")
-        return final_message
-
-    async def run_stream(  # noqa: PLR0915
+    async def _stream_events(  # noqa: PLR0915
         self,
-        *prompts: PromptCompatible,
+        prompts: list[UserContent],
+        *,
+        user_msg: ChatMessage[Any],
+        effective_parent_id: str | None,
         message_id: str | None = None,
+        conversation_id: str | None = None,
+        parent_id: str | None = None,
         input_provider: InputProvider | None = None,
         message_history: MessageHistory | None = None,
-        **kwargs: Any,
+        deps: TDeps | None = None,
+        event_handlers: Sequence[IndividualEventHandler | BuiltinEventHandlerType] | None = None,
+        wait_for_connections: bool | None = None,
+        store_history: bool = True,
     ) -> AsyncIterator[RichAgentStreamEvent[str]]:
-        """Execute prompt with streaming events.
-
-        Sends the prompt to the remote AG-UI agent along with any registered tools.
-        When the agent requests a tool call, the tool is executed locally and the
-        result is sent back in a continuation request.
-
-        Args:
-            prompts: Prompts to send
-            message_id: Optional message ID
-            input_provider: Optional input provider for tool confirmation requests
-            message_history: Optional MessageHistory to use instead of agent's own
-            **kwargs: Additional arguments (ignored for compatibility)
-
-        Yields:
-            Native streaming events converted from AG-UI protocol
-        """
         # Update input provider if provided
         if input_provider is not None:
             self._input_provider = input_provider
+
+        # 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
+
         from ag_ui.core import (
             RunAgentInput,
             TextMessageChunkEvent,
@@ -379,22 +404,27 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
         )
 
         from agentpool.agents.agui_agent.agui_converters import (
-            ToolCallAccumulator,
             agui_to_native_event,
             to_agui_input_content,
             to_agui_tool,
         )
+        from agentpool.agents.tool_call_accumulator import ToolCallAccumulator
 
-        if not self._client or not self._thread_id:
+        if not self._client:
             msg = "Agent not initialized - use async context manager"
             raise RuntimeError(msg)
 
         # Reset cancellation state
         self._cancelled = False
-        self._current_stream_task = asyncio.current_task()
+
+        # Conversation ID initialization handled by BaseAgent
+
+        # Set thread_id from conversation_id (needed for AG-UI protocol)
+        if self._thread_id is None:
+            self._thread_id = self.conversation_id
 
         conversation = message_history if message_history is not None else self.conversation
-        user_msg, processed_prompts, _original_message = await prepare_prompts(*prompts)
+        processed_prompts = prompts
         self._run_id = str(uuid4())  # New run ID for each run
         self._chunk_transformer.reset()  # Reset chunk transformer
         # Track messages in pydantic-ai format: ModelRequest -> ModelResponse -> ModelRequest...
@@ -406,15 +436,15 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
         current_response_parts: list[TextPart | ThinkingPart | ToolCallPart] = []
         text_chunks: list[str] = []  # For final content string
 
+        assert self.conversation_id is not None  # Initialized by BaseAgent.run_stream()
         run_started = RunStartedEvent(
             thread_id=self._thread_id or self.conversation_id,
             run_id=self._run_id or str(uuid4()),
             agent_name=self.name,
         )
-        for handler in self.event_handler._wrapped_handlers:
-            await handler(None, run_started)
-        yield run_started
 
+        await handler(None, run_started)
+        yield run_started
         # Get pending parts from conversation and convert them
         pending_parts = conversation.get_pending_parts()
         pending_content = to_agui_input_content(pending_parts)
@@ -432,6 +462,8 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
         tool_accumulator = ToolCallAccumulator()
         pending_tool_results: list[ToolMessage] = []
         self.log.debug("Sending prompt to AG-UI agent", tool_names=[t.name for t in agui_tools])
+        # Track files modified during this run
+        file_tracker = FileTracker()
         # Loop to handle tool calls - agent may request multiple rounds
         try:
             while True:
@@ -509,25 +541,24 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
 
                                 # Convert to native event and distribute to handlers
                                 if native_event := agui_to_native_event(event):
+                                    # Track file modifications
+                                    file_tracker.process_event(native_event)
                                     # Check for queued custom events first
                                     while not self._event_queue.empty():
                                         try:
                                             custom_event = self._event_queue.get_nowait()
-                                            for handler in self.event_handler._wrapped_handlers:
-                                                await handler(None, custom_event)
+                                            await handler(None, custom_event)
                                             yield custom_event
                                         except asyncio.QueueEmpty:
                                             break
                                     # Distribute to handlers
-                                    for handler in self.event_handler._wrapped_handlers:
-                                        await handler(None, native_event)
+                                    await handler(None, native_event)
                                     yield native_event
 
                         # Flush any pending chunk events at end of stream
                         for event in self._chunk_transformer.flush():
                             if native_event := agui_to_native_event(event):
-                                for handler in self.event_handler._wrapped_handlers:
-                                    await handler(None, native_event)
+                                await handler(None, native_event)
                                 yield native_event
 
                 except httpx.HTTPError:
@@ -592,14 +623,14 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
                 name=self.name,
                 message_id=message_id or str(uuid4()),
                 conversation_id=self.conversation_id,
+                parent_id=user_msg.message_id,
                 messages=model_messages,
                 finish_reason="stop",
+                metadata=file_tracker.get_metadata(),
             )
             complete_event = StreamCompleteEvent(message=final_message)
-            for handler in self.event_handler._wrapped_handlers:
-                await handler(None, complete_event)
+            await handler(None, complete_event)
             yield complete_event
-            self._current_stream_task = None
             return
 
         # Flush any remaining response parts
@@ -610,56 +641,91 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
         while not self._event_queue.empty():
             try:
                 queued_event = self._event_queue.get_nowait()
-                for handler in self.event_handler._wrapped_handlers:
-                    await handler(None, queued_event)
+                await handler(None, queued_event)
                 yield queued_event
             except asyncio.QueueEmpty:
                 break
 
         text_content = "".join(text_chunks)
+
+        # Calculate approximate token usage from what we can observe
+        input_parts = [*processed_prompts, *pending_parts]
+        usage, cost_info = await calculate_usage_from_parts(
+            input_parts=input_parts,
+            response_parts=current_response_parts,
+            text_content=text_content,
+            model_name=self.model_name,
+        )
+
         final_message = ChatMessage[str](
             content=text_content,
             role="assistant",
             name=self.name,
             message_id=message_id or str(uuid4()),
             conversation_id=self.conversation_id,
+            parent_id=user_msg.message_id,
             messages=model_messages,
+            metadata=file_tracker.get_metadata(),
+            usage=usage,
+            cost_info=cost_info,
         )
         complete_event = StreamCompleteEvent(message=final_message)
-        for handler in self.event_handler._wrapped_handlers:
-            await handler(None, complete_event)
-        yield complete_event
-        # Record to conversation history
-        conversation.add_chat_messages([user_msg, final_message])
-
-    async def run_iter(
-        self,
-        *prompt_groups: Sequence[PromptCompatible],
-        message_id: str | None = None,
-        **kwargs: Any,
-    ) -> AsyncIterator[ChatMessage[str]]:
-        """Execute multiple prompt groups sequentially.
-
-        Args:
-            prompt_groups: Groups of prompts to execute
-            message_id: Optional message ID base
-            **kwargs: Additional arguments (ignored for compatibility)
-
-        Yields:
-            ChatMessage for each completed prompt group
-        """
-        for i, prompts in enumerate(prompt_groups):
-            mid = f"{message_id or 'msg'}_{i}" if message_id else None
-            yield await self.run(*prompts, message_id=mid)
+        await handler(None, complete_event)
+        yield complete_event  # Post-processing handled by base class
 
     @property
     def model_name(self) -> str | None:
         """Get model name (AG-UI doesn't expose this)."""
         return None
 
-    async def get_stats(self) -> MessageStats:
-        """Get message statistics for this node."""
-        return MessageStats()
+    async def set_model(self, model: str) -> None:
+        """Set model (no-op for AG-UI as model is controlled by remote server)."""
+        # AG-UI agents don't support model selection - the model is
+        # determined by the remote server configuration
+
+    async def get_available_models(self) -> list[ModelInfo] | None:
+        """Get available models for AG-UI agent.
+
+        AG-UI doesn't expose model information, so returns a placeholder model
+        indicating the model is determined by the remote server.
+
+        Returns:
+            List with a single placeholder ModelInfo
+        """
+        from tokonomics.model_discovery.model_info import ModelInfo
+
+        return [
+            ModelInfo(
+                id="server-determined",
+                name="Determined by server",
+                description="The model is determined by the remote AG-UI server",
+            )
+        ]
+
+    async def get_modes(self) -> list[ModeCategory]:
+        """Get available modes for AG-UI agent.
+
+        AG-UI doesn't expose any configurable modes - model is server-controlled.
+
+        Returns:
+            Empty list - no modes supported
+        """
+        return []
+
+    async def set_mode(self, mode: ModeInfo | str, category_id: str | None = None) -> None:
+        """Set a mode for AG-UI agent.
+
+        AG-UI doesn't support mode switching - model is controlled by remote server.
+
+        Args:
+            mode: The mode to set (not supported)
+            category_id: Category ID (not supported)
+
+        Raises:
+            ValueError: Always - AG-UI doesn't support modes
+        """
+        msg = "AG-UI agent does not support mode switching - model is controlled by remote server"
+        raise ValueError(msg)
 
 
 if __name__ == "__main__":

agentpool/agents/agui_agent/agui_converters.py

@@ -12,7 +12,6 @@ from __future__ import annotations
 import base64
 from typing import TYPE_CHECKING, Any
 
-import anyenv
 from pydantic_ai import (
     AudioUrl,
     BinaryContent,
@@ -291,133 +290,3 @@ def to_agui_tool(tool: Tool) -> AGUITool:
         description=func_schema.get("description", ""),
         parameters=func_schema.get("parameters", {"type": "object", "properties": {}}),
     )
-
-
-def _repair_partial_json(buffer: str) -> str:
-    """Attempt to repair truncated JSON for preview purposes.
-
-    Handles common truncation cases:
-    - Unclosed strings
-    - Missing closing braces/brackets
-    - Trailing commas
-
-    Args:
-        buffer: Potentially incomplete JSON string
-
-    Returns:
-        Repaired JSON string (may still be invalid in edge cases)
-    """
-    if not buffer:
-        return "{}"
-
-    result = buffer.rstrip()
-
-    # Check if we're in the middle of a string by counting unescaped quotes
-    in_string = False
-    i = 0
-    while i < len(result):
-        char = result[i]
-        if char == "\\" and i + 1 < len(result):
-            i += 2  # Skip escaped character
-            continue
-        if char == '"':
-            in_string = not in_string
-        i += 1
-
-    # Close unclosed string
-    if in_string:
-        result += '"'
-
-    # Remove trailing comma (invalid JSON)
-    result = result.rstrip()
-    if result.endswith(","):
-        result = result[:-1]
-
-    # Balance braces and brackets
-    open_braces = result.count("{") - result.count("}")
-    open_brackets = result.count("[") - result.count("]")
-
-    result += "]" * max(0, open_brackets)
-    result += "}" * max(0, open_braces)
-
-    return result
-
-
-class ToolCallAccumulator:
-    """Accumulates streamed tool call arguments.
-
-    AG-UI streams tool call arguments as deltas, this class accumulates them
-    and provides the complete arguments when the tool call ends.
-    """
-
-    def __init__(self) -> None:
-        self._calls: dict[str, dict[str, Any]] = {}
-
-    def start(self, tool_call_id: str, tool_name: str) -> None:
-        """Start tracking a new tool call."""
-        self._calls[tool_call_id] = {"name": tool_name, "args_buffer": ""}
-
-    def add_args(self, tool_call_id: str, delta: str) -> None:
-        """Add argument delta to a tool call."""
-        if tool_call_id in self._calls:
-            self._calls[tool_call_id]["args_buffer"] += delta
-
-    def complete(self, tool_call_id: str) -> tuple[str, dict[str, Any]] | None:
-        """Complete a tool call and return (tool_name, parsed_args).
-
-        Returns:
-            Tuple of (tool_name, args_dict) or None if call not found
-        """
-        if tool_call_id not in self._calls:
-            return None
-
-        call_data = self._calls.pop(tool_call_id)
-        args_str = call_data["args_buffer"]
-        try:
-            args = anyenv.load_json(args_str) if args_str else {}
-        except anyenv.JsonLoadError:
-            args = {"raw": args_str}
-        return call_data["name"], args
-
-    def get_pending(self, tool_call_id: str) -> tuple[str, str] | None:
-        """Get pending call data (tool_name, args_buffer) without completing."""
-        if tool_call_id not in self._calls:
-            return None
-        data = self._calls[tool_call_id]
-        return data["name"], data["args_buffer"]
-
-    def get_partial_args(self, tool_call_id: str) -> dict[str, Any]:
-        """Get best-effort parsed args from incomplete JSON stream.
-
-        Uses heuristics to complete truncated JSON for preview purposes.
-        Handles unclosed strings, missing braces/brackets, and trailing commas.
-
-        Args:
-            tool_call_id: Tool call ID
-
-        Returns:
-            Partially parsed arguments or empty dict
-        """
-        if tool_call_id not in self._calls:
-            return {}
-
-        buffer = self._calls[tool_call_id]["args_buffer"]
-        if not buffer:
-            return {}
-
-        # Try direct parse first
-        try:
-            return anyenv.load_json(buffer)
-        except anyenv.JsonLoadError:
-            pass
-
-        # Try to repair truncated JSON
-        try:
-            repaired = _repair_partial_json(buffer)
-            return anyenv.load_json(repaired)
-        except anyenv.JsonLoadError:
-            return {}
-
-    def clear(self) -> None:
-        """Clear all pending tool calls."""
-        self._calls.clear()

agentpool/agents/agui_agent/helpers.py

@@ -16,8 +16,7 @@ if TYPE_CHECKING:
     from ag_ui.core import Event, ToolMessage
     import httpx
 
-    from agentpool.agents.context import ConfirmationResult
-    from agentpool.messaging.context import NodeContext
+    from agentpool.agents.context import AgentContext, ConfirmationResult
     from agentpool.tools import Tool
     from agentpool.ui.base import InputProvider
     from agentpool_config.nodes import ToolConfirmationMode
@@ -51,7 +50,7 @@ async def _get_tool_confirmation(
     tool: Tool,
     args: dict[str, Any],
     input_provider: InputProvider,
-    context: NodeContext[Any],
+    context: AgentContext[Any],
 ) -> ConfirmationResult:
     """Get confirmation for tool execution.
 
@@ -78,7 +77,7 @@ async def execute_tool_calls(
     *,
     confirmation_mode: ToolConfirmationMode = "never",
     input_provider: InputProvider | None = None,
-    context: NodeContext[Any] | None = None,
+    context: AgentContext[Any] | None = None,
 ) -> list[ToolMessage]:
     """Execute tool calls locally and return results.