scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

Merge pull request #42239 from BerriAI/litellm_agentcore_a2a_message_stream

Browse source 
fix(a2a): send message/stream for Bedrock AgentCore streaming requests
This commit is contained in:
Yassin Kortam 2026-09-21 09:53:41 -05:00 • committed by GitHub
commit 32133a329c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 33 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/a2a_protocol/providers/bedrock_agentcore/handler.py
View file

@ -98,7 +98,7 @@ class BedrockAgentCoreA2AHandler:
request_id=request_id,
params=params,
litellm_params=litellm_params,
method="message/send",
method="message/stream",
stream=True,
agent_extra_headers=agent_extra_headers,
)

32
tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py
View file

@ -605,6 +605,38 @@ class TestNonStreaming:
assert result["error"]["message"] == "Bad request"
class TestStreaming:
"""Streaming requests must ask AgentCore for a stream, not a single send."""
@pytest.mark.asyncio
async def test_streaming_request_uses_message_stream_method_and_yields_sse_events(self, httpx_transport):
from litellm.a2a_protocol.providers.bedrock_agentcore.config import (
BedrockAgentCoreA2AConfig,
)
sse_body = (
'data: {"jsonrpc": "2.0", "id": "req-001", "result": {"kind": "task", "id": "t1"}}\n\n'
'data: {"jsonrpc": "2.0", "id": "req-001", "result": {"kind": "status-update", "final": true}}\n\n'
)
with respx.mock(assert_all_called=True) as router:
route = router.post(url__regex=r".*/invocations.*").mock(
return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, text=sse_body)
)
events = [
event
async for event in BedrockAgentCoreA2AConfig().handle_streaming(
request_id="req-001",
params=SAMPLE_PARAMS,
litellm_params=SAMPLE_LITELLM_PARAMS,
)
]
sent_body = json.loads(route.calls.last.request.content)
assert sent_body["method"] == "message/stream", sent_body
assert sent_body["params"]["message"]["messageId"] == "msg-001"
assert [event["result"]["kind"] for event in events] == ["task", "status-update"]
class TestConfigManager:
"""Test that config manager routes 'bedrock' correctly."""