Conversation
|
|
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (11)
🚧 Files skipped from review as they are similar to previous changes (11)
📝 WalkthroughSummary by CodeRabbit
WalkthroughCohere rerank responses now expose billed search-unit usage. Rerank pricing adds ChangesCohere rerank usage and pricing
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant CohereRerankResponse
participant ToBifrostRerankResponse
participant PricingCatalog
participant computeRerankCost
CohereRerankResponse->>ToBifrostRerankResponse: return billed search units
ToBifrostRerankResponse->>PricingCatalog: provide Usage.NumSearchQueries
PricingCatalog->>computeRerankCost: provide InputCostPerQuery
computeRerankCost->>computeRerankCost: multiply rate by NumSearchQueries
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 golangci-lint (2.12.2)level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies" 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 |
Contributor
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
Adds support for treating input_cost_per_query as the per-query rerank/search-context cost during Entry JSON unmarshalling, and ensures Cohere rerank responses report usage even when only “search units” are present.
Changes:
- Extend
Entry.UnmarshalJSONto foldinput_cost_per_queryintoSearchContextCostPerQuerywhen no explicitsearch_context_cost_per_queryis provided. - Add datasheet unmarshalling tests for per-query pricing precedence behavior.
- Populate Bifrost rerank usage from Cohere “search units” and add tests for that mapping.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| framework/modelcatalog/datasheet/types_test.go | Adds tests covering per-query pricing unmarshalling and precedence. |
| framework/modelcatalog/datasheet/types.go | Implements input_cost_per_query folding into SearchContextCostPerQuery. |
| core/providers/cohere/rerank_test.go | Adds tests ensuring search unit usage is surfaced in Bifrost usage. |
| core/providers/cohere/rerank.go | Populates Usage when search units are present even without token counts. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
|
||
| require.NoError(t, err) | ||
| require.NotNil(t, entry.SearchContextCostPerQuery) | ||
| assert.Equal(t, 0.002, *entry.SearchContextCostPerQuery) |
|
|
||
| require.NoError(t, err) | ||
| require.NotNil(t, entry.SearchContextCostPerQuery) | ||
| assert.Equal(t, 0.01, *entry.SearchContextCostPerQuery) |
Comment on lines
+93
to
+98
| // Rerank entries carry their per-query rate as input_cost_per_query; fold it | ||
| // onto SearchContextCostPerQuery so computeRerankCost can consume it. An | ||
| // explicit search_context_cost_per_query value always wins. | ||
| if p.SearchContextCostPerQuery == nil && raw.InputCostPerQuery != nil { | ||
| p.SearchContextCostPerQuery = raw.InputCostPerQuery | ||
| } |
Contributor
There was a problem hiding this comment.
@kypkk this is a valid comment - we can't blindly associate these fields at the central place. We may have to do it in reranking flow (the mapping)
Contributor
Author
There was a problem hiding this comment.
@akshaydeo I've addressed the review — gated the fold on mode == "rerank" as discussed. Could you take another look when you have a moment? 🙏
Comment on lines
+96
to
+98
| if p.SearchContextCostPerQuery == nil && raw.InputCostPerQuery != nil { | ||
| p.SearchContextCostPerQuery = raw.InputCostPerQuery | ||
| } |
Contributor
Confidence Score: 5/5This looks safe to merge.
Important Files Changed
Reviews (13): Last reviewed commit: "refactor: promote input_cost_per_query t..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
core/providers/cohere/rerank_test.go (2)
105-122: 🧹 Nitpick | 🔵 Trivial | 💤 Low valuePrefer
schemas.Ptr()for pointer literals.The test creates local variables for
searchUnits,inputTokens, andoutputTokensand takes their addresses. Per repository conventions, preferschemas.Ptr(value)directly.♻️ Suggested refactor
func TestCohereRerankResponseToBifrostRerankResponseSearchUnitsWithTokenUsage(t *testing.T) { - searchUnits := 1 - inputTokens := 7 - outputTokens := 3 - response := (&CohereRerankResponse{ ID: "rerank-response-id", Results: []CohereRerankResult{ {Index: 0, RelevanceScore: 0.91}, }, Meta: &CohereRerankMeta{ BilledUnits: &CohereBilledUnits{ - InputTokens: &inputTokens, - OutputTokens: &outputTokens, - SearchUnits: &searchUnits, + InputTokens: schemas.Ptr(7), + OutputTokens: schemas.Ptr(3), + SearchUnits: schemas.Ptr(1), }, }, }).ToBifrostRerankResponse(nil, false)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@core/providers/cohere/rerank_test.go` around lines 105 - 122, The test TestCohereRerankResponseToBifrostRerankResponseSearchUnitsWithTokenUsage constructs pointer fields by taking addresses of local ints; replace those with schemas.Ptr(...) calls instead. Update the CohereRerankResponse.Meta.BilledUnits fields (CohereBilledUnits.InputTokens, OutputTokens, SearchUnits) to use schemas.Ptr(7), schemas.Ptr(3), and schemas.Ptr(1) respectively so the test follows repository pointer conventions; no other logic in CohereRerankResponse or ToBifrostRerankResponse needs changing.Source: Learnings
81-92: 🧹 Nitpick | 🔵 Trivial | 💤 Low valuePrefer
schemas.Ptr()for pointer literals.The test creates a local
searchUnitsvariable and takes its address. Per repository conventions, preferschemas.Ptr(2)directly for cleaner test fixtures.♻️ Suggested refactor
func TestCohereRerankResponseToBifrostRerankResponseSearchUnitsUsage(t *testing.T) { - searchUnits := 2 - response := (&CohereRerankResponse{ ID: "rerank-response-id", Results: []CohereRerankResult{ {Index: 0, RelevanceScore: 0.91}, }, Meta: &CohereRerankMeta{ - BilledUnits: &CohereBilledUnits{SearchUnits: &searchUnits}, + BilledUnits: &CohereBilledUnits{SearchUnits: schemas.Ptr(2)}, }, }).ToBifrostRerankResponse(nil, false)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@core/providers/cohere/rerank_test.go` around lines 81 - 92, Replace the local pointer variable usage in TestCohereRerankResponseToBifrostRerankResponseSearchUnitsUsage: instead of creating searchUnits := 2 and taking &searchUnits for CohereRerankMeta.Costs/CohereBilledUnits.SearchUnits, use schemas.Ptr(2) so the test fixture directly sets Meta: &CohereRerankMeta{BilledUnits: &CohereBilledUnits{SearchUnits: schemas.Ptr(2)}}; update any similar pointer literal usages in this test to follow the same pattern.Source: Learnings
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@core/providers/cohere/rerank_test.go`:
- Around line 105-122: The test
TestCohereRerankResponseToBifrostRerankResponseSearchUnitsWithTokenUsage
constructs pointer fields by taking addresses of local ints; replace those with
schemas.Ptr(...) calls instead. Update the CohereRerankResponse.Meta.BilledUnits
fields (CohereBilledUnits.InputTokens, OutputTokens, SearchUnits) to use
schemas.Ptr(7), schemas.Ptr(3), and schemas.Ptr(1) respectively so the test
follows repository pointer conventions; no other logic in CohereRerankResponse
or ToBifrostRerankResponse needs changing.
- Around line 81-92: Replace the local pointer variable usage in
TestCohereRerankResponseToBifrostRerankResponseSearchUnitsUsage: instead of
creating searchUnits := 2 and taking &searchUnits for
CohereRerankMeta.Costs/CohereBilledUnits.SearchUnits, use schemas.Ptr(2) so the
test fixture directly sets Meta: &CohereRerankMeta{BilledUnits:
&CohereBilledUnits{SearchUnits: schemas.Ptr(2)}}; update any similar pointer
literal usages in this test to follow the same pattern.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: ae936b80-a218-4f95-b57a-4cb5e7ec6988
📒 Files selected for processing (2)
core/providers/cohere/rerank_test.goframework/modelcatalog/datasheet/types_test.go
akshaydeo
requested changes
Jun 11, 2026
Comment on lines
+93
to
+98
| // Rerank entries carry their per-query rate as input_cost_per_query; fold it | ||
| // onto SearchContextCostPerQuery so computeRerankCost can consume it. An | ||
| // explicit search_context_cost_per_query value always wins. | ||
| if p.SearchContextCostPerQuery == nil && raw.InputCostPerQuery != nil { | ||
| p.SearchContextCostPerQuery = raw.InputCostPerQuery | ||
| } |
Contributor
There was a problem hiding this comment.
@kypkk this is a valid comment - we can't blindly associate these fields at the central place. We may have to do it in reranking flow (the mapping)
Contributor
Author
|
@akshaydeo So do you agree that I keep the fold at parse time, but gate it on mode == "rerank" so the rate can only ever attach to rerank entries and can't leak into the web-search pricing path? |
Contributor
|
Yes |
Contributor
Author
|
got it onto it |
Comment on lines
+175
to
+181
| // Rerank is billed by search units rather than tokens, so usage must be | ||
| // populated even when no token counts are present. | ||
| var searchUnits *int | ||
| if response.Meta.BilledUnits != nil && response.Meta.BilledUnits.SearchUnits != nil { | ||
| searchUnits = response.Meta.BilledUnits.SearchUnits | ||
| } | ||
| if hasTokenUsage || searchUnits != nil { |
|
|
||
| require.NoError(t, err) | ||
| require.NotNil(t, entry.SearchContextCostPerQuery) | ||
| assert.Equal(t, 0.002, *entry.SearchContextCostPerQuery) |
Comment on lines
+187
to
+191
| if searchUnits != nil { | ||
| bifrostResponse.Usage.CompletionTokensDetails = &schemas.ChatCompletionTokensDetails{ | ||
| NumSearchQueries: schemas.Ptr(*searchUnits), | ||
| } | ||
| } |
kypkk
changed the title
18 tasks
akshaydeo
force-pushed
the
dev
branch
4 times, most recently
from
June 21, 2026 11:44
fa15f50 to
ca190fc
Compare
kypkk
force-pushed
the
fix/cohere-rerank-search-units
branch
from
June 24, 2026 15:27
3ff01a4 to
b4bd9ba
Compare
Contributor
Author
|
rebased, ready for review |
akshaydeo
force-pushed
the
dev
branch
2 times, most recently
from
July 1, 2026 12:24
ac30a53 to
7c66b20
Compare
kypkk
force-pushed
the
fix/cohere-rerank-search-units
branch
from
July 5, 2026 23:04
64bfbc5 to
510cce8
Compare
kypkk
force-pushed
the
fix/cohere-rerank-search-units
branch
from
July 13, 2026 17:45
92ec7fa to
d07cb9a
Compare
kypkk
force-pushed
the
fix/cohere-rerank-search-units
branch
from
July 13, 2026 18:10
d07cb9a to
92e2c59
Compare
akshaydeo
dismissed
coderabbitai[bot]’s stale review
July 14, 2026 08:11
The merge-base changed after approval.
## Summary Briefly explain the purpose of this PR and the problem it solves. ## Changes - What was changed and why - Any notable design decisions or trade-offs ## Type of change - [ ] Bug fix - [ ] Feature - [ ] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [ ] Core (Go) - [ ] Transports (HTTP) - [ ] Providers/Integrations - [ ] Plugins - [ ] UI (React) - [ ] Docs ## How to test Describe the steps to validate this change. Include commands and expected outcomes. ```sh # Core/Transports go version go test ./... # UI cd ui pnpm i || npm i pnpm test || npm test pnpm build || npm run build ``` If adding new configs or environment variables, document them here. ## Screenshots/Recordings If UI changes, add before/after screenshots or short clips. ## Breaking changes - [ ] Yes - [ ] No If yes, describe impact and migration instructions. ## Related issues Link related issues and discussions. Example: Closes maximhq#123 ## Security considerations Note any security implications (auth, secrets, PII, sandboxing, etc.). ## Checklist - [ ] I read `docs/contributing/README.md` and followed the guidelines - [ ] I added/updated tests where appropriate - [ ] I updated documentation where needed - [ ] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable
…nk pricing Cohere bills rerank by search units, not tokens. The response carries meta.billed_units.search_units and no token counts, so Usage stayed nil and rerank cost was always null in logs/UI/metrics. - populate Usage.CompletionTokensDetails.NumSearchQueries from billed_units.search_units (with a non-nil Usage) so computeRerankCost can consume it; token counts, when present, are still mapped alongside - parse the datasheet rerank per-query rate (input_cost_per_query) and fold it onto SearchContextCostPerQuery, gated on mode == "rerank" so it can never leak into the web-search pricing path; an explicit search_context_cost_per_query value always wins - regression tests: search-units-only usage, combined token+search-unit usage, per-query rate parsing, non-rerank mode gating, and tiered-value precedence (floats asserted with InDelta) Fixes maximhq#4239
Per review: instead of folding the rerank per-query rate onto SearchContextCostPerQuery in the central datasheet parser, give it full custom-pricing wiring (mirroring the fast-mode cache pricing pattern): - Options and TableModelPricing gain InputCostPerQuery, with the add_input_cost_per_query_column configstore migration, pricing sync column, override patching, and Entry<->TableModelPricing conversion - computeRerankCost consumes InputCostPerQuery for the per-query term, falling back to search_context_cost_per_query for entries and overrides that still use the shared field - on rerank rows, an explicit override of the shared search_context_cost_per_query field supersedes the datasheet input_cost_per_query rate in patchPricing, so pre-existing rerank overrides saved via the shared field keep taking effect - custom pricing override UI exposes the new field for rerank - the UnmarshalJSON fold is removed; the rerank-only association now lives in the rerank cost path Also hardens the Cohere rerank usage conversion per review: the search-units read guards its full chain independently, and CompletionTokensDetails is initialized only when missing instead of being replaced.
kypkk
force-pushed
the
fix/cohere-rerank-search-units
branch
from
July 15, 2026 16:23
92e2c59 to
6942574
Compare
1 task
akshaydeo
dismissed
coderabbitai[bot]’s stale review
July 17, 2026 23:27
The merge-base changed after approval.
akshaydeo
force-pushed
the
dev
branch
2 times, most recently
from
July 18, 2026 01:10
44564de to
493bff0
Compare
18 tasks
18 tasks
18 tasks
akshaydeo
force-pushed
the
dev
branch
3 times, most recently
from
August 13, 2026 09:47
244a01d to
ce1b2a6
Compare
18 tasks
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
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
Cohere bills rerank by search units, not tokens. The rerank response carries
meta.billed_units.search_unitsand no token counts, so the usage conversionleft
Usageasnil— every Cohere rerank request showedusage: null,cost: nullin logs/UI, and nobifrost_cost_totalincrement. Additionally,datasheet rerank entries carry their per-query rate under
input_cost_per_query, a key with no corresponding field, so the rate wassilently discarded on unmarshal. This PR fixes both ends of the data flow so
rerank cost is tracked like every other request type.
Changes
ToBifrostRerankResponsenow populatesUsage.CompletionTokensDetails.NumSearchQueriesfrom
meta.billed_units.search_units(with a non-nilUsage), whichcomputeRerankCostalready consumes. Token counts, when present, are stillmapped alongside.
Entry.UnmarshalJSONin the pricing datasheet now parsesinput_cost_per_queryand folds it ontoSearchContextCostPerQuery—reusing the existing cost path and DB column, no schema change needed. An
explicit
search_context_cost_per_queryvalue (Perplexity's tiered object)always takes precedence.
search units + token usage combined, per-query rate parsing, and tiered-value
precedence.
Type of change
Affected areas
How to test
Expected: the two new search-units usage tests and the two new unmarshal tests
pass; full cohere and datasheet suites pass. End-to-end: a
POST /v1/rerankagainst a Cohere provider now returns a non-null
usagewithcompletion_tokens_details.num_search_queries, and rerank cost appears inlogs/UI/metrics once the datasheet rerank entry's per-query rate is loaded.
Screenshots/Recordings
N/A
Breaking changes
Related issues
Fixes #4239
Security considerations
None.
Checklist
docs/contributing/README.mdand followed the guidelines