fix: gemini 原生格式流模式中断请求未计费 - #1272
Conversation
WalkthroughThe Gemini text generation stream handler was updated to accumulate all streamed text parts into a buffer and recalculate usage tokens if completion tokens are zero. Additionally, multiple token counting functions and calls across various relay channels and services were simplified by removing error returns and related error handling, assuming token counting never fails. Some error checks were removed along with error assignments, resulting in more consistent omission of error handling. Changes
Poem
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (5)
🚧 Files skipped from review as they are similar to previous changes (5)
✨ Finishing Touches
🪧 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: 0
🧹 Nitpick comments (1)
relay/channel/gemini/relay-gemini-native.go (1)
95-100: Verify token counting accuracy in streaming context.The token counting logic looks correct for addressing the billing issue. However, consider these potential optimizations:
Performance: Token counting on every streaming chunk adds latency. Consider batching or optimizing if this becomes a bottleneck.
Error handling: Currently errors are only logged. Consider tracking if token counting consistently fails to alert on potential under-billing.
For improved error visibility, consider tracking token counting failures:
+var tokenCountingErrors int // 本地统计completion tokens textTokens, err := service.CountTextToken(part.Text, info.UpstreamModelName) if err != nil { common.LogError(c, "error counting text token: "+err.Error()) + tokenCountingErrors++ } localCompletionTokens += textTokens
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
relay/channel/gemini/relay-gemini-native.go(3 hunks)
🔇 Additional comments (2)
relay/channel/gemini/relay-gemini-native.go (2)
78-79: LGTM: Clean variable initialization.The local completion tokens counter is properly initialized and clearly named.
134-138: Review the interaction with image token counting logic.The fallback logic correctly uses local token counts when the API doesn't provide completion tokens. However, verify the interaction with the existing image counting logic (lines 128-132).
Current flow:
- Image logic sets
CompletionTokens = imageCount * 258ifCompletionTokens == 0- Text fallback only triggers if
CompletionTokensis still 0Potential issue: For responses containing both images and text, only image tokens may be counted, potentially under-billing text content.
#!/bin/bash # Description: Search for test cases or examples that handle both images and text in Gemini responses # Expected: Find test cases that verify mixed content billing # Search for test files related to Gemini functionality fd -e go -t f | grep -i test | xargs rg -l "gemini.*image\|image.*gemini" -A 5 -B 5 # Search for mixed content handling patterns rg -A 10 -B 5 "imageCount.*CompletionTokens\|CompletionTokens.*imageCount"
There was a problem hiding this comment.
Actionable comments posted: 4
🔭 Outside diff range comments (3)
relay/gemini_handler.go (1)
109-114: Remove dead code - error handling for function that no longer returns error.The
getGeminiInputTokensfunction no longer returns an error, making the error handling code on lines 110-112 unreachable dead code.Apply this diff to remove the dead code:
} else { promptTokens := getGeminiInputTokens(req, relayInfo) - if err != nil { - return service.OpenAIErrorWrapperLocal(err, "count_input_tokens_error", http.StatusBadRequest) - } c.Set("prompt_tokens", promptTokens) }service/token_counter.go (2)
174-178: Remove inconsistent error handling for function that no longer returns error.The
CountTokenInputfunction no longer returns an error, making this error handling code unreachable and inconsistent with the refactoring.Apply this diff to remove the inconsistent error handling:
toolTokens := CountTokenInput(countStr, request.Model) - if err != nil { - return 0, err - } tkm += 8
197-202: Remove inconsistent error handling for function that no longer returns error.The
CountTokenInputfunction no longer returns an error, making this error handling code unreachable and inconsistent with the refactoring.Apply this diff to remove the inconsistent error handling:
systemTokens := CountTokenInput(request.System, model) - if err != nil { - return 0, err - } tkm += systemTokens
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (20)
-
relay/audio_handler.go(1 hunks) -
relay/channel/claude/relay-claude.go(3 hunks) -
relay/channel/cloudflare/relay_cloudflare.go(3 hunks) -
relay/channel/cohere/relay-cohere.go(1 hunks) -
relay/channel/coze/relay-coze.go(1 hunks) -
relay/channel/dify/relay-dify.go(1 hunks) -
relay/channel/gemini/relay-gemini-native.go(4 hunks) -
relay/channel/openai/relay-openai.go(6 hunks) -
relay/channel/openai/relay_responses.go(1 hunks) -
relay/channel/palm/adaptor.go(1 hunks) -
relay/channel/palm/relay-palm.go(1 hunks) -
relay/channel/tencent/adaptor.go(1 hunks) -
relay/channel/xai/text.go(1 hunks) -
relay/embedding_handler.go(1 hunks) -
relay/gemini_handler.go(3 hunks) -
relay/relay-text.go(1 hunks) -
relay/rerank_handler.go(1 hunks) -
relay/responses_handler.go(2 hunks) -
service/token_counter.go(9 hunks) -
service/usage_helpr.go(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- relay/channel/openai/relay-openai.go
🚧 Files skipped from review as they are similar to previous changes (1)
- relay/channel/gemini/relay-gemini-native.go
🔇 Additional comments (23)
relay/channel/palm/relay-palm.go (1)
158-158: ```shell
#!/bin/bashDescription: Inspect implementations of getTokenEncoder and getTokenNum to verify they never fail/panic
ast-grep --pattern 'func getTokenEncoder($) $ { $$$ }'
ast-grep --pattern 'func getTokenNum($, $)$_ { $ $$ }'</details> <details> <summary>relay/channel/xai/text.go (1)</summary> `71-71`: **Verify that usage calculation never fails before removing error handling.** The error return from `service.ResponseText2Usage` has been removed. This function is used as a fallback when stream usage information is not available, making error handling particularly important. Please verify that ResponseText2Usage can never fail: ```shell #!/bin/bash # Description: Check the implementation of ResponseText2Usage to verify it handles all edge cases # Expected: Function should gracefully handle all inputs without returning errors ast-grep --pattern 'func ResponseText2Usage($_, $_, $_) $_ { $$$ }'relay/channel/openai/relay_responses.go (1)
113-113: Verify error handling removal is safe for fallback token counting.The error return from
service.CountTextTokenhas been removed in a fallback scenario (whenusage.CompletionTokens == 0). Since this is used when primary token counting fails, robust error handling is especially important here.relay/rerank_handler.go (1)
17-21: Verify that unconditional token accumulation is safe.The error handling logic has been significantly changed. Previously, document tokens were only added to the total if token counting succeeded. Now, all tokens are added unconditionally, which could lead to incorrect billing if token counting fails but returns a default value.
The logic change is more significant than just removing error returns. Please verify that CountTokenInput never fails and always returns accurate counts:
#!/bin/bash # Description: Check CountTokenInput implementation and verify it handles all edge cases safely # Expected: Function should never fail and always return accurate token counts ast-grep --pattern 'func CountTokenInput($_, $_) $_ { $$$ }'relay/embedding_handler.go (1)
18-18: LGTM: Clean removal of error handling from token counting.The change correctly removes error assignment from
service.CountTokenInputcall, aligning with the broader refactor to simplify token counting operations.relay/channel/cohere/relay-cohere.go (1)
165-165: LGTM: Consistent removal of error handling from usage calculation.The change correctly removes error assignment from
service.ResponseText2Usagecall, maintaining consistency with the broader refactor to simplify usage calculation operations.relay/channel/dify/relay-dify.go (1)
253-253: LGTM: Proper simplification of token counting call.The change correctly removes error assignment from
service.CountTextTokencall, consistent with the refactor to streamline token counting operations.relay/channel/claude/relay-claude.go (1)
552-552: LGTM: Consistent usage calculation simplification.Both calls to
service.ResponseText2Usagecorrectly remove error handling, aligning with the broader refactor to streamline usage calculation operations.Also applies to: 561-561
relay/channel/palm/adaptor.go (1)
77-77: I wasn’t able to locate the definition or usages with the previous AST query—let’s fall back to a grep search:#!/bin/bash # Locate the definition of ResponseText2Usage rg -n 'func ResponseText2Usage' # Find all call sites of ResponseText2Usage rg -n 'ResponseText2Usage'relay/channel/tencent/adaptor.go (1)
101-101: LGTM! Consistent with service function refactor.This change matches the pattern seen in other relay channels where error handling is removed from
service.ResponseText2Usagecalls.relay/channel/cloudflare/relay_cloudflare.go (3)
74-74: LGTM! Service function signature simplified.The removal of error handling from
service.ResponseText2Usageis consistent with the broader refactor.
111-111: LGTM! Consistent error handling removal.This change aligns with the pattern of removing error returns from token counting service functions.
153-153: Let’s locate theCountTextTokendefinition and all its usages to confirm its signature change:#!/bin/bash # Find where CountTextToken is defined rg -n "func CountTextToken" . # Find all calls to CountTextToken (qualified or not) rg -n "CountTextToken" .relay/relay-text.go (3)
254-254: LGTM! CountTokenInput simplified for completions mode.The removal of error handling from
service.CountTokenInputis consistent with the service layer simplification.
256-256: LGTM! CountTokenInput simplified for moderations mode.Consistent with the pattern of removing error returns from token counting functions.
258-258: LGTM! CountTokenInput simplified for embeddings mode.This completes the consistent pattern of removing error handling from
service.CountTokenInputacross all relevant relay modes.relay/responses_handler.go (1)
43-47: LGTM! Function signature simplified correctly.The
getInputTokensfunction has been correctly updated to remove error returns, consistent with theservice.CountTokenInputsimplification.service/usage_helpr.go (1)
19-26: LGTM! Clean signature simplification.The function signature change and implementation are consistent with the broader refactoring to remove error returns from token counting functions. The logic remains correct.
relay/gemini_handler.go (1)
62-77: LGTM! Function signature change implemented correctly.The function signature change from returning
(int, error)tointis properly implemented, and the call toservice.CountTokenInputis updated accordingly.service/token_counter.go (4)
400-418: LGTM! Function signature simplified correctly.The removal of error return from
CountTokenInputis implemented correctly. The function logic remains intact and the recursive call maintains consistency.
437-443: LGTM! Function signature simplified correctly.The removal of error return from
CountTTSTokenis implemented correctly with proper return statements.
474-480: LGTM! Function signature simplified correctly.The removal of error return from
CountTextTokenis implemented correctly with proper early return and final return statements.
400-480: I want to inspect how errors from the tokenizer are handled. Let’s locategetTokenNumandgetTokenEncoderinservice/token_counter.go:#!/bin/bash # Find definition of getTokenNum to see if it propagates errors rg -A5 -B5 "func getTokenNum" --type go # Find definition of getTokenEncoder to see if it can ever return nil or error rg -A5 -B5 "func getTokenEncoder" --type go
…mpletion-count-fix fix: gemini 原生格式流模式中断请求未计费
2 participants
Summary by CodeRabbit