agentrun-sdk 0.1.2__py3-none-any.whl

agentrun_sdk/multiagent/a2a/server.py

@@ -0,0 +1,252 @@
+"""A2A-compatible wrapper for Strands Agent.
+
+This module provides the A2AAgent class, which adapts a Strands Agent to the A2A protocol,
+allowing it to be used in A2A-compatible systems.
+"""
+
+import logging
+from typing import Any, Literal
+from urllib.parse import urlparse
+
+import uvicorn
+from a2a.server.apps import A2AFastAPIApplication, A2AStarletteApplication
+from a2a.server.events import QueueManager
+from a2a.server.request_handlers import DefaultRequestHandler
+from a2a.server.tasks import InMemoryTaskStore, PushNotificationConfigStore, PushNotificationSender, TaskStore
+from a2a.types import AgentCapabilities, AgentCard, AgentSkill
+from fastapi import FastAPI
+from starlette.applications import Starlette
+
+from ...agent.agent import Agent as SAAgent
+from .executor import StrandsA2AExecutor
+
+logger = logging.getLogger(__name__)
+
+
+class A2AServer:
+    """A2A-compatible wrapper for Strands Agent."""
+
+    def __init__(
+        self,
+        agent: SAAgent,
+        *,
+        # AgentCard
+        host: str = "0.0.0.0",
+        port: int = 9000,
+        http_url: str | None = None,
+        serve_at_root: bool = False,
+        version: str = "0.0.1",
+        skills: list[AgentSkill] | None = None,
+        # RequestHandler
+        task_store: TaskStore | None = None,
+        queue_manager: QueueManager | None = None,
+        push_config_store: PushNotificationConfigStore | None = None,
+        push_sender: PushNotificationSender | None = None,
+
+    ):
+        """Initialize an A2A-compatible server from a Strands agent.
+
+        Args:
+            agent: The Strands Agent to wrap with A2A compatibility.
+            host: The hostname or IP address to bind the A2A server to. Defaults to "0.0.0.0".
+            port: The port to bind the A2A server to. Defaults to 9000.
+            http_url: The public HTTP URL where this agent will be accessible. If provided,
+                this overrides the generated URL from host/port and enables automatic
+                path-based mounting for load balancer scenarios.
+                Example: "http://my-alb.amazonaws.com/agent1"
+            serve_at_root: If True, forces the server to serve at root path regardless of
+                http_url path component. Use this when your load balancer strips path prefixes.
+                Defaults to False.
+            version: The version of the agent. Defaults to "0.0.1".
+            skills: The list of capabilities or functions the agent can perform.
+            task_store: Custom task store implementation for managing agent tasks. If None,
+                uses InMemoryTaskStore.
+            queue_manager: Custom queue manager for handling message queues. If None,
+                no queue management is used.
+            push_config_store: Custom store for push notification configurations. If None,
+                no push notification configuration is used.
+            push_sender: Custom push notification sender implementation. If None,
+                no push notifications are sent.
+        """
+        self.host = host
+        self.port = port
+        self.version = version
+
+        if http_url:
+            # Parse the provided URL to extract components for mounting
+            self.public_base_url, self.mount_path = self._parse_public_url(http_url)
+            self.http_url = http_url.rstrip("/") + "/"
+
+            # Override mount path if serve_at_root is requested
+            if serve_at_root:
+                self.mount_path = ""
+        else:
+            # Fall back to constructing the URL from host and port
+            self.public_base_url = f"http://{host}:{port}"
+            self.http_url = f"{self.public_base_url}/"
+            self.mount_path = ""
+
+        self.strands_agent = agent
+        self.name = self.strands_agent.name
+        self.description = self.strands_agent.description
+        self.capabilities = AgentCapabilities(streaming=True)
+        self.request_handler = DefaultRequestHandler(
+            agent_executor=StrandsA2AExecutor(self.strands_agent),
+            task_store=task_store or InMemoryTaskStore(),
+            queue_manager=queue_manager,
+            push_config_store=push_config_store,
+            push_sender=push_sender,
+        )
+        self._agent_skills = skills
+        logger.info("Strands' integration with A2A is experimental. Be aware of frequent breaking changes.")
+
+    def _parse_public_url(self, url: str) -> tuple[str, str]:
+        """Parse the public URL into base URL and mount path components.
+
+        Args:
+            url: The full public URL (e.g., "http://my-alb.amazonaws.com/agent1")
+
+        Returns:
+            tuple: (base_url, mount_path) where base_url is the scheme+netloc
+                  and mount_path is the path component
+
+        Example:
+            _parse_public_url("http://my-alb.amazonaws.com/agent1")
+            Returns: ("http://my-alb.amazonaws.com", "/agent1")
+        """
+        parsed = urlparse(url.rstrip("/"))
+        base_url = f"{parsed.scheme}://{parsed.netloc}"
+        mount_path = parsed.path if parsed.path != "/" else ""
+        return base_url, mount_path
+
+    @property
+    def public_agent_card(self) -> AgentCard:
+        """Get the public AgentCard for this agent.
+
+        The AgentCard contains metadata about the agent, including its name,
+        description, URL, version, skills, and capabilities. This information
+        is used by other agents and systems to discover and interact with this agent.
+
+        Returns:
+            AgentCard: The public agent card containing metadata about this agent.
+
+        Raises:
+            ValueError: If name or description is None or empty.
+        """
+        if not self.name:
+            raise ValueError("A2A agent name cannot be None or empty")
+        if not self.description:
+            raise ValueError("A2A agent description cannot be None or empty")
+
+        return AgentCard(
+            name=self.name,
+            description=self.description,
+            url=self.http_url,
+            version=self.version,
+            skills=self.agent_skills,
+            default_input_modes=["text"],
+            default_output_modes=["text"],
+            capabilities=self.capabilities,
+        )
+
+    def _get_skills_from_tools(self) -> list[AgentSkill]:
+        """Get the list of skills from Strands agent tools.
+
+        Skills represent specific capabilities that the agent can perform.
+        Strands agent tools are adapted to A2A skills.
+
+        Returns:
+            list[AgentSkill]: A list of skills this agent provides.
+        """
+        return [
+            AgentSkill(name=config["name"], id=config["name"], description=config["description"], tags=[])
+            for config in self.strands_agent.tool_registry.get_all_tools_config().values()
+        ]
+
+    @property
+    def agent_skills(self) -> list[AgentSkill]:
+        """Get the list of skills this agent provides."""
+        return self._agent_skills if self._agent_skills is not None else self._get_skills_from_tools()
+
+    @agent_skills.setter
+    def agent_skills(self, skills: list[AgentSkill]) -> None:
+        """Set the list of skills this agent provides.
+
+        Args:
+            skills: A list of AgentSkill objects to set for this agent.
+        """
+        self._agent_skills = skills
+
+    def to_starlette_app(self) -> Starlette:
+        """Create a Starlette application for serving this agent via HTTP.
+
+        Automatically handles path-based mounting if a mount path was derived
+        from the http_url parameter.
+
+        Returns:
+            Starlette: A Starlette application configured to serve this agent.
+        """
+        a2a_app = A2AStarletteApplication(agent_card=self.public_agent_card, http_handler=self.request_handler).build()
+
+        if self.mount_path:
+            # Create parent app and mount the A2A app at the specified path
+            parent_app = Starlette()
+            parent_app.mount(self.mount_path, a2a_app)
+            logger.info("Mounting A2A server at path: %s", self.mount_path)
+            return parent_app
+
+        return a2a_app
+
+    def to_fastapi_app(self) -> FastAPI:
+        """Create a FastAPI application for serving this agent via HTTP.
+
+        Automatically handles path-based mounting if a mount path was derived
+        from the http_url parameter.
+
+        Returns:
+            FastAPI: A FastAPI application configured to serve this agent.
+        """
+        a2a_app = A2AFastAPIApplication(agent_card=self.public_agent_card, http_handler=self.request_handler).build()
+
+        if self.mount_path:
+            # Create parent app and mount the A2A app at the specified path
+            parent_app = FastAPI()
+            parent_app.mount(self.mount_path, a2a_app)
+            logger.info("Mounting A2A server at path: %s", self.mount_path)
+            return parent_app
+
+        return a2a_app
+
+    def serve(
+        self,
+        app_type: Literal["fastapi", "starlette"] = "starlette",
+        *,
+        host: str | None = None,
+        port: int | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Start the A2A server with the specified application type.
+
+        This method starts an HTTP server that exposes the agent via the A2A protocol.
+        The server can be implemented using either FastAPI or Starlette, depending on
+        the specified app_type.
+
+        Args:
+            app_type: The type of application to serve, either "fastapi" or "starlette".
+                Defaults to "starlette".
+            host: The host address to bind the server to. Defaults to "0.0.0.0".
+            port: The port number to bind the server to. Defaults to 9000.
+            **kwargs: Additional keyword arguments to pass to uvicorn.run.
+        """
+        try:
+            logger.info("Starting Strands A2A server...")
+            if app_type == "fastapi":
+                uvicorn.run(self.to_fastapi_app(), host=host or self.host, port=port or self.port, **kwargs)
+            else:
+                uvicorn.run(self.to_starlette_app(), host=host or self.host, port=port or self.port, **kwargs)
+        except KeyboardInterrupt:
+            logger.warning("Strands A2A server shutdown requested (KeyboardInterrupt).")
+        except Exception:
+            logger.exception("Strands A2A server encountered exception.")
+        finally:
+            logger.info("Strands A2A server has shutdown.")

agentrun_sdk/multiagent/base.py

@@ -0,0 +1,92 @@
+"""Multi-Agent Base Class.
+
+Provides minimal foundation for multi-agent patterns (Swarm, Graph).
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Union
+
+from ..agent import AgentResult
+from ..types.content import ContentBlock
+from ..types.event_loop import Metrics, Usage
+
+
+class Status(Enum):
+    """Execution status for both graphs and nodes."""
+
+    PENDING = "pending"
+    EXECUTING = "executing"
+    COMPLETED = "completed"
+    FAILED = "failed"
+
+
+@dataclass
+class NodeResult:
+    """Unified result from node execution - handles both Agent and nested MultiAgentBase results.
+
+    The status field represents the semantic outcome of the node's work:
+    - COMPLETED: The node's task was successfully accomplished
+    - FAILED: The node's task failed or produced an error
+    """
+
+    # Core result data - single AgentResult, nested MultiAgentResult, or Exception
+    result: Union[AgentResult, "MultiAgentResult", Exception]
+
+    # Execution metadata
+    execution_time: int = 0
+    status: Status = Status.PENDING
+
+    # Accumulated metrics from this node and all children
+    accumulated_usage: Usage = field(default_factory=lambda: Usage(inputTokens=0, outputTokens=0, totalTokens=0))
+    accumulated_metrics: Metrics = field(default_factory=lambda: Metrics(latencyMs=0))
+    execution_count: int = 0
+
+    def get_agent_results(self) -> list[AgentResult]:
+        """Get all AgentResult objects from this node, flattened if nested."""
+        if isinstance(self.result, Exception):
+            return []  # No agent results for exceptions
+        elif isinstance(self.result, AgentResult):
+            return [self.result]
+        else:
+            # Flatten nested results from MultiAgentResult
+            flattened = []
+            for nested_node_result in self.result.results.values():
+                flattened.extend(nested_node_result.get_agent_results())
+            return flattened
+
+
+@dataclass
+class MultiAgentResult:
+    """Result from multi-agent execution with accumulated metrics.
+
+    The status field represents the outcome of the MultiAgentBase execution:
+    - COMPLETED: The execution was successfully accomplished
+    - FAILED: The execution failed or produced an error
+    """
+
+    status: Status = Status.PENDING
+    results: dict[str, NodeResult] = field(default_factory=lambda: {})
+    accumulated_usage: Usage = field(default_factory=lambda: Usage(inputTokens=0, outputTokens=0, totalTokens=0))
+    accumulated_metrics: Metrics = field(default_factory=lambda: Metrics(latencyMs=0))
+    execution_count: int = 0
+    execution_time: int = 0
+
+
+class MultiAgentBase(ABC):
+    """Base class for multi-agent helpers.
+
+    This class integrates with existing Strands Agent instances and provides
+    multi-agent orchestration capabilities.
+    """
+
+    @abstractmethod
+    async def invoke_async(self, task: str | list[ContentBlock], **kwargs: Any) -> MultiAgentResult:
+        """Invoke asynchronously."""
+        raise NotImplementedError("invoke_async not implemented")
+
+    @abstractmethod
+    def __call__(self, task: str | list[ContentBlock], **kwargs: Any) -> MultiAgentResult:
+        """Invoke synchronously."""
+        raise NotImplementedError("__call__ not implemented")