strands-swarms 0.1.1__py3-none-any.whl → 0.1.2__py3-none-any.whl

strands_swarms/__init__.py

@@ -26,37 +26,35 @@ Example:
         '''Search the web.'''
         return f"Results for: {query}"
 
-    swarm = DynamicSwarm(
-        available_tools={"search_web": search_web},
-        verbose=True,
-    )
+    # Basic usage
+    swarm = DynamicSwarm(available_tools={"search_web": search_web})
     result = swarm.execute("Research AI trends and summarize")
+
+    # Streaming trajectory
+    import asyncio
+
+    async def run():
+        async for event in swarm.stream_async(
+            "Research AI trends and summarize",
+            include_subagent_events=False,
+        ):
+            print(event)
+
+    asyncio.run(run())
 """
 
 # Re-export strands types for convenience
-from strands.hooks import HookProvider, HookRegistry
 from strands.multiagent.base import Status
 
-from .events import (
-    AgentSpawnedEvent,
-    ExecutionCompletedEvent,
-    # Execution events
-    ExecutionStartedEvent,
-    PlanningCompletedEvent,
-    PlanningStartedEvent,
-    # Hook provider
-    PrintingHookProvider,
-    SwarmCompletedEvent,
-    SwarmFailedEvent,
-    # Planning/Orchestration events
-    SwarmStartedEvent,
-    TaskCompletedEvent,
-    TaskCreatedEvent,
-    TaskFailedEvent,
-    TaskStartedEvent,
+from .definition import (
+    AgentDefinition,
+    DynamicSwarmCapabilities,
+    SessionConfig,
+    SwarmDefinition,
+    TaskDefinition,
 )
+from .dynamic_swarm import DynamicSwarm, DynamicSwarmResult, build_swarm
 from .orchestrator import create_orchestrator_agent
-from .swarm import DynamicSwarm, DynamicSwarmResult
 
 __version__ = "0.1.1"
 
@@ -64,25 +62,14 @@ __all__ = [
     # Main API
     "DynamicSwarm",
     "DynamicSwarmResult",
-    # Orchestrator (handles both planning AND completion in same conversation)
+    "build_swarm",
+    # Definition types
+    "DynamicSwarmCapabilities",
+    "SwarmDefinition",
+    "AgentDefinition",
+    "TaskDefinition",
+    "SessionConfig",
+    # Orchestrator
     "create_orchestrator_agent",
-    # Status enum
     "Status",
-    # Events (for custom hooks)
-    "SwarmStartedEvent",
-    "PlanningStartedEvent",
-    "AgentSpawnedEvent",
-    "TaskCreatedEvent",
-    "PlanningCompletedEvent",
-    "ExecutionStartedEvent",
-    "TaskStartedEvent",
-    "TaskCompletedEvent",
-    "TaskFailedEvent",
-    "ExecutionCompletedEvent",
-    "SwarmCompletedEvent",
-    "SwarmFailedEvent",
-    # Hook system
-    "PrintingHookProvider",
-    "HookProvider",
-    "HookRegistry",
 ]

strands_swarms/definition.py

@@ -0,0 +1,169 @@
+"""Swarm definition types - data structures for defining multi-agent workflows."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Callable
+
+from strands.session.file_session_manager import FileSessionManager
+
+# ANSI Color Constants (for consistent agent output coloring)
+AGENT_COLORS: list[str] = [
+    "\033[94m",  # Blue
+    "\033[92m",  # Green
+    "\033[93m",  # Yellow
+    "\033[95m",  # Magenta
+    "\033[96m",  # Cyan
+    "\033[91m",  # Red
+]
+
+if TYPE_CHECKING:
+    from strands.models import Model
+
+
+@dataclass(frozen=True)
+class DynamicSwarmCapabilities:
+    """Immutable configuration for available tools and models."""
+
+    available_tools: dict[str, Callable[..., Any]]
+    available_models: dict[str, "Model"]
+    default_model: str | None = None
+
+    def validate_tools(self, tool_names: list[str]) -> None:
+        """Validate that all tool names exist in available tools."""
+        for tool_name in tool_names:
+            if tool_name not in self.available_tools:
+                available = list(self.available_tools.keys())
+                raise ValueError(
+                    f"Tool '{tool_name}' not in available tools: {available}"
+                )
+
+    def validate_model(self, model_name: str | None) -> None:
+        """Validate that the model name exists in available models."""
+        if model_name and model_name not in self.available_models:
+            available = list(self.available_models.keys())
+            raise ValueError(
+                f"Model '{model_name}' not in available models: {available}"
+            )
+
+    @property
+    def available_tool_names(self) -> list[str]:
+        return list(self.available_tools.keys())
+
+    @property
+    def available_model_names(self) -> list[str]:
+        return list(self.available_models.keys())
+
+
+@dataclass
+class AgentDefinition:
+    """Definition of a sub-agent to be spawned."""
+
+    name: str
+    role: str
+    instructions: str | None = None
+    tools: list[str] = field(default_factory=list)
+    model: str | None = None
+    color: str | None = None
+
+    def build_system_prompt(self) -> str:
+        parts = [f"You are a {self.role}."]
+        if self.instructions:
+            parts.append(f"\n\nInstructions:\n{self.instructions}")
+        return "\n".join(parts)
+
+
+@dataclass
+class TaskDefinition:
+    """Definition of a task to be executed by a sub-agent."""
+
+    name: str
+    agent: str
+    description: str | None = None
+    depends_on: list[str] = field(default_factory=list)
+
+
+@dataclass
+class SessionConfig:
+    """Configuration for session persistence."""
+
+    session_id: str
+    storage_dir: str = "./.swarm_sessions"
+
+    def for_agent(self, agent_name: str) -> FileSessionManager:
+        return FileSessionManager(
+            session_id=f"{self.session_id}-{agent_name}",
+            storage_dir=self.storage_dir,
+        )
+
+    def for_graph(self) -> FileSessionManager:
+        return FileSessionManager(
+            session_id=f"{self.session_id}-graph",
+            storage_dir=self.storage_dir,
+        )
+
+
+class SwarmDefinition:
+    """Per-query definition of sub-agents and tasks created during planning."""
+
+    def __init__(
+        self,
+        capabilities: DynamicSwarmCapabilities,
+    ) -> None:
+        self._capabilities = capabilities
+        self.sub_agents: dict[str, AgentDefinition] = {}
+        self.tasks: dict[str, TaskDefinition] = {}
+        self._color_index = 0
+
+    @property
+    def capabilities(self) -> DynamicSwarmCapabilities:
+        return self._capabilities
+
+    def register_agent(self, definition: AgentDefinition) -> None:
+        """Register a sub-agent definition."""
+        if definition.name in self.sub_agents:
+            raise ValueError(f"Agent '{definition.name}' already exists")
+
+        self._capabilities.validate_tools(definition.tools)
+        self._capabilities.validate_model(definition.model)
+
+        definition.color = AGENT_COLORS[self._color_index % len(AGENT_COLORS)]
+        self._color_index += 1
+
+        self.sub_agents[definition.name] = definition
+
+    def register_task(self, definition: TaskDefinition) -> None:
+        """Register a task definition.
+
+        Note: dependencies must already be registered (create dependency tasks first).
+        """
+        if definition.name in self.tasks:
+            raise ValueError(f"Task '{definition.name}' already exists")
+
+        if definition.agent not in self.sub_agents:
+            available = list(self.sub_agents.keys())
+            raise ValueError(
+                f"Agent '{definition.agent}' not found. Available: {available}"
+            )
+
+        for dep in definition.depends_on:
+            if dep not in self.tasks:
+                available = list(self.tasks.keys())
+                raise ValueError(
+                    f"Dependency '{dep}' not found. Available: {available}"
+                )
+
+        self.tasks[definition.name] = definition
+
+    def get_summary(self) -> str:
+        """Get a summary of registered sub-agents and tasks."""
+        lines = [
+            f"Agents ({len(self.sub_agents)}):",
+            *[f"  - {name}: {d.role}" for name, d in self.sub_agents.items()],
+            f"\nTasks ({len(self.tasks)}):",
+            *[
+                f"  - {name} -> {d.agent}" + (f" (depends: {d.depends_on})" if d.depends_on else "")
+                for name, d in self.tasks.items()
+            ],
+        ]
+        return "\n".join(lines)

