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

agentpool/agents/agui_agent/agui_agent.py

@@ -31,11 +31,10 @@ 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.tools import ToolManager
-from agentpool.utils.streams import FileTracker
 from agentpool.utils.token_breakdown import calculate_usage_from_parts
 
 
@@ -45,25 +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__)
@@ -194,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.
 
@@ -207,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,
@@ -219,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."""
@@ -309,71 +360,38 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
             self._startup_process = None
             self.log.info("AG-UI server stopped")
 
-    async def run(
+    async def _stream_events(  # noqa: PLR0915
         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
-        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,
@@ -392,20 +410,21 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
         )
         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
-        # Get parent_id from last message in history for tree structure
-        last_msg_id = conversation.get_last_message_id()
-        user_msg, processed_prompts, _original_message = await prepare_prompts(
-            *prompts, parent_id=last_msg_id
-        )
+        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...
@@ -417,13 +436,14 @@ 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,
         )
 
-        await self.event_handler(None, run_started)
+        await handler(None, run_started)
         yield run_started
         # Get pending parts from conversation and convert them
         pending_parts = conversation.get_pending_parts()
@@ -527,18 +547,18 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
                                     while not self._event_queue.empty():
                                         try:
                                             custom_event = self._event_queue.get_nowait()
-                                            await self.event_handler(None, custom_event)
+                                            await handler(None, custom_event)
                                             yield custom_event
                                         except asyncio.QueueEmpty:
                                             break
                                     # Distribute to handlers
-                                    await self.event_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):
-                                await self.event_handler(None, native_event)
+                                await handler(None, native_event)
                                 yield native_event
 
                 except httpx.HTTPError:
@@ -609,9 +629,8 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
                 metadata=file_tracker.get_metadata(),
             )
             complete_event = StreamCompleteEvent(message=final_message)
-            await self.event_handler(None, complete_event)
+            await handler(None, complete_event)
             yield complete_event
-            self._current_stream_task = None
             return
 
         # Flush any remaining response parts
@@ -622,7 +641,7 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
         while not self._event_queue.empty():
             try:
                 queued_event = self._event_queue.get_nowait()
-                await self.event_handler(None, queued_event)
+                await handler(None, queued_event)
                 yield queued_event
             except asyncio.QueueEmpty:
                 break
@@ -651,30 +670,8 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
             cost_info=cost_info,
         )
         complete_event = StreamCompleteEvent(message=final_message)
-        await self.event_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:
@@ -705,10 +702,10 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
             )
         ]
 
-    def get_modes(self) -> list[ModeCategory]:
+    async def get_modes(self) -> list[ModeCategory]:
         """Get available modes for AG-UI agent.
 
-        AG-UI doesn't expose mode information, so returns an empty list.
+        AG-UI doesn't expose any configurable modes - model is server-controlled.
 
         Returns:
             Empty list - no modes supported
@@ -718,16 +715,16 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
     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 modes, so this always raises an error.
+        AG-UI doesn't support mode switching - model is controlled by remote server.
 
         Args:
             mode: The mode to set (not supported)
-            category_id: Optional category ID (not supported)
+            category_id: Category ID (not supported)
 
         Raises:
             ValueError: Always - AG-UI doesn't support modes
         """
-        msg = "AG-UI agent does not support modes"
+        msg = "AG-UI agent does not support mode switching - model is controlled by remote server"
         raise ValueError(msg)
 
 

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.