Conversation
Refactored Google provider parsing to streamline logic by removing excessive JSON handling. Improved clarity and reduced redundant code.
|
""" WalkthroughThis change refactors the streaming response handling logic in the chat completions endpoint. It removes provider-specific incremental JSON parsing for Google providers, unifying streaming chunk parsing for all providers by processing lines prefixed with "data: ". It enhances token usage estimation for Google and Anthropic providers, refines finish reason extraction, and adds final usage chunk emission. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~15 minutes Possibly related PRs
Note ⚡️ Unit Test Generation is now available in beta!Learn more here, or try it out under "Finishing Touches" below. 📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous 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). (3)
✨ 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. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. 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: 2
🧹 Nitpick comments (1)
apps/gateway/src/chat/chat.ts (1)
2233-2239: Consider adding more context to error logs.While the error handling is appropriate, consider including the provider name in the error log context for easier debugging of provider-specific issues.
} catch (e) { console.warn("Failed to parse streaming JSON:", { error: e instanceof Error ? e.message : String(e), lineContent: line.substring(0, 100), // First 100 chars for debugging provider: usedProvider, + model: usedModel, + lineNumber: lines.indexOf(line), }); }
📜 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
📓 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
🧠 Learnings (1)
apps/gateway/src/chat/chat.ts (3)
Learnt from: CR
PR: theopenco/llmgateway#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-23T19:33:55.702Z
Learning: Applies to {apps/api,apps/gateway,packages/db}/**/*.ts : Use Drizzle ORM with latest object syntax for database operations
Learnt from: CR
PR: theopenco/llmgateway#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-23T19:33:55.702Z
Learning: Applies to {apps/api,apps/gateway,packages/db}/**/*.ts : For reads, use db().query.<table>.findMany() or db().query.<table>.findFirst()
Learnt from: CR
PR: theopenco/llmgateway#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-23T19:33:55.702Z
Learning: Applies to apps/api/routes/**/*.ts : Run pnpm generate after API route changes to update OpenAPI schemas
⏰ 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). (3)
- GitHub Check: e2e / run
- GitHub Check: build / run
- GitHub Check: autofix
🔇 Additional comments (2)
apps/gateway/src/chat/chat.ts (2)
2047-2054: LGTM! Proper handling of incomplete lines in streaming data.The buffer management logic correctly handles cases where JSON objects might be split across chunks by keeping incomplete lines in the buffer.
2047-2239: Verify Google provider streaming compatibilityWe didn’t find any existing tests or documentation that cover streaming responses for
google-vertexorgoogle-ai-studiousing the new line-based SSE parsing. Please ensure that:
- The chat handler correctly breaks and parses SSE “data: …” lines for both Google providers.
- The
streamGenerateContent?alt=sseendpoint responses are compatible withtransformStreamingChunkToOpenAIFormatandextractContentFromProvider.- Any provider-specific nuances (e.g., partial JSON payloads) are handled as expected.
Consider adding unit or integration tests that simulate streaming chunks from each Google provider to catch any parsing regressions.
|
|
||
| // Check for finish reason | ||
| if (data.candidates && data.candidates[0]?.finishReason) { | ||
| finishReason = data.candidates[0].finishReason; | ||
|
|
||
| // Send final chunk when we get a finish reason | ||
| if (finishReason) { | ||
| await writeSSEAndCache({ | ||
| event: "done", | ||
| data: "[DONE]", | ||
| id: String(eventId++), | ||
| }); | ||
| } | ||
| if (finalCompletionTokens === null) { | ||
| finalCompletionTokens = | ||
| estimateTokensFromContent(fullContent); | ||
| } | ||
|
|
||
| // Extract token usage using helper function | ||
| const usage = extractTokenUsage(data, usedProvider); | ||
| if (usage.promptTokens !== null) { | ||
| promptTokens = usage.promptTokens; | ||
| } | ||
| if (usage.completionTokens !== null) { | ||
| completionTokens = usage.completionTokens; | ||
| } | ||
| if (usage.totalTokens !== null) { | ||
| totalTokens = usage.totalTokens; | ||
| } | ||
| if (usage.reasoningTokens !== null) { | ||
| reasoningTokens = usage.reasoningTokens; | ||
| if (finalTotalTokens === null) { | ||
| finalTotalTokens = | ||
| (finalPromptTokens || 0) + (finalCompletionTokens || 0); | ||
| } | ||
|
|
||
| // For Google AI Studio, if candidatesTokenCount is not provided, | ||
| // we'll calculate it later from the fullContent | ||
| // Send final usage chunk before [DONE] if we have any usage data | ||
| if ( | ||
| (usedProvider === "google-ai-studio" || | ||
| usedProvider === "google-vertex") && | ||
| !usage.completionTokens && |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Token estimation uses inconsistent methods.
The code estimates tokens using simple division by 4, but elsewhere in the file (lines 241-276), there are more sophisticated methods using encodeChat and encode from gpt-tokenizer. Consider using the existing estimateTokens helper function for consistency.
// Estimate missing tokens if needed using helper function
-if (finalPromptTokens === null) {
- finalPromptTokens = Math.round(
- messages.reduce(
- (acc, m) => acc + (m.content?.length || 0),
- 0,
- ) / 4,
- );
-}
-
-if (finalCompletionTokens === null) {
- finalCompletionTokens =
- estimateTokensFromContent(fullContent);
-}
+const { calculatedPromptTokens, calculatedCompletionTokens } = estimateTokens(
+ usedProvider,
+ messages,
+ fullContent,
+ finalPromptTokens,
+ finalCompletionTokens,
+);
+finalPromptTokens = finalPromptTokens ?? calculatedPromptTokens;
+finalCompletionTokens = finalCompletionTokens ?? calculatedCompletionTokens;📝 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.
| if (finalPromptTokens === null) { | |
| finalPromptTokens = Math.round( | |
| messages.reduce( | |
| (acc, m) => acc + (m.content?.length || 0), | |
| 0, | |
| ) / 4, | |
| ); | |
| } | |
| // Check for finish reason | |
| if (data.candidates && data.candidates[0]?.finishReason) { | |
| finishReason = data.candidates[0].finishReason; | |
| // Send final chunk when we get a finish reason | |
| if (finishReason) { | |
| await writeSSEAndCache({ | |
| event: "done", | |
| data: "[DONE]", | |
| id: String(eventId++), | |
| }); | |
| } | |
| if (finalCompletionTokens === null) { | |
| finalCompletionTokens = | |
| estimateTokensFromContent(fullContent); | |
| } | |
| // Extract token usage using helper function | |
| const usage = extractTokenUsage(data, usedProvider); | |
| if (usage.promptTokens !== null) { | |
| promptTokens = usage.promptTokens; | |
| } | |
| if (usage.completionTokens !== null) { | |
| completionTokens = usage.completionTokens; | |
| } | |
| if (usage.totalTokens !== null) { | |
| totalTokens = usage.totalTokens; | |
| } | |
| if (usage.reasoningTokens !== null) { | |
| reasoningTokens = usage.reasoningTokens; | |
| if (finalTotalTokens === null) { | |
| finalTotalTokens = | |
| (finalPromptTokens || 0) + (finalCompletionTokens || 0); | |
| } | |
| // Estimate missing tokens if needed using helper function | |
| const { calculatedPromptTokens, calculatedCompletionTokens } = estimateTokens( | |
| usedProvider, | |
| messages, | |
| fullContent, | |
| finalPromptTokens, | |
| finalCompletionTokens, | |
| ); | |
| finalPromptTokens = finalPromptTokens ?? calculatedPromptTokens; | |
| finalCompletionTokens = finalCompletionTokens ?? calculatedCompletionTokens; | |
| if (finalTotalTokens === null) { | |
| finalTotalTokens = | |
| (finalPromptTokens || 0) + (finalCompletionTokens || 0); | |
| } |
🤖 Prompt for AI Agents
In apps/gateway/src/chat/chat.ts around lines 2064 to 2081, the token estimation
uses a simple division by 4 which is inconsistent with the more accurate methods
used elsewhere in the file. Replace the manual token calculation with the
existing estimateTokens helper function to ensure consistent and accurate token
estimation across the codebase.
| completion_tokens: usage.output_tokens, | ||
| total_tokens: | ||
| estimatedPromptTokens + usage.output_tokens, | ||
| }; | ||
| } | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Duplicate token estimation logic for Anthropic.
This token estimation logic duplicates the code at lines 2219-2224. Consider extracting this into a shared helper function to maintain DRY principles.
// Estimate prompt tokens if not provided
-const estimatedPromptTokens = Math.round(
- messages.reduce(
- (acc, m) => acc + (m.content?.length || 0),
- 0,
- ) / 4,
-);
+const estimatedPromptTokens = promptTokens ?? estimatePromptTokensFromMessages(messages);Where estimatePromptTokensFromMessages would be a new helper function that encapsulates the token estimation logic and uses the more accurate encodeChat method when possible.
Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In apps/gateway/src/chat/chat.ts around lines 2138 to 2143, the token estimation
logic duplicates similar code found at lines 2219 to 2224. To fix this, extract
the token estimation code into a new shared helper function named
estimatePromptTokensFromMessages that uses the encodeChat method for accuracy
when available. Replace both instances with calls to this helper to adhere to
DRY principles and improve maintainability.
Enhanced error handling by adding detailed logs for failed streaming JSON parsing, including line content and provider info. Refactored nested logic for improved clarity and robustness.
Added logic to extract and assign finishReason from responses of Google Vertex and Google AI Studio providers. Improved consistency in provider-specific handling of response finish reasons.
Added fallback logic to estimate missing completion token counts using content length when data is incomplete. Improved token usage calculation for providers, ensuring accurate results even if metadata is partially available. Enhanced error handling for encoding failures with descriptive logs.
Added logic to extract and include token usage metadata (prompt, completion, and total tokens) for Google Vertex and Google AI Studio providers when available. Improved streaming data consistency by integrating usage details into streaming chunks.
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (3)
apps/gateway/src/chat/chat.ts (3)
359-384: Duplicate token estimation logic detected.This token estimation logic duplicates the code added in
parseProviderResponse(lines 166-184). This violates DRY principles and makes the code harder to maintain.Extract the Google provider token estimation logic into a shared helper function that can be used by both
parseProviderResponseandextractTokenUsage:+function estimateGoogleCompletionTokens(content: string | null): number | null { + if (!content) return null; + + try { + return encode(content).length; + } catch (error) { + console.error(`Failed to encode completion text: ${error}`); + return Math.max(1, Math.round(content.length / 4)); + } +} // In parseProviderResponse and extractTokenUsage -if (completionTokens === null && content) { - try { - completionTokens = encode(content).length; - } catch (error) { - console.error(`Failed to encode completion text: ${error}`); - completionTokens = Math.max(1, Math.round(content.length / 4)); - } -} +if (completionTokens === null) { + completionTokens = estimateGoogleCompletionTokens(content); +}
2112-2124: Token estimation uses inconsistent methods.The code estimates tokens using simple division by 4, but elsewhere in the file (lines 241-276), there are more sophisticated methods using
encodeChatandencodefrom gpt-tokenizer. Consider using the existingestimateTokenshelper function for consistency.// Estimate missing tokens if needed using helper function -if (finalPromptTokens === null) { - finalPromptTokens = Math.round( - messages.reduce( - (acc, m) => acc + (m.content?.length || 0), - 0, - ) / 4, - ); -} - -if (finalCompletionTokens === null) { - finalCompletionTokens = - estimateTokensFromContent(fullContent); -} +const { calculatedPromptTokens, calculatedCompletionTokens } = estimateTokens( + usedProvider, + messages, + fullContent, + finalPromptTokens, + finalCompletionTokens, +); +finalPromptTokens = finalPromptTokens ?? calculatedPromptTokens; +finalCompletionTokens = finalCompletionTokens ?? calculatedCompletionTokens;
2187-2207: Duplicate token estimation logic for Anthropic.This token estimation logic duplicates the code at lines 2195-2200. Consider extracting this into a shared helper function to maintain DRY principles.
// Estimate prompt tokens if not provided -const estimatedPromptTokens = Math.round( - messages.reduce( - (acc, m) => acc + (m.content?.length || 0), - 0, - ) / 4, -); +const estimatedPromptTokens = promptTokens ?? estimatePromptTokensFromMessages(messages);Where
estimatePromptTokensFromMessageswould be a new helper function that encapsulates the token estimation logic and uses the more accurateencodeChatmethod when possible.
🧹 Nitpick comments (1)
apps/gateway/src/chat/chat.ts (1)
164-184: Enhanced token estimation for Google providers looks good.The addition of completion token estimation when
candidatesTokenCountis missing is a valuable improvement. The implementation properly uses gpt-tokenizer with fallback estimation and includes appropriate error handling.However, this token estimation logic is duplicated in the
extractTokenUsagefunction (lines 359-384). Consider extracting this into a shared helper function to maintain DRY principles.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
apps/gateway/src/chat/chat.ts(9 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/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
🧠 Learnings (1)
apps/gateway/src/chat/chat.ts (3)
Learnt from: CR
PR: theopenco/llmgateway#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-23T19:33:55.702Z
Learning: Applies to {apps/api,apps/gateway,packages/db}/**/*.ts : Use Drizzle ORM with latest object syntax for database operations
Learnt from: CR
PR: theopenco/llmgateway#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-23T19:33:55.702Z
Learning: Applies to {apps/api,apps/gateway,packages/db}/**/*.ts : For reads, use db().query.<table>.findMany() or db().query.<table>.findFirst()
Learnt from: CR
PR: theopenco/llmgateway#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-23T19:33:55.702Z
Learning: Applies to apps/api/routes/**/*.ts : Run pnpm generate after API route changes to update OpenAPI schemas
⏰ 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). (3)
- GitHub Check: build / run
- GitHub Check: e2e / run
- GitHub Check: autofix
🔇 Additional comments (8)
apps/gateway/src/chat/chat.ts (8)
3-3: LGTM on import additions.The new imports (
ApiKey,shortid,getModelStreamingSupport,hasProviderEnvironmentToken) are appropriately added and used throughout the file.Also applies to: 7-7, 12-12, 43-43
2095-2101: Excellent unification of streaming parsing logic.The change from provider-specific parsing to unified line-based parsing with "data: " prefixes is a significant improvement. The buffer management for incomplete lines is correctly implemented and handles edge cases properly.
2106-2160: Final usage chunk emission logic is well-implemented.The logic to emit final usage chunks when token data is missing or incomplete is comprehensive and handles various scenarios properly. This ensures consistent usage reporting across all providers.
2167-2178: Improved error handling for JSON parsing.The addition of try-catch around JSON parsing with detailed logging is a good improvement. The warning includes helpful context (provider, line content) for debugging malformed streaming responses.
2209-2232: Enhanced Google provider token usage extraction.The integration of the enhanced
extractTokenUsagefunction with thefullContentparameter properly addresses the missing completion token issue for Google providers during streaming.
2255-2283: Provider-specific finish reason extraction is well-structured.The switch statement clearly handles finish reason extraction for each provider type with appropriate fallbacks. The logic correctly handles the various Anthropic event types and Google's finishReason format.
2286-2305: Enhanced token usage extraction with fullContent parameter.The updated call to
extractTokenUsagewith thefullContentparameter enables better token estimation for Google providers during streaming, which aligns with the PR objectives.
2519-2519: Minor formatting improvement.This appears to be a minor whitespace/formatting change that improves code readability without affecting functionality.
Removed excessive logging in token estimation logic for Google streaming. Improved clarity and reduced unnecessary console clutter.
1 participant
Refactored Google provider parsing to streamline logic by removing excessive JSON handling. Improved clarity and reduced redundant code.
Summary by CodeRabbit