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

agentpool/delegation/team.py

@@ -4,12 +4,11 @@ from __future__ import annotations
 
 import asyncio
 from time import perf_counter
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal
 from uuid import uuid4
 
 from anyenv.async_run import as_generated
 import anyio
-from toprompt import to_prompt
 
 from agentpool.common_types import SupportsRunStream
 from agentpool.delegation.base_team import BaseTeam
@@ -33,36 +32,6 @@ if TYPE_CHECKING:
     from agentpool_config.task import Job
 
 
-async def normalize_stream_for_teams(
-    node: MessageNode[Any, Any],
-    *args: Any,
-    **kwargs: Any,
-) -> AsyncIterator[tuple[MessageNode[Any, Any], RichAgentStreamEvent[Any]]]:
-    """Normalize any streaming node to yield (node, event) tuples for team composition.
-
-    Args:
-        node: The streaming node (Agent, Team, etc.)
-        *args: Arguments to pass to run_stream
-        **kwargs: Keyword arguments to pass to run_stream
-
-    Yields:
-        Tuples of (node, event) where node is the MessageNode instance
-        and event is the streaming event from that node.
-    """
-    if not isinstance(node, SupportsRunStream):
-        msg = f"Node {node.name} does not support streaming"
-        raise TypeError(msg)
-
-    stream = node.run_stream(*args, **kwargs)
-    async for item in stream:
-        if isinstance(item, tuple):
-            # Already normalized (from Team or other composite node)
-            yield item
-        else:
-            # Raw event (from Agent) - wrap it with the source node
-            yield (node, item)
-
-
 class Team[TDeps = None](BaseTeam[TDeps, Any]):
     """Group of agents that can execute together."""
 
@@ -77,8 +46,7 @@ class Team[TDeps = None](BaseTeam[TDeps, Any]):
         final_prompt = list(prompts)
         if self.shared_prompt:
             final_prompt.insert(0, self.shared_prompt)
-        combined_prompt = "\n".join([await to_prompt(p) for p in final_prompt])
-        all_nodes = list(await self.pick_agents(combined_prompt))
+        all_nodes = list(self.nodes)
         # Create Talk connections for monitoring this execution
         execution_talks: list[Talk[Any]] = []
         for node in all_nodes:
@@ -130,8 +98,7 @@ class Team[TDeps = None](BaseTeam[TDeps, Any]):
                 await queue.put(None)
 
         # Get nodes to run
-        combined_prompt = "\n".join([await to_prompt(p) for p in prompts])
-        all_nodes = list(await self.pick_agents(combined_prompt))
+        all_nodes = list(self.nodes)
 
         # Start all agents
         tasks = [asyncio.create_task(_run(n), name=f"run_{n.name}") for n in all_nodes]
@@ -163,8 +130,8 @@ class Team[TDeps = None](BaseTeam[TDeps, Any]):
     ) -> ChatMessage[list[Any]]:
         """Run all agents in parallel and return combined message."""
         # Prepare prompts and create user message
-        user_msg, processed_prompts, original_message = await prepare_prompts(*prompts)
-        self.message_received.emit(user_msg)
+        user_msg, processed_prompts = await prepare_prompts(*prompts)
+        await self.message_received.emit(user_msg)
 
         # Execute team logic
         result = await self.execute(*processed_prompts, **kwargs)
@@ -176,6 +143,7 @@ class Team[TDeps = None](BaseTeam[TDeps, Any]):
             name=self.name,
             message_id=message_id,
             conversation_id=user_msg.conversation_id,
+            parent_id=user_msg.message_id,
             metadata={
                 "agent_names": [r.agent_name for r in result],
                 "errors": {name: str(error) for name, error in result.errors.items()},
@@ -194,7 +162,6 @@ class Team[TDeps = None](BaseTeam[TDeps, Any]):
             user_msg,
             self,
             self.connections,
-            original_message,
             wait_for_connections,
         )
 
@@ -202,7 +169,7 @@ class Team[TDeps = None](BaseTeam[TDeps, Any]):
         self,
         *prompts: PromptCompatible,
         **kwargs: Any,
-    ) -> AsyncIterator[tuple[MessageNode[Any, Any], RichAgentStreamEvent[Any]]]:
+    ) -> AsyncIterator[RichAgentStreamEvent[Any]]:
         """Stream responses from all team members in parallel.
 
         Args:
@@ -210,23 +177,53 @@ class Team[TDeps = None](BaseTeam[TDeps, Any]):
             kwargs: Additional arguments passed to each agent
 
         Yields:
-            Tuples of (agent, event) where agent is the Agent instance
-            and event is the streaming event from that agent.
+            RichAgentStreamEvent, with member events wrapped in SubAgentEvent
         """
-        # Get nodes to run
-        combined_prompt = "\n".join([await to_prompt(p) for p in prompts])
-        all_nodes = list(await self.pick_agents(combined_prompt))
+        from agentpool.agents.events import SubAgentEvent
 
-        # Create list of streams that yield (agent, event) tuples
-        agent_streams = [
-            normalize_stream_for_teams(agent, *prompts, **kwargs)
-            for agent in all_nodes
-            if isinstance(agent, SupportsRunStream)
-        ]
-
-        # Merge all agent streams
-        async for agent_event_tuple in as_generated(agent_streams):
-            yield agent_event_tuple
+        # Get nodes to run
+        all_nodes = list(self.nodes)
+
+        # Create list of streams
+        async def wrap_stream(
+            node: MessageNode[Any, Any],
+        ) -> AsyncIterator[RichAgentStreamEvent[Any]]:
+            """Wrap a node's stream events in SubAgentEvent."""
+            if not isinstance(node, SupportsRunStream):
+                return
+            async for event in node.run_stream(*prompts, **kwargs):
+                # Handle already-wrapped SubAgentEvents (nested teams)
+                if isinstance(event, SubAgentEvent):
+                    yield SubAgentEvent(
+                        source_name=event.source_name,
+                        source_type=event.source_type,
+                        event=event.event,
+                        depth=event.depth + 1,
+                    )
+                else:
+                    # Determine source type based on node type
+                    from agentpool.delegation.teamrun import TeamRun
+
+                    if isinstance(node, TeamRun):
+                        source_type: Literal["team_parallel", "team_sequential", "agent"] = (
+                            "team_sequential"
+                        )
+                    elif isinstance(node, BaseTeam):
+                        source_type = "team_parallel"
+                    else:
+                        source_type = "agent"
+
+                    yield SubAgentEvent(
+                        source_name=node.name,
+                        source_type=source_type,
+                        event=event,
+                        depth=1,
+                    )
+
+        streams = [wrap_stream(node) for node in all_nodes]
+        # Merge all streams
+        async for event in as_generated(streams):
+            yield event
 
     async def run_job[TJobResult](
         self,
@@ -310,6 +307,7 @@ if __name__ == "__main__":
 
     async def main() -> None:
         from agentpool import Agent, TeamRun
+        from agentpool.agents.events import SubAgentEvent
 
         agent_a = Agent(name="A", model="test")
         agent_b = Agent(name="B", model="test")
@@ -319,7 +317,10 @@ if __name__ == "__main__":
         outer_team = Team([inner_run, agent_c], name="Parallel")
 
         print("Testing Team containing TeamRun...")
-        async for node, event in outer_team.run_stream("test"):
-            print(f"{node.name}: {type(event).__name__}")
+        async for event in outer_team.run_stream("test"):
+            if isinstance(event, SubAgentEvent):
+                print(f"[depth={event.depth}] {event.source_name}: {type(event.event).__name__}")
+            else:
+                print(f"Event: {type(event).__name__}")
 
     anyio.run(main)

agentpool/delegation/teamrun.py

@@ -9,11 +9,9 @@ from typing import TYPE_CHECKING, Any, Literal, overload
 from uuid import uuid4
 
 import anyio
-from pydantic_ai import PartDeltaEvent, TextPartDelta
 
 from agentpool.common_types import SupportsRunStream
 from agentpool.delegation.base_team import BaseTeam
-from agentpool.delegation.team import normalize_stream_for_teams
 from agentpool.log import get_logger
 from agentpool.messaging import AgentResponse, ChatMessage, TeamResponse
 from agentpool.messaging.processing import finalize_message, prepare_prompts
@@ -27,7 +25,7 @@ if TYPE_CHECKING:
 
     from agentpool import MessageNode
     from agentpool.agents.events import RichAgentStreamEvent
-    from agentpool.common_types import PromptCompatible, SupportsStructuredOutput
+    from agentpool.common_types import PromptCompatible
     from agentpool.delegation import AgentPool
 
 
@@ -65,9 +63,6 @@ class TeamRun[TDeps, TResult](BaseTeam[TDeps, TResult]):
         display_name: str | None = None,
         shared_prompt: str | None = None,
         validator: MessageNode[Any, TResult],
-        picker: SupportsStructuredOutput | None = None,
-        num_picks: int | None = None,
-        pick_prompt: str | None = None,
         agent_pool: AgentPool | None = None,
     ) -> None: ...
 
@@ -81,9 +76,6 @@ class TeamRun[TDeps, TResult](BaseTeam[TDeps, TResult]):
         display_name: str | None = None,
         shared_prompt: str | None = None,
         validator: None = None,
-        picker: SupportsStructuredOutput | None = None,
-        num_picks: int | None = None,
-        pick_prompt: str | None = None,
         agent_pool: AgentPool | None = None,
     ) -> None: ...
 