strands_swarms/dynamic_swarm.py

@@ -0,0 +1,439 @@
+"""DynamicSwarm - automatically construct and execute multi-agent workflows."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Callable
+
+from strands import Agent
+from strands.multiagent.base import MultiAgentResult, Status
+from strands.multiagent.graph import (
+    Graph,
+    GraphBuilder,
+    GraphResult,
+)
+
+from .definition import (
+    AgentDefinition,
+    DynamicSwarmCapabilities,
+    SessionConfig,
+    SwarmDefinition,
+    TaskDefinition,
+)
+
+if TYPE_CHECKING:
+    from strands.models import Model
+
+
+# =============================================================================
+# DynamicSwarm
+# =============================================================================
+
+class DynamicSwarm:
+    """Dynamically construct and execute multi-agent workflows.
+
+    An orchestrator agent analyzes queries and coordinates workflows by:
+    1. Planning: Creating specialized sub-agents with tools
+    2. Execution: Running tasks with dependencies in parallel where possible
+    3. Synthesis: Combining task outputs into a final response
+    """
+
+    def __init__(
+        self,
+        available_tools: dict[str, Callable[..., Any]] | None = None,
+        available_models: dict[str, "Model"] | None = None,
+        *,
+        orchestrator_model: "Model" | None = None,
+        default_agent_model: str | None = None,
+        execution_timeout: float = 900.0,
+        task_timeout: float = 300.0,
+        session_id: str | None = None,
+        session_storage_dir: str = "./.swarm_sessions",
+    ) -> None:
+        self._capabilities = DynamicSwarmCapabilities(
+            available_tools=available_tools or {},
+            available_models=available_models or {},
+            default_model=default_agent_model,
+        )
+        self._orchestrator_model = orchestrator_model
+        self._execution_timeout = execution_timeout
+        self._task_timeout = task_timeout
+        self._session_config = (
+            SessionConfig(session_id=session_id, storage_dir=session_storage_dir)
+            if session_id
+            else None
+        )
+
+    def execute(self, query: str) -> DynamicSwarmResult:
+        """Execute a query synchronously."""
+        return asyncio.run(self.execute_async(query))
+
+    async def execute_async(self, query: str) -> DynamicSwarmResult:
+        """Execute a query asynchronously (non-streaming)."""
+        result: DynamicSwarmResult | None = None
+        async for event in self.stream_async(query):
+            if event.get("type") == "swarm_result":
+                result = event.get("result")
+        if not isinstance(result, DynamicSwarmResult):
+            raise RuntimeError("DynamicSwarm stream ended without producing a result")
+        return result
+
+    async def stream_async(
+        self,
+        query: str,
+        *,
+        include_subagent_events: bool = False,
+    ) -> AsyncIterator[dict[str, Any]]:
+        """Stream execution events, including planning and graph execution.
+
+        This yields:
+        - High-level DynamicSwarm events (planning/execution/synthesis)
+        - Raw `strands` Graph stream events (e.g. multiagent_node_start/stop, multiagent_result)
+
+        By default, per-node agent stream events (type: "multiagent_node_stream") are
+        filtered to avoid interleaved token/tool output when tasks run in parallel.
+        Pass `include_subagent_events=True` to include them.
+        """
+        definition = SwarmDefinition(capabilities=self._capabilities)
+
+        yield {
+            "type": "swarm_started",
+            "query": query,
+            "available_tools": self._capabilities.available_tool_names,
+            "available_models": self._capabilities.available_model_names,
+        }
+
+        # Phase 1: Planning
+        yield {"type": "planning_started"}
+        planning_result = await self._run_planning(query, definition)
+
+        if not planning_result.success:
+            error = planning_result.error or "Planning failed"
+            yield {"type": "planning_failed", "error": error}
+            yield {
+                "type": "swarm_result",
+                "result": DynamicSwarmResult(
+                    status=Status.FAILED,
+                    planning_output=planning_result.output,
+                    agents_spawned=list(definition.sub_agents.values()),
+                    tasks_created=list(definition.tasks.values()),
+                    error=error,
+                ),
+            }
+            return
+
+        if not definition.tasks:
+            error = "No tasks were created"
+            yield {"type": "planning_failed", "error": error}
+            yield {
+                "type": "swarm_result",
+                "result": DynamicSwarmResult(
+                    status=Status.FAILED,
+                    planning_output=planning_result.output,
+                    agents_spawned=list(definition.sub_agents.values()),
+                    tasks_created=[],
+                    error=error,
+                ),
+            }
+            return
+
+        yield {
+            "type": "planning_completed",
+            "planning_output": planning_result.output,
+            "summary": definition.get_summary(),
+            "agents": definition.sub_agents,
+            "tasks": definition.tasks,
+        }
+
+        # Phase 2: Execution
+        graph = build_swarm(
+            definition,
+            execution_timeout=self._execution_timeout,
+            task_timeout=self._task_timeout,
+            session_config=self._session_config,
+        )
+        yield {"type": "execution_started", "tasks": list(definition.tasks.keys())}
+
+        execution_result: GraphResult | None = None
+        execution_error: str | None = None
+        try:
+            async for event in graph.stream_async(query):
+                if event.get("type") == "multiagent_result":
+                    candidate = event.get("result")
+                    if isinstance(candidate, GraphResult):
+                        execution_result = candidate
+                if not include_subagent_events and event.get("type") == "multiagent_node_stream":
+                    continue
+                yield event
+        except Exception as e:
+            execution_error = str(e)
+
+        if execution_error:
+            yield {"type": "execution_failed", "error": execution_error}
+            yield {
+                "type": "swarm_result",
+                "result": DynamicSwarmResult(
+                    status=Status.FAILED,
+                    planning_output=planning_result.output,
+                    execution_result=execution_result,
+                    final_response=None,
+                    agents_spawned=list(definition.sub_agents.values()),
+                    tasks_created=list(definition.tasks.values()),
+                    error=execution_error,
+                ),
+            }
+            return
+
+        if not execution_result:
+            error = "Graph completed without producing a result"
+            yield {"type": "execution_failed", "error": error}
+            yield {
+                "type": "swarm_result",
+                "result": DynamicSwarmResult(
+                    status=Status.FAILED,
+                    planning_output=planning_result.output,
+                    execution_result=None,
+                    final_response=None,
+                    agents_spawned=list(definition.sub_agents.values()),
+                    tasks_created=list(definition.tasks.values()),
+                    error=error,
+                ),
+            }
+            return
+
+        yield {"type": "execution_completed", "status": execution_result.status.value}
+
+        # Phase 3: Synthesis
+        yield {"type": "synthesis_started"}
+        assert planning_result.orchestrator is not None
+        final_response = await self._synthesize_final_response(
+            query, definition, execution_result, planning_result.orchestrator
+        )
+        yield {"type": "synthesis_completed", "final_response": final_response}
+
+        result = DynamicSwarmResult(
+            status=execution_result.status,
+            planning_output=planning_result.output,
+            execution_result=execution_result,
+            final_response=final_response,
+            agents_spawned=list(definition.sub_agents.values()),
+            tasks_created=list(definition.tasks.values()),
+            error=(
+                None
+                if execution_result.status == Status.COMPLETED
+                else f"Execution ended with status {execution_result.status.value}"
+            ),
+        )
+
+        yield {"type": "swarm_result", "result": result}
+
+    async def _run_planning(
+        self, query: str, definition: SwarmDefinition
+    ) -> _PlanningResult:
+        from .orchestrator import create_orchestrator_agent
+
+        orchestrator = create_orchestrator_agent(
+            definition=definition,
+            model=self._orchestrator_model,
+        )
+
+        prompt = PLANNING_PROMPT.format(
+            query=query,
+            available_tools=definition.capabilities.available_tool_names or ["none"],
+            available_models=definition.capabilities.available_model_names or ["default"],
+        )
+
+        try:
+            output = _extract_message_text(orchestrator(prompt))
+            return _PlanningResult(success=True, output=output, orchestrator=orchestrator)
+        except Exception as e:
+            return _PlanningResult(success=False, error=str(e))
+
+    async def _synthesize_final_response(
+        self,
+        query: str,
+        definition: SwarmDefinition,
+        execution_result: MultiAgentResult | None,
+        orchestrator: Agent,
+    ) -> str | None:
+        if not execution_result or not hasattr(execution_result, "results"):
+            return None
+
+        task_outputs: list[str] = []
+        for task_name in definition.tasks:
+            node_result = execution_result.results.get(task_name)
+            if node_result:
+                task_outputs.append(f"[{task_name}]:\n{node_result.result}")
+
+        if not task_outputs:
+            return None
+
+        prompt = SYNTHESIS_PROMPT.format(
+            query=query,
+            task_outputs="\n\n".join(task_outputs),
+        )
+
+        try:
+            return _extract_message_text(orchestrator(prompt))
+        except Exception:
+            return None
+
+
+# =============================================================================
+# Results
+# =============================================================================
+
+
+@dataclass
+class DynamicSwarmResult:
+    """Result from DynamicSwarm execution."""
+
+    status: Status
+    planning_output: str | None = None
+    execution_result: GraphResult | None = None
+    final_response: str | None = None
+    agents_spawned: list[AgentDefinition] = field(default_factory=list)
+    tasks_created: list[TaskDefinition] = field(default_factory=list)
+    error: str | None = None
+
+    @property
+    def agents_spawned_count(self) -> int:
+        return len(self.agents_spawned)
+
+    @property
+    def tasks_created_count(self) -> int:
+        return len(self.tasks_created)
+
+    def get_output(self, task_name: str) -> Any | None:
+        if self.execution_result and hasattr(self.execution_result, "results"):
+            node_result = self.execution_result.results.get(task_name)
+            if node_result:
+                return str(node_result.result)
+        return None
+
+    def __bool__(self) -> bool:
+        return self.status == Status.COMPLETED
+
+
+@dataclass
+class _PlanningResult:
+    success: bool
+    output: str | None = None
+    error: str | None = None
+    orchestrator: Agent | None = None
+
+
+# =============================================================================
+# Graph Building
+# =============================================================================
+
+def build_swarm(
+    definition: SwarmDefinition,
+    *,
+    execution_timeout: float = 900.0,
+    task_timeout: float = 300.0,
+    session_config: SessionConfig | None = None,
+) -> Graph:
+    """Build a strands Graph from a SwarmDefinition."""
+    if not definition.tasks:
+        raise ValueError("No tasks registered - cannot build swarm")
+
+    capabilities = definition.capabilities
+
+    # Build per-agent config from definitions.
+    #
+    # Note: strands Graph requires each node executor to be a unique object
+    # instance, so we use these configs to create a fresh Agent per task.
+    agent_configs: dict[str, dict[str, Any]] = {}
+    for name, agent_def in definition.sub_agents.items():
+        tools = [capabilities.available_tools[t] for t in agent_def.tools]
+        model_name = agent_def.model or capabilities.default_model
+        model = capabilities.available_models.get(model_name) if model_name else None
+        agent_configs[name] = {
+            "system_prompt": agent_def.build_system_prompt(),
+            "model": model,
+            "tools": tools or None,
+        }
+
+    # Build the execution graph
+    builder = GraphBuilder()
+    for task_name, task in definition.tasks.items():
+        agent_config = agent_configs[task.agent]
+        builder.add_node(
+            Agent(
+                name=task.agent,
+                system_prompt=agent_config["system_prompt"],
+                model=agent_config["model"],
+                tools=agent_config["tools"],  # type: ignore[arg-type]
+                callback_handler=None,
+            ),
+            task_name,
+        )
+    for task_name, task in definition.tasks.items():
+        for dep_name in task.depends_on:
+            builder.add_edge(dep_name, task_name)
+
+    builder.set_execution_timeout(execution_timeout)
+    builder.set_node_timeout(task_timeout)
+    if session_config:
+        builder.set_session_manager(session_config.for_graph())
+
+    return builder.build()
+
+
+# =============================================================================
+# Prompts
+# =============================================================================
+
+
+PLANNING_PROMPT = """\
+Analyze this request and design a multi-agent workflow to complete it:
+
+REQUEST: {query}
+
+AVAILABLE TOOLS: {available_tools}
+AVAILABLE MODELS: {available_models}
+
+INSTRUCTIONS:
+1. Break down the request into logical steps
+2. Create specialized agents with appropriate tools using spawn_agent()
+3. Create tasks with dependencies using create_task()
+   - Important: dependencies must already exist (create dependency tasks first)
+4. Call finalize_plan() when done
+
+Keep the workflow simple - only create agents and tasks that are necessary."""
+
+SYNTHESIS_PROMPT = """\
+The workflow has completed. Here are the results from each task:
+
+{task_outputs}
+
+ORIGINAL REQUEST: {query}
+
+Synthesize these results into a final response. \
+Be direct - deliver the answer without mentioning the workflow or agents."""
+
+
+# =============================================================================
+# Helpers
+# =============================================================================
+
+
+def _extract_message_text(result: Any) -> str | None:
+    """Extract the first text block from a strands Agent result, if present."""
+    message = getattr(result, "message", None)
+    if not isinstance(message, dict):
+        return None
+
+    content = message.get("content")
+    if not isinstance(content, list) or not content:
+        return None
+
+    first = content[0]
+    if not isinstance(first, dict):
+        return None
+
+    text = first.get("text")
+    return text if isinstance(text, str) else None