Go: add token usage for orcarouter, baichuan, cohere and novita (#17455)

As title #17284

internal/entity/models/baichuan.go

@@ -49,6 +49,29 @@ func (b *BaichuanModel) Name() string {
 	return "BaiChuan"
 }
 
+type BaiChuanChatResponse struct {
+	ID      string `json:"id"`
+	Choices []struct {
+		FinishReason string `json:"finish_reason"`
+		Index        int    `json:"index"`
+		Message      struct {
+			Content   string           `json:"content"`
+			Role      string           `json:"role"`
+			ToolCalls []map[string]any `json:"tool_calls"`
+		} `json:"message"`
+		Logprobs interface{} `json:"logprobs"`
+	} `json:"choices"`
+	Created int    `json:"created"`
+	Model   string `json:"model"`
+	Object  string `json:"object"`
+	Usage   struct {
+		CompletionTokens int `json:"completion_tokens"`
+		PromptTokens     int `json:"prompt_tokens"`
+		TotalTokens      int `json:"total_tokens"`
+		SearchCount      int `json:"search_count"`
+	} `json:"usage"`
+}
+
 func (b *BaichuanModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) {
 	if err := b.baseModel.APIConfigCheck(apiConfig); err != nil {
 		return nil, err
@@ -96,41 +119,34 @@ func (b *BaichuanModel) ChatWithMessages(ctx context.Context, modelName string,
 	}
 
 	// Parse response
-	var result map[string]interface{}
-	if err := json.Unmarshal(body, &result); err != nil {
-		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
-	}
+	return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, chatConfig *ChatConfig) (chatResponseParts, error) {
+		var result BaiChuanChatResponse
+		if err := json.Unmarshal(body, &result); err != nil {
+			return chatResponseParts{}, fmt.Errorf("failed to unmarshal response: %w", err)
+		}
+		if len(result.Choices) == 0 {
+			return chatResponseParts{}, fmt.Errorf("no choices in response")
+		}
 
-	choices, ok := result["choices"].([]interface{})
-	if !ok {
-		return nil, fmt.Errorf("no choices in response")
-	}
+		choice := &result.Choices[0]
+		content := choice.Message.Content
+		if content == "" && len(choice.Message.ToolCalls) == 0 {
+			return chatResponseParts{}, fmt.Errorf("no message in response")
+		}
 
-	firstChoice, ok := choices[0].(map[string]interface{})
-	if !ok {
-		return nil, fmt.Errorf("no choices in response")
-	}
-
-	messageMap, ok := firstChoice["message"].(map[string]interface{})
-	if !ok {
-		return nil, fmt.Errorf("no message in response")
-	}
-
-	content, ok := messageMap["content"].(string)
-	toolCalls := extractToolCalls(messageMap)
-	if !ok && len(toolCalls) == 0 {
-		return nil, fmt.Errorf("no message in response")
-	}
-
-	// baichuan not support think
-	emptyReason := ""
-	chatResponse := &ChatResponse{
-		Answer:        &content,
-		ReasonContent: &emptyReason,
-		ToolCalls:     toolCalls,
-	}
-
-	return chatResponse, nil
+		emptyReason := ""
+		return chatResponseParts{
+			RequestID:     result.ID,
+			Content:       &content,
+			ReasonContent: &emptyReason,
+			ToolCalls:     choice.Message.ToolCalls,
+			Usage: &TokenUsage{
+				PromptTokens:     result.Usage.PromptTokens,
+				CompletionTokens: result.Usage.CompletionTokens,
+				TotalTokens:      result.Usage.TotalTokens,
+			},
+		}, nil
+	})
 }
 
 func (b *BaichuanModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error {
@@ -183,6 +199,14 @@ func (b *BaichuanModel) ChatStreamlyWithSender(ctx context.Context, modelName st
 	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(modelConfig, modelUsage, tokenUsage)
+		}
+
 		choices, ok := event["choices"].([]interface{})
 		if !ok || len(choices) == 0 {
 			return nil
@@ -279,10 +303,15 @@ func (b *BaichuanModel) Embed(ctx context.Context, modelName *string, texts []st
 	}
 
 	var parsedResponse struct {
+		ID   string `json:"id"`
 		Data []struct {
 			Embedding []float64 `json:"embedding"`
 			Index     int       `json:"index"`
 		} `json:"data"`
+		Usage struct {
+			PromptTokens int `json:"prompt_tokens"`
+			TotalTokens  int `json:"total_tokens"`
+		}
 	}
 
 	if err = json.Unmarshal(body, &parsedResponse); err != nil {
@@ -300,6 +329,10 @@ func (b *BaichuanModel) Embed(ctx context.Context, modelName *string, texts []st
 			Index:     dataElem.Index,
 		})
 	}
+	recordResponseUsage(modelUsage, parsedResponse.ID, &TokenUsage{
+		PromptTokens: parsedResponse.Usage.PromptTokens,
+		TotalTokens:  parsedResponse.Usage.TotalTokens,
+	}, "embedding")
 
 	return embeddings, nil
 }

internal/entity/models/cohere.go

@@ -53,6 +53,86 @@ func (c *CoHereModel) Name() string {
 	return "Cohere"
 }
 
+type CohereChatResponse struct {
+	ID           string `json:"id"`
+	FinishReason string `json:"finish_reason"`
+	Message      struct {
+		Role      string           `json:"role"`
+		ToolCalls []map[string]any `json:"tool_calls"`
+
+		Content []struct {
+			Type     string `json:"type"`
+			Text     string `json:"text"`
+			Thinking string `json:"thinking"`
+		} `json:"content"`
+	} `json:"message"`
+	Usage struct {
+		Tokens struct {
+			InputTokens  float64 `json:"input_tokens"`
+			OutputTokens float64 `json:"output_tokens"`
+		} `json:"tokens"`
+	} `json:"usage"`
+}
+
+func buildCohereTokenUsage(inputTokens, outputTokens float64) *TokenUsage {
+	input := int(inputTokens)
+	output := int(outputTokens)
+	return &TokenUsage{
+		PromptTokens:     input,
+		CompletionTokens: output,
+		TotalTokens:      input + output,
+	}
+}
+
+func decodeCohereStreamUsage(event map[string]interface{}) (*TokenUsage, bool, error) {
+	eventType, _ := event["type"].(string)
+	if eventType != "message-end" {
+		return nil, false, nil
+	}
+	var parsed struct {
+		Delta struct {
+			Usage struct {
+				Tokens struct {
+					InputTokens  float64 `json:"input_tokens"`
+					OutputTokens float64 `json:"output_tokens"`
+				} `json:"tokens"`
+			} `json:"usage"`
+		} `json:"delta"`
+		Usage struct {
+			Tokens struct {
+				InputTokens  float64 `json:"input_tokens"`
+				OutputTokens float64 `json:"output_tokens"`
+			} `json:"tokens"`
+		} `json:"usage"`
+		Response struct {
+			Usage struct {
+				Tokens struct {
+					InputTokens  float64 `json:"input_tokens"`
+					OutputTokens float64 `json:"output_tokens"`
+				} `json:"tokens"`
+			} `json:"usage"`
+		} `json:"response"`
+	}
+	body, err := json.Marshal(event)
+	if err != nil {
+		return nil, false, err
+	}
+	if err := json.Unmarshal(body, &parsed); err != nil {
+		return nil, false, err
+	}
+	tokens := parsed.Response.Usage.Tokens
+	if tokens.InputTokens == 0 && tokens.OutputTokens == 0 {
+		tokens = parsed.Delta.Usage.Tokens
+	}
+	if tokens.InputTokens == 0 && tokens.OutputTokens == 0 {
+		tokens = parsed.Usage.Tokens
+	}
+	if tokens.InputTokens == 0 && tokens.OutputTokens == 0 {
+		return nil, false, nil
+	}
+	return buildCohereTokenUsage(tokens.InputTokens, tokens.OutputTokens), true, nil
+}
+
 func (c *CoHereModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) {
 	if err := c.baseModel.APIConfigCheck(apiConfig); err != nil {
 		return nil, err
@@ -120,44 +200,33 @@ func (c *CoHereModel) ChatWithMessages(ctx context.Context, modelName string, me
 		return nil, fmt.Errorf("Cohere chat API error: %d %s", resp.StatusCode, string(body))
 	}
 
-	// Parse response
-	var result map[string]interface{}
-	if err := json.Unmarshal(body, &result); err != nil {
-		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
-	}
-
-	messageMap, ok := result["message"].(map[string]interface{})
-	if !ok {
-		return nil, fmt.Errorf("no message found in Cohere response: %s", string(body))
-	}
-
-	contentArray, ok := messageMap["content"].([]interface{})
-	if !ok {
-		return nil, fmt.Errorf("content is not an array in Cohere response")
-	}
-
-	var fullContent string
-	var reasonContent string
-	for _, cBlock := range contentArray {
-		cmap, ok := cBlock.(map[string]interface{})
-		if !ok {
-			continue
+	return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, chatConfig *ChatConfig) (chatResponseParts, error) {
+		var result CohereChatResponse
+		if err := json.Unmarshal(body, &result); err != nil {
+			return chatResponseParts{}, fmt.Errorf("failed to unmarshal response: %w", err)
 		}
-		if blockType, ok := cmap["type"].(string); ok && blockType == "thinking" {
-			if thinkingText, ok := cmap["thinking"].(string); ok {
-				reasonContent += thinkingText
+		if len(result.Message.Content) == 0 && len(result.Message.ToolCalls) == 0 {
+			return chatResponseParts{}, fmt.Errorf("content is not an array in Cohere response")
+		}
+
+		var fullContent string
+		var reasonContent string
+		for _, block := range result.Message.Content {
+			if block.Type == "thinking" {
+				reasonContent += block.Thinking
+			} else {
+				fullContent += block.Text
 			}
-		} else if text, ok := cmap["text"].(string); ok {
-			fullContent += text
 		}
-	}
 
-	chatResponse := &ChatResponse{
-		Answer:        &fullContent,
-		ReasonContent: &reasonContent,
-	}
-
-	return chatResponse, nil
+		return chatResponseParts{
+			RequestID:     result.ID,
+			Content:       &fullContent,
+			ReasonContent: &reasonContent,
+			ToolCalls:     result.Message.ToolCalls,
+			Usage:         buildCohereTokenUsage(result.Usage.Tokens.InputTokens, result.Usage.Tokens.OutputTokens),
+		}, nil
+	})
 }
 
 func (c *CoHereModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error {
@@ -231,6 +300,14 @@ func (c *CoHereModel) ChatStreamlyWithSender(ctx context.Context, modelName stri
 			return nil
 		}
 
+		tokenUsage, found, usageErr := decodeCohereStreamUsage(event)
+		if usageErr != nil {
+			return usageErr
+		}
+		if found {
+			applyStreamUsage(modelConfig, modelUsage, tokenUsage)
+		}
+
 		if eventType == "message-end" {
 			sawTerminal = true
 			return nil
@@ -341,9 +418,16 @@ func (c *CoHereModel) Embed(ctx context.Context, modelName *string, texts []stri
 	}
 
 	var result struct {
+		ID         string `json:"id"`
 		Embeddings struct {
 			Float [][]float64 `json:"float"`
 		} `json:"embeddings"`
+		Meta struct {
+			Tokens struct {
+				InputTokens  float64 `json:"input_tokens"`
+				OutputTokens float64 `json:"output_tokens"`
+			}
+		} `json:"meta"`
 	}
 	if err = json.Unmarshal(body, &result); err != nil {
 		return nil, fmt.Errorf("failed to decode response: %w", err)
@@ -360,6 +444,7 @@ func (c *CoHereModel) Embed(ctx context.Context, modelName *string, texts []stri
 			Index:     i,
 		})
 	}
+	recordResponseUsage(modelUsage, result.ID, buildCohereTokenUsage(result.Meta.Tokens.InputTokens, result.Meta.Tokens.OutputTokens), "embedding")
 
 	return embeddings, nil
 }
@@ -426,10 +511,17 @@ func (c *CoHereModel) Rerank(ctx context.Context, modelName *string, query strin
 	}
 
 	var rerankResp struct {
+		ID      string `json:"id"`
 		Results []struct {
 			Index          int     `json:"index"`
 			RelevanceScore float64 `json:"relevance_score"`
 		} `json:"results"`
+		Meta struct {
+			Tokens struct {
+				InputTokens  float64 `json:"input_tokens"`
+				OutputTokens float64 `json:"output_tokens"`
+			}
+		} `json:"meta"`
 	}
 
 	if err := json.Unmarshal(body, &rerankResp); err != nil {
@@ -444,6 +536,7 @@ func (c *CoHereModel) Rerank(ctx context.Context, modelName *string, query strin
 		}
 		rerankResponse.Data = append(rerankResponse.Data, rerankResult)
 	}
+	recordResponseUsage(modelUsage, rerankResp.ID, buildCohereTokenUsage(rerankResp.Meta.Tokens.InputTokens, rerankResp.Meta.Tokens.OutputTokens), "rerank")
 
 	return &rerankResponse, nil
 }

internal/entity/models/cohere_test.go

@@ -0,0 +1,127 @@
+package models
+
+import (
+	"encoding/json"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"ragflow/internal/common"
+	"strings"
+	"testing"
+)
+
+func newCohereForTest(baseURL string) *CoHereModel {
+	return NewCoHereModel(
+		map[string]string{"default": baseURL},
+		URLSuffix{Chat: "v2/chat", Embedding: "v2/embed", Rerank: "v2/rerank"},
+	)
+}
+
+func TestCohereChatAllowsToolCallOnlyResponse(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path != "/v2/chat" {
+			t.Errorf("path=%s", r.URL.Path)
+		}
+		_ = json.NewEncoder(w).Encode(map[string]any{
+			"id": "cohere-chat-tool",
+			"message": map[string]any{
+				"role": "assistant",
+				"tool_calls": []map[string]any{
+					{
+						"id":   "tool-1",
+						"type": "function",
+						"function": map[string]any{
+							"name":      "lookup",
+							"arguments": `{"q":"ragflow"}`,
+						},
+					},
+				},
+			},
+			"usage": map[string]any{
+				"tokens": map[string]any{
+					"input_tokens":  4,
+					"output_tokens": 2,
+				},
+			},
+		})
+	}))
+	defer srv.Close()
+
+	apiKey := "test-key"
+	usage := &common.ModelUsage{}
+	resp, err := newCohereForTest(srv.URL).ChatWithMessages(
+		t.Context(),
+		"command-a",
+		[]Message{{Role: "user", Content: "use a tool"}},
+		&APIConfig{ApiKey: &apiKey},
+		nil,
+		usage,
+	)
+	if err != nil {
+		t.Fatalf("ChatWithMessages: %v", err)
+	}
+	if resp.Answer == nil || *resp.Answer != "" {
+		t.Fatalf("Answer=%v, want empty content", resp.Answer)
+	}
+	if len(resp.ToolCalls) != 1 || resp.ToolCalls[0]["id"] != "tool-1" {
+		t.Fatalf("ToolCalls=%#v, want tool-1", resp.ToolCalls)
+	}
+	assertModelUsage(t, usage, 4, 2, 6)
+}
+
+func TestCohereStreamRecordsDeltaUsage(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path != "/v2/chat" {
+			t.Errorf("path=%s", r.URL.Path)
+		}
+		raw, err := io.ReadAll(r.Body)
+		if err != nil {
+			t.Fatalf("read body: %v", err)
+		}
+		var body map[string]any
+		if err := json.Unmarshal(raw, &body); err != nil {
+			t.Fatalf("unmarshal body: %v", err)
+		}
+		if body["stream"] != true {
+			t.Fatalf("stream=%v, want true", body["stream"])
+		}
+
+		w.Header().Set("Content-Type", "text/event-stream")
+		_, _ = io.WriteString(w,
+			`data: {"type":"content-delta","delta":{"message":{"content":{"text":"hello"}}}}`+"\n"+
+				`data: {"type":"message-end","delta":{"usage":{"tokens":{"input_tokens":4,"output_tokens":2}}}}`+"\n",
+		)
+	}))
+	defer srv.Close()
+
+	apiKey := "test-key"
+	cfg := &ChatConfig{}
+	usage := &common.ModelUsage{}
+	var chunks []string
+	err := newCohereForTest(srv.URL).ChatStreamlyWithSender(
+		t.Context(),
+		"command-a",
+		[]Message{{Role: "user", Content: "hi"}},
+		&APIConfig{ApiKey: &apiKey},
+		cfg,
+		usage,
+		func(content *string, reasoning *string) error {
+			if content != nil {
+				chunks = append(chunks, *content)
+			}
+			return nil
+		},
+	)
+	if err != nil {
+		t.Fatalf("ChatStreamlyWithSender: %v", err)
+	}
+	if strings.Join(chunks, "") != "hello[DONE]" {
+		t.Fatalf("chunks=%q, want hello[DONE]", strings.Join(chunks, ""))
+	}
+	if cfg.UsageResult == nil || cfg.UsageResult.PromptTokens != 4 || cfg.UsageResult.CompletionTokens != 2 || cfg.UsageResult.TotalTokens != 6 {
+		t.Fatalf("UsageResult=%#v, want prompt=4 completion=2 total=6", cfg.UsageResult)
+	}
+	if usage.InputTokens != 4 || usage.OutputTokens != 2 || usage.TotalTokens != 6 {
+		t.Fatalf("model usage=(%d,%d,%d), want (4,2,6)", usage.InputTokens, usage.OutputTokens, usage.TotalTokens)
+	}
+}

