Skip to content

Commit 336d882

Browse files
yaozheng-fangyaozheng-fang
authored
feat: support local a2a hub server and client (#438)
1 parent b2af2d0 commit 336d882

File tree

4 files changed

+271
-0
lines changed

4 files changed

+271
-0
lines changed

‎veadk/a2a/hub/__init__.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.

‎veadk/a2a/hub/a2a_hub_client.py‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import requests
16+
from a2a.types import AgentCard
17+
18+
19+
class A2AHubClient:
20+
def __init__(self, server_host: str, server_port: int):
21+
self.server_host = server_host
22+
self.server_port = server_port
23+
self.health_check()
24+
25+
def health_check(self) -> None:
26+
"""Check the health of the server."""
27+
response = requests.get(f"http://{self.server_host}:{self.server_port}/ping")
28+
assert response.status_code == 200, (
29+
f"unexpected status code from A2A hub server: {response.status_code}"
30+
)
31+
32+
def get_agent_cards(
33+
self, group_id: str, target_agents: list[str] = []
34+
) -> list[dict]:
35+
"""Get the agent cards of the agents in the group."""
36+
ret = []
37+
38+
response = requests.get(
39+
f"http://{self.server_host}:{self.server_port}/group/{group_id}/agents"
40+
).json()
41+
agent_infos = response["agent_infos"]
42+
for agent_info in agent_infos:
43+
agent_id = agent_info["agent_id"]
44+
if target_agents:
45+
if agent_id in target_agents:
46+
ret.append(agent_info)
47+
else:
48+
ret.append(agent_info)
49+
50+
return ret
51+
52+
def register_agent(self, group_id: str, agent_id: str, agent_card: AgentCard):
53+
response = requests.post(
54+
f"http://{self.server_host}:{self.server_port}/register_agent",
55+
json={
56+
"group_id": group_id,
57+
"agent_id": agent_id,
58+
"agent_card": agent_card.model_dump(),
59+
},
60+
)
61+
62+
assert response.status_code == 200, (
63+
f"unexpected status code from A2A hub server: {response.status_code}"
64+
)
65+
66+
def create_group(self, group_id: str):
67+
response = requests.post(
68+
f"http://{self.server_host}:{self.server_port}/create_group",
69+
params={
70+
"group_id": group_id,
71+
},
72+
)
73+
74+
assert response.status_code == 200, (
75+
f"unexpected status code from A2A hub server: {response.status_code}"
76+
)

‎veadk/a2a/hub/a2a_hub_server.py‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import uvicorn
16+
from fastapi import FastAPI
17+
from fastapi.responses import JSONResponse
18+
19+
from veadk.a2a.hub.models import (
20+
AgentInformation,
21+
GetAgentResponse,
22+
GetAgentsResponse,
23+
GetGroupsResponse,
24+
RegisterAgentRequest,
25+
RegisterAgentResponse,
26+
RegisterGroupResponse,
27+
)
28+
29+
30+
class A2AHubServer:
31+
def __init__(self):
32+
self.app = FastAPI()
33+
34+
self.groups: list[str] = []
35+
36+
# group_id -> agent_id -> agent_card
37+
self.agent_cards: dict[str, dict[str, dict]] = {}
38+
39+
@self.app.get("/ping")
40+
def ping() -> JSONResponse:
41+
return JSONResponse(content={"msg": "pong!"})
42+
43+
@self.app.post("/create_group")
44+
def create_group(group_id: str) -> RegisterGroupResponse:
45+
"""Create a group."""
46+
self.groups.append(group_id)
47+
self.agent_cards[group_id] = {}
48+
return RegisterGroupResponse(group_id=group_id)
49+
50+
@self.app.post("/register_agent")
51+
def register_agent(
52+
request: RegisterAgentRequest,
53+
) -> RegisterAgentResponse:
54+
"""Register an agent to a specified group."""
55+
if request.group_id not in self.groups:
56+
return RegisterAgentResponse(
57+
err_code=1, msg=f"group {request.group_id} not exist"
58+
)
59+
if request.agent_id in self.agent_cards[request.group_id]:
60+
return RegisterAgentResponse(
61+
err_code=1, msg=f"agent {request.agent_id} already exist"
62+
)
63+
64+
self.agent_cards[request.group_id][request.agent_id] = request.agent_card
65+
return RegisterAgentResponse(
66+
group_id=request.group_id,
67+
agent_id=request.agent_id,
68+
agent_card=self.agent_cards[request.group_id][request.agent_id],
69+
)
70+
71+
@self.app.get("/group/{group_id}/agents")
72+
def agents(group_id: str) -> GetAgentsResponse:
73+
"""Get all agents in a specified group."""
74+
if group_id not in self.groups:
75+
return GetAgentsResponse(err_code=1, msg=f"group {group_id} not exist")
76+
77+
agent_infos = [
78+
AgentInformation(agent_id=agent_id, agent_card=agent_card)
79+
for agent_id, agent_card in self.agent_cards[group_id].items()
80+
]
81+
return GetAgentsResponse(group_id=group_id, agent_infos=agent_infos)
82+
83+
@self.app.get("/group/{group_id}/agent/{agent_id}")
84+
def agent(group_id: str, agent_id: str) -> GetAgentResponse:
85+
"""Get the agent card of a specified agent in a specified group."""
86+
if group_id not in self.groups:
87+
return GetAgentResponse(err_code=1, msg=f"group {group_id} not exist")
88+
if agent_id not in self.agent_cards[group_id]:
89+
return GetAgentResponse(
90+
err_code=1,
91+
msg=f"agent {agent_id} in group {group_id} not exist",
92+
)
93+
return GetAgentResponse(
94+
agent_id=agent_id,
95+
agent_card=self.agent_cards[group_id][agent_id],
96+
)
97+
98+
@self.app.get("/groups")
99+
def groups() -> GetGroupsResponse:
100+
"""Get all registered groups."""
101+
return GetGroupsResponse(group_ids=self.groups)
102+
103+
def serve(self, **kwargs):
104+
uvicorn.run(self.app, **kwargs)

‎veadk/a2a/hub/models.py‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from pydantic import BaseModel, Field
16+
17+
18+
class BaseResponse(BaseModel):
19+
err_code: int = 0
20+
21+
msg: str = ""
22+
"""The message of the response."""
23+
24+
25+
class RegisterGroupResponse(BaseResponse):
26+
group_id: str = ""
27+
"""The id of the group."""
28+
29+
30+
class RegisterAgentRequest(BaseResponse):
31+
group_id: str
32+
"""Target group id."""
33+
34+
agent_id: str
35+
"""The id of the agent."""
36+
37+
agent_card: dict
38+
"""The agent card of the agent in json format."""
39+
40+
41+
class RegisterAgentResponse(BaseResponse):
42+
group_id: str = ""
43+
"""Target group id."""
44+
45+
agent_id: str = ""
46+
"""The id of the agent."""
47+
48+
agent_card: dict = Field(default_factory=dict)
49+
"""The agent card of the agent in json format."""
50+
51+
52+
class AgentInformation(BaseModel):
53+
agent_id: str = ""
54+
"""The id of the agent."""
55+
56+
agent_card: dict = Field(default_factory=dict)
57+
"""The agent card of the agent in json format."""
58+
59+
60+
class GetAgentsResponse(BaseResponse):
61+
group_id: str = ""
62+
"""Target group id."""
63+
64+
agent_infos: list[AgentInformation] = Field(default_factory=list)
65+
"""The agent cards of the agents in json format."""
66+
67+
68+
class GetAgentResponse(BaseResponse):
69+
agent_id: str = ""
70+
"""The id of the agent."""
71+
72+
agent_card: dict = Field(default_factory=dict)
73+
"""The agent card of the agent in json format."""
74+
75+
76+
class GetGroupsResponse(BaseResponse):
77+
group_ids: list[str] = Field(default_factory=list)
78+
"""The ids of the groups."""

0 commit comments

Comments
 (0)