Repository navigation
fix(agentcore): Convert SSE stream iterator to async for proper streaming support - #16293
4 commits merged into
Conversation
|
@busla is attempting to deploy a commit to the CLERKIEAI Team on Vercel. A member of the Team first needs to authorize it. |
ishaan-jaff
left a comment
There was a problem hiding this comment.
please can you show the before and after of manual testing results
This worked for me when testing
updated the PR with the error before the change |
ishaan-jaff
left a comment
There was a problem hiding this comment.
does this break existing functionality, can you run test_bedrock_agentcore.py with your changes + agents and validate it works
I think your changes would break existing sync completion
|
I tested this against my runtime which contains Strands agents with async stream, so the released LiteLLM feature breaks runtimes that stream responses. Would be great if you could test my changes against your JSON response agentcore runtime. AgentCore Runtime supports two response formats:
From the AWS documentation example: if "text/event-stream" in response.get("contentType", ""):
# Handle streaming response - uses iter_lines()
for line in response["response"].iter_lines():
# Process SSE chunks
elif response.get("contentType") == "application/json":
# Handle standard JSON response
content = json.loads(response_data)Streaming Path (Affected by Changes) File: transformation.py:405-420
Non-Streaming Path (NOT Affected) File: transformation.py:557-623
Important: def _parse_sse_stream(self, response_text: str) -> AgentCoreParsedResponse:
# ...This method takes a string parameter response_text which comes from: response_text = raw_response.text # Synchronous read of full response bodyThe async iterator (AgentCoreSSEStreamIterator) is ONLY used when:
For synchronous responses:
Synchronous responses are NOT affected because they:
|
…ader for runtime oauth support
|
Updated the PR to map |
|
@busla i believe the ask was to just confirm if your pr works. Can you just confirm this PR works for you on streaming in Strands? (Which i assume uses bedrock runtime) |
|
yes, it works on my agentcore agent that streams responses. |
…ming support (BerriAI#16293) * fix(agentcore): support async agentcore runtime streaming * revert: CLAUDE.md * revert: .gitignore * fix: map runtimeUserId to X-Amzn-Bedrock-AgentCore-Runtime-User-Id header for runtime oauth support
Description
Fixes async/sync iterator mismatch in AgentCore SSE streaming that was causing the error:
Problem
The
AgentCoreSSEStreamIteratorwas implemented as a synchronous iterator (__iter__and__next__) but was being used with async HTTP response streams from httpx. When httpx provides an async response in streaming mode, it requires async iteration.Solution
Converted
AgentCoreSSEStreamIteratorfrom sync to async iterator:__iter__()→__aiter__()__next__()→async __anext__()self.response.iter_lines()→self.response.aiter_lines()for line in self.line_iterator:→async for line in self.line_iterator:StopIteration→StopAsyncIterationThe
CustomStreamWrapperalready has proper async handling (checksis_async_iterable()and usesasync for), so it works correctly with the async iterator.Additional Enhancement: Runtime User ID Support
Added support for tracking user identity across AgentCore requests:
_get_runtime_user_id()method to extractruntimeUserIdfrom optional paramstransform_request()to setX-Amzn-Bedrock-AgentCore-Runtime-User-Idheader whenruntimeUserIdis providedruntimeSessionId→X-Amzn-Bedrock-AgentCore-Runtime-Session-IdmappingThis allows AgentCore runtimes to track individual users within sessions for better personalization and analytics.
Changes
Files Modified:
litellm/llms/bedrock/chat/agentcore/sse_iterator.py- Fixed async streaminglitellm/llms/bedrock/chat/agentcore/transformation.py- Added user ID header supportTesting
runtimeSessionIdandruntimeUserIdoptional params are correctly mapped to headersImpact
LibreChat example
LiteLLM example