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) <noreply@anthropic.com> Co-authored-by: Haruko386 <tryeverypossible@163.com> Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
2026-08-01 21:37:33 +08:00 · 2026-07-31 18:32:07 +08:00
parent 29287ef74b
commit fe02c9b95f
7 changed files with 646 additions and 129 deletions
--- 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 {
--- 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
 }
--- 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
--- a/internal/entity/models/parser_config.go
+++ 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,
+}
--- a/internal/entity/models/response_handler.go
+++ 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 <think> 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 ""
+}
--- a/internal/entity/models/usage_parser.go
+++ 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
+}
--- a/internal/entity/models/usage_parser_test.go
+++ 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)
+	}
+}