fix: count tokens request building - #5620
Merged
Merged
Conversation
Contributor
📝 WalkthroughWalkthroughGemini count-token requests now use a dedicated request type, normalized envelopes, and early model validation. Gemini and Vertex response conversion now aggregates modality token details while separating cached read tokens. ChangesGemini Count Tokens
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant GenAIRoute
participant GeminiCountTokens
participant GeminiAPI
GenAIRoute->>GeminiCountTokens: create and populate count-token request
GeminiCountTokens->>GeminiCountTokens: validate model and normalize body
GeminiCountTokens->>GeminiAPI: send countTokens request
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
|
tejas ghatte seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account. You have signed the CLA already but the status is still pending? Let us recheck it. |
Collaborator
Author
This stack of pull requests is managed by Graphite. Learn more about stacking. |
18 tasks
TejasGhatte
marked this pull request as ready for review
July 28, 2026 12:59
coderabbitai
Bot
requested review from
Pratham-Mishra04,
akshaydeo and
roroghost17
July 28, 2026 13:01
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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.
Inline comments:
In `@core/providers/gemini/payload_ordering_test.go`:
- Around line 134-175: Extend TestGeminiCountTokensRequestToGenerationRequest
with a case that sets top-level GeminiCountTokensRequest.Fallbacks, calls
ToGeminiGenerationRequest, and asserts the resulting GeminiGenerationRequest
retains the same fallback chain.
In `@core/providers/gemini/types.go`:
- Around line 2422-2426: Add a top-level Fallbacks field to
GeminiCountTokensRequest in core/providers/gemini/types.go at lines 2422-2426,
copy non-empty request.Fallbacks into generationRequest.Fallbacks in
core/providers/gemini/count_tokens.go at lines 64-69, and update
core/providers/gemini/payload_ordering_test.go at lines 134-175 to verify
top-level fallbacks survive flattening.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 2a575976-bf46-4cd0-b494-c84e14161668
📒 Files selected for processing (8)
core/providers/gemini/count_tokens.gocore/providers/gemini/gemini.gocore/providers/gemini/payload_ordering_test.gocore/providers/gemini/types.gocore/providers/utils/utils.gocore/providers/vertex/count_tokens.gocore/providers/vertex/types.gotransports/bifrost-http/integrations/genai.go
TejasGhatte
force-pushed
the
07-28-fix_count_tokens_request_building
branch
from
July 29, 2026 07:27
c9edcc6 to
6ee7c1f
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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.
Inline comments:
In `@core/providers/gemini/count_tokens.go`:
- Around line 107-108: Update the modality-detail aggregation around
AddModalityTokens so CacheTokensDetails is processed unconditionally, including
when CachedContentTokenCount is nonzero. Keep the top-level cached token count
authoritative while ensuring each cached modality contributes to the text,
image, and audio breakdown.
In `@transports/bifrost-http/integrations/genai.go`:
- Around line 1497-1502: Update the explicit Gemini branch in the
GeminiCountTokensRequest handling to capture and enable the raw request body
only when the configured send_back_raw_request and store_raw_request_response
policies permit it, honoring client.disable_content_logging and both per-request
override gates. Otherwise leave raw-body context unset so later handling cannot
retain or expose the count-token payload.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 5ed75955-e750-407d-983f-71fcca01fc85
📒 Files selected for processing (8)
core/providers/gemini/count_tokens.gocore/providers/gemini/gemini.gocore/providers/gemini/payload_ordering_test.gocore/providers/gemini/types.gocore/providers/utils/utils.gocore/providers/vertex/count_tokens.gocore/providers/vertex/types.gotransports/bifrost-http/integrations/genai.go
🚧 Files skipped from review as they are similar to previous changes (5)
- core/providers/gemini/types.go
- core/providers/vertex/count_tokens.go
- core/providers/utils/utils.go
- core/providers/vertex/types.go
- core/providers/gemini/gemini.go
TejasGhatte
force-pushed
the
07-28-fix_count_tokens_request_building
branch
from
July 29, 2026 07:38
6ee7c1f to
af8a7ee
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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.
Inline comments:
In `@core/providers/gemini/count_tokens.go`:
- Around line 96-109: Update the cached-token handling in the token-count
conversion flow to always pass non-nil entries from resp.CacheTokensDetails
through AddModalityTokens, preserving their modality breakdown. Keep
CachedContentTokenCount authoritative for inputDetails.CachedReadTokens, using
the existing aggregate fallback only when that count is zero.
- Around line 64-73: Update the conversion logic around generationRequest so it
copies request.GenerateContentRequest before assigning Model, IsCountTokens, or
Fallbacks. Preserve the existing default construction when the envelope is nil,
while ensuring the original request envelope remains unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: dc3992ac-89fb-428c-8626-ae41aabbafd1
📒 Files selected for processing (8)
core/providers/gemini/count_tokens.gocore/providers/gemini/gemini.gocore/providers/gemini/payload_ordering_test.gocore/providers/gemini/types.gocore/providers/utils/utils.gocore/providers/vertex/count_tokens.gocore/providers/vertex/types.gotransports/bifrost-http/integrations/genai.go
🚧 Files skipped from review as they are similar to previous changes (7)
- core/providers/vertex/count_tokens.go
- core/providers/gemini/types.go
- core/providers/vertex/types.go
- core/providers/gemini/payload_ordering_test.go
- core/providers/utils/utils.go
- transports/bifrost-http/integrations/genai.go
- core/providers/gemini/gemini.go
18 tasks
TejasGhatte
force-pushed
the
07-28-fix_count_tokens_request_building
branch
from
July 29, 2026 09:02
af8a7ee to
c012275
Compare
8 tasks
Contributor
Merge activity
|
akshaydeo
changed the base branch from
07-28-fix_vertex_count_tokens_unsupported_fields
to
graphite-base/5620
July 29, 2026 11:07
akshaydeo
dismissed
coderabbitai[bot]’s stale review
July 29, 2026 11:07
The base branch was changed.
akshaydeo
pushed a commit
that referenced
this pull request
Jul 29, 2026
## Summary Fixes a regression where `:countTokens` for both Gemini (`generativelanguage`) and Vertex (`aiplatform`) reported only the `contents` token count because `systemInstruction`, `generationConfig`, and `toolConfig` were unconditionally stripped from the request body. In the reported case, a ~7.4k-token system prompt plus tool declarations counted as only 17 tokens on Vertex. The root cause is that the two endpoints require opposite request shapes: - **Gemini**: `systemInstruction`, `tools`, `toolConfig`, and `generationConfig` are rejected as `Unknown name` at the top level and must be wrapped inside a `generateContentRequest` envelope. Top-level `contents`/`model` are silently ignored when the envelope is present. - **Vertex**: These fields are accepted flat; only `toolConfig` (and a few others) must be stripped. There is no envelope. The `/genai` ingress must now parse both shapes and emit whichever the resolved provider requires. The Bifrost-only `fallbacks` routing field must be honoured for routing but stripped before the upstream call. ## Changes - Added E2E test suite **34. Gemini/Vertex countTokens full-prompt accounting** covering: - **34.1 / 34.6** – Contents-only baseline for Gemini and Vertex respectively, recording `totalTokens` as a collection variable for relative assertions in subsequent cases. - **34.2 / 34.7** – Flat `systemInstruction` is counted (not stripped) for both providers. - **34.3** – Flat `tools`, `toolConfig`, and `generationConfig` are wrapped into the Gemini envelope rather than forwarded verbatim (which would produce a hard 400). - **34.8** – Vertex keeps `tools` and `generationConfig` flat but strips `toolConfig`. - **34.4 / 34.9** – A `generateContentRequest` envelope is passed through for Gemini and unwrapped for Vertex. - **34.5** – The Bifrost `fallbacks` field is stripped before the upstream call and does not alter the counted token total. - Assertions are relative (count must exceed the contents-only baseline) rather than absolute, so they remain valid across model version updates. A dropped field collapses the count back to baseline, which is exactly the regression signature these cases catch. - HTTP 400 is intentionally **not** in the infra guard for the tools/toolConfig/fallbacks cases — a 400 from upstream is itself the regression signature for those scenarios. ## Type of change - [x] Bug fix ## Affected areas - [x] Providers/Integrations ## How to test Import the updated `provider-harness.json` collection into Postman or Newman and run group **34** with valid `genaiKey`, `genaiModel`, `vertexModel`, and `baseUrl` collection variables set. ```sh newman run tests/e2e/api/collections/provider-harness.json \ --folder "34. Gemini/Vertex countTokens full-prompt accounting (PR #5620)" \ --env-var baseUrl=<your-base-url> \ --env-var genaiKey=<your-api-key> \ --env-var genaiModel=gemini-2.5-pro \ --env-var vertexModel=gemini-2.5-pro ``` Each case should return HTTP 2xx with `totalTokens` strictly greater than the baseline recorded in 34.1/34.6. Any case that returns an `Unknown name` error or a token count equal to the baseline indicates the regression has re-appeared. ## Breaking changes - [x] No ## Related issues Closes #5620 ## Security considerations No auth, secrets, or PII changes. Test requests use a scoped no-op tool declaration (`probe_context`) with no real side effects. ## Checklist - [x] I read `docs/contributing/README.md` and followed the guidelines - [x] I added/updated tests where appropriate - [ ] I updated documentation where needed - [x] I verified builds succeed (Go and UI) - [x] I verified the CI pipeline passes locally if applicable
akhsaul
pushed a commit
to akhsaul/bifrost
that referenced
this pull request
Aug 27, 2026
## Summary The Gemini `countTokens` endpoint requires all prompt configuration (system instruction, tools, tool config, generation config) to be nested inside a `generateContentRequest` envelope rather than at the top level. Previously, these fields were simply stripped before sending, meaning token counts were incomplete. This PR rewrites the request body into the correct envelope shape so the full prompt is counted accurately. ## Changes - Introduced `wrapGeminiCountTokensBody` which rewrites any flat `generateContent`-style body into the `generateContentRequest` envelope, handles already-enveloped bodies without double-wrapping, qualifies the model name with the `models/` prefix, and strips fields not accepted by `GenerateContentRequest` (e.g. `labels`, `fallbacks`). - Added `GeminiCountTokensRequest` type that accepts both bare `contents` and the `generateContentRequest` envelope, with a `ToGeminiGenerationRequest` converter that resolves precedence between the two. - Replaced the previous approach of deleting `toolConfig`, `generationConfig`, and `systemInstruction` from the top-level body with the envelope wrapping strategy, which preserves those fields inside the envelope instead. - Extracted `AddModalityTokens` as a shared helper that maps Gemini modality strings (`audio`, `image`, `video`, text) to the neutral `ResponsesResponseInputTokens` fields, replacing duplicated inline logic in both `ToBifrostCountTokensResponse` and the Vertex equivalent. - Extended `ToGeminiCountTokensResponse` to round-trip `PromptTokensDetails` per modality and to respect `TotalTokens` when present. - Updated `VertexCountTokensResponse` to include `PromptTokensDetails` and `TotalBillableCharacters`, and wired `AddModalityTokens` into the Vertex count tokens response converter. - Added `SetRawJSONField` to provider utils to insert pre-encoded JSON verbatim without re-marshaling. - Updated the GenAI HTTP transport to parse count tokens requests into `GeminiCountTokensRequest` and route them through the new converter, including raw body passthrough for explicit Gemini requests. - Added comprehensive tests for `wrapGeminiCountTokensBody` and `ToGeminiGenerationRequest` covering envelope wrapping, double-wrap prevention, model prefix normalization, unsupported field stripping, and content fallback behavior. ## Type of change - [ ] Bug fix - [x] Feature - [ ] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [x] Core (Go) - [x] Transports (HTTP) - [x] Providers/Integrations - [ ] Plugins - [ ] UI (React) - [ ] Docs ## How to test ```sh go test ./core/providers/gemini/... ./core/providers/vertex/... ./transports/bifrost-http/... ``` Send a `countTokens` request with a system instruction and tools attached and verify the returned token count reflects the full prompt rather than only the `contents` array. Confirm that a body already wrapped in `generateContentRequest` is not double-wrapped by inspecting the outgoing request payload. ## Screenshots/Recordings N/A ## Breaking changes - [ ] Yes - [x] No ## Related issues N/A ## Security considerations No auth, secrets, or PII implications. The change only affects how request bodies are serialized before being forwarded to the Gemini and Vertex APIs. ## Checklist - [ ] I read `docs/contributing/README.md` and followed the guidelines - [x] I added/updated tests where appropriate - [ ] I updated documentation where needed - [x] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable
akhsaul
pushed a commit
to akhsaul/bifrost
that referenced
this pull request
Aug 27, 2026
## Summary Fixes a regression where `:countTokens` for both Gemini (`generativelanguage`) and Vertex (`aiplatform`) reported only the `contents` token count because `systemInstruction`, `generationConfig`, and `toolConfig` were unconditionally stripped from the request body. In the reported case, a ~7.4k-token system prompt plus tool declarations counted as only 17 tokens on Vertex. The root cause is that the two endpoints require opposite request shapes: - **Gemini**: `systemInstruction`, `tools`, `toolConfig`, and `generationConfig` are rejected as `Unknown name` at the top level and must be wrapped inside a `generateContentRequest` envelope. Top-level `contents`/`model` are silently ignored when the envelope is present. - **Vertex**: These fields are accepted flat; only `toolConfig` (and a few others) must be stripped. There is no envelope. The `/genai` ingress must now parse both shapes and emit whichever the resolved provider requires. The Bifrost-only `fallbacks` routing field must be honoured for routing but stripped before the upstream call. ## Changes - Added E2E test suite **34. Gemini/Vertex countTokens full-prompt accounting** covering: - **34.1 / 34.6** – Contents-only baseline for Gemini and Vertex respectively, recording `totalTokens` as a collection variable for relative assertions in subsequent cases. - **34.2 / 34.7** – Flat `systemInstruction` is counted (not stripped) for both providers. - **34.3** – Flat `tools`, `toolConfig`, and `generationConfig` are wrapped into the Gemini envelope rather than forwarded verbatim (which would produce a hard 400). - **34.8** – Vertex keeps `tools` and `generationConfig` flat but strips `toolConfig`. - **34.4 / 34.9** – A `generateContentRequest` envelope is passed through for Gemini and unwrapped for Vertex. - **34.5** – The Bifrost `fallbacks` field is stripped before the upstream call and does not alter the counted token total. - Assertions are relative (count must exceed the contents-only baseline) rather than absolute, so they remain valid across model version updates. A dropped field collapses the count back to baseline, which is exactly the regression signature these cases catch. - HTTP 400 is intentionally **not** in the infra guard for the tools/toolConfig/fallbacks cases — a 400 from upstream is itself the regression signature for those scenarios. ## Type of change - [x] Bug fix ## Affected areas - [x] Providers/Integrations ## How to test Import the updated `provider-harness.json` collection into Postman or Newman and run group **34** with valid `genaiKey`, `genaiModel`, `vertexModel`, and `baseUrl` collection variables set. ```sh newman run tests/e2e/api/collections/provider-harness.json \ --folder "34. Gemini/Vertex countTokens full-prompt accounting (PR maximhq#5620)" \ --env-var baseUrl=<your-base-url> \ --env-var genaiKey=<your-api-key> \ --env-var genaiModel=gemini-2.5-pro \ --env-var vertexModel=gemini-2.5-pro ``` Each case should return HTTP 2xx with `totalTokens` strictly greater than the baseline recorded in 34.1/34.6. Any case that returns an `Unknown name` error or a token count equal to the baseline indicates the regression has re-appeared. ## Breaking changes - [x] No ## Related issues Closes maximhq#5620 ## Security considerations No auth, secrets, or PII changes. Test requests use a scoped no-op tool declaration (`probe_context`) with no real side effects. ## Checklist - [x] I read `docs/contributing/README.md` and followed the guidelines - [x] I added/updated tests where appropriate - [ ] I updated documentation where needed - [x] I verified builds succeed (Go and UI) - [x] I verified the CI pipeline passes locally if applicable
occcat
pushed a commit
to occcat/bifrost
that referenced
this pull request
Sep 2, 2026
## Summary The Gemini `countTokens` endpoint requires all prompt configuration (system instruction, tools, tool config, generation config) to be nested inside a `generateContentRequest` envelope rather than at the top level. Previously, these fields were simply stripped before sending, meaning token counts were incomplete. This PR rewrites the request body into the correct envelope shape so the full prompt is counted accurately. ## Changes - Introduced `wrapGeminiCountTokensBody` which rewrites any flat `generateContent`-style body into the `generateContentRequest` envelope, handles already-enveloped bodies without double-wrapping, qualifies the model name with the `models/` prefix, and strips fields not accepted by `GenerateContentRequest` (e.g. `labels`, `fallbacks`). - Added `GeminiCountTokensRequest` type that accepts both bare `contents` and the `generateContentRequest` envelope, with a `ToGeminiGenerationRequest` converter that resolves precedence between the two. - Replaced the previous approach of deleting `toolConfig`, `generationConfig`, and `systemInstruction` from the top-level body with the envelope wrapping strategy, which preserves those fields inside the envelope instead. - Extracted `AddModalityTokens` as a shared helper that maps Gemini modality strings (`audio`, `image`, `video`, text) to the neutral `ResponsesResponseInputTokens` fields, replacing duplicated inline logic in both `ToBifrostCountTokensResponse` and the Vertex equivalent. - Extended `ToGeminiCountTokensResponse` to round-trip `PromptTokensDetails` per modality and to respect `TotalTokens` when present. - Updated `VertexCountTokensResponse` to include `PromptTokensDetails` and `TotalBillableCharacters`, and wired `AddModalityTokens` into the Vertex count tokens response converter. - Added `SetRawJSONField` to provider utils to insert pre-encoded JSON verbatim without re-marshaling. - Updated the GenAI HTTP transport to parse count tokens requests into `GeminiCountTokensRequest` and route them through the new converter, including raw body passthrough for explicit Gemini requests. - Added comprehensive tests for `wrapGeminiCountTokensBody` and `ToGeminiGenerationRequest` covering envelope wrapping, double-wrap prevention, model prefix normalization, unsupported field stripping, and content fallback behavior. ## Type of change - [ ] Bug fix - [x] Feature - [ ] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [x] Core (Go) - [x] Transports (HTTP) - [x] Providers/Integrations - [ ] Plugins - [ ] UI (React) - [ ] Docs ## How to test ```sh go test ./core/providers/gemini/... ./core/providers/vertex/... ./transports/bifrost-http/... ``` Send a `countTokens` request with a system instruction and tools attached and verify the returned token count reflects the full prompt rather than only the `contents` array. Confirm that a body already wrapped in `generateContentRequest` is not double-wrapped by inspecting the outgoing request payload. ## Screenshots/Recordings N/A ## Breaking changes - [ ] Yes - [x] No ## Related issues N/A ## Security considerations No auth, secrets, or PII implications. The change only affects how request bodies are serialized before being forwarded to the Gemini and Vertex APIs. ## Checklist - [ ] I read `docs/contributing/README.md` and followed the guidelines - [x] I added/updated tests where appropriate - [ ] I updated documentation where needed - [x] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable
occcat
pushed a commit
to occcat/bifrost
that referenced
this pull request
Sep 2, 2026
## Summary Fixes a regression where `:countTokens` for both Gemini (`generativelanguage`) and Vertex (`aiplatform`) reported only the `contents` token count because `systemInstruction`, `generationConfig`, and `toolConfig` were unconditionally stripped from the request body. In the reported case, a ~7.4k-token system prompt plus tool declarations counted as only 17 tokens on Vertex. The root cause is that the two endpoints require opposite request shapes: - **Gemini**: `systemInstruction`, `tools`, `toolConfig`, and `generationConfig` are rejected as `Unknown name` at the top level and must be wrapped inside a `generateContentRequest` envelope. Top-level `contents`/`model` are silently ignored when the envelope is present. - **Vertex**: These fields are accepted flat; only `toolConfig` (and a few others) must be stripped. There is no envelope. The `/genai` ingress must now parse both shapes and emit whichever the resolved provider requires. The Bifrost-only `fallbacks` routing field must be honoured for routing but stripped before the upstream call. ## Changes - Added E2E test suite **34. Gemini/Vertex countTokens full-prompt accounting** covering: - **34.1 / 34.6** – Contents-only baseline for Gemini and Vertex respectively, recording `totalTokens` as a collection variable for relative assertions in subsequent cases. - **34.2 / 34.7** – Flat `systemInstruction` is counted (not stripped) for both providers. - **34.3** – Flat `tools`, `toolConfig`, and `generationConfig` are wrapped into the Gemini envelope rather than forwarded verbatim (which would produce a hard 400). - **34.8** – Vertex keeps `tools` and `generationConfig` flat but strips `toolConfig`. - **34.4 / 34.9** – A `generateContentRequest` envelope is passed through for Gemini and unwrapped for Vertex. - **34.5** – The Bifrost `fallbacks` field is stripped before the upstream call and does not alter the counted token total. - Assertions are relative (count must exceed the contents-only baseline) rather than absolute, so they remain valid across model version updates. A dropped field collapses the count back to baseline, which is exactly the regression signature these cases catch. - HTTP 400 is intentionally **not** in the infra guard for the tools/toolConfig/fallbacks cases — a 400 from upstream is itself the regression signature for those scenarios. ## Type of change - [x] Bug fix ## Affected areas - [x] Providers/Integrations ## How to test Import the updated `provider-harness.json` collection into Postman or Newman and run group **34** with valid `genaiKey`, `genaiModel`, `vertexModel`, and `baseUrl` collection variables set. ```sh newman run tests/e2e/api/collections/provider-harness.json \ --folder "34. Gemini/Vertex countTokens full-prompt accounting (PR maximhq#5620)" \ --env-var baseUrl=<your-base-url> \ --env-var genaiKey=<your-api-key> \ --env-var genaiModel=gemini-2.5-pro \ --env-var vertexModel=gemini-2.5-pro ``` Each case should return HTTP 2xx with `totalTokens` strictly greater than the baseline recorded in 34.1/34.6. Any case that returns an `Unknown name` error or a token count equal to the baseline indicates the regression has re-appeared. ## Breaking changes - [x] No ## Related issues Closes maximhq#5620 ## Security considerations No auth, secrets, or PII changes. Test requests use a scoped no-op tool declaration (`probe_context`) with no real side effects. ## Checklist - [x] I read `docs/contributing/README.md` and followed the guidelines - [x] I added/updated tests where appropriate - [ ] I updated documentation where needed - [x] I verified builds succeed (Go and UI) - [x] I verified the CI pipeline passes locally if applicable
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The Gemini
countTokensendpoint requires all prompt configuration (system instruction, tools, tool config, generation config) to be nested inside agenerateContentRequestenvelope rather than at the top level. Previously, these fields were simply stripped before sending, meaning token counts were incomplete. This PR rewrites the request body into the correct envelope shape so the full prompt is counted accurately.Changes
wrapGeminiCountTokensBodywhich rewrites any flatgenerateContent-style body into thegenerateContentRequestenvelope, handles already-enveloped bodies without double-wrapping, qualifies the model name with themodels/prefix, and strips fields not accepted byGenerateContentRequest(e.g.labels,fallbacks).GeminiCountTokensRequesttype that accepts both barecontentsand thegenerateContentRequestenvelope, with aToGeminiGenerationRequestconverter that resolves precedence between the two.toolConfig,generationConfig, andsystemInstructionfrom the top-level body with the envelope wrapping strategy, which preserves those fields inside the envelope instead.AddModalityTokensas a shared helper that maps Gemini modality strings (audio,image,video, text) to the neutralResponsesResponseInputTokensfields, replacing duplicated inline logic in bothToBifrostCountTokensResponseand the Vertex equivalent.ToGeminiCountTokensResponseto round-tripPromptTokensDetailsper modality and to respectTotalTokenswhen present.VertexCountTokensResponseto includePromptTokensDetailsandTotalBillableCharacters, and wiredAddModalityTokensinto the Vertex count tokens response converter.SetRawJSONFieldto provider utils to insert pre-encoded JSON verbatim without re-marshaling.GeminiCountTokensRequestand route them through the new converter, including raw body passthrough for explicit Gemini requests.wrapGeminiCountTokensBodyandToGeminiGenerationRequestcovering envelope wrapping, double-wrap prevention, model prefix normalization, unsupported field stripping, and content fallback behavior.Type of change
Affected areas
How to test
go test ./core/providers/gemini/... ./core/providers/vertex/... ./transports/bifrost-http/...Send a
countTokensrequest with a system instruction and tools attached and verify the returned token count reflects the full prompt rather than only thecontentsarray. Confirm that a body already wrapped ingenerateContentRequestis not double-wrapped by inspecting the outgoing request payload.Screenshots/Recordings
N/A
Breaking changes
Related issues
N/A
Security considerations
No auth, secrets, or PII implications. The change only affects how request bodies are serialized before being forwarded to the Gemini and Vertex APIs.
Checklist
docs/contributing/README.mdand followed the guidelines