@@ -97,9 +89,6 @@ class TeamRun[TDeps, TResult](BaseTeam[TDeps, TResult]):
         display_name: str | None = None,
         shared_prompt: str | None = None,
         validator: MessageNode[Any, TResult] | None = None,
-        picker: SupportsStructuredOutput | None = None,
-        num_picks: int | None = None,
-        pick_prompt: str | None = None,
         agent_pool: AgentPool | None = None,
     ) -> None: ...
 
@@ -112,9 +101,6 @@ class TeamRun[TDeps, TResult](BaseTeam[TDeps, TResult]):
         display_name: str | None = None,
         shared_prompt: str | None = None,
         validator: MessageNode[Any, TResult] | None = None,
-        picker: SupportsStructuredOutput | None = None,
-        num_picks: int | None = None,
-        pick_prompt: str | None = None,
         agent_pool: AgentPool | None = None,
         # result_mode: ResultMode = "last",
     ) -> None:
@@ -124,9 +110,6 @@ class TeamRun[TDeps, TResult](BaseTeam[TDeps, TResult]):
             description=description,
             display_name=display_name,
             shared_prompt=shared_prompt,
-            picker=picker,
-            num_picks=num_picks,
-            pick_prompt=pick_prompt,
             agent_pool=agent_pool,
         )
         self.validator = validator
@@ -147,8 +130,8 @@ class TeamRun[TDeps, TResult](BaseTeam[TDeps, TResult]):
     ) -> ChatMessage[TResult]:
         """Run agents sequentially and return combined message."""
         # Prepare prompts and create user message
-        user_msg, processed_prompts, original_message = await prepare_prompts(*prompts)
-        self.message_received.emit(user_msg)
+        user_msg, processed_prompts = await prepare_prompts(*prompts)
+        await self.message_received.emit(user_msg)
         # Execute sequential logic
         message_id = str(uuid4())  # Always generate unique response ID
         result = await self.execute(*processed_prompts, **kwargs)
@@ -172,6 +155,7 @@ class TeamRun[TDeps, TResult](BaseTeam[TDeps, TResult]):
             associated_messages=all_messages,
             message_id=message_id,
             conversation_id=user_msg.conversation_id,
+            parent_id=user_msg.message_id,
             metadata={
                 "execution_order": [r.agent_name for r in result],
                 "start_time": result.start_time.isoformat(),
@@ -186,7 +170,6 @@ class TeamRun[TDeps, TResult](BaseTeam[TDeps, TResult]):
             user_msg,
             self,
             self.connections,
-            original_message,
             wait_for_connections,
         )
 
@@ -223,12 +206,9 @@ class TeamRun[TDeps, TResult](BaseTeam[TDeps, TResult]):
         *prompt: PromptCompatible,
         **kwargs: Any,
     ) -> AsyncIterator[Talk[Any] | AgentResponse[Any]]:
-        from toprompt import to_prompt
-
         connections: list[Talk[Any]] = []
         try:
-            combined_prompt = "\n".join([await to_prompt(p) for p in prompt])
-            all_nodes = list(await self.pick_agents(combined_prompt))
+            all_nodes = list(self.nodes)
             if self.validator:
                 all_nodes.append(self.validator)
             first = all_nodes[0]
@@ -273,7 +253,7 @@ class TeamRun[TDeps, TResult](BaseTeam[TDeps, TResult]):
         *prompts: PromptCompatible,
         require_all: bool = True,
         **kwargs: Any,
-    ) -> AsyncIterator[tuple[MessageNode[Any, Any], RichAgentStreamEvent[Any]]]:
+    ) -> AsyncIterator[RichAgentStreamEvent[Any]]:
         """Stream responses through the chain of team members.
 
         Args:
@@ -283,50 +263,61 @@ class TeamRun[TDeps, TResult](BaseTeam[TDeps, TResult]):
             kwargs: Additional arguments passed to each agent
 
         Yields:
-            Tuples of (agent, event) where agent is the Agent instance
-            and event is the streaming event.
+            RichAgentStreamEvent, with member events wrapped in SubAgentEvent
         """
-        from agentpool.agents.events import StreamCompleteEvent
+        from agentpool.agents.events import StreamCompleteEvent, SubAgentEvent
+        from agentpool.delegation.team import Team
 
         current_message = prompts
-        collected_content = []
-        for agent in self.nodes:
+        for node in self.nodes:
             try:
-                agent_content = []
-
-                # Use wrapper to normalize all streaming nodes to (agent, event) tuples
-                if not isinstance(agent, SupportsRunStream):
-                    msg = f"Agent {agent.name} does not support streaming"
+                if not isinstance(node, SupportsRunStream):
+                    msg = f"Node {node.name} does not support streaming"
                     raise TypeError(msg)  # noqa: TRY301
 
