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

agentpool/messaging/messagenode.py

@@ -5,9 +5,8 @@ from __future__ import annotations
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any, Literal, Self, overload
-from uuid import uuid4
 
-from psygnal import Signal
+from anyenv.signals import Signal
 
 from agentpool.log import get_logger
 from agentpool.messaging import ChatMessage
@@ -20,8 +19,8 @@ if TYPE_CHECKING:
     from datetime import timedelta
     from types import TracebackType
 
-    from evented.configs import EventConfig
     from evented.event_data import EventData
+    from evented_config import EventConfig
 
     from agentpool.common_types import (
         AnyTransformFn,
@@ -34,7 +33,7 @@ if TYPE_CHECKING:
     from agentpool.storage import StorageManager
     from agentpool.talk import Talk, TeamTalk
     from agentpool.talk.stats import AggregatedMessageStats, MessageStats
-    from agentpool.tools.base import Tool
+    from agentpool.tools.base import FunctionTool
     from agentpool_config.forward_targets import ConnectionType
     from agentpool_config.mcp_server import MCPServerConfig
 
@@ -45,10 +44,10 @@ logger = get_logger(__name__)
 class MessageNode[TDeps, TResult](ABC):
     """Base class for all message processing nodes."""
 
-    message_received = Signal(ChatMessage)
+    message_received = Signal[ChatMessage[Any]]()
     """Signal emitted when node receives a message."""
 
-    message_sent = Signal(ChatMessage)
+    message_sent = Signal[ChatMessage[Any]]()
     """Signal emitted when node creates a message."""
 
     def __init__(
@@ -86,18 +85,33 @@ class MessageNode[TDeps, TResult](ABC):
         name_ = f"node_{self._name}"
         self.mcp = MCPManager(name_, servers=mcp_servers, owner=self.name)
         self.enable_db_logging = enable_logging
-        self.conversation_id = str(uuid4())
-        # Connect to the combined signal to capture all messages
-        # TODO: need to check this
-        # node.message_received.connect(self.log_message)
+        self.conversation_id: str | None = None
+        self.conversation_title: str | None = None
 
-    async def __aenter__(self) -> Self:
-        """Initialize base message node."""
-        if self.enable_db_logging and self.storage:
+    def _set_conversation_title(self, title: str) -> None:
+        """Callback for setting conversation title (called by storage manager)."""
+        self.conversation_title = title
+
+    async def log_conversation(self, initial_prompt: str | None = None) -> None:
+        """Log conversation to storage if enabled.
+
+        Should be called at the start of run_stream() after conversation_id is set.
+        For native agents, generate conversation_id first with uuid4().
+        For wrapped agents (Claude Code), set conversation_id from SDK session first.
+
+        Args:
+            initial_prompt: Optional initial prompt to trigger title generation.
+        """
+        if self.enable_db_logging and self.storage and self.conversation_id:
             await self.storage.log_conversation(
                 conversation_id=self.conversation_id,
                 node_name=self.name,
+                initial_prompt=initial_prompt,
+                on_title_generated=self._set_conversation_title,
             )
+
+    async def __aenter__(self) -> Self:
+        """Initialize base message node."""
         try:
             await self._events.__aenter__()
             await self.mcp.__aenter__()
@@ -161,7 +175,7 @@ class MessageNode[TDeps, TResult](ABC):
 
     def to_tool(
         self, *, name: str | None = None, description: str | None = None, **kwargs: Any
-    ) -> Tool[TResult]:
+    ) -> FunctionTool[TResult]:
         """Convert node to a callable tool.
 
         Args:
@@ -354,6 +368,30 @@ class MessageNode[TDeps, TResult](ABC):
     async def run(self, *prompts: Any, **kwargs: Any) -> ChatMessage[TResult]:
         """Execute node with prompts. Implementation-specific run logic."""
 
+    async def run_message(
+        self,
+        message: ChatMessage[Any],
+        **kwargs: Any,
+    ) -> ChatMessage[TResult]:
+        """Run with an incoming ChatMessage (e.g., from Talk routing).
+
+        Extracts content from the message, preserves conversation_id,
+        and sets parent_id to track the message chain.
+
+        Args:
+            message: The incoming ChatMessage to process
+            **kwargs: Additional arguments passed to run()
+
+        Returns:
+            Response ChatMessage with message chain tracked via parent_id
+        """
+        return await self.run(
+            message.content,
+            conversation_id=message.conversation_id,
+            parent_id=message.message_id,
+            **kwargs,
+        )
+
     async def get_message_history(self, limit: int | None = None) -> list[ChatMessage[Any]]:
         """Get message history from storage."""
         from agentpool_config.session import SessionQuery

agentpool/messaging/messages.py

@@ -30,6 +30,7 @@ from agentpool.common_types import MessageRole, SimpleJsonType  # noqa: TC001
 from agentpool.log import get_logger
 from agentpool.utils.inspection import dataclasses_no_defaults_repr
 from agentpool.utils.now import get_now
+from agentpool.utils.pydantic_ai_helpers import safe_args_as_dict
 
 
 if TYPE_CHECKING:
@@ -71,9 +72,7 @@ Metadata:
   {{ key }}: {{ value }}
 {%- endfor %}
 {%- endif %}
-{%- if forwarded_from %}
-Forwarded via: {{ forwarded_from|join(' -> ') }}
-{%- endif %}"""
+"""
 
 MARKDOWN_TEMPLATE = """## {{ name or role.title() }}
 *{{ timestamp.strftime('%Y-%m-%d %H:%M:%S') }}*
@@ -97,10 +96,7 @@ MARKDOWN_TEMPLATE = """## {{ name or role.title() }}
 ```
 {%- endif %}
 
-{% if forwarded_from %}
-
-*Forwarded via: {{ forwarded_from|join(' → ') }}*
-{% endif %}"""
+"""
 
 MESSAGE_TEMPLATES = {
     "simple": SIMPLE_TEMPLATE,
@@ -198,6 +194,14 @@ class ChatMessage[TContent]:
     conversation_id: str | None = None
     """ID of the conversation this message belongs to."""
 
+    parent_id: str | None = None
+    """ID of the parent message for tree-structured conversations.
+
+    This enables branching conversations where each message knows its predecessor.
+    For user messages, this typically points to the previous assistant response.
+    For assistant responses, this points to the user message being responded to.
+    """
+
     response_time: float | None = None
     """Time it took the LLM to respond."""
 
@@ -207,9 +211,6 @@ class ChatMessage[TContent]:
     name: str | None = None
     """Display name for the message sender in UI."""
 
-    forwarded_from: list[str] = field(default_factory=list)
-    """List of agent names (the chain) that forwarded this message to the sender."""
-
     provider_details: dict[str, Any] = field(default_factory=dict)
     """Provider specific metadata / extra information."""
 
@@ -295,12 +296,29 @@ class ChatMessage[TContent]:
         message: TPromptContent,
         conversation_id: str | None = None,
         instructions: str | None = None,
+        parent_id: str | None = None,
     ) -> ChatMessage[TPromptContent]:
-        """Create a user prompt message."""
+        """Create a user prompt message.
+
+        Args:
+            message: The prompt content
+            conversation_id: ID of the conversation
+            instructions: Optional instructions for the model
+            parent_id: ID of the parent message (typically the previous assistant response)
+
+        Returns:
+            A ChatMessage representing the user prompt
+        """
         part = UserPromptPart(content=message)
         request = ModelRequest(parts=[part], instructions=instructions)
         id_ = conversation_id or str(uuid4())
-        return ChatMessage(messages=[request], role="user", content=message, conversation_id=id_)
+        return ChatMessage(
+            messages=[request],
+            role="user",
+            content=message,
+            conversation_id=id_,
+            parent_id=parent_id,
+        )
 
     @classmethod
     def from_pydantic_ai[TContentType](
@@ -309,7 +327,7 @@ class ChatMessage[TContent]:
         message: ModelMessage,
         conversation_id: str | None = None,
         name: str | None = None,
-        forwarded_from: list[str] | None = None,
+        parent_id: str | None = None,
     ) -> ChatMessage[TContentType]:
         """Convert a Pydantic model to a ChatMessage."""
         match message:
@@ -319,9 +337,8 @@ class ChatMessage[TContent]:
                     content=content,
                     role="user",
                     message_id=run_id or str(uuid.uuid4()),
-                    # instructions=instructions,
-                    forwarded_from=forwarded_from or [],
                     name=name,
+                    parent_id=parent_id,
                 )
             case ModelResponse(
                 usage=usage,
@@ -347,7 +364,7 @@ class ChatMessage[TContent]:
                     finish_reason=finish_reason,
                     provider_response_id=provider_response_id,
                     name=name,
-                    forwarded_from=forwarded_from or [],
+                    parent_id=parent_id,
                 )
             case _ as unreachable:
                 assert_never(unreachable)
@@ -360,7 +377,9 @@ class ChatMessage[TContent]:
         agent_name: str | None = None,
         message_id: str | None = None,
         conversation_id: str | None = None,
+        parent_id: str | None = None,
         response_time: float,
+        metadata: SimpleJsonType | None = None,
     ) -> ChatMessage[OutputDataT]:
         """Create a ChatMessage from a PydanticAI run result.
 
@@ -369,19 +388,29 @@ class ChatMessage[TContent]:
             agent_name: Name of the agent that generated this response
             message_id: Unique message identifier
             conversation_id: Conversation identifier
+            parent_id: ID of the parent message (typically the user message)
             response_time: Total time taken for the response
+            metadata: Optional metadata to attach to the message
 
         Returns:
             A ChatMessage with all fields populated from the result
         """
-        # Calculate costs
+        # Calculate costs - prefer provider-reported cost if available
         run_usage = result.usage()
         usage = result.response.usage
-        cost_info = await TokenCost.from_usage(
-            model=result.response.model_name or "",
-            usage=run_usage,
-            provider=result.response.provider_name,
-        )
+        provider_cost = (result.response.provider_details or {}).get("cost")
+        if provider_cost is not None:
+            # Use actual cost from provider (e.g., OpenRouter returns this)
+            cost_info: TokenCost | None = TokenCost(
+                token_usage=run_usage, total_cost=Decimal(str(provider_cost))
+            )
+        else:
+            # Fall back to calculated cost
+            cost_info = await TokenCost.from_usage(
+                model=result.response.model_name or "",
+                usage=run_usage,
+                provider=result.response.provider_name,
+            )
 
         return ChatMessage[OutputDataT](
             content=result.output,
@@ -395,23 +424,13 @@ class ChatMessage[TContent]:
             provider_name=result.response.provider_name,
             message_id=message_id or str(uuid4()),
             conversation_id=conversation_id,
+            parent_id=parent_id,
             cost_info=cost_info,
             response_time=response_time,
             provider_details={},
+            metadata=metadata or {},
         )
 
-    def forwarded(self, previous_message: ChatMessage[Any]) -> Self:
-        """Create new message showing it was forwarded from another message.
-
-        Args:
-            previous_message: The message that led to this one's creation
-
-        Returns:
-            New message with updated chain showing the path through previous message
-        """
-        from_ = [*previous_message.forwarded_from, previous_message.name or "unknown"]
-        return replace(self, forwarded_from=from_)
-
     @property
     def response(self) -> ModelResponse:
         """Return the last response from the message history."""
@@ -509,7 +528,7 @@ class ChatMessage[TContent]:
                 call_part = call_parts[part.tool_call_id]
                 tool_info = ToolCallInfo(
                     tool_name=call_part.tool_name,
-                    args=call_part.args_as_dict(),
+                    args=safe_args_as_dict(call_part),
                     agent_name=agent_name or "UNSET",
                     result=part.content,
                     tool_call_id=call_part.tool_call_id or str(uuid4()),

agentpool/messaging/processing.py

@@ -17,32 +17,27 @@ if TYPE_CHECKING:
 
 
 async def prepare_prompts(
-    *prompt: PromptCompatible | ChatMessage[Any],
-) -> tuple[ChatMessage[Any], list[UserContent], ChatMessage[Any] | None]:
+    *prompt: PromptCompatible,
+    parent_id: str | None = None,
+    conversation_id: str | None = None,
+) -> tuple[ChatMessage[Any], list[UserContent]]:
     """Prepare prompts for processing.
 
-    Extracted from MessageNode.pre_run logic.
-
     Args:
         *prompt: The prompt(s) to prepare.
+        parent_id: Optional ID of the parent message (typically the previous response).
+        conversation_id: Optional conversation ID for the user message.
 
     Returns:
         A tuple of:
-            - Either incoming message, or a constructed incoming message based
-              on the prompt(s).
+            - A ChatMessage representing the user prompt.
             - A list of prompts to be sent to the model.
-            - The original ChatMessage if forwarded, None otherwise
     """
-    if len(prompt) == 1 and isinstance(prompt[0], ChatMessage):
-        original_msg = prompt[0]
-        # Update received message's chain to show it came through its source
-        user_msg = original_msg.forwarded(original_msg).to_request()
-        prompts = await convert_prompts([user_msg.content])
-        # clear cost info to avoid double-counting
-        return user_msg, prompts, original_msg
     prompts = await convert_prompts(prompt)
-    user_msg = ChatMessage.user_prompt(message=prompts)
-    return user_msg, prompts, None
+    user_msg = ChatMessage.user_prompt(
+        message=prompts, parent_id=parent_id, conversation_id=conversation_id
+    )
+    return user_msg, prompts
 
 
 async def finalize_message(
@@ -50,7 +45,6 @@ async def finalize_message(
     previous_message: ChatMessage[Any] | None,
     node: MessageNode[Any, Any],
     connections: ConnectionManager,
-    original_message: ChatMessage[Any] | None,
     wait_for_connections: bool | None = None,
 ) -> ChatMessage[Any]:
     """Handle message finalization and routing.
@@ -60,16 +54,12 @@ async def finalize_message(
         previous_message: The original user message (if any)
         node: The message node that produced the message
         connections: Connection manager for routing
-        original_message: The original ChatMessage if forwarded, None otherwise
         wait_for_connections: Whether to wait for connected nodes
 
     Returns:
         The finalized message
     """
-    # For chain processing, update the response's chain if input was forwarded
-    if original_message:
-        message = message.forwarded(original_message)
-    node.message_sent.emit(message)  # Emit signals
+    await node.message_sent.emit(message)  # Emit signals
     await node.log_message(message)  # Log message if enabled
     # Route to connections
     await connections.route_message(message, wait=wait_for_connections)

agentpool/models/__init__.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from agentpool.models.acp_agents import ACPAgentConfig, ACPAgentConfigTypes, BaseACPAgentConfig
-from agentpool.models.agents import NativeAgentConfig
+from agentpool.models.agents import AnyToolConfig, NativeAgentConfig  # noqa: F401
 from agentpool.models.agui_agents import AGUIAgentConfig
 from agentpool.models.claude_code_agents import ClaudeCodeAgentConfig
 from agentpool.models.manifest import AgentsManifest, AnyAgentConfig

agentpool/models/acp_agents/base.py

@@ -7,13 +7,12 @@ import os
 import tempfile
 from typing import TYPE_CHECKING, Annotated, Any, Literal
 
-from exxec import ExecutionEnvironmentStr, get_environment  # noqa: TC002
-from exxec.configs import (
+from exxec_config import (
     E2bExecutionEnvironmentConfig,
     ExecutionEnvironmentConfig,  # noqa: TC002
+    ExecutionEnvironmentStr,  # noqa: TC002
 )
 from pydantic import ConfigDict, Field
-from tokonomics.model_discovery import ProviderType  # noqa: TC002
 
 from agentpool_config.nodes import NodeConfig
 from agentpool_config.system_prompts import PromptConfig  # noqa: TC001
@@ -224,6 +223,8 @@ class BaseACPAgentConfig(NodeConfig):
 
     def get_execution_environment(self) -> ExecutionEnvironment:
         """Create execution environment from config."""
+        from exxec import get_environment
+
         if isinstance(self.execution_environment, str):
             return get_environment(self.execution_environment)
         return self.execution_environment.get_provider()
@@ -233,21 +234,14 @@ class BaseACPAgentConfig(NodeConfig):
 
         Returns None if not configured (caller should fall back to main env).
         """
+        from exxec import get_environment
+
         if self.client_execution_environment is None:
             return None
         if isinstance(self.client_execution_environment, str):
             return get_environment(self.client_execution_environment)
         return self.client_execution_environment.get_provider()
 
-    @property
-    def model_providers(self) -> list[ProviderType]:
-        """Return the model providers used by this ACP agent.
-
-        Override in subclasses to specify which providers the agent uses.
-        Used for intelligent model discovery and fallback configuration.
-        """
-        return []
-
 
 class ACPAgentConfig(BaseACPAgentConfig):
     """Configuration for a custom ACP agent with explicit command.
@@ -286,18 +280,6 @@ class ACPAgentConfig(BaseACPAgentConfig):
     )
     """Arguments to pass to the command."""
 
-    providers: list[ProviderType] = Field(
-        default_factory=list,
-        title="Providers",
-        examples=[["openai", "anthropic"], ["gemini"]],
-    )
-    """Model providers this agent can use."""
-
-    @property
-    def model_providers(self) -> list[ProviderType]:
-        """Return configured providers for custom ACP agents."""
-        return list(self.providers)
-
     def get_command(self) -> str:
         """Get the command to spawn the ACP server."""
         return self.command