Repository navigation
fix(bedrock): handle toolResult attachment co-location per-model via bedrock_tool_result_colocatable_content - #6098
Conversation
|
There's some related logic here: pydantic-ai/pydantic_ai_slim/pydantic_ai/models/bedrock.py Lines 1297 to 1303 in e2bdafb Worth checking to see if we can make it more consistent / fix it in one place. |
|
I have addressed this! I've updated the logic at the end of |
573a32e to
79b7ab4
Compare
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (5)
🔗 Linked repositories identifiedCodeRabbit considers these linked repositories for cross-repo context during reviews:
✅ Files skipped from review due to trivial changes (1)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughThis PR updates Bedrock message mapping so ChangesRelated issues: Sequence Diagram(s)sequenceDiagram
participant _map_messages
participant BedrockModelProfile
participant BedrockMessages
_map_messages->>BedrockModelProfile: read bedrock_tool_result_colocatable_content
_map_messages->>_map_messages: inspect following user turn content kinds
alt toolResult + disallowed content kind
_map_messages->>BedrockMessages: insert synthetic assistant "."
else allowed content kind
_map_messages->>BedrockMessages: merge into prior user message
end
Poem: 🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
🧹 Nitpick comments (2)
pydantic_ai_slim/pydantic_ai/models/bedrock.py (2)
1192-1207: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConfirm this block is truly needed given it's marked unreachable.
This normalization pass is marked
# pragma: no cover(per the commit, intentionally, since_map_user_promptand the merge logic already guarantee the document+text invariant per-message before this runs). Keeping unreachable defensive code is a reasonable safety net, but worth a second look: if it's truly dead given current invariants, consider removing it to avoid maintaining unreachable logic; if it's meant as future-proofing against new merge paths, a short comment explaining why would help future readers.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pydantic_ai_slim/pydantic_ai/models/bedrock.py` around lines 1192 - 1207, The normalization in bedrock.py that injects a fallback text block for user messages with documents appears unreachable given the existing _map_user_prompt and merge logic, so either remove the dead defensive branch from the processed_messages loop or keep it only if needed for future-proofing and add a short explanatory comment near the role/content check to justify why the invariant is not guaranteed elsewhere.
1128-1148: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate detection logic could be extracted into a helper.
The
any(... for k in ('image', 'document', 'video'))/any('toolResult' in block ...)checks are repeated four times here (and similar logic reappears at Line 1196 and Line 1330-1331). Extracting a small helper (e.g._has_block_kind(content, *keys)) would reduce duplication and improve readability.As per coding guidelines, "Consolidate methods with duplicated logic using helpers, delegation, or type overloads."♻️ Proposed refactor
+ `@staticmethod` + def _has_block_kind(content: list[Any], *keys: str) -> bool: + return any(any(k in block for k in keys) for block in content) + for current_message in bedrock_messages: if ( last_message is not None and current_message['role'] == last_message['role'] and current_message['role'] == 'user' ): - last_has_tool_result = any('toolResult' in block for block in last_message['content']) - current_has_attachment = any( - any(k in block for k in ('image', 'document', 'video')) for block in current_message['content'] - ) - last_has_attachment = any( - any(k in block for k in ('image', 'document', 'video')) for block in last_message['content'] - ) - current_has_tool_result = any('toolResult' in block for block in current_message['content']) + last_has_tool_result = self._has_block_kind(last_message['content'], 'toolResult') + current_has_attachment = self._has_block_kind(current_message['content'], 'image', 'document', 'video') + last_has_attachment = self._has_block_kind(last_message['content'], 'image', 'document', 'video') + current_has_tool_result = self._has_block_kind(current_message['content'], 'toolResult')🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pydantic_ai_slim/pydantic_ai/models/bedrock.py` around lines 1128 - 1148, The duplicate block-type detection in bedrock message processing should be consolidated into a helper to reduce repetition. Extract the repeated `any(...)` content checks used for `toolResult` and attachment kinds into a small method or local helper, such as one used by the message-merging logic in `bedrock.py`, and reuse it wherever the same pattern appears in the message processing flow. Keep the merge decision in the existing code path, but delegate the repeated detection work to the helper so the `processed_messages` assembly stays readable and easier to maintain.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@pydantic_ai_slim/pydantic_ai/models/bedrock.py`:
- Around line 1192-1207: The normalization in bedrock.py that injects a fallback
text block for user messages with documents appears unreachable given the
existing _map_user_prompt and merge logic, so either remove the dead defensive
branch from the processed_messages loop or keep it only if needed for
future-proofing and add a short explanatory comment near the role/content check
to justify why the invariant is not guaranteed elsewhere.
- Around line 1128-1148: The duplicate block-type detection in bedrock message
processing should be consolidated into a helper to reduce repetition. Extract
the repeated `any(...)` content checks used for `toolResult` and attachment
kinds into a small method or local helper, such as one used by the
message-merging logic in `bedrock.py`, and reuse it wherever the same pattern
appears in the message processing flow. Keep the merge decision in the existing
code path, but delegate the repeated detection work to the helper so the
`processed_messages` assembly stays readable and easier to maintain.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 6b3ba337-4fc2-4c00-9b40-1c03c92432fc
📒 Files selected for processing (1)
pydantic_ai_slim/pydantic_ai/models/bedrock.py
🔗 Linked repositories identified
CodeRabbit considers these linked repositories for cross-repo context during reviews:
pydantic/logfire(manual)pydantic/pydantic-ai(manual)pydantic/pydantic(auto-detected)
…l_result and media attachment turns
…ssistant message Bedrock rejects a user message that co-locates a `toolResult` block with an attachment block, and it also requires alternating roles, so a plain split into two `user` turns is not enough. Insert a minimal assistant turn between them, using `.` because Anthropic on Bedrock rejects whitespace-only text. - drop the redundant post-merge document-text guard; `_map_user_prompt` already guarantees a text block alongside any document - add a VCR regression test proving Converse accepts the reshaped history - keep a unit test pinning the exact mapped message sequence Refs pydantic#6081
a376ac7 to
2f1f287
Compare
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
pydantic_ai_slim/pydantic_ai/models/bedrock.py (1)
1003-1035: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winSplit unsupported tool-return attachments into a separate user turn The
elsebranch still appendssibling_contentto the sameusermessage astoolResult, so a tool return with avideo(or any other unsupported media kind) can still produce the forbiddentoolResult+ attachment shape and hit the same Bedrock validation failure. Emit a second user message here instead.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pydantic_ai_slim/pydantic_ai/models/bedrock.py` around lines 1003 - 1035, The Bedrock tool-return assembly in the branch handling is_multi_modal_content still combines the toolResult payload and unsupported attachment sibling_content into one user message, which can trigger the same invalid shape. Update this path in bedrock.py so the toolResult message is appended first, then emit the attachment content as a separate user turn instead of extending the same user_content list; keep the logic around _map_file_to_content_block, tool_result_content, and sibling_content intact.
🧹 Nitpick comments (1)
tests/models/test_bedrock.py (1)
5182-5211: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueSolid regression coverage; only
documentattachment type exercised.The fix's merge-guard checks
image/document/videouniformly, but onlydocumentis covered by these new tests. Given the logic is generic and already unit-pinned, this is a minor coverage gap rather than a real risk.As per coding guidelines, "Provide comprehensive tests covering 100% of code paths, favoring integration tests and real requests."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/models/test_bedrock.py` around lines 5182 - 5211, Add regression coverage for the merge-guard in the Bedrock history reshaping logic by exercising the other attachment types handled alongside `document`, especially `image` and `video`. Update the Bedrock test suite around `test_bedrock_tool_result_followed_by_attachment_accepted` and the related history-mapping helpers so the same tool-result-followed-by-attachment scenario is verified for all supported attachment variants, not just `document`.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@pydantic_ai_slim/pydantic_ai/models/bedrock.py`:
- Around line 1003-1035: The Bedrock tool-return assembly in the branch handling
is_multi_modal_content still combines the toolResult payload and unsupported
attachment sibling_content into one user message, which can trigger the same
invalid shape. Update this path in bedrock.py so the toolResult message is
appended first, then emit the attachment content as a separate user turn instead
of extending the same user_content list; keep the logic around
_map_file_to_content_block, tool_result_content, and sibling_content intact.
---
Nitpick comments:
In `@tests/models/test_bedrock.py`:
- Around line 5182-5211: Add regression coverage for the merge-guard in the
Bedrock history reshaping logic by exercising the other attachment types handled
alongside `document`, especially `image` and `video`. Update the Bedrock test
suite around `test_bedrock_tool_result_followed_by_attachment_accepted` and the
related history-mapping helpers so the same tool-result-followed-by-attachment
scenario is verified for all supported attachment variants, not just `document`.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: d5230da7-77d0-433f-8320-74ae9fddfbac
📒 Files selected for processing (3)
pydantic_ai_slim/pydantic_ai/models/bedrock.pytests/models/cassettes/test_bedrock/test_bedrock_tool_result_followed_by_attachment_accepted.yamltests/models/test_bedrock.py
🔗 Linked repositories identified
CodeRabbit considers these linked repositories for cross-repo context during reviews:
pydantic/logfire(manual)pydantic/pydantic-ai(manual)pydantic/pydantic(auto-detected)
…-result-attachments-conflict
The previous fix split every `toolResult`-then-attachment turn with a synthetic
assistant message for all models, but live Converse testing shows the constraint
is per-model: Anthropic rejects only documents/video beside a `toolResult` (text
and images are fine), Llama and Mistral reject any content sharing the turn, and
Nova has no constraint at all — so the blanket split needlessly polluted Nova
(and every permissive model) with a dummy turn.
Introduce a `bedrock_tool_result_colocatable_content` profile flag listing the
content-block kinds a model accepts alongside a `toolResult`. The merge pass
splits (separating the turns with a synthetic assistant message, since Bedrock
re-merges consecutive same-role turns) only when the combined content isn't
co-locatable for that model. Default is permissive, preserving behavior for
uncovered families. Set to `{text, image}` for Anthropic and `frozenset()` for
Meta and Mistral (the latter reproduce the same defect for plain text, a broader
form of pydantic#6081).
Replace the PDF VCR recording (which Anthropic accepts even without the fix, so
it proved nothing) with `text/plain` recordings that actually reproduce the
rejection, plus Nova (co-located, no split) and Llama (text isolated) cassettes.
Refs pydantic#6081
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/models/test_bedrock.py (1)
5169-5227: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd VCR-exemption rationale to the Nova/Llama unit tests.
test_bedrock_anthropic_tool_result_and_document_alternateexplains why it's a unit test rather than a VCR test (cassette matcher isn't sensitive to request body).test_bedrock_nova_tool_result_and_document_mergeandtest_bedrock_llama_tool_result_isolated_from_textpin the same kind of internal request-payload shape but their docstrings omit this rationale.As per path instructions, "Each unit test should explain why it is not (or cannot be) a VCR test, especially when it pins internal behavior, pre-request guards, or request payload shape that cassettes would not reliably catch."
Source: Path instructions
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 0d65779c-6227-41f3-a628-f0ae36d23f96
📒 Files selected for processing (7)
pydantic_ai_slim/pydantic_ai/models/bedrock.pypydantic_ai_slim/pydantic_ai/providers/bedrock.pytests/models/cassettes/test_bedrock/test_bedrock_anthropic_tool_result_followed_by_document_accepted.yamltests/models/cassettes/test_bedrock/test_bedrock_llama_tool_result_followed_by_text_accepted.yamltests/models/cassettes/test_bedrock/test_bedrock_nova_tool_result_followed_by_document_accepted.yamltests/models/test_bedrock.pytests/providers/test_bedrock.py
🔗 Linked repositories identified
CodeRabbit considers these linked repositories for cross-repo context during reviews:
pydantic/logfire(manual)pydantic/pydantic-ai(manual)pydantic/pydantic(auto-detected)
✅ Files skipped from review due to trivial changes (2)
- tests/models/cassettes/test_bedrock/test_bedrock_anthropic_tool_result_followed_by_document_accepted.yaml
- tests/models/cassettes/test_bedrock/test_bedrock_llama_tool_result_followed_by_text_accepted.yaml
🚧 Files skipped from review as they are similar to previous changes (1)
- pydantic_ai_slim/pydantic_ai/models/bedrock.py
…result_colocatable_content`
The new profile flag appears in the normalized profile snapshots for Anthropic
(`{text, image}`), Meta, and Mistral (`frozenset()`) on Bedrock.
toolResult attachment co-location per-model via bedrock_tool_result_colocatable_content
| """Default: `False`.""" | ||
| bedrock_supported_media_kinds_in_tool_returns: frozenset[str] | ||
| """Default: `frozenset({'image'})`.""" | ||
| bedrock_tool_result_colocatable_content: frozenset[str] |
There was a problem hiding this comment.
Let's use a Literal union with the actual supported values, instead of str. That also goes for supported_media_kinds....
…eral` unions Address review: use `frozenset[Literal[...]]` instead of `frozenset[str]` for `bedrock_tool_result_colocatable_content` and `bedrock_supported_media_kinds_in_tool_returns`.
…flict Resolve additive conflicts from pydantic#6255 (leading-assistant-message flag): keep both that flag and this PR's `bedrock_tool_result_colocatable_content` in `BedrockModelProfile`, the Anthropic profile, and the resolution-matrix / test-bedrock snapshots.
…flict Resolve conflicts from pydantic#6254 (media-kinds fix): union the two `bedrock_meta_model_profile` definitions (both `bedrock_tool_result_colocatable_content` and `bedrock_supported_media_kinds_in_tool_returns`), and keep both flags' values in the Mistral profile and in the resolution-matrix / provider-test assertions.
Closes #6081
Problem
On Bedrock, pydantic-ai merges consecutive user turns into one Converse message. When a tool-result turn is followed by a user turn carrying an attachment, they merge into one user message co-locating a
toolResultblock with the attachment — which some models reject with a misleadingValidationExceptionabout a missingtool_result.The original fix split every such turn with a synthetic assistant message for all models. Live testing against Bedrock (
us-east-1) shows the constraint is per-model, so a blanket split is both incomplete (misses models with a stricter rule) and over-broad (needlessly pollutes models with no rule).Evidence (live Converse, what may share a user message with a
toolResult)toolResult, but accepts text and images. (PDF happens to pass, but the pdf-vs-other-document split is undocumented and brittle, so we conservatively treat all documents as non-co-locatable.)toolResultturn, even plain text — a broader form of the same bug that the attachment-only guard missed.Mechanism notes, also verified: a bare split into two consecutive user turns does not work (Bedrock re-merges same-role turns), so an alternating assistant turn is required; the separator must be non-whitespace (
.works everywhere); block order within the message is irrelevant. AWS does not document thetoolResult+document constraint — it's model behavior — which is why this is driven by empirical per-family testing.Fix
Add a
bedrock_tool_result_colocatable_contentmodel-profile flag listing the content-block kinds a model accepts alongside atoolResult. The merge pass splits the turns (separating them with a synthetic.assistant message) only when the combined content isn't co-locatable for that model; otherwise it merges as before.{'text', 'image'}frozenset()toolResultisolated from all other contentThe default is permissive, so no behavior change for uncovered families.
Tests
text/plaindocument (the original PDF recording is accepted even without the fix, so it proved nothing — a text document actually reproduces the rejection), Nova co-located with no dummy turn, and Llama with text isolated.Follow-ups (separate PRs)
Live testing surfaced two adjacent, distinct defects that are handled in their own PRs rather than widening this one:
bedrock_supported_media_kinds_in_tool_returnshas wrong per-family values (Mistral/Writer acceptdocumentbut rejectimageinside tool results; Qwen accepts neither; Nova rejects binary documents).🤖 Generated with Claude Code