-
Notifications
You must be signed in to change notification settings - Fork 8.4k
fix(bedrock): preserve streaming tool call arguments at contentBlockStop #6150
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Vidit-Ostwal
merged 7 commits into
crewAIInc:main
from
kimnamu:fix/bedrock-streaming-tool-args
Sep 4, 2026
+208
−0
Merged
Changes from all commits
Commits
Show all changes
7 commits
Select commit
Hold shift + click to select a range
2a4ca0e
fix(bedrock): preserve streaming tool call arguments at contentBlockStop
kimnamu a3697df
fix(bedrock): coerce non-dict streaming tool input to empty dict
kimnamu ad6620d
Merge branch 'main' into fix/bedrock-streaming-tool-args
kimnamu 4979052
Merge branch 'main' into fix/bedrock-streaming-tool-args
kimnamu 1074788
Merge remote-tracking branch 'origin/main' into fix/bedrock-streaming…
kimnamu 8134c7c
Merge branch 'main' into fix/bedrock-streaming-tool-args
Vidit-Ostwal 620202e
Merge branch 'main' into fix/bedrock-streaming-tool-args
Vidit-Ostwal File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
182 changes: 182 additions & 0 deletions
182
lib/crewai/tests/llms/bedrock/test_bedrock_streaming_tool_args.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,182 @@ | ||
| """Regression tests for streaming Bedrock tool-call argument handling. | ||
|
|
||
| The streaming Converse handlers deliver tool input as a sequence of JSON | ||
| string deltas (``contentBlockDelta`` -> ``toolUse.input``) that are | ||
| accumulated separately from the tool-use block. These tests assert that the | ||
| accumulated input is folded back into the tool call at ``contentBlockStop``, | ||
| so executed tools receive their real arguments instead of an empty ``{}``. | ||
|
|
||
| This is the streaming counterpart of the non-streaming fix in #5415 | ||
| (issue #4972). | ||
| """ | ||
|
|
||
| import os | ||
| from unittest.mock import MagicMock, patch | ||
|
|
||
| import pytest | ||
|
|
||
| from crewai.llm import LLM | ||
| from crewai.llms.providers.bedrock.completion import BedrockCompletion | ||
|
|
||
|
|
||
| def _make_tool_use_stream() -> list[dict]: | ||
| """Synthetic Converse stream: a single tool call with JSON-chunked input.""" | ||
| # Tool input is delivered as two partial JSON string fragments that only | ||
| # form valid JSON once concatenated: '{"city":' + ' "Paris"}'. | ||
| chunk1 = '{"city":' | ||
| chunk2 = ' "Paris"}' | ||
| return [ | ||
| {"messageStart": {"role": "assistant"}}, | ||
| { | ||
| "contentBlockStart": { | ||
| "start": {"toolUse": {"toolUseId": "tool-1", "name": "get_weather"}}, | ||
| "contentBlockIndex": 0, | ||
| } | ||
| }, | ||
| {"contentBlockDelta": {"delta": {"toolUse": {"input": chunk1}}}}, | ||
| {"contentBlockDelta": {"delta": {"toolUse": {"input": chunk2}}}}, | ||
| {"contentBlockStop": {}}, | ||
| {"messageStop": {"stopReason": "tool_use"}}, | ||
| ] | ||
|
|
||
|
|
||
| def _make_non_dict_tool_use_stream() -> list[dict]: | ||
| """Synthetic Converse stream whose tool input is valid JSON but not an object. | ||
|
|
||
| ``json.loads`` succeeds here (returns a string), so the parsed value must | ||
| still be coerced to a dict before it reaches ``fn(**function_args)``. | ||
| """ | ||
| return [ | ||
| {"messageStart": {"role": "assistant"}}, | ||
| { | ||
| "contentBlockStart": { | ||
| "start": {"toolUse": {"toolUseId": "tool-1", "name": "get_weather"}}, | ||
| "contentBlockIndex": 0, | ||
| } | ||
| }, | ||
| {"contentBlockDelta": {"delta": {"toolUse": {"input": '"oops"'}}}}, | ||
| {"contentBlockStop": {}}, | ||
| {"messageStop": {"stopReason": "tool_use"}}, | ||
| ] | ||
|
|
||
|
|
||
| def _build_completion() -> BedrockCompletion: | ||
| """Build a BedrockCompletion with mocked AWS credentials/session.""" | ||
| with patch.dict( | ||
| os.environ, | ||
| { | ||
| "AWS_ACCESS_KEY_ID": "test-access-key", | ||
| "AWS_SECRET_ACCESS_KEY": "test-secret-key", | ||
| "AWS_DEFAULT_REGION": "us-east-1", | ||
| }, | ||
| ): | ||
| with patch("crewai.llms.providers.bedrock.completion.Session"): | ||
| llm = LLM(model="bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0") | ||
| assert isinstance(llm, BedrockCompletion) | ||
| return llm | ||
|
|
||
|
|
||
| def test_streaming_tool_call_preserves_arguments(): | ||
| """Sync streaming: function_args must carry the streamed tool input.""" | ||
| llm = _build_completion() | ||
|
|
||
| captured: dict = {} | ||
|
|
||
| def capture(function_args, **kwargs): | ||
| captured["args"] = function_args | ||
| return None # returning None stops the recursive _handle_converse call | ||
|
|
||
| mock_client = MagicMock() | ||
| mock_client.converse_stream.return_value = {"stream": _make_tool_use_stream()} | ||
|
|
||
| with ( | ||
| patch.object(llm, "_get_sync_client", return_value=mock_client), | ||
| patch.object(llm, "_handle_tool_execution", side_effect=capture), | ||
| ): | ||
| llm._handle_streaming_converse( | ||
| messages=[{"role": "user", "content": "weather in Paris?"}], | ||
| body={}, | ||
| available_functions={"get_weather": lambda **kw: "sunny"}, | ||
| ) | ||
|
|
||
| assert captured["args"] == {"city": "Paris"} | ||
|
|
||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_async_streaming_tool_call_preserves_arguments(): | ||
| """Async streaming: function_args must carry the streamed tool input.""" | ||
| llm = _build_completion() | ||
|
|
||
| class _AsyncStream: | ||
| def __init__(self, events): | ||
| self._events = events | ||
|
|
||
| def __aiter__(self): | ||
| self._it = iter(self._events) | ||
| return self | ||
|
|
||
| async def __anext__(self): | ||
| try: | ||
| return next(self._it) | ||
| except StopIteration: | ||
| raise StopAsyncIteration | ||
|
|
||
| async def _converse_stream(**kwargs): | ||
| return {"stream": _AsyncStream(_make_tool_use_stream())} | ||
|
|
||
| mock_async_client = MagicMock() | ||
| mock_async_client.converse_stream = _converse_stream | ||
|
|
||
| async def _ensure(*args, **kwargs): | ||
| return mock_async_client | ||
|
|
||
| captured: dict = {} | ||
|
|
||
| def capture(function_args, **kwargs): | ||
| captured["args"] = function_args | ||
| return None | ||
|
|
||
| with ( | ||
| patch.object(llm, "_ensure_async_client", side_effect=_ensure), | ||
| patch.object(llm, "_handle_tool_execution", side_effect=capture), | ||
| ): | ||
| await llm._ahandle_streaming_converse( | ||
| messages=[{"role": "user", "content": "weather in Paris?"}], | ||
| body={}, | ||
| available_functions={"get_weather": lambda **kw: "sunny"}, | ||
| ) | ||
|
|
||
| assert captured["args"] == {"city": "Paris"} | ||
|
|
||
|
|
||
| def test_streaming_non_dict_tool_input_coerced_to_empty_dict(): | ||
| """Valid-but-non-object JSON input must be coerced to ``{}``. | ||
|
|
||
| ``json.loads('"oops"')`` returns a string; passing it on as | ||
| ``fn(**function_args)`` would raise ``TypeError``. The handler must | ||
| guard against this and fall back to an empty dict. | ||
| """ | ||
| llm = _build_completion() | ||
|
|
||
| captured: dict = {} | ||
|
|
||
| def capture(function_args, **kwargs): | ||
| captured["args"] = function_args | ||
| return None | ||
|
|
||
| mock_client = MagicMock() | ||
| mock_client.converse_stream.return_value = { | ||
| "stream": _make_non_dict_tool_use_stream() | ||
| } | ||
|
|
||
| with ( | ||
| patch.object(llm, "_get_sync_client", return_value=mock_client), | ||
| patch.object(llm, "_handle_tool_execution", side_effect=capture), | ||
| ): | ||
| llm._handle_streaming_converse( | ||
| messages=[{"role": "user", "content": "weather in Paris?"}], | ||
| body={}, | ||
| available_functions={"get_weather": lambda **kw: "sunny"}, | ||
| ) | ||
|
|
||
| assert captured["args"] == {} |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.