distributed-a2a 0.2.1__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.2.1/distributed_a2a.egg-info → distributed_a2a-0.2.2}/PKG-INFO +1 -1
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/distributed_a2a/executors.py +4 -32
- distributed_a2a-0.2.2/distributed_a2a/files.py +68 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2/distributed_a2a.egg-info}/PKG-INFO +1 -1
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/distributed_a2a.egg-info/SOURCES.txt +1 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/pyproject.toml +1 -1
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/tests/test_executor_files.py +103 -69
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/tests/test_files.py +145 -65
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/LICENSE +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/MANIFEST.in +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/README.md +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/distributed_a2a/__init__.py +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/distributed_a2a/agent.py +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/distributed_a2a/client.py +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/distributed_a2a/config.py +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/distributed_a2a/model.py +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/distributed_a2a/py.typed +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/distributed_a2a/registry.py +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/distributed_a2a/registry_server/__init__.py +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/distributed_a2a/registry_server/bootstrap.py +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/distributed_a2a/registry_server/dynamo_db.py +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/distributed_a2a/registry_server/in_memory_registry_storage.py +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/distributed_a2a/registry_server/model.py +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/distributed_a2a/registry_server/storage.py +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/distributed_a2a/router.py +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/distributed_a2a/schemas/agent-schema.json +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/distributed_a2a/schemas/router-agent-schema.json +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/distributed_a2a/server.py +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/distributed_a2a.egg-info/dependency_links.txt +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/distributed_a2a.egg-info/requires.txt +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/distributed_a2a.egg-info/top_level.txt +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/requirements.txt +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/setup.cfg +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/setup.py +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/tests/test_app.py +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/tests/test_client.py +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/tests/test_rejection.py +0 -0
- {distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/tests/test_timeout.py +0 -0
|
@@ -4,50 +4,22 @@ 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, FilePart,
|
|
8
|
-
|
|
9
|
-
TaskStatusUpdateEvent)
|
|
7
|
+
from a2a.types import (Artifact, FilePart, Part, TaskArtifactUpdateEvent,
|
|
8
|
+
TaskState, TaskStatus, TaskStatusUpdateEvent)
|
|
10
9
|
from a2a.utils import new_text_artifact
|
|
11
|
-
from langchain_core.messages import BaseMessage, ToolMessage
|
|
12
10
|
from langchain_core.tools import BaseTool
|
|
13
11
|
from langchain_mcp_adapters.client import MultiServerMCPClient
|
|
14
12
|
from langgraph.checkpoint.base import BaseCheckpointSaver
|
|
15
|
-
from mcp.types import BlobResourceContents, EmbeddedResource, ImageContent
|
|
16
13
|
|
|
17
14
|
from .agent import RoutingResponse, StatusAgent, StringResponse
|
|
18
15
|
from .config import settings
|
|
16
|
+
from .files import extract_file_parts
|
|
19
17
|
from .model import AgentConfig, RouterConfig
|
|
20
18
|
from .registry import AgentRegistryLookupClient, McpRegistryLookup
|
|
21
19
|
|
|
22
20
|
logger = logging.getLogger(__name__)
|
|
23
21
|
|
|
24
22
|
|
|
25
|
-
def _extract_file_parts(messages: list[BaseMessage]) -> list[tuple[str, FilePart]]:
|
|
26
|
-
out: list[tuple[str, FilePart]] = []
|
|
27
|
-
for message in messages:
|
|
28
|
-
if not isinstance(message, ToolMessage):
|
|
29
|
-
continue
|
|
30
|
-
artifact = getattr(message, "artifact", None)
|
|
31
|
-
if not artifact:
|
|
32
|
-
continue
|
|
33
|
-
for block in artifact:
|
|
34
|
-
if isinstance(block, EmbeddedResource) and isinstance(block.resource, BlobResourceContents):
|
|
35
|
-
uri_str = str(block.resource.uri)
|
|
36
|
-
name = uri_str.rsplit("/", 1)[-1] or "file.bin"
|
|
37
|
-
out.append((name, FilePart(file=FileWithBytes(
|
|
38
|
-
name=name,
|
|
39
|
-
mime_type=block.resource.mimeType or "application/octet-stream",
|
|
40
|
-
bytes=block.resource.blob,
|
|
41
|
-
))))
|
|
42
|
-
elif isinstance(block, ImageContent):
|
|
43
|
-
out.append(("image", FilePart(file=FileWithBytes(
|
|
44
|
-
name="image",
|
|
45
|
-
mime_type=block.mimeType,
|
|
46
|
-
bytes=block.data,
|
|
47
|
-
))))
|
|
48
|
-
return out
|
|
49
|
-
|
|
50
|
-
|
|
51
23
|
class RoutingFailed(Exception):
|
|
52
24
|
def __init__(self, message: str) -> None:
|
|
53
25
|
super().__init__(message)
|
|
@@ -155,7 +127,7 @@ class RoutingAgentExecutor(AgentExecutor):
|
|
|
155
127
|
artifact = await _route_request_to_matching_agent(self.routing_agent, self.agent_registry, context)
|
|
156
128
|
else:
|
|
157
129
|
logger.info(f"Request with id {context.context_id} was successfully processed by agent.")
|
|
158
|
-
file_parts =
|
|
130
|
+
file_parts = extract_file_parts(invocation.messages)
|
|
159
131
|
artifact = new_text_artifact(
|
|
160
132
|
name='current_result',
|
|
161
133
|
description='Result of request to agent.',
|
|
@@ -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
|
|
@@ -1,6 +1,7 @@
|
|
|
1
1
|
from __future__ import annotations
|
|
2
2
|
|
|
3
3
|
import base64
|
|
4
|
+
from types import SimpleNamespace
|
|
4
5
|
from typing import Any
|
|
5
6
|
|
|
6
7
|
import pytest
|
|
@@ -11,16 +12,12 @@ from a2a.types import Message as A2AMessage
|
|
|
11
12
|
from a2a.types import (MessageSendParams, Part, Role, TaskArtifactUpdateEvent,
|
|
12
13
|
TaskState, TaskStatusUpdateEvent, TextPart)
|
|
13
14
|
from langchain_core.messages import BaseMessage, HumanMessage, ToolMessage
|
|
14
|
-
from mcp.types import BlobResourceContents, EmbeddedResource
|
|
15
|
-
from pydantic import AnyUrl
|
|
16
15
|
|
|
17
16
|
from distributed_a2a.agent import AgentInvocation, StringResponse
|
|
18
17
|
from distributed_a2a.executors import RoutingAgentExecutor
|
|
19
18
|
|
|
20
19
|
|
|
21
20
|
class _StubStatusAgent:
|
|
22
|
-
"""Minimal stand-in for ``StatusAgent`` used in the executor."""
|
|
23
|
-
|
|
24
21
|
def __init__(self, response: StringResponse, messages: list[BaseMessage]) -> None:
|
|
25
22
|
self._response = response
|
|
26
23
|
self._messages = messages
|
|
@@ -47,35 +44,7 @@ def _make_request_context() -> RequestContext:
|
|
|
47
44
|
)
|
|
48
45
|
|
|
49
46
|
|
|
50
|
-
def _make_tool_message_with_docx(b64: str) -> ToolMessage:
|
|
51
|
-
return ToolMessage(
|
|
52
|
-
content='{"filename": "cv-foo.docx"}',
|
|
53
|
-
tool_call_id="call-1",
|
|
54
|
-
artifact=[
|
|
55
|
-
EmbeddedResource(
|
|
56
|
-
type="resource",
|
|
57
|
-
resource=BlobResourceContents(
|
|
58
|
-
uri=AnyUrl("cv://cv-foo.docx"),
|
|
59
|
-
mimeType=(
|
|
60
|
-
"application/vnd.openxmlformats-officedocument."
|
|
61
|
-
"wordprocessingml.document"
|
|
62
|
-
),
|
|
63
|
-
blob=b64,
|
|
64
|
-
),
|
|
65
|
-
),
|
|
66
|
-
],
|
|
67
|
-
)
|
|
68
|
-
|
|
69
|
-
|
|
70
47
|
async def _drain_queue(queue: EventQueue) -> list[Any]:
|
|
71
|
-
"""Dequeue every event currently in ``queue`` without blocking.
|
|
72
|
-
|
|
73
|
-
Uses ``no_wait=True`` and stops on the first empty/closed exception. We
|
|
74
|
-
intentionally avoid ``await queue.close()`` because the graceful close
|
|
75
|
-
waits for ``queue.join()`` — which requires the consumer to call
|
|
76
|
-
``task_done()`` for every enqueued item — and would otherwise deadlock
|
|
77
|
-
the test once we've finished dequeuing.
|
|
78
|
-
"""
|
|
79
48
|
events: list[Any] = []
|
|
80
49
|
while True:
|
|
81
50
|
try:
|
|
@@ -88,17 +57,54 @@ async def _drain_queue(queue: EventQueue) -> list[Any]:
|
|
|
88
57
|
|
|
89
58
|
|
|
90
59
|
@pytest.mark.asyncio
|
|
91
|
-
async def
|
|
92
|
-
|
|
93
|
-
|
|
94
|
-
|
|
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
|
+
|
|
95
85
|
|
|
96
|
-
|
|
97
|
-
|
|
98
|
-
|
|
99
|
-
|
|
100
|
-
|
|
101
|
-
|
|
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
|
+
)
|
|
102
108
|
|
|
103
109
|
executor = RoutingAgentExecutor.__new__(RoutingAgentExecutor)
|
|
104
110
|
executor.agent_config = SimpleNamespace( # type: ignore[assignment]
|
|
@@ -120,50 +126,51 @@ async def test_executor_emits_file_part_event_for_embedded_resource(
|
|
|
120
126
|
await executor.execute(ctx, queue)
|
|
121
127
|
events = await _drain_queue(queue)
|
|
122
128
|
|
|
123
|
-
# Expected sequence: working status, file artifact, text artifact, final status.
|
|
124
129
|
artifact_events = [e for e in events if isinstance(e, TaskArtifactUpdateEvent)]
|
|
125
|
-
status_events = [e for e in events if isinstance(e, TaskStatusUpdateEvent)]
|
|
126
|
-
|
|
127
130
|
assert len(artifact_events) == 2, (
|
|
128
|
-
|
|
129
|
-
|
|
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."
|
|
130
134
|
)
|
|
131
|
-
|
|
132
135
|
file_event, text_event = artifact_events
|
|
133
|
-
# 1) file artifact carries a FilePart, is NOT the final chunk
|
|
134
136
|
assert file_event.last_chunk is False
|
|
135
|
-
assert file_event.artifact.name == "cv-
|
|
136
|
-
assert len(file_event.artifact.parts) == 1
|
|
137
|
+
assert file_event.artifact.name == "cv-bob.docx"
|
|
137
138
|
file_part = file_event.artifact.parts[0].root
|
|
138
139
|
assert isinstance(file_part, FilePart)
|
|
139
140
|
assert isinstance(file_part.file, FileWithBytes)
|
|
140
|
-
assert file_part.file.name == "cv-
|
|
141
|
+
assert file_part.file.name == "cv-bob.docx"
|
|
141
142
|
assert file_part.file.bytes == docx_b64
|
|
142
143
|
|
|
143
|
-
# 2) text artifact carries the LLM-visible summary, IS the final chunk
|
|
144
144
|
assert text_event.last_chunk is True
|
|
145
145
|
assert text_event.artifact.name == "current_result"
|
|
146
|
-
text_part = text_event.artifact.parts[0].root
|
|
147
|
-
assert isinstance(text_part, TextPart)
|
|
148
|
-
assert text_part.text == "*cv-agent*: Here is your CV."
|
|
149
|
-
|
|
150
|
-
# 3) Final status event is `completed`
|
|
151
|
-
final_status = [e for e in status_events if e.final]
|
|
152
|
-
assert len(final_status) == 1
|
|
153
|
-
assert final_status[0].status.state == TaskState.completed
|
|
154
146
|
|
|
155
147
|
|
|
156
148
|
@pytest.mark.asyncio
|
|
157
|
-
async def
|
|
158
|
-
|
|
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
|
+
)
|
|
159
165
|
|
|
160
166
|
executor = RoutingAgentExecutor.__new__(RoutingAgentExecutor)
|
|
161
167
|
executor.agent_config = SimpleNamespace( # type: ignore[assignment]
|
|
162
|
-
agent=SimpleNamespace(card=SimpleNamespace(name="
|
|
168
|
+
agent=SimpleNamespace(card=SimpleNamespace(name="cv-agent")),
|
|
163
169
|
)
|
|
164
170
|
executor.agent = _StubStatusAgent( # type: ignore[assignment]
|
|
165
|
-
StringResponse(status=TaskState.completed,
|
|
166
|
-
|
|
171
|
+
StringResponse(status=TaskState.completed,
|
|
172
|
+
response="Here are your two CVs."),
|
|
173
|
+
[HumanMessage(content="render two CVs"), tool_msg],
|
|
167
174
|
)
|
|
168
175
|
|
|
169
176
|
async def _noop_reinit() -> None:
|
|
@@ -177,6 +184,33 @@ async def test_executor_emits_no_file_event_when_no_artifacts() -> None:
|
|
|
177
184
|
events = await _drain_queue(queue)
|
|
178
185
|
|
|
179
186
|
artifact_events = [e for e in events if isinstance(e, TaskArtifactUpdateEvent)]
|
|
180
|
-
assert len(artifact_events) ==
|
|
181
|
-
|
|
182
|
-
|
|
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
|
|
@@ -9,117 +9,145 @@ from a2a.types import (AgentCapabilities, AgentCard, Artifact, FilePart,
|
|
|
9
9
|
FileWithBytes, FileWithUri, Message, Part, Task,
|
|
10
10
|
TaskState, TaskStatus, TextPart)
|
|
11
11
|
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
|
|
12
|
-
from mcp.types import (BlobResourceContents, EmbeddedResource, ImageContent,
|
|
13
|
-
TextResourceContents)
|
|
14
|
-
from pydantic import AnyUrl
|
|
15
12
|
|
|
16
13
|
from distributed_a2a.client import AgentReply, RemoteAgentConnection
|
|
17
|
-
from distributed_a2a.
|
|
14
|
+
from distributed_a2a.files import extract_file_parts
|
|
18
15
|
|
|
19
16
|
|
|
20
17
|
def _b64(payload: bytes) -> str:
|
|
21
18
|
return base64.b64encode(payload).decode("ascii")
|
|
22
19
|
|
|
23
20
|
|
|
24
|
-
def
|
|
25
|
-
|
|
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
|
+
)
|
|
26
37
|
tool_msg = ToolMessage(
|
|
27
|
-
content=
|
|
28
|
-
|
|
29
|
-
|
|
30
|
-
|
|
31
|
-
|
|
32
|
-
|
|
33
|
-
|
|
34
|
-
|
|
35
|
-
|
|
36
|
-
),
|
|
37
|
-
),
|
|
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"},
|
|
38
47
|
],
|
|
48
|
+
tool_call_id="call-cv",
|
|
49
|
+
artifact={"structured_content": {"filename": "cv-alice.docx"}},
|
|
39
50
|
)
|
|
40
51
|
|
|
41
|
-
parts =
|
|
52
|
+
parts = extract_file_parts([HumanMessage(content="hi"), tool_msg])
|
|
42
53
|
|
|
43
54
|
assert len(parts) == 1
|
|
44
55
|
name, file_part = parts[0]
|
|
45
|
-
assert name == "cv-
|
|
56
|
+
assert name == "cv-alice.docx"
|
|
46
57
|
assert isinstance(file_part, FilePart)
|
|
47
58
|
assert isinstance(file_part.file, FileWithBytes)
|
|
48
|
-
assert file_part.file.name == "cv-
|
|
59
|
+
assert file_part.file.name == "cv-alice.docx"
|
|
49
60
|
assert file_part.file.mime_type == (
|
|
50
61
|
"application/vnd.openxmlformats-officedocument.wordprocessingml.document"
|
|
51
62
|
)
|
|
52
63
|
assert file_part.file.bytes == docx_b64
|
|
53
64
|
|
|
54
65
|
|
|
55
|
-
def
|
|
66
|
+
def test_extract_file_parts_reads_langchain_image_content_block() -> None:
|
|
56
67
|
png_b64 = _b64(b"\x89PNG\r\n\x1a\n fake")
|
|
57
68
|
tool_msg = ToolMessage(
|
|
58
|
-
content=
|
|
59
|
-
|
|
60
|
-
|
|
61
|
-
ImageContent(type="image", data=png_b64, mimeType="image/png"),
|
|
69
|
+
content=[
|
|
70
|
+
{"type": "image", "base64": png_b64,
|
|
71
|
+
"mime_type": "image/png", "id": "lc_img_1"},
|
|
62
72
|
],
|
|
73
|
+
tool_call_id="call-img",
|
|
63
74
|
)
|
|
64
75
|
|
|
65
|
-
parts =
|
|
76
|
+
parts = extract_file_parts([tool_msg])
|
|
66
77
|
|
|
67
78
|
assert len(parts) == 1
|
|
68
79
|
name, file_part = parts[0]
|
|
69
|
-
assert name
|
|
80
|
+
assert name.startswith("image")
|
|
81
|
+
assert name.endswith(".png")
|
|
70
82
|
assert isinstance(file_part.file, FileWithBytes)
|
|
71
83
|
assert file_part.file.mime_type == "image/png"
|
|
72
84
|
assert file_part.file.bytes == png_b64
|
|
73
85
|
|
|
74
86
|
|
|
75
|
-
def
|
|
76
|
-
|
|
77
|
-
|
|
78
|
-
|
|
79
|
-
|
|
80
|
-
|
|
81
|
-
|
|
82
|
-
|
|
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])
|
|
83
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
|
|
84
111
|
|
|
85
|
-
|
|
112
|
+
|
|
113
|
+
def test_extract_file_parts_ignores_url_only_file_block() -> None:
|
|
86
114
|
tool_msg = ToolMessage(
|
|
87
|
-
content=
|
|
88
|
-
|
|
89
|
-
|
|
90
|
-
|
|
91
|
-
|
|
92
|
-
resource=TextResourceContents(
|
|
93
|
-
uri=AnyUrl("text://note.txt"),
|
|
94
|
-
mimeType="text/plain",
|
|
95
|
-
text="some inline text",
|
|
96
|
-
),
|
|
97
|
-
),
|
|
115
|
+
content=[
|
|
116
|
+
{"type": "file",
|
|
117
|
+
"url": "https://example.com/report.pdf",
|
|
118
|
+
"mime_type": "application/pdf",
|
|
119
|
+
"id": "lc_file_1"},
|
|
98
120
|
],
|
|
121
|
+
tool_call_id="call-url-only",
|
|
99
122
|
)
|
|
100
|
-
assert
|
|
123
|
+
assert extract_file_parts([tool_msg]) == []
|
|
101
124
|
|
|
102
125
|
|
|
103
|
-
def
|
|
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")
|
|
104
132
|
tool_msg = ToolMessage(
|
|
105
|
-
content=
|
|
106
|
-
|
|
107
|
-
|
|
108
|
-
|
|
109
|
-
|
|
110
|
-
|
|
111
|
-
|
|
112
|
-
|
|
113
|
-
|
|
114
|
-
),
|
|
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"},
|
|
115
142
|
],
|
|
143
|
+
tool_call_id="call-multi",
|
|
116
144
|
)
|
|
117
|
-
|
|
118
|
-
|
|
119
|
-
|
|
120
|
-
assert name == "
|
|
121
|
-
assert
|
|
122
|
-
assert
|
|
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]
|
|
123
151
|
|
|
124
152
|
|
|
125
153
|
class _StubAgentClient:
|
|
@@ -235,3 +263,55 @@ async def test_remote_agent_connection_handles_file_with_uri(
|
|
|
235
263
|
assert f.mime_type == "application/pdf"
|
|
236
264
|
assert f.bytes_b64 == ""
|
|
237
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
|
+
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{distributed_a2a-0.2.1 → distributed_a2a-0.2.2}/distributed_a2a/registry_server/bootstrap.py
RENAMED
|
File without changes
|
{distributed_a2a-0.2.1 → 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.2.1 → distributed_a2a-0.2.2}/distributed_a2a/schemas/router-agent-schema.json
RENAMED
|
File without changes
|
|
File without changes
|
{distributed_a2a-0.2.1 → 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
|
|
File without changes
|
|
File without changes
|
|
File without changes
|