|
19 | 19 | from a2a.server.apps import A2AStarletteApplication |
20 | 20 | from a2a.server.request_handlers import DefaultRequestHandler |
21 | 21 | from a2a.server.tasks import InMemoryTaskStore |
22 | | -from a2a.types import AgentCapabilities, AgentCard, AgentSkill |
23 | | -from a2ui.extension.a2ui_extension import get_a2ui_agent_extension |
24 | 22 | from agent import RestaurantAgent |
25 | 23 | from agent_executor import RestaurantAgentExecutor |
26 | 24 | from dotenv import load_dotenv |
@@ -49,39 +47,19 @@ def main(host, port): |
49 | 47 | "GEMINI_API_KEY environment variable not set and GOOGLE_GENAI_USE_VERTEXAI is not TRUE." |
50 | 48 | ) |
51 | 49 |
|
52 | | - capabilities = AgentCapabilities( |
53 | | - streaming=True, |
54 | | - extensions=[get_a2ui_agent_extension()], |
55 | | - ) |
56 | | - skill = AgentSkill( |
57 | | - id="find_restaurants", |
58 | | - name="Find Restaurants Tool", |
59 | | - description="Helps find restaurants based on user criteria (e.g., cuisine, location).", |
60 | | - tags=["restaurant", "finder"], |
61 | | - examples=["Find me the top 10 chinese restaurants in the US"], |
62 | | - ) |
63 | | - |
64 | 50 | base_url = f"http://{host}:{port}" |
65 | 51 |
|
66 | | - agent_card = AgentCard( |
67 | | - name="Restaurant Agent", |
68 | | - description="This agent helps find restaurants based on user criteria.", |
69 | | - url=base_url, # <-- Use base_url here |
70 | | - version="1.0.0", |
71 | | - default_input_modes=RestaurantAgent.SUPPORTED_CONTENT_TYPES, |
72 | | - default_output_modes=RestaurantAgent.SUPPORTED_CONTENT_TYPES, |
73 | | - capabilities=capabilities, |
74 | | - skills=[skill], |
75 | | - ) |
| 52 | + ui_agent = RestaurantAgent(base_url=base_url, use_ui=True) |
| 53 | + text_agent = RestaurantAgent(base_url=base_url, use_ui=False) |
76 | 54 |
|
77 | | - agent_executor = RestaurantAgentExecutor(base_url=base_url) |
| 55 | + agent_executor = RestaurantAgentExecutor(ui_agent, text_agent) |
78 | 56 |
|
79 | 57 | request_handler = DefaultRequestHandler( |
80 | 58 | agent_executor=agent_executor, |
81 | 59 | task_store=InMemoryTaskStore(), |
82 | 60 | ) |
83 | 61 | server = A2AStarletteApplication( |
84 | | - agent_card=agent_card, http_handler=request_handler |
| 62 | + agent_card=ui_agent.get_agent_card(), http_handler=request_handler |
85 | 63 | ) |
86 | 64 | import uvicorn |
87 | 65 |
|
|
0 commit comments