-                stream = normalize_stream_for_teams(agent, *current_message, **kwargs)
-
-                async for agent_event_tuple in stream:
-                    actual_agent, event = agent_event_tuple
-                    match event:
-                        case PartDeltaEvent(delta=TextPartDelta(content_delta=delta)):
-                            agent_content.append(delta)
-                            collected_content.append(delta)
-                            yield (actual_agent, event)  # Yield tuple with agent context
-                        case StreamCompleteEvent(message=message):
-                            # Use complete response as input for next agent
-                            current_message = (message.content,)
-                            yield (actual_agent, event)  # Yield tuple with agent context
-                        case _:
-                            yield (actual_agent, event)  # Yield tuple with agent context
+                async for event in node.run_stream(*current_message, **kwargs):
+                    # Handle already-wrapped SubAgentEvents (nested teams)
+                    if isinstance(event, SubAgentEvent):
+                        yield SubAgentEvent(
+                            source_name=event.source_name,
+                            source_type=event.source_type,
+                            event=event.event,
+                            depth=event.depth + 1,
+                        )
+                    else:
+                        # Determine source type based on node type
+                        if isinstance(node, Team):
+                            source_type: Literal["team_parallel", "team_sequential", "agent"] = (
+                                "team_parallel"
+                            )
+                        elif isinstance(node, BaseTeam):
+                            source_type = "team_sequential"
+                        else:
+                            source_type = "agent"
+
+                        yield SubAgentEvent(
+                            source_name=node.name,
+                            source_type=source_type,
+                            event=event,
+                        )
+
+                        # Extract content for next agent in chain
+                        if isinstance(event, StreamCompleteEvent):
+                            current_message = (event.message.content,)
 
             except Exception as e:
                 if require_all:
-                    msg = f"Chain broken at {agent.name}: {e}"
+                    msg = f"Chain broken at {node.name}: {e}"
                     logger.exception(msg)
                     raise ValueError(msg) from e
-                logger.warning("Chain handler failed", name=agent.name, error=e)
+                logger.warning("Chain handler failed", name=node.name, error=e)
 
 
 if __name__ == "__main__":
 
     async def main() -> None:
         from agentpool import Agent, Team
+        from agentpool.agents.events import SubAgentEvent
 
         agent1 = Agent(name="Agent1", model="test")
         agent2 = Agent(name="Agent2", model="test")
@@ -335,8 +326,13 @@ if __name__ == "__main__":
         outer_run = TeamRun([inner_team, agent3], name="Sequential")
         print("Testing TeamRun containing Team...")
         try:
-            async for node, event in outer_run.run_stream("test"):
-                print(f"{node.name}: {type(event).__name__}")
+            async for event in outer_run.run_stream("test"):
+                if isinstance(event, SubAgentEvent):
+                    print(
+                        f"[depth={event.depth}] {event.source_name}: {type(event.event).__name__}"
+                    )
+                else:
+                    print(f"Event: {type(event).__name__}")
         except Exception as e:  # noqa: BLE001
             print(f"Error: {e}")
 

agentpool/diagnostics/__init__.py

@@ -0,0 +1,53 @@
+"""Diagnostics module for code analysis via LSP and CLI tools.
+
+This module provides:
+- LSP server lifecycle management via LSPManager
+- LSP proxy for stdio-based servers exposed via Unix socket
+- Integration with execution environments (local, remote, sandboxed)
+"""
+
+from agentpool.diagnostics.lsp_manager import LSPManager
+from agentpool.diagnostics.lsp_proxy import LSPProxy
+from agentpool.diagnostics.models import (
+    COMPLETION_KIND_MAP,
+    CODE_ACTION_KIND_MAP,
+    SYMBOL_KIND_MAP,
+    CallHierarchyCall,
+    CallHierarchyItem,
+    CodeAction,
+    CompletionItem,
+    Diagnostic,
+    DiagnosticsResult,
+    HoverInfo,
+    Location,
+    LSPServerState,
+    Position,
+    Range,
+    RenameResult,
+    SignatureInfo,
+    SymbolInfo,
+    TypeHierarchyItem,
+)
+
+__all__ = [
+    "CODE_ACTION_KIND_MAP",
+    "COMPLETION_KIND_MAP",
+    "SYMBOL_KIND_MAP",
+    "CallHierarchyCall",
+    "CallHierarchyItem",
+    "CodeAction",
+    "CompletionItem",
+    "Diagnostic",
+    "DiagnosticsResult",
+    "HoverInfo",
+    "LSPManager",
+    "LSPProxy",
+    "LSPServerState",
+    "Location",
+    "Position",
+    "Range",
+    "RenameResult",
+    "SignatureInfo",
+    "SymbolInfo",
+    "TypeHierarchyItem",
+]