fix(chat): improve reasoning token estimation logic - #1155
Conversation
WalkthroughAdds logic to compute reasoning token counts when reasoning content exists but explicit reasoningTokens are missing. Applies encoding-based token measurement with a fallback estimate on failure, and propagates the calculated value into streaming and non‑streaming usage logging, final usage chunks, and cost calculations. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant ChatHandler
participant TokenEncoder
participant Logger
participant CostCalc
Note over ChatHandler: Incoming request (streaming or non-streaming)
Client->>ChatHandler: send chat request (may include reasoningContent, reasoningTokens?)
alt reasoningTokens present
ChatHandler->>Logger: use provided reasoningTokens
ChatHandler->>CostCalc: compute costs with provided reasoningTokens
else reasoningContent present and reasoningTokens missing
ChatHandler->>TokenEncoder: encode reasoningContent
alt encode succeeds
TokenEncoder-->>ChatHandler: tokenCount
else encode fails
TokenEncoder-->>ChatHandler: error -> fallbackEstimate
end
ChatHandler->>Logger: log calculatedReasoningTokens
ChatHandler->>CostCalc: compute costs with calculatedReasoningTokens
end
alt streaming
Note over ChatHandler: propagate calculatedReasoningTokens into stream final usage chunk and log
else non-streaming
Note over ChatHandler: include calculatedReasoningTokens in final response logging
end
ChatHandler-->>Client: stream/return response
Estimated code review effort🎯 2 (Simple) | ⏱️ ~12 minutes
Possibly related PRs
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
apps/gateway/src/chat/chat.ts (2)
3065-3077: Use calculated reasoning tokens for cost calculation.The
calculateCostsfunction is called withreasoningTokens(line 3076), but the log insertion usescalculatedReasoningTokens(line 3149). When reasoning tokens are calculated from content rather than provided by the API, the cost calculation will miss reasoning token costs, leading to inaccurate billing.Apply this diff:
const costs = calculateCosts( usedModel, usedProvider, calculatedPromptTokens, calculatedCompletionTokens, cachedTokens, { prompt: messages.map((m) => m.content).join("\n"), completion: fullContent, toolResults: streamingToolCalls || undefined, }, - reasoningTokens, + calculatedReasoningTokens, );
3614-3626: Use calculated reasoning tokens for cost calculation.The
calculateCostsfunction is called withreasoningTokens(line 3625), but the log insertion usescalculatedReasoningTokens(line 3715). When reasoning tokens are calculated from content rather than provided by the API, the cost calculation will miss reasoning token costs, leading to inaccurate billing.Apply this diff:
const costs = calculateCosts( usedModel, usedProvider, calculatedPromptTokens, calculatedCompletionTokens, cachedTokens, { prompt: messages.map((m) => m.content).join("\n"), completion: content, toolResults: toolResults, }, - reasoningTokens, + calculatedReasoningTokens, );
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
apps/gateway/src/chat/chat.ts(4 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
apps/gateway/src/chat/chat.ts (1)
packages/logger/src/index.ts (2)
error(153-160)logger(181-181)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (10)
- GitHub Check: e2e-shards (5)
- GitHub Check: e2e-shards (3)
- GitHub Check: e2e-shards (4)
- GitHub Check: generate / run
- GitHub Check: lint / run
- GitHub Check: e2e-shards (2)
- GitHub Check: test / run
- GitHub Check: build / run
- GitHub Check: e2e-shards (1)
- GitHub Check: autofix
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
apps/gateway/src/chat/chat.ts (2)
3063-3075: Critical: Use calculated reasoning tokens for cost calculation.The
calculateCostscall usesreasoningTokens(the original value from the provider) instead ofcalculatedReasoningTokens(the estimated value). This means the reasoning token estimation logic added in this PR won't affect cost calculations, which is likely incorrect.Apply this diff:
const costs = calculateCosts( usedModel, usedProvider, calculatedPromptTokens, calculatedCompletionTokens, cachedTokens, { prompt: messages.map((m) => m.content).join("\n"), completion: fullContent, toolResults: streamingToolCalls || undefined, }, - reasoningTokens, + calculatedReasoningTokens, );
3610-3622: Critical: Use calculated reasoning tokens for cost calculation.Same issue as in the streaming path:
calculateCostsusesreasoningTokensinstead ofcalculatedReasoningTokens. The estimated reasoning tokens won't affect cost calculations.Apply this diff:
const costs = calculateCosts( usedModel, usedProvider, calculatedPromptTokens, calculatedCompletionTokens, cachedTokens, { prompt: messages.map((m) => m.content).join("\n"), completion: content, toolResults: toolResults, }, - reasoningTokens, + calculatedReasoningTokens, );
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
apps/gateway/src/chat/chat.ts(4 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
apps/gateway/src/chat/chat.ts (1)
packages/logger/src/index.ts (2)
error(153-160)logger(181-181)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (10)
- GitHub Check: build / run
- GitHub Check: test / run
- GitHub Check: lint / run
- GitHub Check: generate / run
- GitHub Check: e2e-shards (3)
- GitHub Check: e2e-shards (5)
- GitHub Check: e2e-shards (4)
- GitHub Check: e2e-shards (1)
- GitHub Check: e2e-shards (2)
- GitHub Check: autofix
| // Estimate reasoning tokens if not provided but reasoning content exists | ||
| let calculatedReasoningTokens = reasoningTokens; | ||
| if (!reasoningTokens && fullReasoningContent) { | ||
| try { | ||
| calculatedReasoningTokens = encode(fullReasoningContent).length; | ||
| } catch (error) { | ||
| // Fallback to simple estimation if encoding fails | ||
| logger.error( | ||
| "Failed to encode reasoning text in streaming", | ||
| error instanceof Error ? error : new Error(String(error)), | ||
| ); | ||
| calculatedReasoningTokens = | ||
| estimateTokensFromContent(fullReasoningContent); | ||
| } | ||
| } |
There was a problem hiding this comment.
Use explicit null check to preserve provider-supplied zero values.
The condition !reasoningTokens evaluates to true when reasoningTokens is 0, which means the code will recalculate even if a provider explicitly returns 0 to indicate no reasoning tokens were used.
Apply this diff to use a more explicit check:
// Estimate reasoning tokens if not provided but reasoning content exists
let calculatedReasoningTokens = reasoningTokens;
-if (!reasoningTokens && fullReasoningContent) {
+if (reasoningTokens == null && fullReasoningContent) {
try {
calculatedReasoningTokens = encode(fullReasoningContent).length;
} catch (error) {
// Fallback to simple estimation if encoding fails
logger.error(
"Failed to encode reasoning text in streaming",
error instanceof Error ? error : new Error(String(error)),
);
calculatedReasoningTokens =
estimateTokensFromContent(fullReasoningContent);
}
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| // Estimate reasoning tokens if not provided but reasoning content exists | |
| let calculatedReasoningTokens = reasoningTokens; | |
| if (!reasoningTokens && fullReasoningContent) { | |
| try { | |
| calculatedReasoningTokens = encode(fullReasoningContent).length; | |
| } catch (error) { | |
| // Fallback to simple estimation if encoding fails | |
| logger.error( | |
| "Failed to encode reasoning text in streaming", | |
| error instanceof Error ? error : new Error(String(error)), | |
| ); | |
| calculatedReasoningTokens = | |
| estimateTokensFromContent(fullReasoningContent); | |
| } | |
| } | |
| // Estimate reasoning tokens if not provided but reasoning content exists | |
| let calculatedReasoningTokens = reasoningTokens; | |
| if (reasoningTokens == null && fullReasoningContent) { | |
| try { | |
| calculatedReasoningTokens = encode(fullReasoningContent).length; | |
| } catch (error) { | |
| // Fallback to simple estimation if encoding fails | |
| logger.error( | |
| "Failed to encode reasoning text in streaming", | |
| error instanceof Error ? error : new Error(String(error)), | |
| ); | |
| calculatedReasoningTokens = | |
| estimateTokensFromContent(fullReasoningContent); | |
| } | |
| } |
🤖 Prompt for AI Agents
In apps/gateway/src/chat/chat.ts around lines 2917 to 2931, the check `if
(!reasoningTokens && fullReasoningContent)` treats a provider-supplied 0 as
missing and will recalculate; change the condition to explicitly test for
null/undefined (e.g., reasoningTokens === null || reasoningTokens === undefined)
combined with fullReasoningContent so that an explicit zero is preserved, then
proceed with the existing try/catch encoding and fallback estimation logic.
| // Estimate reasoning tokens if not provided but reasoning content exists | ||
| let calculatedReasoningTokens = reasoningTokens; | ||
| if (!reasoningTokens && reasoningContent) { | ||
| try { | ||
| calculatedReasoningTokens = encode(reasoningContent).length; | ||
| } catch (error) { | ||
| // Fallback to simple estimation if encoding fails | ||
| logger.error( | ||
| "Failed to encode reasoning text", | ||
| error instanceof Error ? error : new Error(String(error)), | ||
| ); | ||
| calculatedReasoningTokens = estimateTokensFromContent(reasoningContent); | ||
| } | ||
| } |
There was a problem hiding this comment.
Use explicit null check to preserve provider-supplied zero values.
Same issue as in the streaming path: the condition !reasoningTokens will be true when reasoningTokens is 0, causing recalculation even when a provider explicitly returns 0.
Apply this diff:
// Estimate reasoning tokens if not provided but reasoning content exists
let calculatedReasoningTokens = reasoningTokens;
-if (!reasoningTokens && reasoningContent) {
+if (reasoningTokens == null && reasoningContent) {
try {
calculatedReasoningTokens = encode(reasoningContent).length;
} catch (error) {
// Fallback to simple estimation if encoding fails
logger.error(
"Failed to encode reasoning text",
error instanceof Error ? error : new Error(String(error)),
);
calculatedReasoningTokens = estimateTokensFromContent(reasoningContent);
}
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| // Estimate reasoning tokens if not provided but reasoning content exists | |
| let calculatedReasoningTokens = reasoningTokens; | |
| if (!reasoningTokens && reasoningContent) { | |
| try { | |
| calculatedReasoningTokens = encode(reasoningContent).length; | |
| } catch (error) { | |
| // Fallback to simple estimation if encoding fails | |
| logger.error( | |
| "Failed to encode reasoning text", | |
| error instanceof Error ? error : new Error(String(error)), | |
| ); | |
| calculatedReasoningTokens = estimateTokensFromContent(reasoningContent); | |
| } | |
| } | |
| // Estimate reasoning tokens if not provided but reasoning content exists | |
| let calculatedReasoningTokens = reasoningTokens; | |
| if (reasoningTokens == null && reasoningContent) { | |
| try { | |
| calculatedReasoningTokens = encode(reasoningContent).length; | |
| } catch (error) { | |
| // Fallback to simple estimation if encoding fails | |
| logger.error( | |
| "Failed to encode reasoning text", | |
| error instanceof Error ? error : new Error(String(error)), | |
| ); | |
| calculatedReasoningTokens = estimateTokensFromContent(reasoningContent); | |
| } | |
| } |
🤖 Prompt for AI Agents
In apps/gateway/src/chat/chat.ts around lines 3596 to 3609, the code uses a
falsy check (`!reasoningTokens`) which treats a valid provider-supplied zero as
missing; change the condition to an explicit null/undefined check (e.g.,
`reasoningTokens == null` or `reasoningTokens === undefined || reasoningTokens
=== null`) so zero is preserved, keep the existing try/catch and fallback logic
intact, and mirror the same explicit check used in the streaming path if
present.
1 participant
Summary by CodeRabbit