internal/entity/models/novita.go

@@ -52,6 +52,29 @@ func (n *NovitaModel) Name() string {
 	return "NovitaAI"
 }
 
+type NovitaChatResponse struct {
+	ID      string `json:"id"`
+	Choices []struct {
+		FinishReason string `json:"finish_reason"`
+		Index        int    `json:"index"`
+		Message      struct {
+			Content          string           `json:"content"`
+			ReasoningContent string           `json:"reasoning_content"`
+			Role             string           `json:"role"`
+			ToolCalls        []map[string]any `json:"tool_calls"`
+		} `json:"message"`
+		Logprobs interface{} `json:"logprobs"`
+	} `json:"choices"`
+	Created int    `json:"created"`
+	Model   string `json:"model"`
+	Object  string `json:"object"`
+	Usage   struct {
+		CompletionTokens int `json:"completion_tokens"`
+		PromptTokens     int `json:"prompt_tokens"`
+		TotalTokens      int `json:"total_tokens"`
+	} `json:"usage"`
+}
+
 const (
 	novitaThinkOpen  = "<think>"
 	novitaThinkClose = "</think>"
@@ -235,58 +258,50 @@ func (n *NovitaModel) ChatWithMessages(ctx context.Context, modelName string, me
 		return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
 	}
 
-	var result map[string]interface{}
-	if err = json.Unmarshal(body, &result); err != nil {
-		return nil, fmt.Errorf("failed to parse response: %w", err)
-	}
+	return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, chatConfig *ChatConfig) (chatResponseParts, error) {
+		var result NovitaChatResponse
+		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")
+		}
 
