mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(a2a): honor agent protocolVersion in the chat completions bridge
This commit is contained in:
Devin AI
parent
60729f733e
commit
367be6b0ad
3 changed files with 173 additions and 10 deletions
|
|
@ -19,6 +19,8 @@ from ..common_utils import (
|
|||
)
|
||||
from .streaming_iterator import A2AModelResponseIterator
|
||||
|
||||
A2A_PROTOCOL_VERSION_PARAM = "a2a_protocol_version"
|
||||
|
||||
|
||||
class A2AConfig(BaseConfig):
|
||||
"""
|
||||
|
|
@ -66,9 +68,10 @@ class A2AConfig(BaseConfig):
|
|||
|
||||
agent = global_agent_registry.get_agent_by_name(agent_name)
|
||||
if agent:
|
||||
# Get api_base from agent card URL
|
||||
if api_base is None and agent.agent_card_params:
|
||||
api_base = agent.agent_card_params.get("url")
|
||||
# Get api_base from the agent card
|
||||
if agent.agent_card_params:
|
||||
if api_base is None:
|
||||
api_base = agent.agent_card_params.get("url")
|
||||
|
||||
# Get api_key, headers, and other params from litellm_params
|
||||
if agent.litellm_params:
|
||||
|
|
@ -225,21 +228,85 @@ class A2AConfig(BaseConfig):
|
|||
"messageId": str(uuid.uuid4()),
|
||||
}
|
||||
|
||||
# Build JSON-RPC 2.0 request
|
||||
# For A2A protocol, the method is "message/send" for non-streaming
|
||||
# and "message/stream" for streaming
|
||||
stream = optional_params.get("stream", False)
|
||||
method = "message/stream" if stream else "message/send"
|
||||
protocol_version = (
|
||||
optional_params.get(A2A_PROTOCOL_VERSION_PARAM)
|
||||
or litellm_params.get(A2A_PROTOCOL_VERSION_PARAM)
|
||||
or self._pinned_protocol_version_from_registry(model)
|
||||
)
|
||||
|
||||
if str(protocol_version) == "1.0":
|
||||
method = "SendStreamingMessage" if stream else "SendMessage"
|
||||
params = self._build_v1_send_params(a2a_message)
|
||||
else:
|
||||
method = "message/stream" if stream else "message/send"
|
||||
params = {"message": a2a_message}
|
||||
|
||||
request_data = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"method": method,
|
||||
"params": {"message": a2a_message},
|
||||
"params": params,
|
||||
}
|
||||
|
||||
return request_data
|
||||
|
||||
@staticmethod
|
||||
def _pinned_protocol_version_from_registry(model: str) -> Optional[str]:
|
||||
"""Look up the agent's pinned ``protocolVersion`` from the proxy registry.
|
||||
|
||||
The completion bridge injects the agent's URL as ``api_base`` and passes the
|
||||
bare agent name as ``model`` (the ``a2a/`` prefix is stripped upstream), so
|
||||
the pinned version is not otherwise threaded through to request building.
|
||||
The registry is the single source of truth; it is only importable in a proxy
|
||||
context, so this returns ``None`` when running as a plain SDK.
|
||||
"""
|
||||
agent_name = model[len("a2a/") :] if model.startswith("a2a/") else model
|
||||
try:
|
||||
from litellm.proxy.agent_endpoints.agent_registry import (
|
||||
global_agent_registry,
|
||||
)
|
||||
except ImportError:
|
||||
return None
|
||||
|
||||
agent = global_agent_registry.get_agent_by_name(agent_name)
|
||||
if agent is None or not agent.agent_card_params:
|
||||
return None
|
||||
version = agent.agent_card_params.get("protocolVersion")
|
||||
return str(version) if version is not None else None
|
||||
|
||||
@staticmethod
|
||||
def _build_v1_send_params(a2a_message: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Build A2A 1.0 JSON-RPC ``params`` (protobuf-JSON) from a 0.3 message dict.
|
||||
|
||||
Mirrors how the a2a-sdk 1.x JSON-RPC transport serializes a
|
||||
``SendMessageRequest``, so 1.0-only agents receive the envelope they expect.
|
||||
"""
|
||||
try:
|
||||
from a2a.compat.v0_3.conversions import (
|
||||
MessageToDict,
|
||||
to_core_send_message_request,
|
||||
)
|
||||
from a2a.compat.v0_3.types import MessageSendParams, SendMessageRequest
|
||||
except ImportError as e:
|
||||
raise A2AError(
|
||||
status_code=500,
|
||||
message=(
|
||||
"The 'a2a-sdk' package is required to call an A2A agent pinned to "
|
||||
"protocolVersion '1.0'. Install it with: pip install a2a-sdk"
|
||||
),
|
||||
) from e
|
||||
|
||||
request = SendMessageRequest(
|
||||
id=str(uuid.uuid4()),
|
||||
params=MessageSendParams.model_validate({"message": a2a_message}),
|
||||
)
|
||||
params: Dict[str, Any] = MessageToDict(
|
||||
to_core_send_message_request(request),
|
||||
preserving_proto_field_name=False,
|
||||
)
|
||||
return params
|
||||
|
||||
def transform_response(
|
||||
self,
|
||||
model: str,
|
||||
|
|
|
|||
|
|
@ -80,7 +80,8 @@ def extract_text_from_a2a_message(message: Dict[str, Any], depth: int = 0, max_d
|
|||
text_parts: List[str] = []
|
||||
|
||||
for part in parts:
|
||||
if part.get("kind") == "text":
|
||||
kind = part.get("kind")
|
||||
if kind == "text" or (kind is None and "text" in part):
|
||||
text_parts.append(part.get("text", ""))
|
||||
# Handle nested parts if they exist
|
||||
elif "parts" in part:
|
||||
|
|
|
|||
|
|
@ -2,9 +2,11 @@
|
|||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.llms.a2a.chat.transformation import A2AConfig
|
||||
from litellm.llms.a2a.chat.transformation import A2A_PROTOCOL_VERSION_PARAM, A2AConfig
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
_MESSAGES = [{"role": "user", "content": "hi there agent"}]
|
||||
|
||||
|
||||
def _raw_response(text: str) -> MagicMock:
|
||||
raw = MagicMock()
|
||||
|
|
@ -21,6 +23,99 @@ def _raw_response(text: str) -> MagicMock:
|
|||
return raw
|
||||
|
||||
|
||||
def _v1_raw_response(text: str) -> MagicMock:
|
||||
"""A2A 1.0 (a2a-sdk 1.x) protobuf-JSON send response: no ``kind`` fields."""
|
||||
raw = MagicMock()
|
||||
raw.status_code = 200
|
||||
raw.headers = {}
|
||||
raw.json.return_value = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": "resp-1",
|
||||
"result": {"message": {"messageId": "m1", "role": "ROLE_AGENT", "parts": [{"text": text}]}},
|
||||
}
|
||||
return raw
|
||||
|
||||
|
||||
def test_transform_request_defaults_to_0_3():
|
||||
"""Without a pinned protocol version the bridge emits the legacy 0.3 method."""
|
||||
request = A2AConfig().transform_request("a2a/agent", _MESSAGES, {}, {}, {})
|
||||
assert request["method"] == "message/send"
|
||||
assert request["params"]["message"]["parts"][0]["kind"] == "text"
|
||||
|
||||
|
||||
def test_transform_request_v1_uses_send_message():
|
||||
"""Regression for #32609: a 1.0-pinned agent must get the 1.0 method and the
|
||||
protobuf-JSON envelope (uppercase role enum, parts without ``kind``)."""
|
||||
request = A2AConfig().transform_request(
|
||||
"a2a/agent", _MESSAGES, {A2A_PROTOCOL_VERSION_PARAM: "1.0"}, {}, {}
|
||||
)
|
||||
assert request["method"] == "SendMessage"
|
||||
message = request["params"]["message"]
|
||||
assert message["role"] == "ROLE_USER"
|
||||
assert message["parts"] == [{"text": "user: hi there agent"}]
|
||||
assert "kind" not in message["parts"][0]
|
||||
|
||||
|
||||
def test_transform_request_v1_streaming_uses_send_streaming_message():
|
||||
request = A2AConfig().transform_request(
|
||||
"a2a/agent", _MESSAGES, {A2A_PROTOCOL_VERSION_PARAM: "1.0", "stream": True}, {}, {}
|
||||
)
|
||||
assert request["method"] == "SendStreamingMessage"
|
||||
|
||||
|
||||
def test_transform_request_v1_from_litellm_params():
|
||||
"""The pinned version may arrive via litellm_params (direct SDK usage)."""
|
||||
request = A2AConfig().transform_request(
|
||||
"a2a/agent", _MESSAGES, {}, {A2A_PROTOCOL_VERSION_PARAM: "1.0"}, {}
|
||||
)
|
||||
assert request["method"] == "SendMessage"
|
||||
|
||||
|
||||
def test_transform_request_reads_pinned_version_from_registry():
|
||||
"""Regression for #32609: the completion bridge strips the ``a2a/`` prefix and
|
||||
injects only api_base, so transform_request must resolve the pinned version from
|
||||
the registry (keyed by the bare agent name) rather than defaulting to 0.3."""
|
||||
try:
|
||||
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
|
||||
from litellm.types.agents import AgentResponse
|
||||
except ImportError:
|
||||
import pytest
|
||||
|
||||
pytest.skip("Registry not available (not in proxy context)")
|
||||
|
||||
agent = AgentResponse(
|
||||
agent_id="v1-id",
|
||||
agent_name="v1-agent",
|
||||
agent_card_params={"url": "http://agent.example.com:9999", "protocolVersion": "1.0"},
|
||||
litellm_params={},
|
||||
)
|
||||
original_agents = global_agent_registry.agent_list.copy()
|
||||
global_agent_registry.register_agent(agent)
|
||||
try:
|
||||
request = A2AConfig().transform_request("v1-agent", _MESSAGES, {}, {}, {})
|
||||
assert request["method"] == "SendMessage"
|
||||
assert "kind" not in request["params"]["message"]["parts"][0]
|
||||
finally:
|
||||
global_agent_registry.agent_list = original_agents
|
||||
|
||||
|
||||
def test_transform_response_extracts_text_from_v1_message():
|
||||
"""Regression for #32609: text must be extracted from a 1.0 protobuf-JSON
|
||||
message result whose parts carry no ``kind`` field."""
|
||||
result = A2AConfig().transform_response(
|
||||
model="a2a/test-agent",
|
||||
raw_response=_v1_raw_response("hello from a 1.0 agent"),
|
||||
model_response=ModelResponse(),
|
||||
logging_obj=MagicMock(),
|
||||
request_data={},
|
||||
messages=_MESSAGES,
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
assert result.choices[0].message.content == "hello from a 1.0 agent"
|
||||
|
||||
|
||||
def test_transform_response_sets_usage():
|
||||
"""Regression: A2AConfig.transform_response must populate usage so per-token
|
||||
pricing computes real cost and callers don't get usage 0/0/0."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue