fix: a2a example (#709)

Co-authored-by: ds-sebastianchwilczynski <[email protected]>

Author: mhordynski

.github/scripts/run_examples.sh

@@ -13,6 +13,11 @@ export LOGFIRE_TOKEN
 GIT_URL_TEMPLATE="git+https://github.com/deepsense-ai/ragbits.git@%s#subdirectory=packages/%s"
 PR_BRANCH="${PR_BRANCH:-main}"
 
+has_script_section() {
+  local file="$1"
+  grep -q "^# /// script" "$file"
+}
+
 patch_dependencies() {
   local file="$1"
   local branch="$2"
@@ -51,8 +56,14 @@ patch_dependencies() {
 }
 
 export -f patch_dependencies
+export -f has_script_section
 
 find examples -name '*.py' | while read -r file; do
+  if ! has_script_section "$file"; then
+    echo "Skipping $file (no script section found)"
+    continue
+  fi
+
   echo "Running the script: $file"
   patch_dependencies "$file" "$PR_BRANCH"
 

examples/agents/a2a/agent_orchestrator.py

@@ -195,10 +195,6 @@ def _list_remote_agents(self) -> list[dict]:
                 "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/hotel_agent.py

@@ -51,7 +51,7 @@ def get_hotel_recommendations(city: str) -> str:
 class HotelPromptInput(BaseModel):
     """Defines the structured input for the hotel recommendation prompt."""
 
-    city: str
+    input: str
 
 
 class HotelPrompt(Prompt[HotelPromptInput]):
@@ -62,7 +62,7 @@ class HotelPrompt(Prompt[HotelPromptInput]):
     """
 
     user_prompt = """
-    I'm planning a trip to {{ city }}. Can you recommend some hotels?
+    {{ input }}
     """
 
 
@@ -71,4 +71,4 @@ class HotelPrompt(Prompt[HotelPromptInput]):
     use_structured_output=True,
 )
 
-agent = Agent(llm=llm, prompt=HotelPrompt, tools=[get_hotel_recommendations])
+hotel_agent = Agent(llm=llm, prompt=HotelPrompt, tools=[get_hotel_recommendations])

examples/agents/a2a/run_orchestrator.py

@@ -9,7 +9,7 @@
 
 To execute this script simply run:
 ```bash
-uv run examples/agents/a2a/orchestrator_client.py
+uv run examples/agents/a2a/run_orchestrator.py
 ```
 """
 
@@ -24,10 +24,8 @@
 import threading
 
 from agent_orchestrator import AgentOrchestrator, ResultsSumarizationPromptInput, RoutingPromptInput
-from flight_agent import FlightPromptInput
-from flight_agent import agent as flight_agent
-from hotel_agent import HotelPromptInput
-from hotel_agent import agent as hotel_agent
+from flight_agent import FlightPromptInput, flight_agent
+from hotel_agent import HotelPromptInput, hotel_agent
 from uvicorn import Server
 
 from ragbits.agents.a2a.server import create_agent_server
@@ -49,7 +47,7 @@ class RoutingPrompt(Prompt[RoutingPromptInput]):
     Each task must include:
     - agent_url: The base URL of the agent that should handle the task.
     - tool: The name of the skill/tool to call on that agent.
-    - parameters: A dictionary of arguments needed for that tool.
+    - parameters: A dictionary with input argument which is just the task for the agent to answer.
 
     Currently available agents and their tools:
     {{ agents }}
@@ -66,12 +64,12 @@ class RoutingPrompt(Prompt[RoutingPromptInput]):
     {
         "agent_url": "http://weather-agent:8000",
         "tool": "get_weather",
-        "parameters": { "location": "Paris" }
+        "parameters": { "input": "What is the weather in Paris?" }
     },
     {
         "agent_url": "http://news-agent:8001",
         "tool": "get_news_summary",
-        "parameters": { "region": "France" }
+        "parameters": { "input": "What are the news in France?" }
     }
     ]
     """

examples/agents/mcp/sse.py

@@ -57,6 +57,7 @@ async def main() -> None:
         params={
             "url": "http://localhost:8000/sse",
         },
+        client_session_timeout_seconds=20,
     ) as server:
         llm = LiteLLM(model_name="gpt-4o-2024-08-06")
         agent = Agent(llm=llm, mcp_servers=[server])

examples/agents/mcp/streamable_http.py

@@ -57,6 +57,7 @@ async def main() -> None:
         params={
             "url": "http://localhost:8000/mcp",
         },
+        client_session_timeout_seconds=20,
     ) as server:
         llm = LiteLLM(model_name="gpt-4o-2024-08-06")
         agent = Agent(llm=llm, mcp_servers=[server])