fix: gemini thought signature handling #1142
Conversation
📝 WalkthroughSummary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings. WalkthroughAdds Gemini thought-signature handling across ingestion, conversion, and streaming: base64 encode/decode paths, custom Part JSON unmarshaling, expanded content→message conversions with consumed-index tracking and signature propagation, a Bedrock schema extraction tweak, test additions/renumbering for Google/Gemini, and version/changelog bumps. Changes
Sequence Diagram(s)mermaid Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes
Poem
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
🧪 Test Suite AvailableThis PR can be tested by a repository admin. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
tests/integrations/tests/test_google.py (3)
117-117: Consider handling missing API key explicitly.
os.environ.get("GEMINI_API_KEY")returnsNoneif the environment variable is not set. This could lead to unclear errors downstream when the client attempts authentication.Proposed fix
- api_key = os.environ.get("GEMINI_API_KEY") + api_key = os.environ.get("GEMINI_API_KEY") + if not api_key: + raise ValueError("GEMINI_API_KEY environment variable is required")
132-133: Hardcoded timeout ignores the configured value.The code checks if a timeout is configured but then uses a hardcoded
30000instead of the actual config value. This seems unintentional.Proposed fix
if api_config.get("timeout"): - http_options_kwargs["timeout"] = 30000 + http_options_kwargs["timeout"] = api_config.get("timeout")
1447-1447: f-string without placeholders.This f-string has no interpolation and can be a regular string.
Proposed fix
- print(f"\n=== Token Usage ===") + print("\n=== Token Usage ===")
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
core/providers/bedrock/utils.go(1 hunks)core/providers/gemini/chat.go(1 hunks)core/providers/gemini/responses.go(8 hunks)core/providers/gemini/types.go(1 hunks)core/providers/gemini/utils.go(2 hunks)tests/integrations/tests/test_google.py(14 hunks)
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
core/providers/gemini/utils.gocore/providers/gemini/types.gocore/providers/gemini/chat.gocore/providers/gemini/responses.gotests/integrations/tests/test_google.pycore/providers/bedrock/utils.go
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
core/providers/gemini/utils.gocore/providers/gemini/types.gocore/providers/gemini/chat.gocore/providers/gemini/responses.gocore/providers/bedrock/utils.go
📚 Learning: 2025-12-19T09:26:54.961Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/providers/utils/utils.go:1050-1051
Timestamp: 2025-12-19T09:26:54.961Z
Learning: Update streaming end-marker handling so HuggingFace is treated as a non-[DONE] provider for backends that do not emit a DONE marker (e.g., meta llama on novita). In core/providers/utils/utils.go, adjust ProviderSendsDoneMarker() (or related logic) to detect providers that may not emit DONE and avoid relying on DONE as the sole end signal. Add tests to cover both DONE-emitting and non-DONE backends, with clear documentation in code comments explaining the rationale and any fallback behavior.
Applied to files:
core/providers/gemini/utils.gocore/providers/gemini/types.gocore/providers/gemini/chat.gocore/providers/gemini/responses.gocore/providers/bedrock/utils.go
🧬 Code graph analysis (3)
core/providers/gemini/chat.go (1)
core/schemas/chatcompletions.go (2)
ChatAssistantMessageToolCall(710-716)ChatToolChoiceTypeFunction(382-382)
core/providers/gemini/responses.go (3)
core/providers/gemini/types.go (5)
Type(782-782)Role(17-17)FunctionCall(1197-1207)Content(977-985)Part(991-1015)core/schemas/responses.go (6)
ResponsesReasoning(741-744)ResponsesMessage(321-334)ResponsesToolMessage(472-492)ResponsesToolMessageOutputStruct(541-545)ResponsesMessageContent(346-351)ResponsesMessageContentBlock(406-421)core/schemas/utils.go (1)
Ptr(16-18)
tests/integrations/tests/test_google.py (3)
tests/integrations/tests/utils/common.py (1)
skip_if_no_api_key(1813-1824)core/providers/gemini/types.go (5)
Tool(604-631)Content(977-985)Part(991-1015)FunctionCall(1197-1207)FunctionResponse(1210-1232)tests/integrations/tests/utils/parametrize.py (1)
get_cross_provider_params_for_scenario(12-47)
🪛 Ruff (0.14.8)
tests/integrations/tests/test_google.py
1238-1238: Unused method argument: test_config
(ARG002)
1447-1447: f-string without any placeholders
Remove extraneous f prefix
(F541)
1455-1455: Unused method argument: test_config
(ARG002)
1537-1537: Unused method argument: test_config
(ARG002)
1591-1591: Unused method argument: test_config
(ARG002)
1638-1638: Unused method argument: test_config
(ARG002)
1682-1682: Unused method argument: test_config
(ARG002)
1727-1727: Unused method argument: test_config
(ARG002)
1787-1787: Unused method argument: test_config
(ARG002)
1821-1821: Unused method argument: test_config
(ARG002)
1841-1841: Unused method argument: test_config
(ARG002)
1881-1881: Unused method argument: test_config
(ARG002)
1924-1924: Unused method argument: test_config
(ARG002)
🔇 Additional comments (12)
core/providers/gemini/chat.go (1)
114-130: LGTM! Clean extraction of tool call construction with thought signature handling.The refactoring to build the tool call in a temporary variable before appending allows for clean conditional attachment of the
ThoughtSignature. The base64 encoding and nested structure undergoogle.thought_signatureis consistent with the decoding logic inutils.go.core/providers/gemini/utils.go (1)
613-617: LGTM! Proper base64 decoding of thought signatures.The decoding logic correctly reverses the encoding performed in
chat.go. Silently skipping on decode failure is appropriate here since a malformed signature shouldn't break the entire request. The structure navigation pathgoogle.thought_signatureis consistent with the encoding side.core/providers/gemini/types.go (1)
1017-1066: LGTM! Robust URL-safe base64 handling for thought signatures.The custom
UnmarshalJSONcorrectly handles thought signatures that may arrive as URL-safe base64 from the Google GenAI SDK. The conversion from URL-safe to standard base64 (replacing-with+and_with/) and the padding logic are both correct. Returning an error on decode failure is appropriate for unmarshaling since it indicates malformed input.core/providers/gemini/responses.go (5)
174-176: LGTM! Well-structured look-ahead pattern for thought signature consumption.The
consumedIndicesmap elegantly solves the problem of pairing function call messages with their subsequent reasoning messages containing thought signatures. The bounds checking at line 240 is correct, and the skip logic at lines 288-290 properly prevents double-processing.Also applies to: 238-251, 287-290
829-842: LGTM! Properly propagates thought signatures in streaming responses.The streaming path correctly base64-encodes the
ThoughtSignatureand attaches it to the stream response'sSignaturefield, maintaining consistency with the non-streaming conversion paths.
1370-1538: LGTM! Comprehensive restructuring with proper thought signature propagation.The rewritten
convertGeminiContentsToResponsesMessagescorrectly:
- Tracks function call IDs by name for correlation with responses
- Emits separate reasoning messages for thought signatures attached to function calls (lines 1412-1424)
- Attaches signatures directly to text content blocks when present (lines 1502-1506)
- Handles the new InlineData and FileData part types
The pattern of emitting a reasoning message immediately after a function call message aligns with the consumption logic in
ToGeminiResponsesResponse.
1728-1732: LGTM! Consistent signature attachment on output text blocks.Properly attaches the base64-encoded thought signature to the text content block that was just created, mirroring the pattern used in
convertGeminiContentsToResponsesMessages.
2387-2396: LGTM! Correctly decodes signatures when converting content blocks to Gemini parts.The signature decoding in
convertContentBlockToGeminiPartproperly reverses the encoding performed elsewhere, and the silent failure on decode error is appropriate since the part can still be used without the signature.tests/integrations/tests/test_google.py (4)
1291-1301: LGTM on the multi-turn thought signature test structure.The test comprehensively validates Gemini 3 Pro thought signature handling across multiple turns:
- Turn 1: Basic thinking with thought signatures
- Turn 2: Tool calls with thought signatures
- Turn 3: Multi-turn conversation preserving thought signatures
- Turn 4: Parallel tool calls
The conversation construction at lines 1358-1386 correctly propagates the
thought_signaturefrom the previous response to maintain context.
1536-1589: Test renumbering looks correct.The file API tests have been correctly renumbered from test_26-29 to test_30-33 to accommodate the new test cases. The test logic and cleanup handling remain proper.
1726-1789: Batch API test renumbering is correct.The batch API tests have been correctly renumbered to maintain sequential ordering after the new test cases were inserted.
1923-1925: E2E batch test renumbering is correct.The end-to-end batch workflow test has been renumbered appropriately and maintains its comprehensive workflow validation: upload → create batch → poll → verify in list → cleanup.
TejasGhatte
force-pushed
the
12-20-fix_gemini_thought_signature_handling
branch
from
ff6b422 to
12987f9
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
core/providers/gemini/responses.go (1)
444-448: Potential double-encoding of thought signatures in streaming reverse pathIn
ToGeminiResponsesStreamResponse, theOutputItemAddedcase currently does:if bifrostResp.Item != nil && bifrostResp.Item.ResponsesReasoning != nil && bifrostResp.Item.EncryptedContent != nil { candidate.Content.Parts = append(candidate.Content.Parts, &Part{ ThoughtSignature: []byte(*bifrostResp.Item.ResponsesReasoning.EncryptedContent), }) }But
EncryptedContentis already a base64 string produced earlier (e.g., viabase64.StdEncoding.EncodeToString(part.ThoughtSignature)). Treating that string as raw bytes means that when thisGenerateContentResponseis later marshaled,ThoughtSignaturewill be base64‑encoded again, effectively double-encoding the signature.To keep the streaming reverse path consistent with the non-streaming conversions and with
convertContentBlockToGeminiPart, you likely want to base64‑decode here instead:- ThoughtSignature: []byte(*bifrostResp.Item.ResponsesReasoning.EncryptedContent), + ThoughtSignature: decoded,after a
base64.StdEncoding.DecodeStringwith appropriate error handling.This isn’t introduced by the new changes, but it becomes more visible now that thought signatures are first-class citizens in the streaming APIs.
🧹 Nitpick comments (12)
plugins/maxim/changelog.md (1)
1-1: Changelog entry should describe the actual improvements, not just version bumps.The changelog entry currently only documents version increments with a "chore:" prefix. However, the PR is classified as a "Bug fix and feature" and includes significant improvements (Gemini 3 Pro thought signature support, multi-turn conversation handling, Bedrock JSON schema fixes). Consider updating the entry to briefly describe what was actually fixed/added, and verify the prefix aligns with the PR type. For reference, the PR objectives mention:
- Adds support for Gemini 3 Pro thought signatures in tool calls and responses
- Improves handling of structured outputs across providers
- Fixes JSON schema handling in Bedrock structured output tools
- Improves handling of thought signatures in multi-turn conversations
tests/integrations/tests/test_google.py (3)
1445-1449: Resolve Ruff F541: f-string without placeholders
print(f"\n=== Token Usage ===")is flagged by Ruff (F541) because it doesn’t interpolate anything. You can silence this by dropping thefprefix:- print(f"\n=== Token Usage ===") + print("\n=== Token Usage ===")
1537-1995: Renumbered Files/Batch tests: behavior unchanged, buttest_configis unusedThe Files API and Batch tests (now
test_30_file_uploadthroughtest_39_batch_e2e_file_api) appear to be pure renames/renumbering plus docstring updates; the core logic (upload/list/retrieve/delete/batch create/list/retrieve/cancel/E2E) is unchanged and still looks correct, including cleanup infinallyblocks.Static analysis (Ruff ARG002) correctly points out that
test_configis unused in these parametrized tests. If you don’t plan to use it, consider dropping the parameter or renaming it to_test_configto keep the fixture wiring explicit while silencing the warning:- def test_30_file_upload(self, test_config, provider, model): + def test_30_file_upload(self, _test_config, provider, model):Same applies to the other file/batch tests where
test_configis unused.
1238-1924: Unusedtest_configfixture across several tests (Ruff ARG002)Beyond the Files/Batch tests,
test_configis also injected but unused in:
test_28_gemini_3_pro_thought_signatures_multi_turntest_29_structured_output_with_thinking- And other parametrized tests highlighted by Ruff (ARG002) in this file.
If these are only present for future expansion or interface consistency, you can either:
- Remove the parameter where not needed, or
- Rename it to
_test_configto make the intent explicit and silence the linter.This is cosmetic but will keep the test module clean of ARG002 noise.
core/providers/gemini/chat.go (1)
96-131: Thought signature propagation on tool calls looks correctBuilding the
toolCallstruct first and then attaching a base64-encodedThoughtSignatureintoExtraContent["google"]["thought_signature"]cleanly exposes Gemini’s thought signatures on tool calls while keeping existing fields (Index,Type,ID,Function) intact. This aligns with howReasoningDetails.Signatureis encoded below and withExtraContent’s purpose as a provider-specific extension field.Minor note: here you use
ChatToolChoiceTypeFunctionforType, while the streaming path usesChatToolTypeFunction; both stringify to"function", so behavior is consistent.core/providers/gemini/types.go (1)
1017-1066: CustomPart.UnmarshalJSONcorrectly decodes URL-safe base64 thought signaturesThe alias-based Unmarshal:
- Preserves all existing fields (
VideoMetadata,Thought,InlineData,FileData,CodeExecutionResult,ExecutableCode,FunctionCall,FunctionResponse,Text).- Treats
thoughtSignatureas a string, converts from URL-safe to standard base64 with padding, and decodes into the[]byteThoughtSignaturefield.- Fails fast with a clear error if decoding breaks, which is appropriate for malformed provider payloads.
Given
Blob.UnmarshalJSONperforms the same URL-safe→standard base64 conversion, you might consider extracting a small helper for the “urlsafe+padding → []byte” path and using%win the error to keep stack context, but that’s purely a readability/DRY improvement.core/providers/gemini/responses.go (6)
174-311: Pairing reasoning messages with tool calls to recover thought signatures looks correctThe additions in
ToGeminiResponsesResponse:
- Track which
ResponsesMessageindices have been “consumed” as carriers of encrypted thought signatures.- When converting a
FunctionCallmessage back into a GeminiPart, look ahead one message for aReasoningmessage withEncryptedContent, base64‑decode it intoPart.ThoughtSignature, and mark that index as consumed.- Skip reasoning messages whose indices are in
consumedIndicesto avoid producing duplicate thought-signature-only parts.- Still handle standalone reasoning messages (summary text and encrypted_content) when not associated with a tool call.
This gives you a clean round-trip for Gemini 3 Pro tool-call thought signatures while preserving backward compatibility for other reasoning messages.
772-846: Streaming text deltas now correctly surface thought signatures viaSignatureIn
processGeminiTextPart, emittingoutput_text.deltaevents now:
- Reuses a single text item and content index (0) via
TextOutputIndexandTextBuffer.- Attaches a
Signaturefield whenpart.ThoughtSignatureis non-empty, encoding it as standard base64.This is a good way to stream encrypted reasoning signatures alongside text without changing the existing delta structure. Downstream consumers of
BifrostResponsesStreamResponsecan now associate a signature with each text span.Just be aware that signatures are only on the delta events, not on the final
output_text.done; make sure any consumers expecting signatures read them from the delta path.
1369-1537: Content→Responses conversion handles per-part messages and thought signatures coherentlyThe reworked
convertGeminiContentsToResponsesMessagesnow:
- Normalizes Gemini roles to Responses roles (
"model"→assistant,"user"→user, defaulting unknown roles to user).- Emits one
ResponsesMessageperPart, making function calls, function responses, reasoning text, regular text, inline data, and file data each first-class items.- For function calls:
- Serializes
Argsto JSON, derives acall_id(ID or name), and tracks it infunctionCallIDsfor later matching with responses.- Emits a
FunctionCallmessage withCallID,Name, andArguments.- If
ThoughtSignatureis present on the part, also emits a siblingReasoningmessage whoseEncryptedContentis the base64-encoded signature (so it can be carried in the Responses schema and reattached later).- For function responses:
- Reconstructs
CallIDby preferringFunctionResponse.ID, falling back to the last-seen call ID for the same function name, then finally to the function name itself.- Serializes the
Responsemap via theoutputkey (or full JSON) intoResponsesToolMessage.Output.- For normal text:
- Uses
output_textvsinput_textblock types based oncontent.Role.- If a
ThoughtSignatureexists, pushes it into the block’sSignaturefield (to be decoded back intoPart.ThoughtSignaturelater).This design gives you enough structure to rehydrate full
Contentobjects on the Gemini side, including encrypted thoughts, while keeping the Responses representation granular and tool-friendly. ThefunctionCallIDsmap is only used as a fallback when response IDs are absent; if you ever have multiple concurrent tool calls with the same name and missing IDs, their responses could be mis-associated, but that’s a niche edge case.
1685-1777: Candidate→Responses conversion now preserves thought signatures for both text and tool callsIn
convertGeminiCandidatesToResponsesOutput:
- Regular assistant text (
part.Textwith nopart.Thought) is converted into anoutput_textcontent block, and whenpart.ThoughtSignatureis present, its base64 encoding is stored inblock.Signature. This aligns with the reverse path inconvertContentBlockToGeminiPart, which decodesSignatureback intoPart.ThoughtSignature.- For
FunctionCallparts, you:
- Serialize
Argsto JSON forResponsesToolMessage.Arguments.- Derive a
CallID(preferringID, falling back toName).- Emit a
FunctionCallmessage withCallIDandName.- If
ThoughtSignatureexists, emit a separateReasoningmessage whoseEncryptedContentcarries the base64-encoded signature.Together with the lookahead logic in
ToGeminiResponsesResponse, this lets you:
- Represent tool-call signatures cleanly in the Responses API.
- Reconstruct them on the Gemini side without inventing a new top-level field on the tool-call message itself.
The switch structure ensures that the “catch-all
case part.ThoughtSignature != nil” is only used for signature-only parts, not for function calls or text that already handle signatures explicitly.
2324-2368: Function-call output conversion now correctly prefersOutputand falls back toContentStrThe
ResponsesMessageTypeFunctionCallOutputbranch inconvertResponsesMessagesToGeminiContents:
- Builds
responseMapfrommsg.ResponsesToolMessage.Output.ResponsesToolCallOutputStrwhen present.- Falls back to
msg.Content.ContentStrfor backward compatibility when the structuredOutputwrapper isn’t populated.- Derives the Gemini function name by preferring
ResponsesToolMessage.Nameand falling back toCallIDwhen the name is missing, which matches how Gemini associates call outputs with tools.- Keeps
FunctionResponse.IDin sync withCallID.This should make the Responses↔Gemini conversion more robust across older and newer payload shapes without changing behavior for existing clients.
2387-2510: DecodingSignatureintoPart.ThoughtSignaturecompletes the round-tripIn
convertContentBlockToGeminiPart, for text blocks:
- You now construct a
PartwithTextand, whenblock.Signatureis present, base64-decode it intoThoughtSignature.- Decode failures are deliberately ignored (you just omit the signature), which is a reasonable fallback for malformed clients.
Combined with the earlier logic that populates
Signaturewhen converting from Gemini to Responses, this gives you a clean, symmetric path for encrypted thought signatures on text content. Other block types (reasoning, refusal, image/audio/file) remain unaffected and continue to behave as before.
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (28)
core/changelog.md(1 hunks)core/providers/bedrock/utils.go(1 hunks)core/providers/gemini/chat.go(1 hunks)core/providers/gemini/responses.go(8 hunks)core/providers/gemini/types.go(1 hunks)core/providers/gemini/utils.go(2 hunks)core/version(1 hunks)framework/changelog.md(1 hunks)framework/version(1 hunks)plugins/governance/changelog.md(1 hunks)plugins/governance/version(1 hunks)plugins/jsonparser/changelog.md(1 hunks)plugins/jsonparser/version(1 hunks)plugins/logging/changelog.md(1 hunks)plugins/logging/version(1 hunks)plugins/maxim/changelog.md(1 hunks)plugins/maxim/version(1 hunks)plugins/mocker/changelog.md(1 hunks)plugins/mocker/version(1 hunks)plugins/otel/changelog.md(1 hunks)plugins/otel/version(1 hunks)plugins/semanticcache/changelog.md(1 hunks)plugins/semanticcache/version(1 hunks)plugins/telemetry/changelog.md(1 hunks)plugins/telemetry/version(1 hunks)tests/integrations/tests/test_google.py(14 hunks)transports/changelog.md(1 hunks)transports/version(1 hunks)
✅ Files skipped from review due to trivial changes (19)
- plugins/otel/changelog.md
- transports/changelog.md
- core/changelog.md
- framework/changelog.md
- plugins/governance/version
- plugins/semanticcache/version
- plugins/semanticcache/changelog.md
- plugins/jsonparser/changelog.md
- plugins/logging/version
- core/version
- plugins/logging/changelog.md
- plugins/jsonparser/version
- transports/version
- framework/version
- plugins/mocker/version
- plugins/governance/changelog.md
- plugins/telemetry/version
- plugins/telemetry/changelog.md
- plugins/otel/version
🚧 Files skipped from review as they are similar to previous changes (2)
- core/providers/bedrock/utils.go
- core/providers/gemini/utils.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
plugins/mocker/changelog.mdplugins/maxim/changelog.mdcore/providers/gemini/chat.gocore/providers/gemini/responses.gocore/providers/gemini/types.goplugins/maxim/versiontests/integrations/tests/test_google.py
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
core/providers/gemini/chat.gocore/providers/gemini/responses.gocore/providers/gemini/types.go
📚 Learning: 2025-12-19T09:26:54.961Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/providers/utils/utils.go:1050-1051
Timestamp: 2025-12-19T09:26:54.961Z
Learning: Update streaming end-marker handling so HuggingFace is treated as a non-[DONE] provider for backends that do not emit a DONE marker (e.g., meta llama on novita). In core/providers/utils/utils.go, adjust ProviderSendsDoneMarker() (or related logic) to detect providers that may not emit DONE and avoid relying on DONE as the sole end signal. Add tests to cover both DONE-emitting and non-DONE backends, with clear documentation in code comments explaining the rationale and any fallback behavior.
Applied to files:
core/providers/gemini/chat.gocore/providers/gemini/responses.gocore/providers/gemini/types.go
🧬 Code graph analysis (2)
core/providers/gemini/chat.go (1)
core/schemas/chatcompletions.go (2)
ChatAssistantMessageToolCall(710-716)ChatToolChoiceTypeFunction(382-382)
core/providers/gemini/responses.go (3)
core/providers/gemini/types.go (7)
Type(782-782)Role(17-17)FunctionCall(1197-1207)FunctionResponse(1210-1232)Content(977-985)FileData(1186-1194)Part(991-1015)core/schemas/responses.go (9)
ResponsesReasoning(741-744)ResponsesMessageRoleType(336-336)ResponsesMessage(321-334)ResponsesToolMessage(472-492)ResponsesReasoningSummary(755-758)ResponsesToolMessageOutputStruct(541-545)ResponsesMessageContent(346-351)ResponsesMessageContentBlock(406-421)ResponsesMessageContentBlockType(392-392)core/schemas/utils.go (1)
Ptr(16-18)
🪛 Ruff (0.14.8)
tests/integrations/tests/test_google.py
1238-1238: Unused method argument: test_config
(ARG002)
1447-1447: f-string without any placeholders
Remove extraneous f prefix
(F541)
1455-1455: Unused method argument: test_config
(ARG002)
1537-1537: Unused method argument: test_config
(ARG002)
1591-1591: Unused method argument: test_config
(ARG002)
1638-1638: Unused method argument: test_config
(ARG002)
1682-1682: Unused method argument: test_config
(ARG002)
1727-1727: Unused method argument: test_config
(ARG002)
1787-1787: Unused method argument: test_config
(ARG002)
1821-1821: Unused method argument: test_config
(ARG002)
1841-1841: Unused method argument: test_config
(ARG002)
1881-1881: Unused method argument: test_config
(ARG002)
1924-1924: Unused method argument: test_config
(ARG002)
🔇 Additional comments (6)
plugins/maxim/version (1)
1-1: LGTM!The patch version bump (1.4.51 → 1.4.52) is appropriate for the bug fix and feature enhancements related to Gemini thought-signature handling introduced in this PR.
plugins/mocker/changelog.md (1)
1-1: Changelog entry is consistent with version bumpsThe note clearly documents the core/framework version upgrades and aligns with the PR scope; no issues here.
tests/integrations/tests/test_google.py (4)
36-49: Scenario list update matches new testsThe docstring enumeration now reflects the new Gemini 3 Pro thought-signature test (28) and the renumbered structured output and Files/Batch API tests (29–39); numbering and descriptions look consistent with the test names below.
1455-1530: Structured-output + thinking test looks solidThe
test_29_structured_output_with_thinkingtest thoroughly validates:
- That
response.parsedis populated and is aMathProbleminstance.- Field presence and types, including
reasoning_stepscontent.- Thinking token accounting for Gemini via
usage_metadata.thoughts_token_count.This should give good coverage for the structured-output + thinking path; no changes needed.
2017-2038: Helperextract_google_function_callsis defensive and appropriateThe helper correctly:
- Short-circuits when
response.function_callsis missing/empty.- Checks for
nameandargson each function-call object.- Safely converts
argsinto adict, catchingAttributeError/TypeErrorand logging a warning without failing the whole extraction.This is a reasonable, defensive adapter for Google GenAI function-call responses.
113-135: Verify GEMINI_API_KEY usage and timeout configuration inconsistenciesThe function has two alignment issues with other provider implementations:
API key retrieval: Uses
os.environ.get("GEMINI_API_KEY")directly instead of delegating toget_api_key()likeget_provider_openai_clientandget_provider_anthropic_clientdo. This creates fragile coupling withskip_if_no_api_key("google"), which expects the same mapping.Timeout handling: Hard-codes
timeout=30000whenapi_config.get("timeout")is truthy, ignoring the actual configured value. The OpenAI implementation correctly usestimeout=api_config.get("timeout", 300)to respect configuration. This prevents respecting per-test timeout customization.Consider using
get_api_key("google")andapi_config.get("timeout", 30000)to align with existing patterns.
TejasGhatte
force-pushed
the
12-20-fix_gemini_thought_signature_handling
branch
from
12987f9 to
d516b42
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
core/providers/gemini/responses.go (1)
443-449: Streaming:OutputItemAddedreasoning branch likely never fires due toEncryptedContentcheckIn
ToGeminiResponsesStreamResponse’sResponsesStreamResponseTypeOutputItemAddedcase:if bifrostResp.Item != nil && bifrostResp.Item.ResponsesReasoning != nil && bifrostResp.Item.EncryptedContent != nil { candidate.Content.Parts = append(candidate.Content.Parts, &Part{ ThoughtSignature: []byte(*bifrostResp.Item.ResponsesReasoning.EncryptedContent), }) }But elsewhere in this file (e.g.
processGeminiThoughtSignaturePartand the non‑streaming conversions), encrypted reasoning content is stored only underResponsesReasoning.EncryptedContent, not on theResponsesMessageitself.That means
bifrostResp.Item.EncryptedContentis almost certainly alwaysnil, so this branch never runs and thought signatures encoded as reasoning items won’t be surfaced back into Gemini streaming responses.You probably intended:
if bifrostResp.Item != nil && bifrostResp.Item.ResponsesReasoning != nil && bifrostResp.Item.ResponsesReasoning.EncryptedContent != nil { candidate.Content.Parts = append(candidate.Content.Parts, &Part{ ThoughtSignature: []byte(*bifrostResp.Item.ResponsesReasoning.EncryptedContent), }) }Fixing this would ensure that reasoning items carrying signatures from the Bifrost stream are correctly converted back into
Part.ThoughtSignaturefor streaming clients using the Gemini format.
🧹 Nitpick comments (3)
core/providers/gemini/types.go (1)
1017-1066: CustomPart.UnmarshalJSONcorrectly handles URL-safe base64, but is brittle to future field additionsThe decoding logic and URL-safe→standard base64 normalization look solid and will handle Gemini’s URL-safe, unpadded signatures correctly. The only real concern is maintainability:
PartAliasre-lists every field ofPartwith a different type forThoughtSignature. Any future field added toPartmust also be added here or it won’t be deserialized.- The URL-safe base64 normalization is now duplicated with
Blob.UnmarshalJSON.Consider a small refactor (non-blocking):
Use an embedded alias to automatically keep all other fields in sync, e.g.:
type rawPart Part aux := struct { ThoughtSignature string `json:"thoughtSignature,omitempty"` *rawPart }{rawPart: (*rawPart)(p)}Then only post-process
aux.ThoughtSignatureintop.ThoughtSignature.Optionally extract a tiny helper like
decodeURLSafeBase64(s string) ([]byte, error)and reuse it here and inBlob.UnmarshalJSON.core/providers/gemini/responses.go (1)
1369-1537: ExpandedconvertGeminiContentsToResponsesMessagescorrectly covers function calls, outputs, thoughts, signatures, and rich mediaThe new per‑
Parthandling here is doing quite a bit, and overall it’s coherent:
- Maps
roleonce perContent(model→assistant,user→user, otherwise defaults to user).- For
FunctionCallparts:
- Marshals
Argsto JSON.- Derives a
callID(preferringID, falling back toName) and tracks it infunctionCallIDsfor later response correlation.- Emits a
ResponsesMessageTypeFunctionCallwithResponsesToolMessage.- Optionally emits a separate
Reasoningmessage carrying the base64‑encodedThoughtSignature.- For
FunctionResponseparts:
- Derives
responseIDfromFunctionResponse.ID, falling back to the tracked call ID or name.- Extracts
"output"from the response map (string if present, otherwise marshalled JSON).- Produces a
FunctionCallOutputmessage withCallID,Output, andName(if non‑blank).- For thought text (
Thought && Text), regular text (with optionalSignature), inline data, and file data it constructs the appropriateResponsesMessage+ content block forms.Two minor notes (non‑blocking):
- The function call and function response messages don’t set
Role; elsewhere in this file, function-call output is typically tagged withassistant. If downstream assumes a role here, you might want to setRole: role(or explicitly assistant for responses) for consistency.- In scenarios with multiple concurrent calls using the same function name, the
functionCallIDs[name]map may be ambiguous. Relying onFunctionResponse.IDwhen present is correct; the name‑based fallback is best‑effort and probably acceptable given Gemini semantics but worth keeping in mind.Overall, the mapping and ThoughtSignature propagation here look correct.
tests/integrations/tests/test_google.py (1)
147-149:test_configfixture is unused in many tests (Ruff ARG002)Ruff correctly notes that
test_configis accepted but never used in a number of tests (e.g.,test_28_gemini_3_pro_thought_signatures_multi_turn,test_29_structured_output_with_thinking, and the file/batch tests 30–39).If you want to quiet ARG002 without changing behavior, you can either:
- Rename the parameter to
_test_configin tests where it’s not used, or- Drop the parameter entirely from those tests if you don’t plan to use the fixture there.
Both options are safe and keep the signatures aligned with how the fixtures are actually consumed.
Also applies to: 1237-1239, 1468-1469, 1550-1552, 1604-1606, 1651-1653, 1695-1697, 1740-1742, 1800-1802, 1834-1836, 1854-1856, 1894-1896, 1937-1939
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (28)
core/changelog.md(1 hunks)core/providers/bedrock/utils.go(1 hunks)core/providers/gemini/chat.go(1 hunks)core/providers/gemini/responses.go(8 hunks)core/providers/gemini/types.go(1 hunks)core/providers/gemini/utils.go(2 hunks)core/version(1 hunks)framework/changelog.md(1 hunks)framework/version(1 hunks)plugins/governance/changelog.md(1 hunks)plugins/governance/version(1 hunks)plugins/jsonparser/changelog.md(1 hunks)plugins/jsonparser/version(1 hunks)plugins/logging/changelog.md(1 hunks)plugins/logging/version(1 hunks)plugins/maxim/changelog.md(1 hunks)plugins/maxim/version(1 hunks)plugins/mocker/changelog.md(1 hunks)plugins/mocker/version(1 hunks)plugins/otel/changelog.md(1 hunks)plugins/otel/version(1 hunks)plugins/semanticcache/changelog.md(1 hunks)plugins/semanticcache/version(1 hunks)plugins/telemetry/changelog.md(1 hunks)plugins/telemetry/version(1 hunks)tests/integrations/tests/test_google.py(14 hunks)transports/changelog.md(1 hunks)transports/version(1 hunks)
✅ Files skipped from review due to trivial changes (3)
- plugins/maxim/changelog.md
- plugins/maxim/version
- plugins/governance/changelog.md
🚧 Files skipped from review as they are similar to previous changes (21)
- core/providers/bedrock/utils.go
- core/changelog.md
- plugins/logging/changelog.md
- plugins/jsonparser/version
- plugins/semanticcache/changelog.md
- core/providers/gemini/chat.go
- transports/changelog.md
- plugins/jsonparser/changelog.md
- plugins/mocker/version
- plugins/governance/version
- core/providers/gemini/utils.go
- plugins/otel/changelog.md
- plugins/semanticcache/version
- plugins/otel/version
- framework/version
- transports/version
- plugins/telemetry/changelog.md
- plugins/logging/version
- plugins/telemetry/version
- framework/changelog.md
- plugins/mocker/changelog.md
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
core/providers/gemini/types.gocore/providers/gemini/responses.gotests/integrations/tests/test_google.pycore/version
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
core/providers/gemini/types.gocore/providers/gemini/responses.go
📚 Learning: 2025-12-19T09:26:54.961Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/providers/utils/utils.go:1050-1051
Timestamp: 2025-12-19T09:26:54.961Z
Learning: Update streaming end-marker handling so HuggingFace is treated as a non-[DONE] provider for backends that do not emit a DONE marker (e.g., meta llama on novita). In core/providers/utils/utils.go, adjust ProviderSendsDoneMarker() (or related logic) to detect providers that may not emit DONE and avoid relying on DONE as the sole end signal. Add tests to cover both DONE-emitting and non-DONE backends, with clear documentation in code comments explaining the rationale and any fallback behavior.
Applied to files:
core/providers/gemini/types.gocore/providers/gemini/responses.go
🧬 Code graph analysis (1)
tests/integrations/tests/test_google.py (3)
tests/integrations/tests/utils/common.py (1)
skip_if_no_api_key(1813-1824)core/providers/gemini/types.go (3)
Part(991-1015)FunctionCall(1197-1207)FunctionResponse(1210-1232)tests/integrations/tests/utils/parametrize.py (1)
get_cross_provider_params_for_scenario(12-47)
🪛 Ruff (0.14.8)
tests/integrations/tests/test_google.py
1238-1238: Unused method argument: test_config
(ARG002)
1460-1460: f-string without any placeholders
Remove extraneous f prefix
(F541)
1468-1468: Unused method argument: test_config
(ARG002)
1550-1550: Unused method argument: test_config
(ARG002)
1604-1604: Unused method argument: test_config
(ARG002)
1651-1651: Unused method argument: test_config
(ARG002)
1695-1695: Unused method argument: test_config
(ARG002)
1740-1740: Unused method argument: test_config
(ARG002)
1800-1800: Unused method argument: test_config
(ARG002)
1834-1834: Unused method argument: test_config
(ARG002)
1854-1854: Unused method argument: test_config
(ARG002)
1894-1894: Unused method argument: test_config
(ARG002)
1937-1937: Unused method argument: test_config
(ARG002)
🔇 Additional comments (9)
core/version (1)
1-1: LGTM! Version bump aligns with the coordinated release.The patch increment from 1.2.40 to 1.2.41 is appropriate for the Gemini thought signature handling fixes and improvements. This version bump is consistent with the documented changelog entries and the framework's upgrade to core 1.2.41.
core/providers/gemini/responses.go (6)
175-312: Reasoning→thought-signature pairing inToGeminiResponsesResponselooks correctThe
consumedIndicesapproach that:
- Looks ahead from a
FunctionCallmessage to a followingReasoningmessage withEncryptedContent,- Decodes the base64 signature into
Part.ThoughtSignature, and- Marks that reasoning index as consumed so it isn't also turned into a separate thought/ThoughtSignature part
is a clean way to align Bifrost’s reasoning messages with Gemini function-call parts. Failure to decode just falls back to the original reasoning path, which is safe.
No issues from a correctness perspective here.
773-846: Streaming: attachingThoughtSignaturetooutput_text.deltaviaSignaturefield is sound
processGeminiTextPartnow:
- Accumulates text in
TextBufferas before.- When emitting
ResponsesStreamResponseTypeOutputTextDelta, base64-encodespart.ThoughtSignatureand wires it intostreamResponse.Signature.This matches the new responses schema and allows downstream consumers to see signatures inline with text deltas. No behavioral regressions visible.
1545-1624: InlineData/FileData → content-block helpers are consistent with the rest of the adapter
convertGeminiInlineDataToContentBlockandconvertGeminiFileDataToContentBlock:
- Respect MIME type, using
isImageMimeTypefor images.- Encode inline data as:
data:<mime>;base64,...URLs for images.{format, data}blocks for audio, stripping the"audio/"prefix.- File blocks with base64 payload + filename for other blobs.
- Map file URIs to image/file blocks, defaulting MIME type to
application/octet-streamwhen missing.These helpers align with the other conversions in this file and make the media handling symmetric across directions. No issues to flag.
1685-1923: Candidate→Responses output: ThoughtSignature propagation and richer part handling look correctIn
convertGeminiCandidatesToResponsesOutputthe new logic:
- For text parts, attaches
ThoughtSignatureas a base64Signatureon the last text content block, mirroring whatconvertGeminiContentsToResponsesMessagesdoes.- For
FunctionCallparts, continues to emit aFunctionCallmessage and now, ifThoughtSignatureis present, emits a separateReasoningmessage withEncryptedContentfor better scalability.- Adds a final
case part.ThoughtSignature != nilto emit a pure‑signatureReasoningmessage when the part has no text/call/etc.- Leaves the existing handling for thoughts, function responses, inline data, file data, code execution results, and executable code intact.
Control flow is safe (the switch is first‑match, so
FunctionCall+signature won’t double‑emit via the trailingThoughtSignaturecase), and the encoding is consistent (standard base64 everywhere).This matches the PR’s goal of surfacing thought signatures in both tool calls and plain text responses.
2225-2379: Responses→Gemini contents: function outputs and thought signatures are wired appropriatelyThe updates in
convertResponsesMessagesToGeminiContentsfor tool messages:
- For
FunctionCall:
- Reconstruct
FunctionCallwithName,Argsdecoded from JSON.- Preserve
CallIDasFunctionCall.ID.- Look ahead to
messages[i+1]for aReasoningmessage withEncryptedContent, base64‑decode it, and attach asPart.ThoughtSignature. This mirrors the forward path where a function call plus separate reasoning message are emitted.- For
FunctionCallOutput:
- Prefer
ResponsesToolMessage.Output.ResponsesToolCallOutputStras"output", with a fallback toContentStrfor compatibility.- Derive
funcNamefromNameif present, otherwise fall back toCallID.- Set
FunctionResponse.NameandFunctionResponse.ID(fromCallID) and attach theresponseMap.Error handling returns contextual errors if argument or content‑block decoding fails, which is helpful.
This should allow multi‑turn and tool‑calling flows (including thought signatures) to round‑trip cleanly.
2381-2510: Text-blockSignature→Part.ThoughtSignaturemapping closes the round-tripIn
convertContentBlockToGeminiPartfor text blocks:
- You now construct
Part{Text: *block.Text}and, ifblock.Signatureis present, base64‑decode it and setPart.ThoughtSignature.- Decode failures are silently ignored, leaving a plain text part, which is reasonable for robustness.
- Other content types (reasoning, refusal, image, audio, file) are unchanged and consistent with the rest of the adapter.
Together with the earlier changes that write
Signaturewhen going Gemini→Responses, this gives you a full round-trip path for thought signatures on text content.No issues here.
tests/integrations/tests/test_google.py (2)
113-137: GEMINI client configuration: direct env key + fixed timeout look reasonableUsing
GEMINI_API_KEYdirectly and wiringtimeout=30000intoHttpOptionswhen a timeout is configured keeps the Google client setup simple and isolates cross-provider routing to thex-model-providerheader. As long as your config loader and CI setGEMINI_API_KEYconsistently, this looks fine.No functional issues spotted here.
1237-1465: Multi-turn Gemini 3 Pro thought-signature test is thorough and avoids prior pitfallsThe new
test_28_gemini_3_pro_thought_signatures_multi_turndoes a good job of validating the end-to-end behavior you just implemented:
- Turn 1 asserts presence of either thought or text content and
thought_signature_count > 0, while guarding all debug prints withgetattr(part, "text", None)before slicing.- Turn 2 asserts at least one function call and at least one
thought_signatureattached to a tool call.- Turn 3 reconstructs a multi-turn conversation, reusing the captured
tool_thought_signatures[0], and asserts that the model produces non-empty text, again guarding thought/text prints.- Turn 4 drives multiple tools in one request and, if any function calls exist, asserts that at least one has a thought signature even at
thinking_level="low".- Finally, it checks
usage_metadata.thoughts_token_count > 0when available.This turns the scenario into a real regression test for thought-signature propagation, not just logging, and aligns tightly with the adapter changes in
core/providers/gemini.
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (4)
tests/integrations/tests/test_google.py (2)
1238-1238: Optional: Remove unusedtest_configfixture parameter.The
test_configparameter is injected as a pytest fixture but not used in multiple test functions. Consider removing it to clean up the signature, or add a comment if it's kept for future use or consistency.This is flagged by static analysis (Ruff ARG002) but is a minor code quality issue.
Also applies to: 1468-1468, 1550-1550, 1604-1604, 1651-1651, 1695-1695, 1740-1740, 1800-1800, 1834-1834, 1854-1854, 1894-1894, 1937-1937
1460-1460: Optional: Remove unnecessary f-string prefix.Line 1460 uses an f-string without any placeholders. You can use a regular string instead.
- print(f"\n=== Token Usage ===") + print("\n=== Token Usage ===")core/providers/gemini/responses.go (2)
238-251: Document the Gemini API convention for thought signatures following function calls.The look-ahead logic assumes that a reasoning message with encrypted content immediately follows a function call message. While the bounds check is correct, adding a comment would clarify this Gemini API convention and help future maintainers understand why this specific pattern is expected.
🔎 Suggested documentation
part := &Part{ FunctionCall: functionCall, } + // Gemini API convention: encrypted thought signatures for function calls + // appear as a separate reasoning message immediately following the function call // Look ahead to see if the next message is a reasoning message with encrypted content // (thought signature for this function call) if i+1 < len(bifrostResp.Output) {
2343-2350: Consider handling the case when both output paths are nil.The extraction logic tries two paths but doesn't explicitly handle when both are absent. While an empty
responseMapmay be acceptable, consider adding a comment or validation to clarify the expected behavior in this edge case.🔎 Suggested enhancement
// Extract output from ResponsesToolMessage.Output if msg.ResponsesToolMessage.Output != nil && msg.ResponsesToolMessage.Output.ResponsesToolCallOutputStr != nil { responseMap["output"] = *msg.ResponsesToolMessage.Output.ResponsesToolCallOutputStr } else if msg.Content != nil && msg.Content.ContentStr != nil { // Fallback to Content.ContentStr for backward compatibility responseMap["output"] = *msg.Content.ContentStr + } else { + // No output available - empty responseMap is acceptable for Gemini }
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (28)
core/changelog.md(1 hunks)core/providers/bedrock/utils.go(1 hunks)core/providers/gemini/chat.go(1 hunks)core/providers/gemini/responses.go(8 hunks)core/providers/gemini/types.go(1 hunks)core/providers/gemini/utils.go(2 hunks)core/version(1 hunks)framework/changelog.md(1 hunks)framework/version(1 hunks)plugins/governance/changelog.md(1 hunks)plugins/governance/version(1 hunks)plugins/jsonparser/changelog.md(1 hunks)plugins/jsonparser/version(1 hunks)plugins/logging/changelog.md(1 hunks)plugins/logging/version(1 hunks)plugins/maxim/changelog.md(1 hunks)plugins/maxim/version(1 hunks)plugins/mocker/changelog.md(1 hunks)plugins/mocker/version(1 hunks)plugins/otel/changelog.md(1 hunks)plugins/otel/version(1 hunks)plugins/semanticcache/changelog.md(1 hunks)plugins/semanticcache/version(1 hunks)plugins/telemetry/changelog.md(1 hunks)plugins/telemetry/version(1 hunks)tests/integrations/tests/test_google.py(14 hunks)transports/changelog.md(1 hunks)transports/version(1 hunks)
✅ Files skipped from review due to trivial changes (4)
- plugins/otel/changelog.md
- framework/version
- framework/changelog.md
- plugins/mocker/changelog.md
🚧 Files skipped from review as they are similar to previous changes (15)
- core/providers/gemini/utils.go
- plugins/telemetry/version
- plugins/maxim/version
- plugins/semanticcache/version
- plugins/maxim/changelog.md
- core/changelog.md
- core/providers/bedrock/utils.go
- plugins/otel/version
- plugins/jsonparser/changelog.md
- transports/version
- core/version
- plugins/governance/changelog.md
- plugins/logging/version
- plugins/mocker/version
- plugins/telemetry/changelog.md
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
transports/changelog.mdcore/providers/gemini/chat.goplugins/governance/versioncore/providers/gemini/responses.gocore/providers/gemini/types.goplugins/semanticcache/changelog.mdplugins/jsonparser/versiontests/integrations/tests/test_google.pyplugins/logging/changelog.md
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
core/providers/gemini/chat.gocore/providers/gemini/responses.gocore/providers/gemini/types.go
📚 Learning: 2025-12-19T09:26:54.961Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/providers/utils/utils.go:1050-1051
Timestamp: 2025-12-19T09:26:54.961Z
Learning: Update streaming end-marker handling so HuggingFace is treated as a non-[DONE] provider for backends that do not emit a DONE marker (e.g., meta llama on novita). In core/providers/utils/utils.go, adjust ProviderSendsDoneMarker() (or related logic) to detect providers that may not emit DONE and avoid relying on DONE as the sole end signal. Add tests to cover both DONE-emitting and non-DONE backends, with clear documentation in code comments explaining the rationale and any fallback behavior.
Applied to files:
core/providers/gemini/chat.gocore/providers/gemini/responses.gocore/providers/gemini/types.go
🧬 Code graph analysis (2)
core/providers/gemini/chat.go (3)
core/schemas/chatcompletions.go (2)
ChatAssistantMessageToolCall(710-716)ChatToolChoiceTypeFunction(382-382)core/providers/gemini/types.go (1)
Type(782-782)core/schemas/utils.go (1)
Ptr(16-18)
tests/integrations/tests/test_google.py (3)
tests/integrations/tests/utils/common.py (1)
skip_if_no_api_key(1813-1824)core/providers/gemini/types.go (5)
Tool(604-631)Content(977-985)Part(991-1015)FunctionCall(1197-1207)FunctionResponse(1210-1232)tests/integrations/tests/utils/parametrize.py (1)
get_cross_provider_params_for_scenario(12-47)
🪛 Ruff (0.14.8)
tests/integrations/tests/test_google.py
1238-1238: Unused method argument: test_config
(ARG002)
1460-1460: f-string without any placeholders
Remove extraneous f prefix
(F541)
1468-1468: Unused method argument: test_config
(ARG002)
1550-1550: Unused method argument: test_config
(ARG002)
1604-1604: Unused method argument: test_config
(ARG002)
1651-1651: Unused method argument: test_config
(ARG002)
1695-1695: Unused method argument: test_config
(ARG002)
1740-1740: Unused method argument: test_config
(ARG002)
1800-1800: Unused method argument: test_config
(ARG002)
1834-1834: Unused method argument: test_config
(ARG002)
1854-1854: Unused method argument: test_config
(ARG002)
1894-1894: Unused method argument: test_config
(ARG002)
1937-1937: Unused method argument: test_config
(ARG002)
🔇 Additional comments (15)
tests/integrations/tests/test_google.py (3)
1237-1465: Well-structured multi-turn thought signature test.The new
test_28_gemini_3_pro_thought_signatures_multi_turntest provides comprehensive coverage of thought signature handling across multiple scenarios:
- Turn 1: Initial thinking validation
- Turn 2: Tool calls with thought signatures
- Turn 3: Multi-turn conversation with tool results
- Turn 4: Multiple parallel tool calls
The assertions added (lines 1328, 1359, 1456) properly validate the presence of thought signatures, addressing previous review feedback. Good work on the comprehensive test coverage!
1415-1416: Guard against None when slicing part.text (duplicate pattern).Same issue as Line 1316:
part_textmight beNone, causing aTypeErrorwhen slicing.🔎 Proposed fix
part_text = getattr(part, 'text', None) - print(f" [THOUGHT] {part_text[:100] if part_text else '(no text)'}...") + preview = part_text[:100] if part_text else '(no text)' + print(f" [THOUGHT] {preview}...")Likely an incorrect or invalid review comment.
1315-1316: Guard against None when slicing part.text.The code uses
getattr(part, 'text', None)to safely retrieve the text attribute, but then immediately slices it withpart_text[:100]. Ifpart_textisNone, this will raise aTypeError.🔎 Proposed fix to add None check
part_text = getattr(part, 'text', None) - print(f" [THOUGHT] {part_text[:100] if part_text else '(no text)'}...") + print(f" [THOUGHT] {part_text[:100] if part_text and len(part_text) > 0 else '(no text)'}...")Or more simply:
part_text = getattr(part, 'text', None) - print(f" [THOUGHT] {part_text[:100] if part_text else '(no text)'}...") + preview = part_text[:100] if part_text else '(no text)' + print(f" [THOUGHT] {preview}...")Likely an incorrect or invalid review comment.
core/providers/gemini/chat.go (1)
114-130: LGTM: Thought signature correctly attached to tool calls.The implementation properly:
- Builds the tool call with standard fields (Index, Type, ID, Function)
- Base64-encodes the binary
ThoughtSignaturefor safe transmission- Attaches it to
ExtraContentunder thegoogle.thought_signaturekey, following the provider-specific extension pattern- Maintains clean separation between standard and provider-specific fields
The nested structure in
ExtraContentis consistent with the field documentation incore/schemas/chatcompletions.go(line 714) which describes it as "Provider-specific fields (e.g., thought_signature for Gemini)".core/providers/gemini/types.go (1)
1017-1066: LGTM: Correct URL-safe base64 decoding for thought signatures.The custom
UnmarshalJSONimplementation correctly handles thethoughtSignaturefield sent as a base64-encoded string from the Google GenAI SDK:
- URL-safe to standard conversion: Properly replaces URL-safe characters (
_→/,-→+)- Padding logic: Correctly adds padding based on length modulo 4
- Error handling: Returns descriptive errors if decoding fails
- Consistency: Mirrors the pattern used in
Blob.UnmarshalJSON(lines 1079-1114)This implementation enables proper round-tripping of encrypted thought signatures through JSON serialization/deserialization.
core/providers/gemini/responses.go (10)
174-176: LGTM! Well-designed state tracking for consumed messages.The
consumedIndicesmap cleanly tracks which reasoning messages have been consumed as thought signatures, preventing duplicate processing later in the loop. This pairs well with the look-ahead logic.
287-290: LGTM! Proper skip logic for consumed messages.The skip logic correctly prevents double-processing of reasoning messages that were already consumed as thought signatures attached to function calls.
829-842: LGTM! Proper signature attachment in streaming deltas.The logic correctly encodes the ThoughtSignature to base64 and attaches it to the stream response delta, maintaining consistency with the signature handling pattern throughout the file.
1728-1732: LGTM! Symmetric signature handling across conversion paths.The signature handling uses two patterns consistently:
- Text content: signature embedded in the content block
- Function calls: separate reasoning message with encrypted content
This mirrors the decoding logic in
ToGeminiResponsesResponse(lines 238-251), ensuring proper round-trip conversion.Also applies to: 1763-1776
2387-2396: LGTM! Proper signature decoding with graceful error handling.The signature decoding correctly handles base64 decoding errors by gracefully degrading (text is preserved even if signature decode fails). This is consistent with the error handling pattern used throughout the file.
244-244: LGTM! Consistent base64 encoding/decoding throughout the file.All ThoughtSignature encoding/decoding operations consistently use
base64.StdEncodingwith proper error handling. The symmetric encode-to-string and decode-from-string pattern ensures correct round-trip conversion.Also applies to: 304-304, 838-838, 1414-1414, 1504-1504, 1730-1730, 1766-1766, 1908-1908, 2330-2330, 2391-2391
303-310: LGTM! Proper handling of standalone reasoning signatures.This correctly processes standalone reasoning messages with encrypted content, which are distinct from thought signatures attached to function calls (those are skipped at lines 287-290).
244-246: LGTM! Robust error handling with graceful degradation.The error handling consistently treats thought signatures as optional enhancements. Base64 decode failures and JSON marshal errors are handled gracefully without disrupting the core message conversion flow. This approach ensures system resilience even when signature data is malformed.
Also applies to: 304-306, 2330-2333
174-176: Verify integration test coverage for thought signature handling includes all multi-turn patterns.Thought signatures appear in content parts (text or functionCall), are returned with all function call responses, appear on the first part of parallel function calls, appear on each part of sequential function calls, and appear in the last part of responses without function calls. The integration test should confirm coverage for:
- Function calls with attached thought signatures (look-ahead pattern)
- Text messages with inline signatures
- Streaming responses with signatures in deltas
- Round-trip conversion (Bifrost → Gemini → Bifrost)
1370-1542: The multiple-message-per-Content refactor is architecturally sound.The change generates multiple
ResponsesMessageentries perContent(one perPart, plus additional reasoning messages when signatures are present). This is handled properly throughout the stack:
bifrostReq.Inputis a slice ([]ResponsesMessage), naturally accommodating variable message counts- The reverse conversion
convertResponsesMessagesToGeminiContentscorrectly reconstructs individualContentobjects from the message stream using lookahead to match reasoning messages with function calls (line 2232 skips orphan reasoning messages; lines 2101-2109 match them with preceding function calls)- No callers assume 1:1
Content-to-Messagemapping
TejasGhatte
force-pushed
the
12-20-fix_gemini_thought_signature_handling
branch
from
d516b42 to
31f1ad5
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
core/providers/gemini/responses.go (1)
738-771: Streaming: FunctionCall parts with ThoughtSignature are dropping the signature and will cause Gemini 3 Pro validation failures.The switch statement in
processGeminiPart(lines 742-767) processes parts as mutually exclusive cases. When a part contains bothFunctionCallandThoughtSignature(which Gemini 3 Pro always includes on the first function call part, and it is mandatory to return that part), only theFunctionCallcase executes.
processGeminiFunctionCallPart(lines 978-1044) does not extract or emit theThoughtSignature. This means Gemini 3 Pro requires thoughtSignature to be passed back during function calling, otherwise validation will fail with a 400 error.Fix by extracting
part.ThoughtSignaturewithinprocessGeminiFunctionCallPartand emitting a parallel reasoning item (similar to howprocessGeminiThoughtSignaturePartdoes), or restructure the switch logic to handle parts with multiple content types.
♻️ Duplicate comments (2)
tests/integrations/tests/test_google.py (2)
113-135: Don’t ignore the configured timeout inget_provider_google_client.The block
if api_config.get("timeout"): http_options_kwargs["timeout"] = 30000drops the configured timeout and hardcodes
30000, reducing flexibility and regressing prior behavior. Prefer honoring the config value with a sensible default, e.g.:Proposed fix
- if api_config.get("timeout"): - http_options_kwargs["timeout"] = 30000 + timeout = api_config.get("timeout") + if timeout: + http_options_kwargs["timeout"] = timeout
1458-1462: Minor: remove unnecessary f-string around static log line.
print(f"\n=== Token Usage ===")is an f-string without interpolation (Ruff F541). It can be a plain string:- if hasattr(response_1.usage_metadata, 'thoughts_token_count'): - print(f"\n=== Token Usage ===") + if hasattr(response_1.usage_metadata, 'thoughts_token_count'): + print("\n=== Token Usage ===")
🧹 Nitpick comments (4)
core/providers/gemini/types.go (1)
1017-1066: Part.UnmarshalJSON correctly normalizes URL-safe thought signatures; consider a small shared helper.The custom unmarshal cleanly handles URL-safe base64 (replacing
-/_and padding) and populatesThoughtSignatureas bytes. This aligns withBlob.UnmarshalJSONand should interoperate correctly with the rest of the pipeline.To avoid duplication and keep future changes in one place, you could extract the URL-safe base64 →
[]bytelogic into a small helper reused by bothPartandBlobunmarshaling.tests/integrations/tests/test_google.py (2)
1238-1238: Clean up unusedtest_configparameters to satisfy Ruff (ARG002).Several tests accept
test_configbut never use it (Ruff ARG002 on these definitions). If you only need the fixture for consistency, renaming to_test_configis enough; otherwise, you can drop the parameter.Example:
Example signature tweak
- def test_28_gemini_3_pro_thought_signatures_multi_turn(self, test_config): + def test_28_gemini_3_pro_thought_signatures_multi_turn(self, _test_config):Apply the same pattern to the other flagged tests to keep the suite lint-clean.
Also applies to: 1468-1468, 1550-1550, 1604-1604, 1651-1651, 1695-1695, 1740-1740, 1800-1800, 1834-1834, 1854-1854, 1894-1894, 1937-1937
2030-2051:extract_google_function_callshelper is currently unused in this file.If you don’t end up wiring this into any tests, consider removing it or adding a narrow unit test that exercises it to avoid dead code in the integration suite.
core/providers/gemini/chat.go (1)
114-131: Tool-call ExtraContent correctly carries ThoughtSignature; keep merge semantics in mind for future fields.Encoding
part.ThoughtSignatureand attaching it undertoolCall.ExtraContent["google"]["thought_signature"]preserves the encrypted signature alongside the tool call without disturbing the core schema, which is exactly what you want here.Right now
toolCallis freshly constructed so assigningExtraContent = map[...]is safe; if you later add other metadata intoExtraContentbefore this block, you’ll want to merge into an existing map rather than overwrite it.Also applies to: 145-152
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (28)
core/changelog.md(1 hunks)core/providers/bedrock/utils.go(1 hunks)core/providers/gemini/chat.go(1 hunks)core/providers/gemini/responses.go(8 hunks)core/providers/gemini/types.go(1 hunks)core/providers/gemini/utils.go(2 hunks)core/version(1 hunks)framework/changelog.md(1 hunks)framework/version(1 hunks)plugins/governance/changelog.md(1 hunks)plugins/governance/version(1 hunks)plugins/jsonparser/changelog.md(1 hunks)plugins/jsonparser/version(1 hunks)plugins/logging/changelog.md(1 hunks)plugins/logging/version(1 hunks)plugins/maxim/changelog.md(1 hunks)plugins/maxim/version(1 hunks)plugins/mocker/changelog.md(1 hunks)plugins/mocker/version(1 hunks)plugins/otel/changelog.md(1 hunks)plugins/otel/version(1 hunks)plugins/semanticcache/changelog.md(1 hunks)plugins/semanticcache/version(1 hunks)plugins/telemetry/changelog.md(1 hunks)plugins/telemetry/version(1 hunks)tests/integrations/tests/test_google.py(13 hunks)transports/changelog.md(1 hunks)transports/version(1 hunks)
✅ Files skipped from review due to trivial changes (6)
- plugins/jsonparser/changelog.md
- plugins/telemetry/changelog.md
- plugins/governance/changelog.md
- plugins/otel/version
- plugins/telemetry/version
- transports/changelog.md
🚧 Files skipped from review as they are similar to previous changes (15)
- plugins/logging/version
- core/providers/bedrock/utils.go
- plugins/semanticcache/version
- plugins/jsonparser/version
- framework/version
- core/version
- plugins/maxim/version
- framework/changelog.md
- plugins/mocker/version
- plugins/governance/version
- core/providers/gemini/utils.go
- plugins/mocker/changelog.md
- plugins/otel/changelog.md
- plugins/logging/changelog.md
- transports/version
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
core/changelog.mdplugins/maxim/changelog.mdcore/providers/gemini/chat.gocore/providers/gemini/responses.gocore/providers/gemini/types.gotests/integrations/tests/test_google.pyplugins/semanticcache/changelog.md
🧠 Learnings (2)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
core/providers/gemini/chat.gocore/providers/gemini/responses.gocore/providers/gemini/types.go
📚 Learning: 2025-12-19T09:26:54.961Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/providers/utils/utils.go:1050-1051
Timestamp: 2025-12-19T09:26:54.961Z
Learning: Update streaming end-marker handling so HuggingFace is treated as a non-[DONE] provider for backends that do not emit a DONE marker (e.g., meta llama on novita). In core/providers/utils/utils.go, adjust ProviderSendsDoneMarker() (or related logic) to detect providers that may not emit DONE and avoid relying on DONE as the sole end signal. Add tests to cover both DONE-emitting and non-DONE backends, with clear documentation in code comments explaining the rationale and any fallback behavior.
Applied to files:
core/providers/gemini/chat.gocore/providers/gemini/responses.gocore/providers/gemini/types.go
🧬 Code graph analysis (1)
tests/integrations/tests/test_google.py (3)
tests/integrations/tests/utils/common.py (1)
skip_if_no_api_key(1813-1824)core/providers/gemini/types.go (5)
Tool(604-631)Content(977-985)Part(991-1015)FunctionCall(1197-1207)FunctionResponse(1210-1232)tests/integrations/tests/utils/parametrize.py (1)
get_cross_provider_params_for_scenario(12-47)
🪛 Ruff (0.14.8)
tests/integrations/tests/test_google.py
1238-1238: Unused method argument: test_config
(ARG002)
1460-1460: f-string without any placeholders
Remove extraneous f prefix
(F541)
1468-1468: Unused method argument: test_config
(ARG002)
1550-1550: Unused method argument: test_config
(ARG002)
1604-1604: Unused method argument: test_config
(ARG002)
1651-1651: Unused method argument: test_config
(ARG002)
1695-1695: Unused method argument: test_config
(ARG002)
1740-1740: Unused method argument: test_config
(ARG002)
1800-1800: Unused method argument: test_config
(ARG002)
1834-1834: Unused method argument: test_config
(ARG002)
1854-1854: Unused method argument: test_config
(ARG002)
1894-1894: Unused method argument: test_config
(ARG002)
1937-1937: Unused method argument: test_config
(ARG002)
🔇 Additional comments (4)
plugins/maxim/changelog.md (1)
1-1: Changelog entry is accurate and properly formatted.The version upgrades (core to 1.2.41 and framework to 1.1.51) align with the PR objectives and are appropriately documented for the plugin's dependency changes.
core/changelog.md (1)
1-1: Changelog entry accurately reflects the Gemini thought-signature fix.Documentation-only change; content is clear and consistent with the code updates elsewhere.
tests/integrations/tests/test_google.py (1)
1237-1465: Gemini 3 Pro thought-signature test exercises the right multi-turn/tool-call paths.The new
test_28_gemini_3_pro_thought_signatures_multi_turnnow:
- Positively asserts presence of thought signatures (not just logs) across turns.
- Reuses tool-call thought signatures in Turn 3 safely (
tool_thought_signaturesis asserted non-empty before indexing).- Guards
part.textslicing viagetattr, avoidingTypeErroron thought-only parts.Given this leans on specific Gemini 3 Pro behavior, just be aware that future model changes around when signatures are emitted could make this test brittle; if that starts happening, you may want to gate some assertions on explicit model metadata or relax counts slightly rather than disabling the test.
core/providers/gemini/responses.go (1)
174-257: Non-stream Responses ↔ Gemini thought-signature round-tripping looks consistent.The updated conversions correctly wire encrypted thought signatures through the Responses API:
ToGeminiResponsesResponseandconvertResponsesMessagesToGeminiContentslook ahead to a following reasoning message, decode itsEncryptedContent, and attach it asPart.ThoughtSignatureon the associated function call, skipping the consumed reasoning message.convertGeminiContentsToResponsesMessages/convertGeminiCandidatesToResponsesOutputemit:
- Function-call messages plus a separate reasoning message when
Part.ThoughtSignatureis present.- Text blocks with
Signatureset when a text part carriesThoughtSignature.convertContentBlockToGeminiPartdecodesblock.Signatureback intoPart.ThoughtSignature, closing the loop.This gives you a coherent, lossless mapping for thought signatures across the non-streaming Responses surface.
Also applies to: 285-311, 1364-1542, 1685-1923, 2382-2511
Merge activity
|
Support for Gemini 3 Pro thought signatures and structured output improvements
This PR adds support for Gemini 3 Pro's thought signatures in tool calls and improves handling of structured outputs across providers.
Changes
Type of change
Affected areas
How to test
Test the Gemini 3 Pro thought signature handling:
Test the Bedrock structured output improvements:
Breaking changes
Related issues
Adds support for Gemini 3 Pro thought signatures which are needed for advanced reasoning capabilities.
Security considerations
The PR handles encrypted thought signatures which are used by Google's Gemini models for advanced reasoning.
Checklist