feat(rerank): emit UsageEvent on /v1/rerank 200 (#405) - #428
Merged
Conversation
Pre-#405, /v1/rerank dropped the UsageEvent entirely. Customers using Cohere or Voyage rerank had spend invisible to cp-api's budget ledger and customer-facing /logs analytics. This PR mirrors PR #402 (embeddings) — rerank, like embeddings, has no completion side, no streaming, no reasoning tokens; the extractor handles the three known wire shapes: - OpenAI-compat: `usage.prompt_tokens` (or `usage.input_tokens`) - Jina: `usage.total_tokens` - Cohere: `meta.billed_units.input_tokens` All three end up surfaced as `UsageEvent.prompt_tokens` because cp-api's `dpmgr_usage_events` has no rerank-specific column; the value is what gets multiplied by the model's per-token price. Architecture: dispatch now parses the upstream body bytes once for usage extraction, then forwards the raw bytes verbatim downstream so any provider-specific fields (Cohere `meta`, Jina extras) round-trip without re-formatting. Emit semantics (lessons baked in from #425 audit MEDIUM-1/-2): - 200 with recognisable usage field → emit - 200 without recognisable usage field → no emit (avoids zero-everything noise rows) - 4xx / 5xx → no emit (negative pinning test) Tests (4 new): - `emits_usage_event_on_200_openai_compat_issue_405` — OpenAI/Jina `usage.prompt_tokens` shape - `emits_usage_event_on_cohere_wire_shape_issue_405` — Cohere `meta.billed_units.input_tokens` shape - `skips_usage_event_when_upstream_lacks_usage_fields` — no recognisable shape → no emit - `upstream_5xx_does_not_emit_usage_event` — negative pinning References: - Parent: #226 - Sibling MVPs: #402 (embeddings), #425 (responses), #426 (completions) - Cohere spec: <https://docs.cohere.com/reference/rerank> - Jina spec: <https://api.jina.ai/v1/rerank>
ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Free Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughThe ChangesRerank Usage Event Emission
Sequence DiagramsequenceDiagram
participant Client
participant RerankHandler
participant UpstreamProvider
participant UsageSink
Client->>RerankHandler: POST /v1/rerank
RerankHandler->>UpstreamProvider: dispatch request
UpstreamProvider-->>RerankHandler: 200 response with usage
RerankHandler->>RerankHandler: parse response body as JSON
RerankHandler->>RerankHandler: extract_rerank_usage (provider-specific)
RerankHandler->>UsageSink: emit_usage_event (prompt_tokens)
RerankHandler-->>Client: rerank response
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Note 🎁 Summarized by CodeRabbit FreeYour organization has reached its limit of developer seats under the Pro Plan. For new users, CodeRabbit will generate a high-level summary and a walkthrough for each pull request. For a comprehensive line-by-line review, please add seats to your subscription by visiting https://app.coderabbit.ai/login.If you believe this is a mistake and have available seats, please assign one to the pull request author through the subscription management page using the link above. Comment |
PR #428 audit raised 1 HIGH + 1 MEDIUM. Both addressed: HIGH — Silent parse failure dropped billing. The previous code did `serde_json::from_slice(&bytes).ok().and_then(...)` which made an upstream that returned 200 + malformed body produce zero billing with zero visibility. Operators couldn't see this in dashboards because the failure surfaced only as missing UsageEvents (no log line, no metric). Fixed: log a `tracing::warn!` with the request_id, model name, and parse error so the failure is operator-visible. MEDIUM — Jina's wire shape uses `usage.total_tokens` only (no `prompt_tokens` or `input_tokens` field). The extractor's precedence chain has the right fallback, but no test exercised the Jina-only path with a real emit assertion — the existing `jina_provider_dispatches_to_upstream_with_bearer_auth` test doesn't wire `usage_sink`, so a refactor breaking the `total_tokens` arm would silently zero every Jina-backed billing row. Added `emits_usage_event_on_jina_total_tokens_only_shape_audit_m1` asserting `event.prompt_tokens == 19` for the Jina-only payload.
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
1 participant
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
Fixes #405. Sibling of PR #402 (embeddings), PR #425 (responses), PR #426 (completions).
Pre-fix, `/v1/rerank` dropped the `UsageEvent` entirely. Customers using Cohere / Voyage rerank had spend invisible to cp-api's budget ledger and customer-facing /logs analytics. This PR closes the gap for all three supported upstreams (OpenAI-compat, Cohere, Jina).
Wire-shape coverage
The three rerank-supporting providers each surface tokens differently. The extractor handles all three:
All three surface as `UsageEvent.prompt_tokens` — rerank has no completion side, and cp-api's `dpmgr_usage_events` table has no rerank-specific columns. The single counter is what gets multiplied by the model's per-token price for billing.
Architecture note
Dispatch now parses the upstream body bytes once for usage extraction, then forwards the raw bytes verbatim downstream. This preserves provider-specific fields (Cohere `meta.api_version`, Jina extras) that a JSON round-trip would reformat.
Emit semantics
Lessons baked in from PR #425 audit MEDIUM-1 / MEDIUM-2:
Test plan
References
Summary by CodeRabbit