-	choices, ok := result["choices"].([]interface{})
-	if !ok || len(choices) == 0 {
-		return nil, fmt.Errorf("no choices in response")
-	}
+		choice := &result.Choices[0]
+		if choice.Message.Content == "" && len(choice.Message.ToolCalls) == 0 {
+			return chatResponseParts{}, fmt.Errorf("invalid content format")
+		}
 
-	firstChoice, ok := choices[0].(map[string]interface{})
-	if !ok {
-		return nil, fmt.Errorf("invalid choice format")
-	}
-
-	messageMap, ok := firstChoice["message"].(map[string]interface{})
-	if !ok {
-		return nil, fmt.Errorf("invalid message format")
-	}
-
-	rawContent, ok := messageMap["content"].(string)
-	toolCalls := extractToolCalls(messageMap)
-	if !ok && len(toolCalls) == 0 {
-		return nil, fmt.Errorf("invalid content format")
-	}
-
-	// Novita emits chain-of-thought in two different shapes depending
-	// on the model and on enable_thinking:
-	//   - qwen3-* and other inline-style models: chain-of-thought is
-	//     embedded inside content as <think>...</think> tags.
-	//   - deepseek-v3.1 / glm-4.5 (and any model with separate
-	//     reasoning enabled): chain-of-thought arrives in a separate
-	//     `reasoning_content` field, with `content` already cleaned.
-	// Handle both so the visible Answer is always tag-free and any
-	// reasoning the upstream supplied is preserved.
-	visible, reasoning := "", ""
-	if ok {
-		visible, reasoning = splitNovitaThink(rawContent)
-		if r, ok := messageMap["reasoning_content"].(string); ok && r != "" {
+		// Novita emits chain-of-thought in two different shapes depending
+		// on the model and on enable_thinking:
+		//   - qwen3-* and other inline-style models: chain-of-thought is
+		//     embedded inside content as <think>...</think> tags.
+		//   - deepseek-v3.1 / glm-4.5 (and any model with separate
+		//     reasoning enabled): chain-of-thought arrives in a separate
+		//     `reasoning_content` field, with `content` already cleaned.
+		// Handle both so the visible Answer is always tag-free and any
+		// reasoning the upstream supplied is preserved.
+		visible, reasoning := splitNovitaThink(choice.Message.Content)
+		if choice.Message.ReasoningContent != "" {
 			if reasoning != "" {
-				reasoning += "\n" + r
+				reasoning += "\n" + choice.Message.ReasoningContent
 			} else {
-				reasoning = r
+				reasoning = choice.Message.ReasoningContent
 			}
 		}
-	}
 
-	return &ChatResponse{
-		Answer:        &visible,
-		ReasonContent: &reasoning,
-		ToolCalls:     toolCalls,
-	}, nil
+		return chatResponseParts{
+			RequestID:     result.ID,
+			Content:       &visible,
+			ReasonContent: &reasoning,
+			ToolCalls:     choice.Message.ToolCalls,
+			Usage: &TokenUsage{
+				PromptTokens:     result.Usage.PromptTokens,
+				CompletionTokens: result.Usage.CompletionTokens,
+				TotalTokens:      result.Usage.TotalTokens,
+			},
+		}, nil
+	})
 }
 
 // ChatStreamlyWithSender sends messages and streams the response via
