Fix for Bedrock Parallel Tool Call User message issue for certain models #1656
Conversation
…o a single user message for Bedrock models. NOTE: Was previously breaking when the LLM was attempting to call multiple tools in one request, because the code was appending multiple user messages in a row. Which breaks the subsequent user/assistant messages required for nova and claude
DouweM
changed the title
There was a problem hiding this comment.
@pandersen-uncharted Thanks Paul for picking up the older PR! We're almost there, just a couple of nitpicky notes on the new code.
| last_message is None | ||
| or current_message['role'] != last_message['role'] | ||
| or current_message['role'] != 'user' | ||
| ): # Remove the "user" test to also allow merging sequential "assistant" messages |
There was a problem hiding this comment.
Should we also merge assistant messages? If not, can you please remove the comment? If so, the test? :)
There was a problem hiding this comment.
I also think this code will be a little easier to follow organized like this:
for current_message in bedrock_messages:
if last_message and ...:
# special behavior
continue
# standard behaviorThere was a problem hiding this comment.
Comment removed, code reworked.
|
|
||
| # Merge together sequential user messages. | ||
| processed_messages: list[MessageUnionTypeDef] = [] | ||
| last_message: dict[str, Any] | None = None |
There was a problem hiding this comment.
Can we type this as MessageUnionTypeDef instead of an arbitrary dict?
There was a problem hiding this comment.
Not easily. When I make that change I get the error below. Attempting to fix that error leads to other errors and in the end I don't know how deep the rabbit hole will be nor how complicated the code will get. This solution keeps the code compact and reasonably easy to follow.
Could not assign item in TypedDict
"list[ContentBlockUnionTypeDef]" is not assignable to "list[ContentBlockOutputTypeDef]"
Type parameter "_T@list" is invariant, but "ContentBlockUnionTypeDef" is not the same as "ContentBlockOutputTypeDef"
Consider switching from "list" to "Sequence" which is covariant
tests/models/test_bedrock.py
Outdated
| import datetime | ||
|
|
||
| from pydantic_ai.messages import ModelRequest, ToolReturnPart | ||
| from pydantic_ai.models.bedrock import BedrockConverseModel |
There was a problem hiding this comment.
Can we move these imports to the top of the file please?
tests/models/test_bedrock.py
Outdated
| ] | ||
| ) | ||
| # Call the mapping function directly | ||
| _, bedrock_messages = await model._map_messages([req]) # type: ignore[reportPrivateUsage] |
There was a problem hiding this comment.
Can we add a ModelResponse and some other parts as well, to ensure this correctly affects only subsequent User parts and not anything else?
Please see the comments on PR #1600 and #1645. This PR is based on seanw7:bedrock_multi_tool_support with DouweM's requested changes.