feat: use model-aware prompt counting in Gemini paths

Preserve Gemini prompt token text during normalization and remove the hardcoded DeepSeek model from native Gemini usage helpers.
This commit is contained in:
shern-point
2026-04-30 00:46:05 +08:00
parent d3018c281b
commit 4b4f097006
3 changed files with 18 additions and 17 deletions

View File

@@ -36,16 +36,17 @@ func normalizeGeminiRequest(store ConfigReader, routeModel string, req map[strin
passThrough := collectGeminiPassThrough(req)
return promptcompat.StandardRequest{
Surface: "google_gemini",
RequestedModel: requestedModel,
ResolvedModel: resolvedModel,
ResponseModel: requestedModel,
Messages: messagesRaw,
FinalPrompt: finalPrompt,
ToolNames: toolNames,
Stream: stream,
Thinking: thinkingEnabled,
Search: searchEnabled,
PassThrough: passThrough,
Surface: "google_gemini",
RequestedModel: requestedModel,
ResolvedModel: resolvedModel,
ResponseModel: requestedModel,
Messages: messagesRaw,
PromptTokenText: finalPrompt,
FinalPrompt: finalPrompt,
ToolNames: toolNames,
Stream: stream,
Thinking: thinkingEnabled,
Search: searchEnabled,
PassThrough: passThrough,
}, nil
}

View File

@@ -227,7 +227,7 @@ func (h *Handler) handleNonStreamGenerateContent(w http.ResponseWriter, resp *ht
//nolint:unused // retained for native Gemini non-stream handling path.
func buildGeminiGenerateContentResponse(model, finalPrompt, finalThinking, finalText string, toolNames []string) map[string]any {
parts := buildGeminiPartsFromFinal(finalText, finalThinking, toolNames)
usage := buildGeminiUsage(finalPrompt, finalThinking, finalText)
usage := buildGeminiUsage(model, finalPrompt, finalThinking, finalText)
return map[string]any{
"candidates": []map[string]any{
{
@@ -245,10 +245,10 @@ func buildGeminiGenerateContentResponse(model, finalPrompt, finalThinking, final
}
//nolint:unused // retained for native Gemini non-stream handling path.
func buildGeminiUsage(finalPrompt, finalThinking, finalText string) map[string]any {
promptTokens := util.EstimateTokens(finalPrompt)
reasoningTokens := util.EstimateTokens(finalThinking)
completionTokens := util.EstimateTokens(finalText)
func buildGeminiUsage(model, finalPrompt, finalThinking, finalText string) map[string]any {
promptTokens := util.CountPromptTokens(finalPrompt, model)
reasoningTokens := util.CountOutputTokens(finalThinking, model)
completionTokens := util.CountOutputTokens(finalText, model)
return map[string]any{
"promptTokenCount": promptTokens,
"candidatesTokenCount": reasoningTokens + completionTokens,

View File

@@ -194,6 +194,6 @@ func (s *geminiStreamRuntime) finalize() {
},
},
"modelVersion": s.model,
"usageMetadata": buildGeminiUsage(s.finalPrompt, finalThinking, finalText),
"usageMetadata": buildGeminiUsage(s.model, s.finalPrompt, finalThinking, finalText),
})
}