feat(completions): ensure prompt tokens are always calculated - #588
Conversation
Co-authored-by: contact <contact@polarlights.llc>
|
Cursor Agent can help with this pull request. Just |
WalkthroughThis change set introduces robust fallback logic to ensure that prompt token counts in chat completion usage reporting are never zero, even if the underlying provider returns zero. It updates chat handling logic, enhances the mock server for testing, and adds comprehensive tests for both streaming and non-streaming scenarios, as well as for the internal token calculation logic. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Gateway
participant Provider
Client->>Gateway: Send chat completion request (may trigger ZERO_TOKENS)
Gateway->>Provider: Forward request
Provider-->>Gateway: Respond with usage (possibly zero prompt_tokens)
alt prompt_tokens is zero or missing
Gateway->>Gateway: Estimate prompt_tokens via fallback logic
Gateway->>Gateway: Ensure prompt_tokens >= 1
end
Gateway-->>Client: Respond with usage (prompt_tokens >= 1)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Poem
Note 🔌 MCP (Model Context Protocol) integration is now available in Early Access!Pro users can now connect to remote MCP servers under the Integrations page to get reviews and chat conversations that understand additional development context. 📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
💤 Files with no reviewable changes (1)
⏰ 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). (2)
✨ 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 comments)
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: 5
♻️ Duplicate comments (1)
apps/gateway/src/chat/chat.ts (1)
1013-1045: Same critical issues in duplicated code blockThis block has the same undefined
messagesvariable issue and needs the same fixes as mentioned above.
🧹 Nitpick comments (3)
apps/gateway/src/chat/chat.ts (3)
651-651: Simplify redundant fallback in prompt_tokens calculationThe expression
Math.max(1, promptTokens || 1)is redundant. IfpromptTokensis falsy, it defaults to 1, making theMath.max(1, 1)unnecessary.- prompt_tokens: Math.max(1, promptTokens || 1), + prompt_tokens: Math.max(1, promptTokens || 0),
692-692: Simplify redundant fallback in prompt_tokens calculation (Anthropic)Same issue as above - the expression
Math.max(1, promptTokens || 1)has redundant fallback logic.- prompt_tokens: Math.max(1, promptTokens || 1), + prompt_tokens: Math.max(1, promptTokens || 0),
3032-3040: Improve readability of complex token calculationsThe nested ternary operators make this code hard to read and maintain.
+ const finalPromptTokens = (promptTokens && promptTokens > 0) + ? promptTokens + : (calculatedPromptTokens || 1); + const finalCompletionTokens = completionTokens || calculatedCompletionTokens || 0; + const finalTotalTokens = totalTokens || calculatedTotalTokens || finalPromptTokens; + usage: { - prompt_tokens: Math.max(1, Math.round( - (promptTokens && promptTokens > 0) ? promptTokens : calculatedPromptTokens || 1, - )), - completion_tokens: Math.round( - completionTokens || calculatedCompletionTokens || 0, - ), - total_tokens: Math.round( - totalTokens || calculatedTotalTokens || Math.max(1, (promptTokens && promptTokens > 0) ? promptTokens : calculatedPromptTokens || 1), - ), + prompt_tokens: Math.max(1, Math.round(finalPromptTokens)), + completion_tokens: Math.round(finalCompletionTokens), + total_tokens: Math.max(1, Math.round(finalTotalTokens)),
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
apps/gateway/src/api.e2e.ts(1 hunks)apps/gateway/src/chat/chat.ts(8 hunks)apps/gateway/src/lib/prompt-tokens.spec.ts(1 hunks)apps/gateway/src/test-utils/mock-openai-server.ts(2 hunks)
🧰 Additional context used
📓 Path-based instructions (4)
**/*.{js,jsx,ts,tsx}
📄 CodeRabbit Inference Engine (.github/copilot-instructions.md)
Use localStorage instead of cookies for client-side data persistence
Files:
apps/gateway/src/test-utils/mock-openai-server.tsapps/gateway/src/chat/chat.tsapps/gateway/src/lib/prompt-tokens.spec.tsapps/gateway/src/api.e2e.ts
**/*.{js,ts}
📄 CodeRabbit Inference Engine (.github/copilot-instructions.md)
**/*.{js,ts}: Use drizzle with the latest object syntax for database operations
For read queries, always usedb().query.<table>.findMany()ordb().query.<table>.findFirst()
Files:
apps/gateway/src/test-utils/mock-openai-server.tsapps/gateway/src/chat/chat.tsapps/gateway/src/lib/prompt-tokens.spec.tsapps/gateway/src/api.e2e.ts
{apps/api,apps/gateway,packages/db}/**/*.ts
📄 CodeRabbit Inference Engine (CLAUDE.md)
{apps/api,apps/gateway,packages/db}/**/*.ts: Use Drizzle ORM with latest object syntax for database operations
For reads, usedb().query.<table>.findMany()ordb().query.<table>.findFirst()
Files:
apps/gateway/src/test-utils/mock-openai-server.tsapps/gateway/src/chat/chat.tsapps/gateway/src/lib/prompt-tokens.spec.tsapps/gateway/src/api.e2e.ts
**/*.{ts,tsx}
📄 CodeRabbit Inference Engine (.cursor/rules/general.mdc)
Never use
as anyor: anyin TypeScript files.
Files:
apps/gateway/src/test-utils/mock-openai-server.tsapps/gateway/src/chat/chat.tsapps/gateway/src/lib/prompt-tokens.spec.tsapps/gateway/src/api.e2e.ts
🧠 Learnings (2)
📚 Learning: after adding features, make sure that the tests pass using `pnpm test:unit`...
Learnt from: CR
PR: theopenco/llmgateway#0
File: .github/copilot-instructions.md:0-0
Timestamp: 2025-07-23T19:33:40.639Z
Learning: After adding features, make sure that the tests pass using `pnpm test:unit`
Applied to files:
apps/gateway/src/lib/prompt-tokens.spec.tsapps/gateway/src/api.e2e.ts
📚 Learning: run `pnpm test:unit` and `pnpm test:e2e` after adding features...
Learnt from: CR
PR: theopenco/llmgateway#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-23T19:33:55.702Z
Learning: Run `pnpm test:unit` and `pnpm test:e2e` after adding features
Applied to files:
apps/gateway/src/api.e2e.ts
🧬 Code Graph Analysis (1)
apps/gateway/src/chat/chat.ts (1)
apps/ui/src/hooks/useChats.ts (1)
ChatMessage(16-22)
🪛 GitHub Check: generate / run
apps/gateway/src/chat/chat.ts
[warning] 973-973:
'error' is defined but never used
[failure] 973-973:
'error' is defined but never used. Allowed unused caught errors must match /^_/u
[warning] 956-956:
'error' is defined but never used
[failure] 956-956:
'error' is defined but never used. Allowed unused caught errors must match /^_/u
[failure] 1041-1041:
'error' is defined but never used. Allowed unused caught errors must match /^_/u
[warning] 1024-1024:
'error' is defined but never used
[failure] 1024-1024:
'error' is defined but never used. Allowed unused caught errors must match /^_/u
apps/gateway/src/lib/prompt-tokens.spec.ts
[failure] 31-31:
'error' is defined but never used. Allowed unused caught errors must match /^_/u
🪛 GitHub Check: lint / run
apps/gateway/src/chat/chat.ts
[warning] 973-973:
'error' is defined but never used
[failure] 973-973:
'error' is defined but never used. Allowed unused caught errors must match /^_/u
[warning] 956-956:
'error' is defined but never used
[failure] 956-956:
'error' is defined but never used. Allowed unused caught errors must match /^_/u
[failure] 1041-1041:
'error' is defined but never used. Allowed unused caught errors must match /^_/u
[warning] 1024-1024:
'error' is defined but never used
[failure] 1024-1024:
'error' is defined but never used. Allowed unused caught errors must match /^_/u
apps/gateway/src/lib/prompt-tokens.spec.ts
[failure] 31-31:
'error' is defined but never used. Allowed unused caught errors must match /^_/u
🪛 GitHub Check: autofix
apps/gateway/src/chat/chat.ts
[warning] 973-973:
'error' is defined but never used
[failure] 973-973:
'error' is defined but never used. Allowed unused caught errors must match /^_/u
[warning] 956-956:
'error' is defined but never used
[failure] 956-956:
'error' is defined but never used. Allowed unused caught errors must match /^_/u
[failure] 1041-1041:
'error' is defined but never used. Allowed unused caught errors must match /^_/u
[warning] 1024-1024:
'error' is defined but never used
[failure] 1024-1024:
'error' is defined but never used. Allowed unused caught errors must match /^_/u
apps/gateway/src/lib/prompt-tokens.spec.ts
[failure] 31-31:
'error' is defined but never used. Allowed unused caught errors must match /^_/u
🪛 ESLint
apps/gateway/src/chat/chat.ts
[error] 956-956: 'error' is defined but never used. Allowed unused caught errors must match /^_/u.
(no-unused-vars)
[error] 973-973: 'error' is defined but never used. Allowed unused caught errors must match /^_/u.
(no-unused-vars)
[error] 1024-1024: 'error' is defined but never used. Allowed unused caught errors must match /^_/u.
(no-unused-vars)
[error] 1041-1041: 'error' is defined but never used. Allowed unused caught errors must match /^_/u.
(no-unused-vars)
apps/gateway/src/lib/prompt-tokens.spec.ts
[error] 31-31: 'error' is defined but never used. Allowed unused caught errors must match /^_/u.
(no-unused-vars)
🪛 GitHub Actions: ci
apps/gateway/src/chat/chat.ts
[warning] 327-327: ESLint: '_e' is defined but never used (unused-imports/no-unused-vars)
[error] 956-956: ESLint: 'error' is defined but never used. Allowed unused caught errors must match /^_/u (no-unused-vars)
⏰ 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). (1)
- GitHub Check: e2e / run
🔇 Additional comments (7)
apps/gateway/src/test-utils/mock-openai-server.ts (2)
56-59: LGTM! Consistent pattern for test scenario detection.The implementation follows the same pattern as the existing error trigger logic and correctly detects messages containing "ZERO_TOKENS" for testing the zero prompt tokens scenario.
77-81: LGTM! Proper usage override for zero tokens testing.The conditional usage override correctly simulates the scenario where a provider returns zero prompt tokens while maintaining realistic completion token counts. This supports the e2e tests for verifying fallback token calculation logic.
apps/gateway/src/lib/prompt-tokens.spec.ts (2)
3-44: LGTM! Comprehensive test coverage for prompt token calculation.The test suite thoroughly covers the fallback logic scenarios:
- Calculates tokens when provider returns 0
- Preserves existing non-zero token counts
- Uses realistic token estimation (character length / 4)
- Ensures minimum fallback of 1 token
The test implementation correctly mirrors the expected production logic.
46-78: LGTM! Excellent edge case coverage.The tests properly handle:
- Minimum token guarantee regardless of input
- Type safety with null/undefined values
- Empty message arrays and content
- Consistent Math.max(1, ...) pattern for minimum enforcement
This comprehensive coverage ensures robust token calculation behavior.
apps/gateway/src/api.e2e.ts (2)
994-1035: LGTM! Comprehensive non-streaming zero tokens test.The test correctly:
- Uses mock-token to trigger mock server behavior
- Sends "ZERO_TOKENS" message to simulate provider returning 0 tokens
- Verifies all usage fields exist and have correct types
- Ensures prompt_tokens > 0 despite provider returning 0
- Validates total_tokens consistency
This validates the core requirement that prompt tokens are never zero.
1037-1069: LGTM! Streaming zero tokens validation.The streaming test properly:
- Uses streaming request with ZERO_TOKENS trigger
- Analyzes stream chunks to find usage information
- Verifies prompt tokens are calculated and > 0 in streaming responses
- Ensures type safety
This complements the non-streaming test to ensure complete coverage of both response modes.
apps/gateway/src/chat/chat.ts (1)
649-658: Good fix for ensuring non-zero prompt tokensThe additional check for
finalPromptTokens === 0ensures that zero prompt tokens trigger the fallback calculation, preventing zero values in usage reporting. This aligns well with the PR objectives.Also applies to: 2688-2690
| usage: data.usageMetadata | ||
| ? { | ||
| prompt_tokens: data.usageMetadata.promptTokenCount || 0, |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Refactor repeated token normalization logic
The same token normalization pattern is repeated across multiple providers. Consider extracting this into a helper function to follow DRY principles.
Create a helper function:
function normalizeTokenUsage(
promptTokens: number | null,
completionTokens: number | null,
totalTokens: number | null
) {
const normalizedPrompt = Math.max(1, promptTokens || 0);
const normalizedCompletion = completionTokens || 0;
const normalizedTotal = Math.max(1, totalTokens || normalizedPrompt);
return {
prompt_tokens: normalizedPrompt,
completion_tokens: normalizedCompletion,
total_tokens: normalizedTotal,
};
}Then use it in all provider cases:
usage: {
- prompt_tokens: Math.max(1, promptTokens || 1),
- completion_tokens: completionTokens || 0,
- total_tokens: Math.max(1, totalTokens || Math.max(1, promptTokens || 1)),
+ ...normalizeTokenUsage(promptTokens, completionTokens, totalTokens),
// ... other fields
},🤖 Prompt for AI Agents
In apps/gateway/src/chat/chat.ts around lines 730 to 732, the token
normalization logic is repeated for multiple providers. Extract this repeated
logic into a single helper function named normalizeTokenUsage that takes
promptTokens, completionTokens, and totalTokens as parameters and returns an
object with normalized prompt_tokens, completion_tokens, and total_tokens
values. Replace all instances of the repeated normalization code with calls to
this helper function to adhere to DRY principles and improve maintainability.
| (data.usageMetadata.promptTokenCount > 0 | ||
| ? data.usageMetadata.promptTokenCount | ||
| : (() => { | ||
| try { |
There was a problem hiding this comment.
Fix linting errors for unused error variables
Static analysis correctly identified unused error variables that don't match the required pattern /^_/u.
-} catch (error) {
+} catch (_error) {Apply this fix to all four occurrences in the streaming chunk transformation logic.
Also applies to: 973-976, 1024-1027, 1041-1044
🧰 Tools
🪛 GitHub Check: generate / run
[warning] 956-956:
'error' is defined but never used
[failure] 956-956:
'error' is defined but never used. Allowed unused caught errors must match /^_/u
🪛 GitHub Check: lint / run
[warning] 956-956:
'error' is defined but never used
[failure] 956-956:
'error' is defined but never used. Allowed unused caught errors must match /^_/u
🪛 GitHub Check: autofix
[warning] 956-956:
'error' is defined but never used
[failure] 956-956:
'error' is defined but never used. Allowed unused caught errors must match /^_/u
🪛 ESLint
[error] 956-956: 'error' is defined but never used. Allowed unused caught errors must match /^_/u.
(no-unused-vars)
🪛 GitHub Actions: ci
[error] 956-956: ESLint: 'error' is defined but never used. Allowed unused caught errors must match /^_/u (no-unused-vars)
🤖 Prompt for AI Agents
In apps/gateway/src/chat/chat.ts at lines 956-959 and similarly at 973-976,
1024-1027, and 1041-1044, the catch blocks declare an error variable that is
unused, causing linting errors. To fix this, rename each unused error variable
to start with an underscore (e.g., change "error" to "_error") to comply with
the linting rule for unused variables. This change should be applied
consistently to all four specified catch blocks.
… missing - Fixed Google providers to calculate prompt tokens when missing or 0 - Added helper function calculatePromptTokensFromMessages for better code reuse - Updated transformToOpenAIFormat to ensure Math.max(1, promptTokens || 1) - Fixed streaming responses to calculate tokens when needed - Added comprehensive test coverage for token calculation logic - Fixed TypeScript compilation errors with proper type annotations - Tests verify tokens are always > 0 even when provider returns 0
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
apps/gateway/src/chat/chat.ts (1)
651-653: Token normalization logic still needs refactoringThe same token normalization pattern is repeated across multiple providers. A previous reviewer already suggested extracting this into a helper function to follow DRY principles.
Also applies to: 692-694, 730-732
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
apps/gateway/src/api.e2e.ts(1 hunks)apps/gateway/src/chat/chat.ts(10 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- apps/gateway/src/api.e2e.ts
🧰 Additional context used
📓 Path-based instructions (4)
**/*.{js,jsx,ts,tsx}
📄 CodeRabbit Inference Engine (.github/copilot-instructions.md)
Use localStorage instead of cookies for client-side data persistence
Files:
apps/gateway/src/chat/chat.ts
**/*.{js,ts}
📄 CodeRabbit Inference Engine (.github/copilot-instructions.md)
**/*.{js,ts}: Use drizzle with the latest object syntax for database operations
For read queries, always usedb().query.<table>.findMany()ordb().query.<table>.findFirst()
Files:
apps/gateway/src/chat/chat.ts
{apps/api,apps/gateway,packages/db}/**/*.ts
📄 CodeRabbit Inference Engine (CLAUDE.md)
{apps/api,apps/gateway,packages/db}/**/*.ts: Use Drizzle ORM with latest object syntax for database operations
For reads, usedb().query.<table>.findMany()ordb().query.<table>.findFirst()
Files:
apps/gateway/src/chat/chat.ts
**/*.{ts,tsx}
📄 CodeRabbit Inference Engine (.cursor/rules/general.mdc)
Never use
as anyor: anyin TypeScript files.
Files:
apps/gateway/src/chat/chat.ts
🪛 ESLint
apps/gateway/src/chat/chat.ts
[error] 758-758: 'error' is defined but never used. Allowed unused caught errors must match /^_/u.
(no-unused-vars)
🪛 GitHub Check: generate / run
apps/gateway/src/chat/chat.ts
[warning] 758-758:
'error' is defined but never used
[failure] 758-758:
'error' is defined but never used. Allowed unused caught errors must match /^_/u
🪛 GitHub Check: lint / run
apps/gateway/src/chat/chat.ts
[warning] 758-758:
'error' is defined but never used
[failure] 758-758:
'error' is defined but never used. Allowed unused caught errors must match /^_/u
🪛 GitHub Actions: ci
apps/gateway/src/chat/chat.ts
[warning] 327-327: ESLint: '_e' is defined but never used (unused-imports/no-unused-vars)
[error] 758-758: ESLint: 'error' is defined but never used. Allowed unused caught errors must match /^_/u (no-unused-vars)
⏰ 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). (1)
- GitHub Check: e2e / run
🔇 Additional comments (5)
apps/gateway/src/chat/chat.ts (5)
767-767: Good fix for the undefined messages variableAdding the
messagesparameter to the function signature correctly resolves the runtime error that would have occurred when referencing the undefinedmessagesvariable.
960-962: Excellent fallback logic for Google provider token handlingThe implementation correctly uses the helper function to calculate prompt tokens when the provider returns zero or missing values, ensuring prompt tokens are never zero in streaming responses.
Also applies to: 966-969, 1005-1007, 1011-1014
2618-2618: Proper zero token handling in streamingThe condition change to include
|| finalPromptTokens === 0ensures estimation is triggered even when providers return zero tokens, and theMath.max(1, ...)calls guarantee minimum token values in the final usage chunk.Also applies to: 2657-2657, 2659-2659
3002-3004: Comprehensive token normalization with proper fallbacksThe logic correctly preserves original prompt tokens when valid (> 0) and falls back to calculated values when zero or missing, ensuring minimum token values are always maintained.
Also applies to: 3009-3009
2707-2707: Correct parameter passing for updated function signatureThe
messagesparameter is properly passed to match the updated function signature.
steebchen
changed the title
Updated authorization tokens in e2e tests to use real tokens instead of mock tokens for better accuracy during testing.
Cleaned up unused `test: "skip"` flags from Zai model providers, improving readability and maintaining consistent configuration.
2 participants
This pull request contains changes generated by Cursor background composer.
Summary by CodeRabbit
Bug Fixes
Tests
Chores
Chores