@@ -352,6 +367,14 @@ func (n *NovitaModel) ChatStreamlyWithSender(ctx context.Context, modelName stri
 	extractor := &novitaThinkExtractor{}
 	sawTerminal := false
 	done, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error {
+		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
@@ -493,9 +516,14 @@ type novitaEmbeddingData struct {
 }
 
 type novitaEmbeddingResponse struct {
+	ID     string                `json:"id"`
 	Data   []novitaEmbeddingData `json:"data"`
 	Model  string                `json:"model"`
 	Object string                `json:"object"`
+	Usage  struct {
+		PromptTokens int `json:"prompt_tokens"`
+		TotalTokens  int `json:"total_tokens"`
+	} `json:"usage"`
 }
 
 // Embed turns a list of texts into embedding vectors using the Novita
@@ -585,6 +613,10 @@ func (n *NovitaModel) Embed(ctx context.Context, modelName *string, texts []stri
 			return nil, fmt.Errorf("novita: missing embedding for input index %d", i)
 		}
 	}
+	recordResponseUsage(modelUsage, parsed.ID, &TokenUsage{
+		PromptTokens: parsed.Usage.PromptTokens,
+		TotalTokens:  parsed.Usage.TotalTokens,
+	}, "embedding")
 
 	return embeddings, nil
 }
