From fe02c9b95f8e04d6dc9d94b4bbc2cae804f4570c Mon Sep 17 00:00:00 2001 From: jay77721 <164177721+jay77721@users.noreply.github.com> Date: Fri, 31 Jul 2026 18:32:07 +0800 Subject: [PATCH] feat(go-models): unify OpenAI-compatible token usage extraction (#17634) ## Summary Relate to #17284 . - Add `usage_parser.go`: `extractOpenAIUsage` / `extractOpenAIStreamUsage` with multi-field fallback (`prompt_tokens`/`input_tokens`, `completion_tokens`/`output_tokens`). - Add `parser_config.go`: `ParserConfig` struct and `OpenAIParserConfig`. - Add `response_handler.go`: `HandleNonStreamingResponse` / `HandleStreamingResponse` as single entry points that extract usage, populate `chatConfig.UsageResult`, and emit a `StreamUsage` log. - Extend `base_model.go` with `doRequest` / `doStreamRequest`. - Migrate `deepseek` driver to shared handlers, cutting ~100 lines. - Add `usage_parser_test.go`. --------- Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: Haruko386 Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> --- internal/entity/models/base_model.go | 80 +++++++ internal/entity/models/common_test.go | 62 +++++- internal/entity/models/deepseek.go | 132 +---------- internal/entity/models/parser_config.go | 36 +++ internal/entity/models/response_handler.go | 235 ++++++++++++++++++++ internal/entity/models/usage_parser.go | 77 +++++++ internal/entity/models/usage_parser_test.go | 153 +++++++++++++ 7 files changed, 646 insertions(+), 129 deletions(-) create mode 100644 internal/entity/models/parser_config.go create mode 100644 internal/entity/models/response_handler.go create mode 100644 internal/entity/models/usage_parser.go create mode 100644 internal/entity/models/usage_parser_test.go diff --git a/internal/entity/models/base_model.go b/internal/entity/models/base_model.go index aa84385a6e..ea154153f7 100644 --- a/internal/entity/models/base_model.go +++ b/internal/entity/models/base_model.go @@ -148,6 +148,86 @@ func (b *BaseModel) APIConfigCheck(apiConfig *APIConfig) error { return nil } +func newJSONPostRequest(ctx context.Context, url string, apiConfig *APIConfig, reqBody map[string]any) (*http.Request, error) { + jsonData, err := json.Marshal(reqBody) + if err != nil { + return nil, fmt.Errorf("failed to marshal request: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData)) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + if auth := BearerAuth(apiConfig); auth != "" { + req.Header.Set("Authorization", auth) + } + + return req, nil +} + +// doRequest sends a JSON POST request and returns the response body. +func (b *BaseModel) doRequest(ctx context.Context, url string, apiConfig *APIConfig, reqBody map[string]any, timeout time.Duration) ([]byte, error) { + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + req, err := newJSONPostRequest(ctx, url, apiConfig, reqBody) + if err != nil { + return nil, err + } + + resp, err := b.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("failed to send request: %w", err) + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("failed to read response: %w", err) + } + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) + } + + return body, nil +} + +// mustMarshal marshals v to JSON, panicking on error. +func mustMarshal(v any) []byte { + b, err := json.Marshal(v) + if err != nil { + panic(fmt.Sprintf("failed to marshal: %v", err)) + } + return b +} + +// doStreamRequest sends a JSON POST request and calls handler with the response body. +func (b *BaseModel) doStreamRequest(ctx context.Context, url string, apiConfig *APIConfig, reqBody map[string]any, timeout time.Duration, handler func(io.ReadCloser) error) error { + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + req, err := newJSONPostRequest(ctx, url, apiConfig, reqBody) + if err != nil { + return err + } + + resp, err := b.httpClient.Do(req) + if err != nil { + return fmt.Errorf("failed to send request: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) + } + + return handler(resp.Body) +} + // BearerAuth returns the Bearer token for Authorization header, // or empty string if apiConfig or its ApiKey is nil/empty. func BearerAuth(apiConfig *APIConfig) string { diff --git a/internal/entity/models/common_test.go b/internal/entity/models/common_test.go index 09d5ca50b3..22fbf5b74e 100644 --- a/internal/entity/models/common_test.go +++ b/internal/entity/models/common_test.go @@ -1,6 +1,14 @@ package models -import "testing" +import ( + "context" + "fmt" + "io" + "net/http" + "net/http/httptest" + "testing" + "time" +) func TestNormalizeModelFamily(t *testing.T) { tests := []struct { @@ -97,6 +105,58 @@ func TestGetThinkingAndAnswerHandlesNilInputs(t *testing.T) { } } +func TestBaseModelDoRequestAuthorizationHeader(t *testing.T) { + tests := []struct { + name string + apiConfig *APIConfig + wantHeader string + }{ + {name: "nil api config"}, + {name: "nil api key", apiConfig: &APIConfig{}}, + {name: "blank api key", apiConfig: &APIConfig{ApiKey: modelFamilyTestString(" ")}}, + {name: "api key", apiConfig: &APIConfig{ApiKey: modelFamilyTestString(" key-1 ")}, wantHeader: "Bearer key-1"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get("Authorization"); got != tt.wantHeader { + t.Fatalf("Authorization = %q, want %q", got, tt.wantHeader) + } + if got := r.Header.Get("Content-Type"); got != "application/json" { + t.Fatalf("Content-Type = %q, want application/json", got) + } + fmt.Fprint(w, `{"ok":true}`) + })) + defer server.Close() + + model := &BaseModel{httpClient: server.Client(), AllowEmptyAPIKey: true} + if _, err := model.doRequest(context.Background(), server.URL, tt.apiConfig, map[string]any{"ok": true}, time.Second); err != nil { + t.Fatalf("doRequest() error = %v", err) + } + }) + } +} + +func TestBaseModelDoStreamRequestAllowsMissingAPIKey(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get("Authorization"); got != "" { + t.Fatalf("Authorization = %q, want empty", got) + } + fmt.Fprint(w, `data: {"ok":true}`) + })) + defer server.Close() + + model := &BaseModel{httpClient: server.Client(), AllowEmptyAPIKey: true} + err := model.doStreamRequest(context.Background(), server.URL, nil, map[string]any{"ok": true}, time.Second, func(body io.ReadCloser) error { + _, err := io.ReadAll(body) + return err + }) + if err != nil { + t.Fatalf("doStreamRequest() error = %v", err) + } +} + func modelFamilyTestString(value string) *string { return &value } diff --git a/internal/entity/models/deepseek.go b/internal/entity/models/deepseek.go index 2ab0212bfe..bb8b0dde14 100644 --- a/internal/entity/models/deepseek.go +++ b/internal/entity/models/deepseek.go @@ -100,7 +100,6 @@ func (d *DeepSeekModel) ChatWithMessages(ctx context.Context, modelName string, } url := fmt.Sprintf("%s/%s", resolvedBaseURL, d.baseModel.URLSuffix.Chat) - // Build request body reqBody := buildRequestBody(chatModelConfig, modelName, messages, false) if chatModelConfig != nil && chatModelConfig.Thinking != nil && *chatModelConfig.Thinking { var thinkingFlag string @@ -111,25 +110,19 @@ func (d *DeepSeekModel) ChatWithMessages(ctx context.Context, modelName string, switch effort { case "none": thinkingFlag = "disabled" - break case "low": thinkingFlag = "disabled" - break case "medium": thinkingFlag = "disabled" - break case "high": thinkingFlag = "enabled" reqBody["reasoning_effort"] = "high" - break case "default": thinkingFlag = "enabled" reqBody["reasoning_effort"] = "high" - break case "max": thinkingFlag = "enabled" reqBody["reasoning_effort"] = "max" - break } reqBody["thinking"] = map[string]interface{}{ "type": thinkingFlag, @@ -140,67 +133,12 @@ func (d *DeepSeekModel) ChatWithMessages(ctx context.Context, modelName string, } } - jsonData, err := json.Marshal(reqBody) + body, err := d.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout) if err != nil { - return nil, fmt.Errorf("failed to marshal request: %w", err) + return nil, err } - ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout) - defer cancel() - - req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData)) - if err != nil { - return nil, fmt.Errorf("failed to create request: %w", err) - } - - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) - - resp, err := d.baseModel.httpClient.Do(req) - if err != nil { - return nil, fmt.Errorf("failed to send request: %w", err) - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - if err != nil { - return nil, fmt.Errorf("failed to read response: %w", err) - } - - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) - } - - return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, _ *ChatConfig) (chatResponseParts, error) { - var result DeepSeekChatResponse - if err := json.Unmarshal(body, &result); err != nil { - return chatResponseParts{}, fmt.Errorf("failed to parse response: %w", err) - } - - if len(result.Choices) == 0 { - return chatResponseParts{}, fmt.Errorf("no choices in response") - } - - choice := &result.Choices[0] - reasonContent := choice.Message.ReasoningContent - // DeepSeek may prefix reasoning content with a newline in non-streaming - // responses; keep the existing provider behavior of removing that prefix. - if reasonContent != "" && reasonContent[0] == '\n' { - reasonContent = reasonContent[1:] - } - - return chatResponseParts{ - RequestID: result.ID, - Content: &choice.Message.Content, - ReasonContent: &reasonContent, - ToolCalls: choice.Message.ToolCalls, - Usage: &TokenUsage{ - PromptTokens: result.Usage.PromptTokens, - CompletionTokens: result.Usage.CompletionTokens, - TotalTokens: result.Usage.TotalTokens, - }, - }, nil - }) + return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel) @@ -219,7 +157,6 @@ func (d *DeepSeekModel) ChatStreamlyWithSender(ctx context.Context, modelName st } url := fmt.Sprintf("%s/chat/completions", resolvedBaseURL) - // Build request body with streaming enabled reqBody := buildRequestBody(chatModelConfig, modelName, messages, true) if chatModelConfig != nil && chatModelConfig.Thinking != nil && *chatModelConfig.Thinking { var thinkingFlag string @@ -282,68 +219,7 @@ func (d *DeepSeekModel) ChatStreamlyWithSender(ctx context.Context, modelName st return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) } - sawTerminal := false - accumulatedToolCalls := make(map[int]map[string]interface{}) - done, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error { - common.Info(fmt.Sprintf("%v", event)) - - tokenUsage, found, usageErr := decodeOpenAICompatibleStreamUsage(event) - if usageErr != nil { - return usageErr - } - if found { - applyStreamUsage(chatModelConfig, modelUsage, tokenUsage) - } - - choices, ok := event["choices"].([]interface{}) - if !ok || len(choices) == 0 { - return nil - } - - firstChoice, ok := choices[0].(map[string]interface{}) - if !ok { - return nil - } - - delta, ok := firstChoice["delta"].(map[string]interface{}) - if !ok { - return nil - } - - accumulateToolCallDeltas(delta, accumulatedToolCalls) - - content, ok := delta["content"].(string) - if ok && content != "" { - if err := sender(&content, nil); err != nil { - return err - } - } - - reasoningContent, ok := delta["reasoning_content"].(string) - if ok && reasoningContent != "" { - if err := sender(nil, &reasoningContent); err != nil { - return err - } - } - - finishReason, ok := firstChoice["finish_reason"].(string) - if ok && finishReason != "" { - sawTerminal = true - } - return nil - }) - if err != nil { - return fmt.Errorf("failed to scan response body: %w", err) - } - if !done && !sawTerminal { - return fmt.Errorf("deepseek: stream ended before [DONE] or finish_reason") - } - - setSortedToolCallsResult(chatModelConfig, accumulatedToolCalls) - - // Send [DONE] marker for OpenAI compatibility - endOfStream := "[DONE]" - return sender(&endOfStream, nil) + return HandleStreamingResponse(resp.Body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) } // Embed embeds a list of texts into embeddings diff --git a/internal/entity/models/parser_config.go b/internal/entity/models/parser_config.go new file mode 100644 index 0000000000..28f7530212 --- /dev/null +++ b/internal/entity/models/parser_config.go @@ -0,0 +1,36 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package models + +// ParserConfig maps a protocol to its usage parsers. Drivers select a +// ParserConfig instead of implementing usage extraction individually. +type ParserConfig struct { + // Protocol is the protocol identifier (e.g. "openai", "claude"). + Protocol string + // ResponseParser extracts usage from a non-streaming response body. + ResponseParser func(map[string]any) (*TokenUsage, bool) + // StreamParser extracts usage from one streaming event. + StreamParser func(map[string]any) (*TokenUsage, bool) +} + +// OpenAIParserConfig is the ParserConfig for OpenAI-compatible APIs +// (NVIDIA NIM, DeepSeek, Aliyun, Moonshot, xAI, OpenRouter, ...). +var OpenAIParserConfig = &ParserConfig{ + Protocol: "openai", + ResponseParser: extractOpenAIUsage, + StreamParser: extractOpenAIStreamUsage, +} diff --git a/internal/entity/models/response_handler.go b/internal/entity/models/response_handler.go new file mode 100644 index 0000000000..f21e93aa83 --- /dev/null +++ b/internal/entity/models/response_handler.go @@ -0,0 +1,235 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package models + +import ( + "encoding/json" + "fmt" + "io" + "ragflow/internal/common" + + "go.uber.org/zap" +) + +// HandleNonStreamingResponse processes a complete non-streaming chat +// response using the ParserConfig's ResponseParser to extract usage. +func HandleNonStreamingResponse( + body []byte, + modelUsage *common.ModelUsage, + chatConfig *ChatConfig, + cfg *ParserConfig, +) (*ChatResponse, error) { + var result map[string]any + if err := json.Unmarshal(body, &result); err != nil { + return nil, fmt.Errorf("failed to parse response: %w", err) + } + + // Extract usage via the protocol-specific parser. + var usage *TokenUsage + if u, ok := cfg.ResponseParser(result); ok { + usage = u + } + + // Extract content / reasoning_content / tool_calls. + content, reasonContent, toolCalls := extractContentAndChoices(result) + + if content == nil && len(toolCalls) == 0 { + return nil, fmt.Errorf("no choices in response") + } + + if usage != nil { + recordResponseUsage(modelUsage, extractRequestID(result), usage, "chat") + if chatConfig != nil { + chatConfig.UsageResult = usage + model, _ := result["model"].(string) + common.Info("StreamUsage", zap.String("model", model), zap.Int("prompt", usage.PromptTokens), zap.Int("completion", usage.CompletionTokens), zap.Int("total", usage.TotalTokens)) + } + } + + return &ChatResponse{ + Answer: content, + ReasonContent: reasonContent, + ToolCalls: toolCalls, + Usage: usage, + }, nil +} + +// HandleStreamingResponse processes a streaming chat response using the +// ParserConfig's StreamParser to extract usage from each event. +func HandleStreamingResponse( + body io.Reader, + modelUsage *common.ModelUsage, + chatConfig *ChatConfig, + cfg *ParserConfig, + sender func(*string, *string) error, +) error { + if sender == nil { + return fmt.Errorf("sender is required") + } + + var streamUsage *TokenUsage + accumulatedToolCalls := make(map[int]map[string]any) + sawTerminal := false + + var streamModel string + done, err := ParseSSEStream[map[string]any](body, func(event map[string]any) error { + if u, ok := cfg.StreamParser(event); ok { + streamUsage = u + } + if m, ok := event["model"].(string); ok { + streamModel = m + } + + if apiErr, ok := event["error"]; ok && apiErr != nil { + return fmt.Errorf("upstream stream error: %v", apiErr) + } + + choices, ok := event["choices"].([]any) + if !ok || len(choices) == 0 { + return nil + } + + firstChoice, ok := choices[0].(map[string]any) + if !ok { + return nil + } + + delta, ok := firstChoice["delta"].(map[string]any) + if !ok { + return nil + } + + accumulateToolCallDeltas(delta, accumulatedToolCalls) + + if reasoningContent, ok := delta["reasoning_content"].(string); ok && reasoningContent != "" { + if err := sender(nil, &reasoningContent); err != nil { + return err + } + } + + if content, ok := delta["content"].(string); ok && content != "" { + if err := sender(&content, nil); err != nil { + return err + } + } + + if finishReason, ok := firstChoice["finish_reason"].(string); ok && finishReason != "" { + sawTerminal = true + } + + return nil + }) + if err != nil { + return fmt.Errorf("failed to scan response body: %w", err) + } + + if chatConfig != nil { + setSortedToolCallsResult(chatConfig, accumulatedToolCalls) + } + + if !done && !sawTerminal { + return fmt.Errorf("stream ended before [DONE] or finish_reason") + } + + if streamUsage != nil { + recordResponseUsage(modelUsage, "", streamUsage, "chat") + if chatConfig != nil { + chatConfig.UsageResult = streamUsage + common.Info("StreamUsage", zap.String("model", streamModel), zap.Int("prompt", streamUsage.PromptTokens), zap.Int("completion", streamUsage.CompletionTokens), zap.Int("total", streamUsage.TotalTokens)) + } + } + + endOfStream := "[DONE]" + return sender(&endOfStream, nil) +} + +// extractContentAndChoices extracts content, reasoning_content, and tool_calls +// from a parsed non-streaming response. +func extractContentAndChoices(result map[string]any) (*string, *string, []map[string]any) { + choices, ok := result["choices"].([]any) + if !ok || len(choices) == 0 { + return nil, nil, nil + } + + firstChoice, ok := choices[0].(map[string]any) + if !ok { + return nil, nil, nil + } + + messageMap, ok := firstChoice["message"].(map[string]any) + if !ok { + return nil, nil, nil + } + + var content *string + if c, ok := messageMap["content"].(string); ok { + cc := c + content = &cc + } else if _, ok := messageMap["tool_calls"].([]any); ok { + // content may be nil when the response only carries tool calls. + // Return an empty-string pointer so callers can rely on a non-nil Answer. + empty := "" + content = &empty + } else if rc, ok := messageMap["reasoning_content"].(string); ok && rc != "" { + // content may be nil when the response only carries reasoning_content. + empty := "" + content = &empty + } + + var reasonContent *string + if rc, ok := messageMap["reasoning_content"].(string); ok { + reason := rc + if reason != "" && reason[0] == '\n' { + reason = reason[1:] + } + reasonContent = &reason + } else { + // Always return a non-nil pointer so callers can rely on it. + empty := "" + reasonContent = &empty + } + + // Some providers (e.g. SiliconFlow) embed thinking inline via tags. + // Extract reasoning into ReasonContent and strip it from Answer. + if content != nil { + if reasoning, answer := extractThinkContent(content); reasoning != nil { + rc := *reasoning + reasonContent = &rc + a := *answer + content = &a + } + } + + var toolCalls []map[string]any + if tcs, ok := messageMap["tool_calls"].([]any); ok { + for _, tc := range tcs { + if tcMap, ok := tc.(map[string]any); ok { + toolCalls = append(toolCalls, tcMap) + } + } + } + + return content, reasonContent, toolCalls +} + +// extractRequestID extracts the request ID from a parsed response. +func extractRequestID(result map[string]any) string { + if id, ok := result["id"].(string); ok { + return id + } + return "" +} diff --git a/internal/entity/models/usage_parser.go b/internal/entity/models/usage_parser.go new file mode 100644 index 0000000000..6282e21773 --- /dev/null +++ b/internal/entity/models/usage_parser.go @@ -0,0 +1,77 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package models + +import "encoding/json" + +// extractOpenAIUsage extracts token usage from an OpenAI-compatible +// non-streaming response body. Returns (nil, false) when the response +// carries no usable usage block. +func extractOpenAIUsage(body map[string]any) (*TokenUsage, bool) { + rawUsage, ok := body["usage"].(map[string]any) + if !ok { + return nil, false + } + + return tokenUsageFromRaw(rawUsage), true +} + +// extractOpenAIStreamUsage extracts token usage from one OpenAI-compatible +// streaming event. A missing or null usage field is not an error. +func extractOpenAIStreamUsage(event map[string]any) (*TokenUsage, bool) { + rawUsage, ok := event["usage"].(map[string]any) + if !ok || rawUsage == nil { + return nil, false + } + + return tokenUsageFromRaw(rawUsage), true +} + +func tokenUsageFromRaw(rawUsage map[string]any) *TokenUsage { + usage := &TokenUsage{} + usage.PromptTokens = extractToken(rawUsage, "prompt_tokens", "input_tokens") + usage.CompletionTokens = extractToken(rawUsage, "completion_tokens", "output_tokens") + usage.TotalTokens = extractToken(rawUsage, "total_tokens") + + if usage.TotalTokens == 0 { + usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens + } + + return usage +} + +// extractToken reads a numeric field from a map, trying each key in order. +// Returns 0 when no key is present or the value is not a number. +func extractToken(m map[string]any, keys ...string) int { + for _, k := range keys { + if v, ok := m[k]; ok { + switch val := v.(type) { + case float64: + return int(val) + case int: + return val + case int64: + return int(val) + case json.Number: + if n, err := val.Int64(); err == nil { + return int(n) + } + } + } + } + return 0 +} diff --git a/internal/entity/models/usage_parser_test.go b/internal/entity/models/usage_parser_test.go new file mode 100644 index 0000000000..674eca7709 --- /dev/null +++ b/internal/entity/models/usage_parser_test.go @@ -0,0 +1,153 @@ +package models + +import ( + "testing" +) + +func TestExtractOpenAIUsage(t *testing.T) { + cases := []struct { + name string + body map[string]any + want *TokenUsage + wantOK bool + }{ + { + name: "standard openai", + body: map[string]any{ + "usage": map[string]any{ + "prompt_tokens": float64(10), + "completion_tokens": float64(20), + "total_tokens": float64(30), + }, + }, + want: &TokenUsage{PromptTokens: 10, CompletionTokens: 20, TotalTokens: 30}, + wantOK: true, + }, + { + name: "total missing falls back to sum", + body: map[string]any{ + "usage": map[string]any{ + "prompt_tokens": float64(10), + "completion_tokens": float64(20), + }, + }, + want: &TokenUsage{PromptTokens: 10, CompletionTokens: 20, TotalTokens: 30}, + wantOK: true, + }, + { + name: "no usage key", + body: map[string]any{ + "choices": []any{map[string]any{"message": map[string]any{"content": "hi"}}}, + }, + want: nil, + wantOK: false, + }, + { + name: "cohere style input_tokens/output_tokens", + body: map[string]any{ + "usage": map[string]any{ + "input_tokens": float64(100), + "output_tokens": float64(50), + }, + }, + want: &TokenUsage{PromptTokens: 100, CompletionTokens: 50, TotalTokens: 150}, + wantOK: true, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got, ok := extractOpenAIUsage(tc.body) + if ok != tc.wantOK { + t.Fatalf("extractOpenAIUsage() ok = %v, want %v", ok, tc.wantOK) + } + if !tc.wantOK { + return + } + if got.PromptTokens != tc.want.PromptTokens || + got.CompletionTokens != tc.want.CompletionTokens || + got.TotalTokens != tc.want.TotalTokens { + t.Fatalf("extractOpenAIUsage() = %+v, want %+v", got, tc.want) + } + }) + } +} + +func TestExtractOpenAIStreamUsage(t *testing.T) { + cases := []struct { + name string + event map[string]any + want *TokenUsage + wantOK bool + }{ + { + name: "standard streaming event", + event: map[string]any{ + "usage": map[string]any{ + "prompt_tokens": float64(100), + "completion_tokens": float64(50), + "total_tokens": float64(150), + }, + }, + want: &TokenUsage{PromptTokens: 100, CompletionTokens: 50, TotalTokens: 150}, + wantOK: true, + }, + { + name: "null usage", + event: map[string]any{ + "usage": nil, + }, + want: nil, + wantOK: false, + }, + { + name: "no usage key", + event: map[string]any{ + "choices": []any{map[string]any{"delta": map[string]any{"content": "hi"}}}, + }, + want: nil, + wantOK: false, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got, ok := extractOpenAIStreamUsage(tc.event) + if ok != tc.wantOK { + t.Fatalf("extractOpenAIStreamUsage() ok = %v, want %v", ok, tc.wantOK) + } + if !tc.wantOK { + return + } + if got.PromptTokens != tc.want.PromptTokens || + got.CompletionTokens != tc.want.CompletionTokens || + got.TotalTokens != tc.want.TotalTokens { + t.Fatalf("extractOpenAIStreamUsage() = %+v, want %+v", got, tc.want) + } + }) + } +} + +func TestExtractToken(t *testing.T) { + m := map[string]any{ + "a": float64(1), + "b": int(2), + "c": int64(3), + } + + if got := extractToken(m, "a"); got != 1 { + t.Fatalf("extractToken(a) = %d, want 1", got) + } + if got := extractToken(m, "b"); got != 2 { + t.Fatalf("extractToken(b) = %d, want 2", got) + } + if got := extractToken(m, "c"); got != 3 { + t.Fatalf("extractToken(c) = %d, want 3", got) + } + if got := extractToken(m, "missing"); got != 0 { + t.Fatalf("extractToken(missing) = %d, want 0", got) + } + if got := extractToken(m, "missing", "a"); got != 1 { + t.Fatalf("extractToken(missing, a) = %d, want 1", got) + } +}