distributed-a2a 0.1.26__tar.gz → 0.2.2__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {distributed_a2a-0.1.26/distributed_a2a.egg-info → distributed_a2a-0.2.2}/PKG-INFO +1 -1
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/__init__.py +3 -1
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/agent.py +21 -3
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/client.py +67 -10
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/executors.py +24 -7
- distributed_a2a-0.2.2/distributed_a2a/files.py +68 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2/distributed_a2a.egg-info}/PKG-INFO +1 -1
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a.egg-info/SOURCES.txt +3 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/pyproject.toml +1 -1
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/tests/test_app.py +7 -4
- distributed_a2a-0.2.2/tests/test_executor_files.py +216 -0
- distributed_a2a-0.2.2/tests/test_files.py +317 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/tests/test_rejection.py +11 -10
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/tests/test_timeout.py +13 -7
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/LICENSE +0 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/MANIFEST.in +0 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/README.md +0 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/config.py +0 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/model.py +0 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/py.typed +0 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/registry.py +0 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/registry_server/__init__.py +0 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/registry_server/bootstrap.py +0 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/registry_server/dynamo_db.py +0 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/registry_server/in_memory_registry_storage.py +0 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/registry_server/model.py +0 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/registry_server/storage.py +0 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/router.py +0 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/schemas/agent-schema.json +0 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/schemas/router-agent-schema.json +0 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/server.py +0 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a.egg-info/dependency_links.txt +0 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a.egg-info/requires.txt +0 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a.egg-info/top_level.txt +0 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/requirements.txt +0 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/setup.cfg +0 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/setup.py +0 -0
- {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/tests/test_client.py +0 -0
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
from .client import A2ATimeoutError, RoutingA2AClient
|
|
1
|
+
from .client import A2ATimeoutError, AgentReply, FileRef, RoutingA2AClient
|
|
2
2
|
from .model import (AgentConfig, AgentItem, CardConfig, LLMConfig,
|
|
3
3
|
RegistryConfig, RegistryItemConfig, RouterConfig,
|
|
4
4
|
RouterItem, SkillConfig)
|
|
@@ -15,6 +15,8 @@ __all__ = [
|
|
|
15
15
|
"load_router",
|
|
16
16
|
"RoutingA2AClient",
|
|
17
17
|
"A2ATimeoutError",
|
|
18
|
+
"AgentReply",
|
|
19
|
+
"FileRef",
|
|
18
20
|
"load_registry",
|
|
19
21
|
"AgentConfig",
|
|
20
22
|
"SkillConfig",
|
|
@@ -3,7 +3,7 @@ from typing import Any, Literal, cast
|
|
|
3
3
|
|
|
4
4
|
from a2a.types import TaskState
|
|
5
5
|
from langchain.agents import create_agent
|
|
6
|
-
from langchain_core.messages import HumanMessage
|
|
6
|
+
from langchain_core.messages import BaseMessage, HumanMessage
|
|
7
7
|
from langchain_core.runnables import RunnableConfig
|
|
8
8
|
from langchain_core.tools import BaseTool
|
|
9
9
|
from langgraph.checkpoint.base import BaseCheckpointSaver
|
|
@@ -38,6 +38,21 @@ class StringResponse(AgentResponse):
|
|
|
38
38
|
response: str = Field(description="The main response to be returned to the user")
|
|
39
39
|
|
|
40
40
|
|
|
41
|
+
class AgentInvocation[ResponseT: AgentResponse](BaseModel):
|
|
42
|
+
"""Result of invoking a :class:`StatusAgent`.
|
|
43
|
+
|
|
44
|
+
Exposes both the structured response (LLM-visible content the agent decided
|
|
45
|
+
to return) and the raw message list produced by the underlying LangGraph
|
|
46
|
+
agent, so callers can mine ``ToolMessage.artifact`` for out-of-band binary
|
|
47
|
+
payloads (e.g. files returned by MCP tools) without those bytes ever passing
|
|
48
|
+
through the LLM context.
|
|
49
|
+
"""
|
|
50
|
+
|
|
51
|
+
structured: ResponseT
|
|
52
|
+
messages: list[BaseMessage]
|
|
53
|
+
model_config = {"arbitrary_types_allowed": True}
|
|
54
|
+
|
|
55
|
+
|
|
41
56
|
class StatusAgent[ResponseT: AgentResponse]:
|
|
42
57
|
|
|
43
58
|
def __init__(self,
|
|
@@ -77,7 +92,7 @@ class StatusAgent[ResponseT: AgentResponse]:
|
|
|
77
92
|
|
|
78
93
|
async def __call__(self,
|
|
79
94
|
message: str,
|
|
80
|
-
context_id: str | None = None) -> ResponseT:
|
|
95
|
+
context_id: str | None = None) -> AgentInvocation[ResponseT]:
|
|
81
96
|
config: RunnableConfig = RunnableConfig(
|
|
82
97
|
configurable={'thread_id': context_id}
|
|
83
98
|
)
|
|
@@ -86,4 +101,7 @@ class StatusAgent[ResponseT: AgentResponse]:
|
|
|
86
101
|
config
|
|
87
102
|
)
|
|
88
103
|
logging.info("agent response: %s", response)
|
|
89
|
-
return
|
|
104
|
+
return AgentInvocation[ResponseT](
|
|
105
|
+
structured=cast(ResponseT, response['structured_response']),
|
|
106
|
+
messages=list(response.get('messages', [])),
|
|
107
|
+
)
|
|
@@ -1,18 +1,48 @@
|
|
|
1
1
|
import asyncio
|
|
2
2
|
import json
|
|
3
3
|
import time
|
|
4
|
+
from dataclasses import dataclass, field
|
|
4
5
|
from uuid import uuid4
|
|
5
6
|
|
|
6
7
|
import httpx
|
|
7
8
|
from a2a.client import (A2ACardResolver, ClientConfig, ClientEvent,
|
|
8
9
|
ClientFactory, create_text_message_object)
|
|
9
|
-
from a2a.types import (AgentCard,
|
|
10
|
-
TaskState,
|
|
10
|
+
from a2a.types import (AgentCard, FilePart, FileWithBytes, FileWithUri,
|
|
11
|
+
Message, Part, Task, TaskQueryParams, TaskState,
|
|
12
|
+
TextPart)
|
|
11
13
|
|
|
12
14
|
DEFAULT_MAX_POLLS = 50
|
|
13
15
|
DEFAULT_POLL_INTERVAL = 1.0
|
|
14
16
|
|
|
15
17
|
|
|
18
|
+
@dataclass
|
|
19
|
+
class FileRef:
|
|
20
|
+
"""A file payload received as part of an A2A agent reply.
|
|
21
|
+
|
|
22
|
+
Exactly one of ``bytes_b64`` (for ``FileWithBytes``) or ``uri`` (for
|
|
23
|
+
``FileWithUri``) is populated. ``bytes_b64`` is the raw base64 string
|
|
24
|
+
delivered over the wire by the A2A SDK — the caller is responsible for
|
|
25
|
+
decoding before forwarding the bytes.
|
|
26
|
+
"""
|
|
27
|
+
|
|
28
|
+
name: str
|
|
29
|
+
mime_type: str
|
|
30
|
+
bytes_b64: str = ""
|
|
31
|
+
uri: str | None = None
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
@dataclass
|
|
35
|
+
class AgentReply:
|
|
36
|
+
"""Structured reply from a routing-aware agent.
|
|
37
|
+
|
|
38
|
+
Carries the user-visible text (if any) plus zero-or-more files that the
|
|
39
|
+
agent emitted out-of-band as ``FilePart`` artifacts.
|
|
40
|
+
"""
|
|
41
|
+
|
|
42
|
+
text: str | None = None
|
|
43
|
+
files: list[FileRef] = field(default_factory=list)
|
|
44
|
+
|
|
45
|
+
|
|
16
46
|
class A2ATimeoutError(Exception):
|
|
17
47
|
"""Raised when polling a remote agent task exceeds max_polls."""
|
|
18
48
|
|
|
@@ -69,8 +99,10 @@ class RemoteAgentConnection:
|
|
|
69
99
|
response: Task = await self.agent_client.get_task(query_params)
|
|
70
100
|
return response
|
|
71
101
|
|
|
72
|
-
async def send_message(self,
|
|
73
|
-
|
|
102
|
+
async def send_message(self,
|
|
103
|
+
message_to_send: str,
|
|
104
|
+
context_id: str,
|
|
105
|
+
task_id: None | str = None) -> AgentReply | AgentCard | TaskState:
|
|
74
106
|
message: Message = create_text_message_object(content=message_to_send)
|
|
75
107
|
message.message_id = str(uuid4())
|
|
76
108
|
message.context_id = context_id
|
|
@@ -112,13 +144,36 @@ class RemoteAgentConnection:
|
|
|
112
144
|
for artifact in response.artifacts or []:
|
|
113
145
|
match artifact.name, artifact.parts:
|
|
114
146
|
case 'routing_error', [Part(root=TextPart(text=error_msg)), *_]:
|
|
115
|
-
return error_msg
|
|
147
|
+
return AgentReply(text=error_msg)
|
|
116
148
|
case 'rejected', [Part(root=TextPart()), *_]:
|
|
117
149
|
return TaskState.rejected
|
|
118
150
|
case 'target_agent', [Part(root=TextPart(text=agent_card_str)), *_]:
|
|
119
151
|
return AgentCard(**json.loads(agent_card_str))
|
|
120
|
-
|
|
121
|
-
|
|
152
|
+
|
|
153
|
+
text_out: str | None = None
|
|
154
|
+
files_out: list[FileRef] = []
|
|
155
|
+
for artifact in response.artifacts or []:
|
|
156
|
+
for part in artifact.parts or []:
|
|
157
|
+
root = getattr(part, "root", None)
|
|
158
|
+
if isinstance(root, TextPart) and artifact.name == "current_result":
|
|
159
|
+
text_out = root.text
|
|
160
|
+
elif isinstance(root, FilePart):
|
|
161
|
+
root_file = root.file
|
|
162
|
+
if isinstance(root_file, FileWithBytes):
|
|
163
|
+
files_out.append(FileRef(
|
|
164
|
+
name=root_file.name or artifact.name or "file.bin",
|
|
165
|
+
mime_type=root_file.mime_type or "application/octet-stream",
|
|
166
|
+
bytes_b64=root_file.bytes,
|
|
167
|
+
))
|
|
168
|
+
elif isinstance(root_file, FileWithUri):
|
|
169
|
+
files_out.append(FileRef(
|
|
170
|
+
name=root_file.name or artifact.name or "file.bin",
|
|
171
|
+
mime_type=root_file.mime_type or "application/octet-stream",
|
|
172
|
+
uri=root_file.uri,
|
|
173
|
+
))
|
|
174
|
+
|
|
175
|
+
if text_out is not None or files_out:
|
|
176
|
+
return AgentReply(text=text_out, files=files_out)
|
|
122
177
|
|
|
123
178
|
if task_state == TaskState.rejected:
|
|
124
179
|
return TaskState.rejected
|
|
@@ -157,8 +212,10 @@ class RoutingA2AClient:
|
|
|
157
212
|
await card_resolver.get_agent_card()
|
|
158
213
|
)
|
|
159
214
|
|
|
160
|
-
async def send_message(self, message: str,
|
|
161
|
-
|
|
215
|
+
async def send_message(self, message: str,
|
|
216
|
+
context_id: str,
|
|
217
|
+
depth: int = 0,
|
|
218
|
+
rejected_agents: list[str] | None = None) -> AgentReply:
|
|
162
219
|
if depth > MAX_RECURSION_DEPTH:
|
|
163
220
|
raise Exception("Maximum recursion depth exceeded. This is likely due to an infinite loop in your agent.")
|
|
164
221
|
|
|
@@ -182,7 +239,7 @@ class RoutingA2AClient:
|
|
|
182
239
|
if rejection_msg not in message:
|
|
183
240
|
message_to_send = f"{message}\n\n{rejection_msg}"
|
|
184
241
|
|
|
185
|
-
agent_response:
|
|
242
|
+
agent_response: AgentReply | AgentCard | TaskState = await agent_connection.send_message(message_to_send, context_id)
|
|
186
243
|
if isinstance(agent_response, AgentCard):
|
|
187
244
|
if agent_response.url == self.current_card.url:
|
|
188
245
|
raise Exception("Agent redirected to itself.")
|
|
@@ -4,8 +4,8 @@ from typing import Any
|
|
|
4
4
|
|
|
5
5
|
from a2a.server.agent_execution import AgentExecutor, RequestContext
|
|
6
6
|
from a2a.server.events import EventQueue
|
|
7
|
-
from a2a.types import (Artifact,
|
|
8
|
-
TaskStatus, TaskStatusUpdateEvent)
|
|
7
|
+
from a2a.types import (Artifact, FilePart, Part, TaskArtifactUpdateEvent,
|
|
8
|
+
TaskState, TaskStatus, TaskStatusUpdateEvent)
|
|
9
9
|
from a2a.utils import new_text_artifact
|
|
10
10
|
from langchain_core.tools import BaseTool
|
|
11
11
|
from langchain_mcp_adapters.client import MultiServerMCPClient
|
|
@@ -13,6 +13,7 @@ from langgraph.checkpoint.base import BaseCheckpointSaver
|
|
|
13
13
|
|
|
14
14
|
from .agent import RoutingResponse, StatusAgent, StringResponse
|
|
15
15
|
from .config import settings
|
|
16
|
+
from .files import extract_file_parts
|
|
16
17
|
from .model import AgentConfig, RouterConfig
|
|
17
18
|
from .registry import AgentRegistryLookupClient, McpRegistryLookup
|
|
18
19
|
|
|
@@ -116,21 +117,36 @@ class RoutingAgentExecutor(AgentExecutor):
|
|
|
116
117
|
task_id=context.task_id
|
|
117
118
|
))
|
|
118
119
|
await self.reinitialize_agent_with_tools()
|
|
119
|
-
|
|
120
|
-
|
|
120
|
+
invocation = await self.agent(message=context.get_user_input(),
|
|
121
|
+
context_id=context.context_id)
|
|
122
|
+
agent_response: StringResponse = invocation.structured
|
|
121
123
|
|
|
122
124
|
artifact: Artifact
|
|
125
|
+
file_parts: list[tuple[str, FilePart]] = []
|
|
123
126
|
if agent_response.status == TaskState.rejected:
|
|
124
127
|
artifact = await _route_request_to_matching_agent(self.routing_agent, self.agent_registry, context)
|
|
125
128
|
else:
|
|
126
129
|
logger.info(f"Request with id {context.context_id} was successfully processed by agent.")
|
|
130
|
+
file_parts = extract_file_parts(invocation.messages)
|
|
127
131
|
artifact = new_text_artifact(
|
|
128
132
|
name='current_result',
|
|
129
133
|
description='Result of request to agent.',
|
|
130
134
|
text=f"*{self.agent_config.agent.card.name}*: {agent_response.response}"
|
|
131
135
|
)
|
|
132
136
|
|
|
133
|
-
|
|
137
|
+
for name, file_part in file_parts:
|
|
138
|
+
await event_queue.enqueue_event(TaskArtifactUpdateEvent(
|
|
139
|
+
append=False,
|
|
140
|
+
last_chunk=False,
|
|
141
|
+
context_id=context.context_id,
|
|
142
|
+
task_id=context.task_id,
|
|
143
|
+
artifact=Artifact(
|
|
144
|
+
artifact_id=name,
|
|
145
|
+
name=name,
|
|
146
|
+
parts=[Part(root=file_part)],
|
|
147
|
+
),
|
|
148
|
+
))
|
|
149
|
+
|
|
134
150
|
await event_queue.enqueue_event(TaskArtifactUpdateEvent(
|
|
135
151
|
append=False,
|
|
136
152
|
context_id=context.context_id,
|
|
@@ -297,8 +313,9 @@ class RoutingExecutor(AgentExecutor):
|
|
|
297
313
|
async def _route_request_to_matching_agent(routing_agent: StatusAgent[RoutingResponse],
|
|
298
314
|
agent_registry: AgentRegistryLookupClient,
|
|
299
315
|
context: RequestContext) -> Artifact:
|
|
300
|
-
|
|
301
|
-
|
|
316
|
+
invocation = await routing_agent(message=context.get_user_input(),
|
|
317
|
+
context_id=context.context_id)
|
|
318
|
+
routing_agent_response: RoutingResponse = invocation.structured
|
|
302
319
|
agent_name: str | None = routing_agent_response.agent_name
|
|
303
320
|
logger.info(f"routing response received: {routing_agent_response}")
|
|
304
321
|
if agent_name is None:
|
|
@@ -0,0 +1,68 @@
|
|
|
1
|
+
import json
|
|
2
|
+
import mimetypes
|
|
3
|
+
from typing import Any
|
|
4
|
+
|
|
5
|
+
from a2a.types import FilePart, FileWithBytes
|
|
6
|
+
from langchain_core.messages import BaseMessage, ToolMessage
|
|
7
|
+
|
|
8
|
+
_LANGCHAIN_BINARY_BLOCK_TYPES: dict[str, str] = {
|
|
9
|
+
"file": "attachment",
|
|
10
|
+
"image": "image",
|
|
11
|
+
}
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
def _filename_from_text_block(block: dict[str, Any]) -> str | None:
|
|
15
|
+
text = block.get("text")
|
|
16
|
+
if not isinstance(text, str):
|
|
17
|
+
return None
|
|
18
|
+
try:
|
|
19
|
+
payload = json.loads(text)
|
|
20
|
+
except (ValueError, TypeError):
|
|
21
|
+
return None
|
|
22
|
+
if isinstance(payload, dict):
|
|
23
|
+
name = payload.get("filename")
|
|
24
|
+
if isinstance(name, str) and name:
|
|
25
|
+
return name
|
|
26
|
+
return None
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def extract_file_parts(messages: list[BaseMessage]) -> list[tuple[str, FilePart]]:
|
|
30
|
+
parts: list[tuple[str, FilePart]] = []
|
|
31
|
+
for message in messages:
|
|
32
|
+
if not (isinstance(message, ToolMessage)
|
|
33
|
+
and isinstance(message.content, list)):
|
|
34
|
+
continue
|
|
35
|
+
pending_name: str | None = None
|
|
36
|
+
counters: dict[str, int] = {"file": 0, "image": 0}
|
|
37
|
+
for block in message.content:
|
|
38
|
+
if not isinstance(block, dict):
|
|
39
|
+
continue
|
|
40
|
+
block_type = block.get("type")
|
|
41
|
+
if not isinstance(block_type, str):
|
|
42
|
+
continue
|
|
43
|
+
if block_type == "text":
|
|
44
|
+
hint = _filename_from_text_block(block)
|
|
45
|
+
if hint:
|
|
46
|
+
pending_name = hint
|
|
47
|
+
continue
|
|
48
|
+
kind = _LANGCHAIN_BINARY_BLOCK_TYPES.get(block_type)
|
|
49
|
+
if kind is None:
|
|
50
|
+
continue
|
|
51
|
+
b64 = block.get("base64")
|
|
52
|
+
if not isinstance(b64, str) or not b64:
|
|
53
|
+
continue
|
|
54
|
+
mime_type = block.get("mime_type") or "application/octet-stream"
|
|
55
|
+
if pending_name is not None:
|
|
56
|
+
name = pending_name
|
|
57
|
+
pending_name = None
|
|
58
|
+
else:
|
|
59
|
+
index = counters[block_type]
|
|
60
|
+
counters[block_type] = index + 1
|
|
61
|
+
guessed_ext = mimetypes.guess_extension(mime_type)
|
|
62
|
+
ext = f"-{guessed_ext}" if guessed_ext is not None else ""
|
|
63
|
+
suffix = f"-{index}" if index > 0 else ""
|
|
64
|
+
name = f"{kind}{suffix}{ext}"
|
|
65
|
+
parts.append((name, FilePart(file=FileWithBytes(
|
|
66
|
+
name=name, mime_type=mime_type, bytes=b64,
|
|
67
|
+
))))
|
|
68
|
+
return parts
|
|
@@ -9,6 +9,7 @@ distributed_a2a/agent.py
|
|
|
9
9
|
distributed_a2a/client.py
|
|
10
10
|
distributed_a2a/config.py
|
|
11
11
|
distributed_a2a/executors.py
|
|
12
|
+
distributed_a2a/files.py
|
|
12
13
|
distributed_a2a/model.py
|
|
13
14
|
distributed_a2a/py.typed
|
|
14
15
|
distributed_a2a/registry.py
|
|
@@ -29,5 +30,7 @@ distributed_a2a/schemas/agent-schema.json
|
|
|
29
30
|
distributed_a2a/schemas/router-agent-schema.json
|
|
30
31
|
tests/test_app.py
|
|
31
32
|
tests/test_client.py
|
|
33
|
+
tests/test_executor_files.py
|
|
34
|
+
tests/test_files.py
|
|
32
35
|
tests/test_rejection.py
|
|
33
36
|
tests/test_timeout.py
|
|
@@ -59,10 +59,12 @@ async def test_app_completed_path(fake_registry_server: str, fake_completed_llm:
|
|
|
59
59
|
with FakeAgent(fake_registry_server, fake_completed_llm, "test-agent") as agent:
|
|
60
60
|
# When
|
|
61
61
|
client = RoutingA2AClient(initial_url=f"http://127.0.0.1:{agent.app_port}/{agent.name}")
|
|
62
|
-
|
|
62
|
+
reply = await client.send_message(message="Hello", context_id="test-context")
|
|
63
63
|
|
|
64
64
|
# Then: Check the response
|
|
65
|
-
assert
|
|
65
|
+
assert reply.text is not None
|
|
66
|
+
assert "This is a mock response from the fake OpenAI server." in reply.text
|
|
67
|
+
assert reply.files == []
|
|
66
68
|
|
|
67
69
|
|
|
68
70
|
@pytest.mark.asyncio
|
|
@@ -75,7 +77,8 @@ async def test_app_redirect_path(fake_registry_server: str, fake_completed_llm:
|
|
|
75
77
|
client = RoutingA2AClient(initial_url=f"http://127.0.0.1:{first_agent.app_port}")
|
|
76
78
|
|
|
77
79
|
# When
|
|
78
|
-
|
|
80
|
+
reply = await client.send_message(message="Hello", context_id="test-context")
|
|
79
81
|
|
|
80
82
|
# Then
|
|
81
|
-
assert
|
|
83
|
+
assert reply.text is not None
|
|
84
|
+
assert FINAL_RESPONSE in reply.text
|
|
@@ -0,0 +1,216 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import base64
|
|
4
|
+
from types import SimpleNamespace
|
|
5
|
+
from typing import Any
|
|
6
|
+
|
|
7
|
+
import pytest
|
|
8
|
+
from a2a.server.agent_execution import RequestContext
|
|
9
|
+
from a2a.server.events import EventQueue
|
|
10
|
+
from a2a.types import FilePart, FileWithBytes
|
|
11
|
+
from a2a.types import Message as A2AMessage
|
|
12
|
+
from a2a.types import (MessageSendParams, Part, Role, TaskArtifactUpdateEvent,
|
|
13
|
+
TaskState, TaskStatusUpdateEvent, TextPart)
|
|
14
|
+
from langchain_core.messages import BaseMessage, HumanMessage, ToolMessage
|
|
15
|
+
|
|
16
|
+
from distributed_a2a.agent import AgentInvocation, StringResponse
|
|
17
|
+
from distributed_a2a.executors import RoutingAgentExecutor
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
class _StubStatusAgent:
|
|
21
|
+
def __init__(self, response: StringResponse, messages: list[BaseMessage]) -> None:
|
|
22
|
+
self._response = response
|
|
23
|
+
self._messages = messages
|
|
24
|
+
|
|
25
|
+
async def __call__(self, message: str, context_id: str | None = None) -> AgentInvocation[StringResponse]:
|
|
26
|
+
return AgentInvocation[StringResponse](
|
|
27
|
+
structured=self._response,
|
|
28
|
+
messages=self._messages,
|
|
29
|
+
)
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def _make_request_context() -> RequestContext:
|
|
33
|
+
msg = A2AMessage(
|
|
34
|
+
message_id="m-1",
|
|
35
|
+
context_id="ctx-1",
|
|
36
|
+
task_id="task-1",
|
|
37
|
+
role=Role.user,
|
|
38
|
+
parts=[Part(root=TextPart(text="render a CV please"))],
|
|
39
|
+
)
|
|
40
|
+
return RequestContext(
|
|
41
|
+
request=MessageSendParams(message=msg),
|
|
42
|
+
task_id="task-1",
|
|
43
|
+
context_id="ctx-1",
|
|
44
|
+
)
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
async def _drain_queue(queue: EventQueue) -> list[Any]:
|
|
48
|
+
events: list[Any] = []
|
|
49
|
+
while True:
|
|
50
|
+
try:
|
|
51
|
+
evt = await queue.dequeue_event(no_wait=True)
|
|
52
|
+
except Exception:
|
|
53
|
+
break
|
|
54
|
+
events.append(evt)
|
|
55
|
+
queue.task_done()
|
|
56
|
+
return events
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
@pytest.mark.asyncio
|
|
60
|
+
async def test_executor_emits_no_file_event_when_no_artifacts() -> None:
|
|
61
|
+
executor = RoutingAgentExecutor.__new__(RoutingAgentExecutor)
|
|
62
|
+
executor.agent_config = SimpleNamespace( # type: ignore[assignment]
|
|
63
|
+
agent=SimpleNamespace(card=SimpleNamespace(name="plain-agent")),
|
|
64
|
+
)
|
|
65
|
+
executor.agent = _StubStatusAgent( # type: ignore[assignment]
|
|
66
|
+
StringResponse(status=TaskState.completed, response="No files here."),
|
|
67
|
+
[HumanMessage(content="hi")],
|
|
68
|
+
)
|
|
69
|
+
|
|
70
|
+
async def _noop_reinit() -> None:
|
|
71
|
+
return None
|
|
72
|
+
|
|
73
|
+
executor.reinitialize_agent_with_tools = _noop_reinit # type: ignore[method-assign]
|
|
74
|
+
|
|
75
|
+
ctx = _make_request_context()
|
|
76
|
+
queue = EventQueue()
|
|
77
|
+
await executor.execute(ctx, queue)
|
|
78
|
+
events = await _drain_queue(queue)
|
|
79
|
+
|
|
80
|
+
artifact_events = [e for e in events if isinstance(e, TaskArtifactUpdateEvent)]
|
|
81
|
+
assert len(artifact_events) == 1
|
|
82
|
+
assert artifact_events[0].artifact.name == "current_result"
|
|
83
|
+
assert artifact_events[0].last_chunk is True
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
@pytest.mark.asyncio
|
|
87
|
+
async def test_executor_emits_file_part_from_langchain_content_block() -> None:
|
|
88
|
+
docx_b64 = base64.b64encode(b"PK\x03\x04 fake docx").decode("ascii")
|
|
89
|
+
summary_json = (
|
|
90
|
+
'{"filename": "cv-bob.docx", '
|
|
91
|
+
'"mime_type": "application/vnd.openxmlformats-officedocument.'
|
|
92
|
+
'wordprocessingml.document"}'
|
|
93
|
+
)
|
|
94
|
+
tool_msg = ToolMessage(
|
|
95
|
+
content=[
|
|
96
|
+
{"type": "text", "text": summary_json, "id": "lc_text_1"},
|
|
97
|
+
{"type": "file",
|
|
98
|
+
"base64": docx_b64,
|
|
99
|
+
"mime_type": (
|
|
100
|
+
"application/vnd.openxmlformats-officedocument."
|
|
101
|
+
"wordprocessingml.document"
|
|
102
|
+
),
|
|
103
|
+
"id": "lc_file_1"},
|
|
104
|
+
],
|
|
105
|
+
tool_call_id="call-cv",
|
|
106
|
+
artifact={"structured_content": {"filename": "cv-bob.docx"}},
|
|
107
|
+
)
|
|
108
|
+
|
|
109
|
+
executor = RoutingAgentExecutor.__new__(RoutingAgentExecutor)
|
|
110
|
+
executor.agent_config = SimpleNamespace( # type: ignore[assignment]
|
|
111
|
+
agent=SimpleNamespace(card=SimpleNamespace(name="cv-agent")),
|
|
112
|
+
)
|
|
113
|
+
executor.agent = _StubStatusAgent( # type: ignore[assignment]
|
|
114
|
+
StringResponse(status=TaskState.completed,
|
|
115
|
+
response="Here is your CV."),
|
|
116
|
+
[HumanMessage(content="render a CV please"), tool_msg],
|
|
117
|
+
)
|
|
118
|
+
|
|
119
|
+
async def _noop_reinit() -> None:
|
|
120
|
+
return None
|
|
121
|
+
|
|
122
|
+
executor.reinitialize_agent_with_tools = _noop_reinit # type: ignore[method-assign]
|
|
123
|
+
|
|
124
|
+
ctx = _make_request_context()
|
|
125
|
+
queue = EventQueue()
|
|
126
|
+
await executor.execute(ctx, queue)
|
|
127
|
+
events = await _drain_queue(queue)
|
|
128
|
+
|
|
129
|
+
artifact_events = [e for e in events if isinstance(e, TaskArtifactUpdateEvent)]
|
|
130
|
+
assert len(artifact_events) == 2, (
|
|
131
|
+
"Expected the executor to emit a file artifact followed by the text "
|
|
132
|
+
"artifact when the tool response uses the production LangChain "
|
|
133
|
+
"content-block shape."
|
|
134
|
+
)
|
|
135
|
+
file_event, text_event = artifact_events
|
|
136
|
+
assert file_event.last_chunk is False
|
|
137
|
+
assert file_event.artifact.name == "cv-bob.docx"
|
|
138
|
+
file_part = file_event.artifact.parts[0].root
|
|
139
|
+
assert isinstance(file_part, FilePart)
|
|
140
|
+
assert isinstance(file_part.file, FileWithBytes)
|
|
141
|
+
assert file_part.file.name == "cv-bob.docx"
|
|
142
|
+
assert file_part.file.bytes == docx_b64
|
|
143
|
+
|
|
144
|
+
assert text_event.last_chunk is True
|
|
145
|
+
assert text_event.artifact.name == "current_result"
|
|
146
|
+
|
|
147
|
+
|
|
148
|
+
@pytest.mark.asyncio
|
|
149
|
+
async def test_executor_emits_one_file_event_per_file_block() -> None:
|
|
150
|
+
b64_a = base64.b64encode(b"aaa docx bytes").decode("ascii")
|
|
151
|
+
b64_b = base64.b64encode(b"bbb docx bytes").decode("ascii")
|
|
152
|
+
docx_mime = (
|
|
153
|
+
"application/vnd.openxmlformats-officedocument."
|
|
154
|
+
"wordprocessingml.document"
|
|
155
|
+
)
|
|
156
|
+
tool_msg = ToolMessage(
|
|
157
|
+
content=[
|
|
158
|
+
{"type": "text", "text": '{"filename": "cv-a.docx"}', "id": "t1"},
|
|
159
|
+
{"type": "file", "base64": b64_a, "mime_type": docx_mime, "id": "f1"},
|
|
160
|
+
{"type": "text", "text": '{"filename": "cv-b.docx"}', "id": "t2"},
|
|
161
|
+
{"type": "file", "base64": b64_b, "mime_type": docx_mime, "id": "f2"},
|
|
162
|
+
],
|
|
163
|
+
tool_call_id="call-multi",
|
|
164
|
+
)
|
|
165
|
+
|
|
166
|
+
executor = RoutingAgentExecutor.__new__(RoutingAgentExecutor)
|
|
167
|
+
executor.agent_config = SimpleNamespace( # type: ignore[assignment]
|
|
168
|
+
agent=SimpleNamespace(card=SimpleNamespace(name="cv-agent")),
|
|
169
|
+
)
|
|
170
|
+
executor.agent = _StubStatusAgent( # type: ignore[assignment]
|
|
171
|
+
StringResponse(status=TaskState.completed,
|
|
172
|
+
response="Here are your two CVs."),
|
|
173
|
+
[HumanMessage(content="render two CVs"), tool_msg],
|
|
174
|
+
)
|
|
175
|
+
|
|
176
|
+
async def _noop_reinit() -> None:
|
|
177
|
+
return None
|
|
178
|
+
|
|
179
|
+
executor.reinitialize_agent_with_tools = _noop_reinit # type: ignore[method-assign]
|
|
180
|
+
|
|
181
|
+
ctx = _make_request_context()
|
|
182
|
+
queue = EventQueue()
|
|
183
|
+
await executor.execute(ctx, queue)
|
|
184
|
+
events = await _drain_queue(queue)
|
|
185
|
+
|
|
186
|
+
artifact_events = [e for e in events if isinstance(e, TaskArtifactUpdateEvent)]
|
|
187
|
+
assert len(artifact_events) == 3, (
|
|
188
|
+
"Expected 2 file artifacts (last_chunk=False) + 1 text artifact "
|
|
189
|
+
f"(last_chunk=True). Got: "
|
|
190
|
+
f"{[(e.artifact.name, e.last_chunk) for e in artifact_events]}"
|
|
191
|
+
)
|
|
192
|
+
file_a_event, file_b_event, text_event = artifact_events
|
|
193
|
+
|
|
194
|
+
assert file_a_event.last_chunk is False
|
|
195
|
+
assert file_a_event.artifact.name == "cv-a.docx"
|
|
196
|
+
file_a_part = file_a_event.artifact.parts[0].root
|
|
197
|
+
assert isinstance(file_a_part, FilePart)
|
|
198
|
+
assert isinstance(file_a_part.file, FileWithBytes)
|
|
199
|
+
assert file_a_part.file.name == "cv-a.docx"
|
|
200
|
+
assert file_a_part.file.bytes == b64_a
|
|
201
|
+
|
|
202
|
+
assert file_b_event.last_chunk is False
|
|
203
|
+
assert file_b_event.artifact.name == "cv-b.docx"
|
|
204
|
+
file_b_part = file_b_event.artifact.parts[0].root
|
|
205
|
+
assert isinstance(file_b_part, FilePart)
|
|
206
|
+
assert isinstance(file_b_part.file, FileWithBytes)
|
|
207
|
+
assert file_b_part.file.name == "cv-b.docx"
|
|
208
|
+
assert file_b_part.file.bytes == b64_b
|
|
209
|
+
|
|
210
|
+
assert text_event.last_chunk is True
|
|
211
|
+
assert text_event.artifact.name == "current_result"
|
|
212
|
+
|
|
213
|
+
final_status = [e for e in events
|
|
214
|
+
if isinstance(e, TaskStatusUpdateEvent) and e.final]
|
|
215
|
+
assert len(final_status) == 1
|
|
216
|
+
assert final_status[0].status.state == TaskState.completed
|
|
@@ -0,0 +1,317 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import base64
|
|
4
|
+
from collections.abc import AsyncGenerator
|
|
5
|
+
|
|
6
|
+
import httpx
|
|
7
|
+
import pytest
|
|
8
|
+
from a2a.types import (AgentCapabilities, AgentCard, Artifact, FilePart,
|
|
9
|
+
FileWithBytes, FileWithUri, Message, Part, Task,
|
|
10
|
+
TaskState, TaskStatus, TextPart)
|
|
11
|
+
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
|
|
12
|
+
|
|
13
|
+
from distributed_a2a.client import AgentReply, RemoteAgentConnection
|
|
14
|
+
from distributed_a2a.files import extract_file_parts
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
def _b64(payload: bytes) -> str:
|
|
18
|
+
return base64.b64encode(payload).decode("ascii")
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def test_extract_file_parts_ignores_non_tool_messages_and_string_content() -> None:
|
|
22
|
+
messages = [
|
|
23
|
+
HumanMessage(content="hi"),
|
|
24
|
+
AIMessage(content="hello"),
|
|
25
|
+
ToolMessage(content="just text", tool_call_id="c1"),
|
|
26
|
+
]
|
|
27
|
+
assert extract_file_parts(messages) == []
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def test_extract_file_parts_reads_langchain_file_content_block() -> None:
|
|
31
|
+
docx_b64 = _b64(b"PK\x03\x04 real docx bytes")
|
|
32
|
+
summary = (
|
|
33
|
+
'{"filename": "cv-alice.docx", '
|
|
34
|
+
'"mime_type": "application/vnd.openxmlformats-officedocument.'
|
|
35
|
+
'wordprocessingml.document", "size_bytes": 4321}'
|
|
36
|
+
)
|
|
37
|
+
tool_msg = ToolMessage(
|
|
38
|
+
content=[
|
|
39
|
+
{"type": "text", "text": summary, "id": "lc_text_1"},
|
|
40
|
+
{"type": "file",
|
|
41
|
+
"base64": docx_b64,
|
|
42
|
+
"mime_type": (
|
|
43
|
+
"application/vnd.openxmlformats-officedocument."
|
|
44
|
+
"wordprocessingml.document"
|
|
45
|
+
),
|
|
46
|
+
"id": "lc_file_1"},
|
|
47
|
+
],
|
|
48
|
+
tool_call_id="call-cv",
|
|
49
|
+
artifact={"structured_content": {"filename": "cv-alice.docx"}},
|
|
50
|
+
)
|
|
51
|
+
|
|
52
|
+
parts = extract_file_parts([HumanMessage(content="hi"), tool_msg])
|
|
53
|
+
|
|
54
|
+
assert len(parts) == 1
|
|
55
|
+
name, file_part = parts[0]
|
|
56
|
+
assert name == "cv-alice.docx"
|
|
57
|
+
assert isinstance(file_part, FilePart)
|
|
58
|
+
assert isinstance(file_part.file, FileWithBytes)
|
|
59
|
+
assert file_part.file.name == "cv-alice.docx"
|
|
60
|
+
assert file_part.file.mime_type == (
|
|
61
|
+
"application/vnd.openxmlformats-officedocument.wordprocessingml.document"
|
|
62
|
+
)
|
|
63
|
+
assert file_part.file.bytes == docx_b64
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
def test_extract_file_parts_reads_langchain_image_content_block() -> None:
|
|
67
|
+
png_b64 = _b64(b"\x89PNG\r\n\x1a\n fake")
|
|
68
|
+
tool_msg = ToolMessage(
|
|
69
|
+
content=[
|
|
70
|
+
{"type": "image", "base64": png_b64,
|
|
71
|
+
"mime_type": "image/png", "id": "lc_img_1"},
|
|
72
|
+
],
|
|
73
|
+
tool_call_id="call-img",
|
|
74
|
+
)
|
|
75
|
+
|
|
76
|
+
parts = extract_file_parts([tool_msg])
|
|
77
|
+
|
|
78
|
+
assert len(parts) == 1
|
|
79
|
+
name, file_part = parts[0]
|
|
80
|
+
assert name.startswith("image")
|
|
81
|
+
assert name.endswith(".png")
|
|
82
|
+
assert isinstance(file_part.file, FileWithBytes)
|
|
83
|
+
assert file_part.file.mime_type == "image/png"
|
|
84
|
+
assert file_part.file.bytes == png_b64
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
def test_extract_file_parts_synthesises_name_without_filename_hint() -> None:
|
|
88
|
+
docx_b64 = _b64(b"PK\x03\x04 bytes")
|
|
89
|
+
tool_msg = ToolMessage(
|
|
90
|
+
content=[
|
|
91
|
+
{"type": "text", "text": "no json here", "id": "lc_text_1"},
|
|
92
|
+
{"type": "file",
|
|
93
|
+
"base64": docx_b64,
|
|
94
|
+
"mime_type": (
|
|
95
|
+
"application/vnd.openxmlformats-officedocument."
|
|
96
|
+
"wordprocessingml.document"
|
|
97
|
+
),
|
|
98
|
+
"id": "lc_file_1"},
|
|
99
|
+
],
|
|
100
|
+
tool_call_id="call-cv",
|
|
101
|
+
)
|
|
102
|
+
|
|
103
|
+
parts = extract_file_parts([tool_msg])
|
|
104
|
+
|
|
105
|
+
assert len(parts) == 1
|
|
106
|
+
name, file_part = parts[0]
|
|
107
|
+
assert name.startswith("attachment")
|
|
108
|
+
assert name.endswith(".docx")
|
|
109
|
+
assert isinstance(file_part.file, FileWithBytes)
|
|
110
|
+
assert file_part.file.bytes == docx_b64
|
|
111
|
+
|
|
112
|
+
|
|
113
|
+
def test_extract_file_parts_ignores_url_only_file_block() -> None:
|
|
114
|
+
tool_msg = ToolMessage(
|
|
115
|
+
content=[
|
|
116
|
+
{"type": "file",
|
|
117
|
+
"url": "https://example.com/report.pdf",
|
|
118
|
+
"mime_type": "application/pdf",
|
|
119
|
+
"id": "lc_file_1"},
|
|
120
|
+
],
|
|
121
|
+
tool_call_id="call-url-only",
|
|
122
|
+
)
|
|
123
|
+
assert extract_file_parts([tool_msg]) == []
|
|
124
|
+
|
|
125
|
+
|
|
126
|
+
def test_extract_file_parts_matches_multiple_filenames_by_order() -> None:
|
|
127
|
+
"""Two [text-summary, file] pairs in the same ToolMessage — each
|
|
128
|
+
text-summary's ``filename`` is consumed by the immediately following
|
|
129
|
+
file block."""
|
|
130
|
+
b64_a = _b64(b"aaa")
|
|
131
|
+
b64_b = _b64(b"bbb")
|
|
132
|
+
tool_msg = ToolMessage(
|
|
133
|
+
content=[
|
|
134
|
+
{"type": "text", "text": '{"filename": "a.bin"}',
|
|
135
|
+
"id": "lc_text_1"},
|
|
136
|
+
{"type": "file", "base64": b64_a,
|
|
137
|
+
"mime_type": "application/octet-stream", "id": "lc_file_1"},
|
|
138
|
+
{"type": "text", "text": '{"filename": "b.bin"}',
|
|
139
|
+
"id": "lc_text_2"},
|
|
140
|
+
{"type": "file", "base64": b64_b,
|
|
141
|
+
"mime_type": "application/octet-stream", "id": "lc_file_2"},
|
|
142
|
+
],
|
|
143
|
+
tool_call_id="call-multi",
|
|
144
|
+
)
|
|
145
|
+
|
|
146
|
+
parts = extract_file_parts([tool_msg])
|
|
147
|
+
|
|
148
|
+
assert [name for name, _ in parts] == ["a.bin", "b.bin"]
|
|
149
|
+
assert parts[0][1].file.bytes == b64_a # type: ignore[union-attr]
|
|
150
|
+
assert parts[1][1].file.bytes == b64_b # type: ignore[union-attr]
|
|
151
|
+
|
|
152
|
+
|
|
153
|
+
class _StubAgentClient:
|
|
154
|
+
def __init__(self, task: Task):
|
|
155
|
+
self._task = task
|
|
156
|
+
|
|
157
|
+
async def send_message(self, _message: Message) -> AsyncGenerator[tuple[Task, None], None]:
|
|
158
|
+
yield self._task, None
|
|
159
|
+
|
|
160
|
+
|
|
161
|
+
class _StubClientFactory:
|
|
162
|
+
def __init__(self, _config: object, task: Task):
|
|
163
|
+
self._task = task
|
|
164
|
+
|
|
165
|
+
def create(self, _agent_card: AgentCard) -> _StubAgentClient:
|
|
166
|
+
return _StubAgentClient(self._task)
|
|
167
|
+
|
|
168
|
+
|
|
169
|
+
def _patch_client_factory(monkeypatch: pytest.MonkeyPatch, task: Task) -> None:
|
|
170
|
+
def _ctor(config: object) -> _StubClientFactory:
|
|
171
|
+
return _StubClientFactory(config, task)
|
|
172
|
+
monkeypatch.setattr("distributed_a2a.client.ClientFactory", _ctor)
|
|
173
|
+
|
|
174
|
+
|
|
175
|
+
def _agent_card() -> AgentCard:
|
|
176
|
+
return AgentCard(
|
|
177
|
+
name="stub",
|
|
178
|
+
description="stub",
|
|
179
|
+
url="http://127.0.0.1:0",
|
|
180
|
+
version="1.0.0",
|
|
181
|
+
default_input_modes=["text"],
|
|
182
|
+
default_output_modes=["text"],
|
|
183
|
+
capabilities=AgentCapabilities(streaming=False, push_notifications=False),
|
|
184
|
+
skills=[],
|
|
185
|
+
preferred_transport="JSONRPC",
|
|
186
|
+
)
|
|
187
|
+
|
|
188
|
+
|
|
189
|
+
@pytest.mark.asyncio
|
|
190
|
+
async def test_remote_agent_connection_returns_text_and_files_in_agent_reply(
|
|
191
|
+
monkeypatch: pytest.MonkeyPatch) -> None:
|
|
192
|
+
docx_b64 = _b64(b"PK\x03\x04 docx-bytes")
|
|
193
|
+
completed_task = Task(
|
|
194
|
+
id="task-files",
|
|
195
|
+
context_id="ctx-files",
|
|
196
|
+
status=TaskStatus(state=TaskState.completed),
|
|
197
|
+
artifacts=[
|
|
198
|
+
Artifact(
|
|
199
|
+
artifact_id="cv-foo.docx",
|
|
200
|
+
name="cv-foo.docx",
|
|
201
|
+
parts=[Part(root=FilePart(file=FileWithBytes(
|
|
202
|
+
name="cv-foo.docx",
|
|
203
|
+
mime_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
|
204
|
+
bytes=docx_b64,
|
|
205
|
+
)))],
|
|
206
|
+
),
|
|
207
|
+
Artifact(
|
|
208
|
+
artifact_id="current_result",
|
|
209
|
+
name="current_result",
|
|
210
|
+
parts=[Part(root=TextPart(text="Here is your CV."))],
|
|
211
|
+
),
|
|
212
|
+
],
|
|
213
|
+
)
|
|
214
|
+
factory_task = completed_task
|
|
215
|
+
_patch_client_factory(monkeypatch, factory_task)
|
|
216
|
+
|
|
217
|
+
async with httpx.AsyncClient() as http_client:
|
|
218
|
+
conn = RemoteAgentConnection(_agent_card(), http_client)
|
|
219
|
+
reply = await conn.send_message("hello", "ctx-files")
|
|
220
|
+
|
|
221
|
+
assert isinstance(reply, AgentReply)
|
|
222
|
+
assert reply.text == "Here is your CV."
|
|
223
|
+
assert len(reply.files) == 1
|
|
224
|
+
f = reply.files[0]
|
|
225
|
+
assert f.name == "cv-foo.docx"
|
|
226
|
+
assert f.mime_type == (
|
|
227
|
+
"application/vnd.openxmlformats-officedocument.wordprocessingml.document"
|
|
228
|
+
)
|
|
229
|
+
assert f.bytes_b64 == docx_b64
|
|
230
|
+
assert f.uri is None
|
|
231
|
+
|
|
232
|
+
|
|
233
|
+
@pytest.mark.asyncio
|
|
234
|
+
async def test_remote_agent_connection_handles_file_with_uri(
|
|
235
|
+
monkeypatch: pytest.MonkeyPatch) -> None:
|
|
236
|
+
completed_task = Task(
|
|
237
|
+
id="task-uri",
|
|
238
|
+
context_id="ctx-uri",
|
|
239
|
+
status=TaskStatus(state=TaskState.completed),
|
|
240
|
+
artifacts=[
|
|
241
|
+
Artifact(
|
|
242
|
+
artifact_id="report",
|
|
243
|
+
name="report",
|
|
244
|
+
parts=[Part(root=FilePart(file=FileWithUri(
|
|
245
|
+
name="report.pdf",
|
|
246
|
+
mime_type="application/pdf",
|
|
247
|
+
uri="https://example.com/report.pdf",
|
|
248
|
+
)))],
|
|
249
|
+
),
|
|
250
|
+
],
|
|
251
|
+
)
|
|
252
|
+
_patch_client_factory(monkeypatch, completed_task)
|
|
253
|
+
|
|
254
|
+
async with httpx.AsyncClient() as http_client:
|
|
255
|
+
conn = RemoteAgentConnection(_agent_card(), http_client)
|
|
256
|
+
reply = await conn.send_message("hello", "ctx-uri")
|
|
257
|
+
|
|
258
|
+
assert isinstance(reply, AgentReply)
|
|
259
|
+
assert reply.text is None
|
|
260
|
+
assert len(reply.files) == 1
|
|
261
|
+
f = reply.files[0]
|
|
262
|
+
assert f.name == "report.pdf"
|
|
263
|
+
assert f.mime_type == "application/pdf"
|
|
264
|
+
assert f.bytes_b64 == ""
|
|
265
|
+
assert f.uri == "https://example.com/report.pdf"
|
|
266
|
+
|
|
267
|
+
|
|
268
|
+
@pytest.mark.asyncio
|
|
269
|
+
async def test_remote_agent_connection_gathers_multiple_files_in_order(
|
|
270
|
+
monkeypatch: pytest.MonkeyPatch) -> None:
|
|
271
|
+
b64_a = _b64(b"aaa")
|
|
272
|
+
b64_b = _b64(b"bbb")
|
|
273
|
+
completed_task = Task(
|
|
274
|
+
id="task-multi",
|
|
275
|
+
context_id="ctx-multi",
|
|
276
|
+
status=TaskStatus(state=TaskState.completed),
|
|
277
|
+
artifacts=[
|
|
278
|
+
Artifact(
|
|
279
|
+
artifact_id="a.bin",
|
|
280
|
+
name="a.bin",
|
|
281
|
+
parts=[Part(root=FilePart(file=FileWithBytes(
|
|
282
|
+
name="a.bin",
|
|
283
|
+
mime_type="application/octet-stream",
|
|
284
|
+
bytes=b64_a,
|
|
285
|
+
)))],
|
|
286
|
+
),
|
|
287
|
+
Artifact(
|
|
288
|
+
artifact_id="b.bin",
|
|
289
|
+
name="b.bin",
|
|
290
|
+
parts=[Part(root=FilePart(file=FileWithBytes(
|
|
291
|
+
name="b.bin",
|
|
292
|
+
mime_type="application/octet-stream",
|
|
293
|
+
bytes=b64_b,
|
|
294
|
+
)))],
|
|
295
|
+
),
|
|
296
|
+
Artifact(
|
|
297
|
+
artifact_id="current_result",
|
|
298
|
+
name="current_result",
|
|
299
|
+
parts=[Part(root=TextPart(text="Two files ready."))],
|
|
300
|
+
),
|
|
301
|
+
],
|
|
302
|
+
)
|
|
303
|
+
_patch_client_factory(monkeypatch, completed_task)
|
|
304
|
+
|
|
305
|
+
async with httpx.AsyncClient() as http_client:
|
|
306
|
+
conn = RemoteAgentConnection(_agent_card(), http_client)
|
|
307
|
+
reply = await conn.send_message("hello", "ctx-multi")
|
|
308
|
+
|
|
309
|
+
assert isinstance(reply, AgentReply)
|
|
310
|
+
assert reply.text == "Two files ready."
|
|
311
|
+
assert [f.name for f in reply.files] == ["a.bin", "b.bin"]
|
|
312
|
+
assert [f.mime_type for f in reply.files] == [
|
|
313
|
+
"application/octet-stream", "application/octet-stream",
|
|
314
|
+
]
|
|
315
|
+
assert [f.bytes_b64 for f in reply.files] == [b64_a, b64_b]
|
|
316
|
+
assert all(f.uri is None for f in reply.files)
|
|
317
|
+
|
|
@@ -1,7 +1,7 @@
|
|
|
1
1
|
import pytest
|
|
2
2
|
from a2a.types import AgentCapabilities, AgentCard, TaskState
|
|
3
3
|
|
|
4
|
-
from distributed_a2a.client import RoutingA2AClient
|
|
4
|
+
from distributed_a2a.client import AgentReply, RoutingA2AClient
|
|
5
5
|
|
|
6
6
|
|
|
7
7
|
def build_card(name: str, url: str) -> AgentCard:
|
|
@@ -38,7 +38,7 @@ async def test_rejection_triggers_automated_rerouting(monkeypatch: pytest.Monkey
|
|
|
38
38
|
def __init__(self, agent_card: AgentCard, _client: object, **_kwargs: object) -> None:
|
|
39
39
|
self.agent_card = agent_card
|
|
40
40
|
|
|
41
|
-
async def send_message(self, message_to_send: str, _context_id: str) ->
|
|
41
|
+
async def send_message(self, message_to_send: str, _context_id: str) -> AgentReply | AgentCard | TaskState:
|
|
42
42
|
if self.agent_card.name == "Router":
|
|
43
43
|
router_messages.append(message_to_send)
|
|
44
44
|
if "Please exclude the following agents from routing: rejecting-agent" in message_to_send:
|
|
@@ -49,7 +49,7 @@ async def test_rejection_triggers_automated_rerouting(monkeypatch: pytest.Monkey
|
|
|
49
49
|
return TaskState.rejected
|
|
50
50
|
|
|
51
51
|
if self.agent_card.name == "success-agent":
|
|
52
|
-
return "final answer"
|
|
52
|
+
return AgentReply(text="final answer")
|
|
53
53
|
|
|
54
54
|
raise AssertionError(f"Unexpected agent {self.agent_card.name}")
|
|
55
55
|
|
|
@@ -57,7 +57,8 @@ async def test_rejection_triggers_automated_rerouting(monkeypatch: pytest.Monkey
|
|
|
57
57
|
|
|
58
58
|
result = await client.send_message("Hello", context_id="ctx-1")
|
|
59
59
|
|
|
60
|
-
assert result
|
|
60
|
+
assert isinstance(result, AgentReply)
|
|
61
|
+
assert result.text == "final answer"
|
|
61
62
|
assert len(router_messages) == 2
|
|
62
63
|
assert "Please exclude the following agents from routing" not in router_messages[0]
|
|
63
64
|
assert "Please exclude the following agents from routing: rejecting-agent" in router_messages[1]
|
|
@@ -88,7 +89,7 @@ async def test_rejected_agents_reset_between_calls(monkeypatch: pytest.MonkeyPat
|
|
|
88
89
|
def __init__(self, agent_card: AgentCard, _client: object, **_kwargs: object) -> None:
|
|
89
90
|
self.agent_card = agent_card
|
|
90
91
|
|
|
91
|
-
async def send_message(self, message_to_send: str, _context_id: str) ->
|
|
92
|
+
async def send_message(self, message_to_send: str, _context_id: str) -> AgentReply | AgentCard | TaskState:
|
|
92
93
|
idx = next_step()
|
|
93
94
|
match idx:
|
|
94
95
|
case 0:
|
|
@@ -104,7 +105,7 @@ async def test_rejected_agents_reset_between_calls(monkeypatch: pytest.MonkeyPat
|
|
|
104
105
|
return success_card
|
|
105
106
|
case 3:
|
|
106
107
|
assert self.agent_card.name == "success-agent"
|
|
107
|
-
return "first response"
|
|
108
|
+
return AgentReply(text="first response")
|
|
108
109
|
case 4:
|
|
109
110
|
assert self.agent_card.name == "success-agent"
|
|
110
111
|
return TaskState.rejected
|
|
@@ -115,7 +116,7 @@ async def test_rejected_agents_reset_between_calls(monkeypatch: pytest.MonkeyPat
|
|
|
115
116
|
return rejecting_card
|
|
116
117
|
case 6:
|
|
117
118
|
assert self.agent_card.name == "rejecting-agent"
|
|
118
|
-
return "second response"
|
|
119
|
+
return AgentReply(text="second response")
|
|
119
120
|
case _:
|
|
120
121
|
raise AssertionError(f"Unexpected step {idx}")
|
|
121
122
|
|
|
@@ -124,8 +125,8 @@ async def test_rejected_agents_reset_between_calls(monkeypatch: pytest.MonkeyPat
|
|
|
124
125
|
first = await client.send_message("First", context_id="ctx-1")
|
|
125
126
|
second = await client.send_message("Second", context_id="ctx-2")
|
|
126
127
|
|
|
127
|
-
assert first == "first response"
|
|
128
|
-
assert second == "second response"
|
|
128
|
+
assert isinstance(first, AgentReply) and first.text == "first response"
|
|
129
|
+
assert isinstance(second, AgentReply) and second.text == "second response"
|
|
129
130
|
|
|
130
131
|
|
|
131
132
|
@pytest.mark.asyncio
|
|
@@ -147,7 +148,7 @@ async def test_fails_when_router_returns_already_rejected_agent(monkeypatch: pyt
|
|
|
147
148
|
def __init__(self, agent_card: AgentCard, _client: object, **_kwargs: object) -> None:
|
|
148
149
|
self.agent_card = agent_card
|
|
149
150
|
|
|
150
|
-
async def send_message(self, message_to_send: str, _context_id: str) ->
|
|
151
|
+
async def send_message(self, message_to_send: str, _context_id: str) -> AgentReply | AgentCard | TaskState:
|
|
151
152
|
idx = step["idx"]
|
|
152
153
|
step["idx"] += 1
|
|
153
154
|
if idx == 0:
|
|
@@ -8,8 +8,8 @@ import pytest
|
|
|
8
8
|
from a2a.types import (AgentCapabilities, AgentCard, Artifact, Message, Part,
|
|
9
9
|
Role, Task, TaskState, TaskStatus, TextPart)
|
|
10
10
|
|
|
11
|
-
from distributed_a2a.client import (A2ATimeoutError,
|
|
12
|
-
RoutingA2AClient)
|
|
11
|
+
from distributed_a2a.client import (A2ATimeoutError, AgentReply,
|
|
12
|
+
RemoteAgentConnection, RoutingA2AClient)
|
|
13
13
|
|
|
14
14
|
|
|
15
15
|
def _agent_card() -> AgentCard:
|
|
@@ -124,7 +124,9 @@ async def test_send_message_honors_completion_on_final_allowed_poll(
|
|
|
124
124
|
)
|
|
125
125
|
result = await connection.send_message("hi", "ctx-x")
|
|
126
126
|
|
|
127
|
-
assert result
|
|
127
|
+
assert isinstance(result, AgentReply)
|
|
128
|
+
assert result.text == "all done"
|
|
129
|
+
assert result.files == []
|
|
128
130
|
assert scripted.get_task_calls == 3
|
|
129
131
|
assert scripted._get_states == []
|
|
130
132
|
|
|
@@ -156,7 +158,9 @@ async def test_send_message_returns_immediately_if_already_completed(
|
|
|
156
158
|
)
|
|
157
159
|
result = await connection.send_message("hi", "ctx-x")
|
|
158
160
|
|
|
159
|
-
assert result
|
|
161
|
+
assert isinstance(result, AgentReply)
|
|
162
|
+
assert result.text == "all done"
|
|
163
|
+
assert result.files == []
|
|
160
164
|
|
|
161
165
|
|
|
162
166
|
def test_a2a_timeout_error_message_includes_diagnostics() -> None:
|
|
@@ -182,8 +186,8 @@ async def test_routing_a2a_client_propagates_poll_kwargs(
|
|
|
182
186
|
captured["poll_interval"] = poll_interval
|
|
183
187
|
self.agent_card = agent_card
|
|
184
188
|
|
|
185
|
-
async def send_message(self, _message: str, _context_id: str) ->
|
|
186
|
-
return "done"
|
|
189
|
+
async def send_message(self, _message: str, _context_id: str) -> AgentReply:
|
|
190
|
+
return AgentReply(text="done")
|
|
187
191
|
|
|
188
192
|
monkeypatch.setattr("distributed_a2a.client.RemoteAgentConnection", _Capturing)
|
|
189
193
|
|
|
@@ -197,5 +201,7 @@ async def test_routing_a2a_client_propagates_poll_kwargs(
|
|
|
197
201
|
monkeypatch.setattr(client, "fetch_initial_card", _fetch)
|
|
198
202
|
|
|
199
203
|
result = await client.send_message("hi", context_id="ctx-1")
|
|
200
|
-
assert result
|
|
204
|
+
assert isinstance(result, AgentReply)
|
|
205
|
+
assert result.text == "done"
|
|
206
|
+
assert result.files == []
|
|
201
207
|
assert captured == {"max_polls": 7, "poll_interval": 0.25}
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/registry_server/__init__.py
RENAMED
|
File without changes
|
{distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/registry_server/bootstrap.py
RENAMED
|
File without changes
|
{distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/registry_server/dynamo_db.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/schemas/router-agent-schema.json
RENAMED
|
File without changes
|
|
File without changes
|
{distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a.egg-info/dependency_links.txt
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|