mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(agentcore): end a stream on messageStop when the runtime sends no metadata or final message
This commit is contained in:
yucheng-berri
parent
04f3765e1a
commit
0ea2557ffe
2 changed files with 106 additions and 1 deletions
|
|
@ -540,6 +540,17 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
|
|||
|
||||
return AgentCoreParsedResponse(content=content, usage=usage_data, final_message=final_message)
|
||||
|
||||
@staticmethod
|
||||
def _message_stop_chunk(model: str) -> ModelResponseStream:
|
||||
chunk: Final = ModelResponseStream(
|
||||
id=f"chatcmpl-{uuid.uuid4()}",
|
||||
created=0,
|
||||
model=model,
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
chunk.choices = [StreamingChoices(finish_reason="stop", index=0, delta=Delta())]
|
||||
return chunk
|
||||
|
||||
def _stream_agentcore_response_sync(
|
||||
self,
|
||||
response: httpx.Response,
|
||||
|
|
@ -549,6 +560,8 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
|
|||
Internal sync generator that parses SSE and yields ModelResponse chunks.
|
||||
"""
|
||||
buffer = ""
|
||||
sent_finish = False
|
||||
message_stopped = False # runtimes that send no metadata or final message end on messageStop
|
||||
for text_chunk in response.iter_text():
|
||||
buffer += text_chunk
|
||||
|
||||
|
|
@ -620,7 +633,9 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
|
|||
total_tokens=usage_data.get("totalTokens", 0),
|
||||
),
|
||||
)
|
||||
sent_finish = True
|
||||
yield chunk
|
||||
message_stopped |= "messageStop" in event_payload
|
||||
|
||||
# Process final message
|
||||
if "message" in data_obj and isinstance(data_obj["message"], dict):
|
||||
|
|
@ -637,12 +652,16 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
|
|||
delta=Delta(),
|
||||
)
|
||||
]
|
||||
sent_finish = True
|
||||
yield chunk
|
||||
|
||||
except json.JSONDecodeError:
|
||||
verbose_logger.debug("Skipping non-JSON SSE line: %s", line[:100])
|
||||
continue
|
||||
|
||||
if message_stopped and not sent_finish:
|
||||
yield self._message_stop_chunk(model)
|
||||
|
||||
def get_sync_custom_stream_wrapper(
|
||||
self,
|
||||
model: str,
|
||||
|
|
@ -771,6 +790,8 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
|
|||
Internal async generator that parses SSE and yields ModelResponse chunks.
|
||||
"""
|
||||
buffer = ""
|
||||
sent_finish = False
|
||||
message_stopped = False # runtimes that send no metadata or final message end on messageStop
|
||||
async for text_chunk in response.aiter_text():
|
||||
buffer += text_chunk
|
||||
|
||||
|
|
@ -842,7 +863,9 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
|
|||
total_tokens=usage_data.get("totalTokens", 0),
|
||||
),
|
||||
)
|
||||
sent_finish = True
|
||||
yield chunk
|
||||
message_stopped |= "messageStop" in event_payload
|
||||
|
||||
# Process final message
|
||||
if "message" in data_obj and isinstance(data_obj["message"], dict):
|
||||
|
|
@ -859,12 +882,16 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
|
|||
delta=Delta(),
|
||||
)
|
||||
]
|
||||
sent_finish = True
|
||||
yield chunk
|
||||
|
||||
except json.JSONDecodeError:
|
||||
verbose_logger.debug("Skipping non-JSON SSE line: %s", line[:100])
|
||||
continue
|
||||
|
||||
if message_stopped and not sent_finish:
|
||||
yield self._message_stop_chunk(model)
|
||||
|
||||
async def get_async_custom_stream_wrapper(
|
||||
self,
|
||||
model: str,
|
||||
|
|
|
|||
|
|
@ -19,8 +19,9 @@ from typing import Final
|
|||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import litellm
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
|
||||
|
||||
class TestAgentCoreAcceptHeader:
|
||||
|
|
@ -1269,3 +1270,80 @@ def test_agentcore_runtime_response_becomes_assistant_message(
|
|||
assert json.loads(route.calls[0].request.content)["prompt"] == "Explain machine learning in simple terms"
|
||||
assert response.choices[0].message.content == "Machine learning lets computers learn from data."
|
||||
assert response.choices[0].finish_reason == "stop"
|
||||
|
||||
|
||||
_AGENTCORE_RUNTIME_URL: Final = (
|
||||
"https://bedrock-agentcore.us-west-2.amazonaws.com/runtimes/"
|
||||
"arn%3Aaws%3Abedrock-agentcore%3Aus-west-2%3A888602223428%3Aruntime%2Fhosted_agent_13sf6-cALnp38iZD/invocations"
|
||||
)
|
||||
_AGENTCORE_DELTAS: Final = (
|
||||
'data: {"event":{"contentBlockDelta":{"delta":{"text":"Machine learning lets "}}}}\n\n'
|
||||
'data: {"event":{"contentBlockDelta":{"delta":{"text":"computers learn from data."}}}}\n\n'
|
||||
)
|
||||
|
||||
|
||||
async def _stream_agentcore(body: str, sync_mode: bool) -> list[tuple[str | None, str | None]]:
|
||||
"""Stream `body` from the runtime as SSE and return each chunk's (content, finish_reason)."""
|
||||
kwargs: Final = {
|
||||
"model": "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/"
|
||||
"hosted_agent_13sf6-cALnp38iZD",
|
||||
"messages": [{"role": "user", "content": "Explain machine learning in simple terms"}],
|
||||
"aws_access_key_id": "AKIAAGENTCOREKEY",
|
||||
"aws_secret_access_key": "agentcore-secret",
|
||||
"stream": True,
|
||||
}
|
||||
if sync_mode:
|
||||
return [
|
||||
(chunk.choices[0].delta.content, chunk.choices[0].finish_reason)
|
||||
for chunk in litellm.completion(**kwargs, client=HTTPHandler())
|
||||
]
|
||||
return [
|
||||
(chunk.choices[0].delta.content, chunk.choices[0].finish_reason)
|
||||
async for chunk in await litellm.acompletion(**kwargs, client=AsyncHTTPHandler())
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ending",
|
||||
[
|
||||
pytest.param('data: {"event":{"messageStop":{"stopReason":"end_turn"}}}\n\n', id="message_stop"),
|
||||
pytest.param(
|
||||
'data: {"event":{"messageStop":{"stopReason":"end_turn"}}}\n\n'
|
||||
'data: {"event":{"metadata":{"usage":{"inputTokens":9,"outputTokens":7,"totalTokens":16}}}}\n\n'
|
||||
'data: {"message":{"role":"assistant","content":[{"text":"Machine learning lets computers learn from data."}]}}\n\n',
|
||||
id="strands_full",
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@pytest.mark.respx(assert_all_called=True)
|
||||
async def test_agentcore_stream_that_signals_its_end_completes(
|
||||
ending: str, sync_mode: bool, respx_mock: respx.Router, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
respx_mock.post(_AGENTCORE_RUNTIME_URL).mock(
|
||||
return_value=httpx.Response(
|
||||
200, content=(_AGENTCORE_DELTAS + ending).encode(), headers={"content-type": "text/event-stream"}
|
||||
)
|
||||
)
|
||||
|
||||
chunks: Final = await _stream_agentcore(_AGENTCORE_DELTAS + ending, sync_mode)
|
||||
|
||||
assert "".join(content for content, _ in chunks if content) == "Machine learning lets computers learn from data."
|
||||
assert [finish for _, finish in chunks if finish] == ["stop"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@pytest.mark.respx(assert_all_called=True)
|
||||
async def test_agentcore_stream_cut_before_it_signals_its_end_raises(
|
||||
sync_mode: bool, respx_mock: respx.Router, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
respx_mock.post(_AGENTCORE_RUNTIME_URL).mock(
|
||||
return_value=httpx.Response(
|
||||
200, content=_AGENTCORE_DELTAS.encode(), headers={"content-type": "text/event-stream"}
|
||||
)
|
||||
)
|
||||
|
||||
with pytest.raises(MidStreamFallbackError, match="Stream ended before a finish_reason was received"):
|
||||
await _stream_agentcore(_AGENTCORE_DELTAS, sync_mode)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue