fix: prompt calculation - #1606
Conversation
User will correctly get estimated prompt usage when upstream returns either zero or nothing.
WalkthroughAdds prompt token storage to relayInfo during request handling and revises OpenAI response handling to recompute token usage when PromptTokens is zero. Introduces a usageModified flag and conditions formatting to run only when forced or when usage was updated, otherwise forwarding the upstream payload unchanged. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Relay
participant Tokenizer
participant PriceHelper
Client->>Relay: Request
Relay->>Tokenizer: Compute prompt tokens
Tokenizer-->>Relay: tokens
Relay->>Relay: relayInfo.SetPromptTokens(tokens)
Relay->>PriceHelper: Calculate price(tokens, ...)
PriceHelper-->>Relay: price
Relay-->>Client: Response
sequenceDiagram
participant Upstream as OpenAI Upstream
participant Handler as OpenAI Handler
participant Client
Upstream-->>Handler: Response (body, usage)
alt PromptTokens == 0
Handler->>Handler: Initialize completionTokens
opt completionTokens == 0
Handler->>Handler: Sum tokens across choices
end
Handler->>Handler: Update usage and set usageModified = true
end
alt forceFormat || usageModified
Handler->>Handler: Reformat simpleResponse and marshal
Handler-->>Client: Formatted response with updated usage
else
Handler-->>Client: Forward original payload unchanged
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. ✨ Finishing Touches
🧪 Generate unit tests
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
relay/channel/openai/relay-openai.go (3)
200-215: Avoid overwriting entire Usage; update in place to preserve detailsWhen recomputing usage you replace the whole dto.Usage, which can clobber upstream fields like PromptTokensDetails, CompletionTokenDetails, Input/OutputTokens, Cost, etc. Prefer updating only the specific fields you fix (PromptTokens, CompletionTokens, TotalTokens).
Apply this diff within the current block:
-usageModified := false +usageModified := false if simpleResponse.Usage.PromptTokens == 0 { completionTokens := simpleResponse.Usage.CompletionTokens if completionTokens == 0 { for _, choice := range simpleResponse.Choices { ctkm := service.CountTextToken(choice.Message.StringContent()+choice.Message.ReasoningContent+choice.Message.Reasoning, info.UpstreamModelName) completionTokens += ctkm } } - simpleResponse.Usage = dto.Usage{ - PromptTokens: info.PromptTokens, - CompletionTokens: completionTokens, - TotalTokens: info.PromptTokens + completionTokens, - } + // Preserve any existing usage details from upstream; only fill the missing core fields. + simpleResponse.Usage.PromptTokens = info.PromptTokens + simpleResponse.Usage.CompletionTokens = completionTokens + simpleResponse.Usage.TotalTokens = simpleResponse.Usage.PromptTokens + simpleResponse.Usage.CompletionTokens usageModified = true }
203-208: Consider counting tool_calls to reduce underestimationWhen upstream omits usage entirely, completions that are mostly tool_call arguments may be undercounted. Optionally include ToolCalls JSON in the token estimate.
Here’s a minimal, low-risk addition:
- for _, choice := range simpleResponse.Choices { - ctkm := service.CountTextToken(choice.Message.StringContent()+choice.Message.ReasoningContent+choice.Message.Reasoning, info.UpstreamModelName) - completionTokens += ctkm - } + for _, choice := range simpleResponse.Choices { + ctkm := service.CountTextToken( + choice.Message.StringContent()+choice.Message.ReasoningContent+choice.Message.Reasoning, + info.UpstreamModelName, + ) + completionTokens += ctkm + if len(choice.Message.ToolCalls) > 0 { + completionTokens += service.CountTextToken(string(choice.Message.ToolCalls), info.UpstreamModelName) + } + }
219-227: Formatting gate is correct; add a clarifying comment to the breakThe “break” exits the switch, preserving the original upstream response body when no formatting/usage updates are needed. Consider a short inline comment to avoid misreads.
- if forceFormat || usageModified { + if forceFormat || usageModified { responseBody, err = common.Marshal(simpleResponse) if err != nil { return nil, types.NewError(err, types.ErrorCodeBadResponseBody) } } else { - break + // Keep upstream payload as-is; responseBody remains the original bytes + break }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (2)
controller/relay.go(1 hunks)relay/channel/openai/relay-openai.go(1 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-06-21T03:37:41.726Z
Learnt from: 9Ninety
PR: QuantumNous/new-api#1273
File: relay/channel/gemini/relay-gemini.go:97-116
Timestamp: 2025-06-21T03:37:41.726Z
Learning: In relay/channel/gemini/relay-gemini.go, the thinking budget calculation logic (including the MaxOutputTokens multiplication) was introduced in PR #1247. PR #1273 focused specifically on decoupling the thoughts summary feature from thinking budget settings and did not modify the existing thinking budget behavior.
Applied to files:
controller/relay.go
🧬 Code Graph Analysis (1)
relay/channel/openai/relay-openai.go (4)
dto/openai_response.go (1)
Usage(217-230)service/token_counter.go (1)
CountTextToken(641-647)dto/openai_request.go (2)
Message(247-258)Reasoning(844-847)types/relay_format.go (2)
RelayFormat(3-3)RelayFormatOpenAI(6-6)
🔇 Additional comments (1)
controller/relay.go (1)
131-132: Confirmed PromptTokens propagation across codebaseI ran the suggested grep and confirmed that
relayInfo.SetPromptTokens(tokens)(controller/relay.go:131) fires before any downstream reference toPromptTokens. All key handlers—usage helpers, quota logic, Convert services, and individual channel adapters—now derive their prompt-token counts fromrelayInfo.PromptTokens, ensuring no zero/missing values slip through. No further changes are needed.
fix: prompt calculation
2 participants
closes #1602
User will correctly get estimated prompt usage when upstream returns either zero or nothing.
Summary by CodeRabbit
Bug Fixes
Improvements