@@ -598,7 +630,13 @@ type novitaRerankResult struct {
 }
 
 type novitaRerankResponse struct {
+	ID      string               `json:"id"`
 	Results []novitaRerankResult `json:"results"`
+	Usage   struct {
+		CompletionTokens int `json:"completion_tokens"`
+		PromptTokens     int `json:"prompt_tokens"`
+		TotalTokens      int `json:"total_tokens"`
+	} `json:"usage"`
 }
 
 // Rerank scores documents against the query using the Novita
@@ -690,6 +728,11 @@ func (n *NovitaModel) Rerank(ctx context.Context, modelName *string, query strin
 		})
 		seen[item.Index] = true
 	}
+	recordResponseUsage(modelUsage, parsed.ID, &TokenUsage{
+		PromptTokens:     parsed.Usage.PromptTokens,
+		CompletionTokens: parsed.Usage.CompletionTokens,
+		TotalTokens:      parsed.Usage.TotalTokens,
+	}, "rerank")
 
 	return &rerankResponse, nil
 }

internal/entity/models/orcarouter.go

@@ -48,6 +48,28 @@ func (o *OrcaRouterModel) Name() string {
 	return "orcarouter"
 }
 
+type OrcaRouterChatResponse struct {
+	ID      string `json:"id"`
+	Choices []struct {
+		FinishReason string `json:"finish_reason"`
+		Index        int    `json:"index"`
+		Message      struct {
+			Content   string           `json:"content"`
+			Role      string           `json:"role"`
+			ToolCalls []map[string]any `json:"tool_calls"`
+		} `json:"message"`
+		Logprobs interface{} `json:"logprobs"`
+	} `json:"choices"`
+	Created int    `json:"created"`
+	Model   string `json:"model"`
+	Object  string `json:"object"`
+	Usage   struct {
+		CompletionTokens int `json:"completion_tokens"`
+		PromptTokens     int `json:"prompt_tokens"`
+		TotalTokens      int `json:"total_tokens"`
+	} `json:"usage"`
+}
+
 func (o *OrcaRouterModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) {
 	if err := o.baseModel.APIConfigCheck(apiConfig); err != nil {
 		return nil, err
@@ -101,39 +123,34 @@ func (o *OrcaRouterModel) ChatWithMessages(ctx context.Context, modelName string
 	}
 
 	// Parse response
-	var result map[string]interface{}
-	if err := json.Unmarshal(body, &result); err != nil {
-		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
-	}
+	return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, chatConfig *ChatConfig) (chatResponseParts, error) {
+		var result OrcaRouterChatResponse
+		if err := json.Unmarshal(body, &result); err != nil {
+			return chatResponseParts{}, fmt.Errorf("failed to unmarshal response: %w", err)
+		}
+		if len(result.Choices) == 0 {
+			return chatResponseParts{}, fmt.Errorf("no choices in response")
+		}
 
-	choices, ok := result["choices"].([]interface{})
-	if !ok {
-		return nil, fmt.Errorf("no choices in response")
-	}
+		choice := &result.Choices[0]
+		content := choice.Message.Content
+		if content == "" && len(choice.Message.ToolCalls) == 0 {
+			return chatResponseParts{}, fmt.Errorf("no message in response")
+		}
 
-	firstChoice, ok := choices[0].(map[string]interface{})
-	if !ok {
-		return nil, fmt.Errorf("no choices in response")
-	}
-
-	messageMap, ok := firstChoice["message"].(map[string]interface{})
-	if !ok {
-		return nil, fmt.Errorf("no message in response")
-	}
-
-	content, ok := messageMap["content"].(string)
-	if !ok {
-		return nil, fmt.Errorf("no message in response")
-	}
-
-	// baichuan not support think
-	emptyReason := ""
-	chatResponse := &ChatResponse{
-		Answer:        &content,
-		ReasonContent: &emptyReason,
-	}
-
-	return chatResponse, nil
+		emptyReason := ""
+		return chatResponseParts{
+			RequestID:     result.ID,
+			Content:       &content,
+			ReasonContent: &emptyReason,
+			ToolCalls:     choice.Message.ToolCalls,
+			Usage: &TokenUsage{
+				PromptTokens:     result.Usage.PromptTokens,
+				CompletionTokens: result.Usage.CompletionTokens,
+				TotalTokens:      result.Usage.TotalTokens,
+			},
+		}, nil
+	})
 }
 
 func (o *OrcaRouterModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error {
@@ -192,6 +209,14 @@ func (o *OrcaRouterModel) ChatStreamlyWithSender(ctx context.Context, modelName
 	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(modelConfig, modelUsage, tokenUsage)
+		}
+
 		choices, ok := event["choices"].([]interface{})
 		if !ok || len(choices) == 0 {
 			return nil
@@ -374,7 +399,6 @@ func (o *OrcaRouterModel) ListModels(ctx context.Context, apiConfig *APIConfig)
 		return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
 	}
 
-	// Parse response
 	// Parse response
 	var modelList ModelList
 	if err = json.Unmarshal(body, &modelList); err != nil {