feat: support A2A protocol (#657)

Author: akotyla

.libraries-whitelist.txt

@@ -11,3 +11,4 @@ psycopg
 pytest-postgresql
 python-bidi
 griffe
+types-requests

examples/agents/a2a/agent_orchestrator.py

@@ -0,0 +1,204 @@
+import json
+from typing import TypeVar
+
+import requests
+from pydantic import BaseModel
+
+from ragbits.agents import Agent, AgentOptions, AgentResult, ToolCallResult
+from ragbits.core.llms import LiteLLM
+from ragbits.core.options import Options
+from ragbits.core.prompt import ChatFormat, Prompt
+
+OptionsT = TypeVar("OptionsT", bound=Options)
+
+
+class RoutingPromptInput(BaseModel):
+    """Represents a routing prompt input."""
+
+    message: str
+    agents: list
+
+
+class ResultsSumarizationPromptInput(BaseModel):
+    """Represents a results summarization prompt input."""
+
+    message: str
+    agent_results: list
+
+
+class RemoteAgentTask(BaseModel):
+    """Model representing a task for a remote agent"""
+
+    agent_url: str
+    parameters: dict
+
+
+class AgentOrchestrator(Agent):
+    """
+    Coordinates querying and aggregating responses from multiple remote agents
+    using tools for routing and task execution.
+    """
+
+    def __init__(
+        self,
+        llm: LiteLLM,
+        routing_prompt: type[Prompt[RoutingPromptInput, list[dict]]],
+        results_summarization_prompt: type[Prompt[ResultsSumarizationPromptInput, list]],
+        timeout: float = 20.0,
+        *,
+        history: ChatFormat | None = None,
+        keep_history: bool = False,
+        default_options: AgentOptions[OptionsT] | None = None,
+    ):
+        """
+        Initialize the orchestrator with tools for agent coordination.
+
+        Args:
+            llm: The LLM to run the agent.
+            routing_prompt: Prompt template for routing messages to agents.
+            results_summarization_prompt: Prompt template for summarizing agent results.
+            timeout: Timeout in seconds for the HTTP request.
+            history: The history of the agent.
+            keep_history: Whether to keep the history of the agent.
+            default_options: The default options for the agent run.
+
+        """
+        super().__init__(
+            llm=llm,
+            prompt=None,
+            history=history,
+            keep_history=keep_history,
+            tools=[self.create_agent_tasks, self.execute_agent_task, self.summarize_agent_results],
+            default_options=default_options,
+        )
+
+        self._timeout = timeout
+
+        self._routing_prompt = routing_prompt
+        self._results_summarization_prompt = results_summarization_prompt
+
+        self._remote_agents: dict[str, dict] = {}
+
+        self._current_tasks: list[RemoteAgentTask] = []
+        self._current_results: list[AgentResult] = []
+
+    def add_remote_agent(self, host: str, port: int, protocol: str = "http") -> None:
+        """
+        Discovers and registers a remote agent by fetching its agent card metadata.
+
+        Args:
+            host: The hostname or IP address of the remote agent.
+            port: The port on which the remote agent server is running.
+            protocol: The communication protocol (http or https). Defaults to "http".
+        """
+        url = f"{protocol}://{host}:{port}"
+        if url not in self._remote_agents:
+            agent_card = requests.get(f"{url}/.well-known/agent.json", timeout=self._timeout)
+            self._remote_agents[url] = agent_card.json()
+
+    async def create_agent_tasks(self, message: str) -> str:
+        """
+        Creates tasks for remote agents based on the input message.
+
+        Args:
+            message: The user query to route to agents
+
+        Returns:
+            JSON string of created tasks
+        """
+        prompt_input = RoutingPromptInput(message=message, agents=self._list_remote_agents())
+        prompt = self._routing_prompt(prompt_input)
+        response = await self.llm.generate(prompt)
+
+        tasks = json.loads(response)
+        self._current_tasks = [RemoteAgentTask(**task) for task in tasks]
+
+        return json.dumps(
+            {
+                "status": "success",
+                "task_count": len(self._current_tasks),
+                "tasks": [task.dict() for task in self._current_tasks],
+            }
+        )
+
+    async def execute_agent_task(self, task_index: int) -> str:
+        """
+        Executes a specific task from the current task list.
+
+        Args:
+            task_index: Index of the task to execute
+
+        Returns:
+            JSON string of the execution result
+        """
+        if not self._current_tasks or task_index >= len(self._current_tasks):
+            return json.dumps({"status": "error", "message": "Invalid task index"})
+
+        task = self._current_tasks[task_index]
+        result = self._execute_single_task(task.agent_url, task.parameters)
+        self._current_results.append(result)
+
+        tool_calls = None
+        if result.tool_calls:
+            tool_calls = [{"name": tc.name, "arguments": tc.arguments, "output": tc.output} for tc in result.tool_calls]
+
+        return json.dumps(
+            {
+                "status": "success",
+                "agent_url": task.agent_url,
+                "result": {"content": result.content, "metadata": result.metadata, "tool_calls": tool_calls},
+            }
+        )
+
+    async def summarize_agent_results(self, message: str) -> str:
+        """
+        Summarizes all collected agent results.
+
+        Args:
+            message: The original user message
+
+        Returns:
+            The summarized response
+        """
+        if not self._current_results:
+            return "No results to summarize"
+
+        input_data = ResultsSumarizationPromptInput(message=message, agent_results=self._current_results)
+        prompt = self._results_summarization_prompt(input_data=input_data)
+        return await self.llm.generate(prompt)
+
+    def _execute_single_task(self, agent_url: str, params: dict) -> AgentResult:
+        payload = {"params": params}
+        raw_response = requests.post(agent_url, json=payload, timeout=self._timeout)
+        raw_response.raise_for_status()
+
+        response = raw_response.json()
+        result_data = response["result"]
+
+        return AgentResult(
+            content=result_data["content"],
+            metadata=result_data.get("metadata", {}),
+            history=result_data["history"],
+            tool_calls=[ToolCallResult(**call) for call in result_data.get("tool_calls", [])] or None,
+        )
+
+    def _list_remote_agents(self) -> list[dict]:
+        """
+        Lists metadata of all registered remote agents in a format suitable for routing.
+
+        Returns:
+            A list of dictionaries describing each remote agent's name, URL,
+            description, and skills.
+        """
+        return [
+            {
+                "name": data.get("name"),
+                "agent_url": url,
+                "description": data.get("description"),
+                "skills": [
+                    {"id": skill.get("id"), "name": skill.get("name"), "description": skill.get("description")}
+                    for skill in data.get("skills", [])
+                ],
+            }
+            for url, data in self._remote_agents.items()
+        ]

examples/agents/a2a/flight_agent.py

@@ -0,0 +1,78 @@
+import json
+
+from pydantic import BaseModel
+
+from ragbits.agents import Agent
+from ragbits.agents.a2a.server import create_agent_app, create_agent_server
+from ragbits.core.llms import LiteLLM
+from ragbits.core.prompt import Prompt
+
+
+def get_flight_info(departure: str, arrival: str) -> str:
+    """
+    Returns flight information between two locations.
+
+    Args:
+        departure: The departure city.
+        arrival: The arrival city.
+
+    Returns:
+        A JSON string with mock flight details.
+    """
+    if "new york" in departure.lower() and "paris" in arrival.lower():
+        return json.dumps(
+            {
+                "from": "New York",
+                "to": "Paris",
+                "flights": [
+                    {"airline": "British Airways", "departure": "10:00 AM", "arrival": "10:00 PM"},
+                    {"airline": "Delta", "departure": "1:00 PM", "arrival": "1:00 AM"},
+                ],
+            }
+        )
+    elif "los angeles" in departure.lower() and "tokyo" in arrival.lower():
+        return json.dumps(
+            {
+                "from": "Los Angeles",
+                "to": "Tokyo",
+                "flights": [
+                    {"airline": "ANA", "departure": "8:00 AM", "arrival": "12:00 PM"},
+                    {"airline": "JAL", "departure": "4:00 PM", "arrival": "8:00 AM"},
+                ],
+            }
+        )
+    else:
+        return json.dumps({"from": departure, "to": arrival, "flights": "No flight data available"})
+
+
+class FlightPromptInput(BaseModel):
+    """Defines the structured input schema for the flight search prompt."""
+
+    departure: str
+    arrival: str
+
+
+class FlightPrompt(Prompt[FlightPromptInput]):
+    """Prompt for a flight search assistant."""
+
+    system_prompt = """
+    You are a helpful travel assistant that finds available flights between two cities.
+    """
+
+    user_prompt = """
+    I need to fly from {{ departure }} to {{ arrival }}. What flights are available?
+    """
+
+
+llm = LiteLLM(
+    model_name="gpt-4o-2024-08-06",
+    use_structured_output=True,
+)
+agent = Agent(llm=llm, prompt=FlightPrompt, tools=[get_flight_info])
+agent_card = agent.get_agent_card(
+    name="Flight Info Agent",
+    description="Provides available flight information between two cities.",
+)
+
+app = create_agent_app(agent, agent_card, FlightPromptInput)
+server = create_agent_server(agent, agent_card, FlightPromptInput)

examples/agents/a2a/hotel_agent.py

@@ -0,0 +1,82 @@
+import json
+
+from pydantic import BaseModel
+
+from ragbits.agents import Agent
+from ragbits.agents.a2a.server import create_agent_app, create_agent_server
+from ragbits.core.llms import LiteLLM
+from ragbits.core.prompt import Prompt
+
+
+def get_hotel_recommendations(city: str) -> str:
+    """
+    Returns hotel recommendations for a given city and date range.
+
+    Args:
+        city: The destination city.
+
+    Returns:
+        A JSON string with mock hotel details.
+    """
+    city_lower = city.lower()
+
+    if "paris" in city_lower:
+        return json.dumps(
+            {
+                "city": "Paris",
+                "hotels": [
+                    {"name": "Hotel Le Meurice", "rating": 5, "price_per_night": 450},
+                    {"name": "Hotel Regina Louvre", "rating": 4, "price_per_night": 300},
+                ],
+            }
+        )
+    elif "rome" in city_lower:
+        return json.dumps(
+            {
+                "city": "Rome",
+                "hotels": [
+                    {"name": "Hotel Eden", "rating": 5, "price_per_night": 400},
+                    {"name": "Hotel Artemide", "rating": 4, "price_per_night": 250},
+                ],
+            }
+        )
+    else:
+        return json.dumps(
+            {
+                "city": city,
+                "hotels": "No hotel data available",
+            }
+        )
+
+
+class HotelPromptInput(BaseModel):
+    """Defines the structured input for the hotel recommendation prompt."""
+
+    city: str
+
+
+class HotelPrompt(Prompt[HotelPromptInput]):
+    """Prompt for a hotel recommendation assistant."""
+
+    system_prompt = """
+    You are a helpful travel assistant that recommends hotels based on the city and travel dates.
+    """
+
+    user_prompt = """
+    I'm planning a trip to {{ city }}. Can you recommend some hotels?
+    """
+
+
+llm = LiteLLM(
+    model_name="gpt-4o-2024-08-06",
+    use_structured_output=True,
+)
+
+agent = Agent(llm=llm, prompt=HotelPrompt, tools=[get_hotel_recommendations])
+
+agent_card = agent.get_agent_card(
+    name="Hotel Recommendation Agent", description="Recommends hotels for a given city and travel dates.", port="8001"
+)
+
+app = create_agent_app(agent, agent_card, HotelPromptInput)
+server = create_agent_server(agent, agent_card, HotelPromptInput)