diff --git a/internal/entity/models/baichuan.go b/internal/entity/models/baichuan.go
index f51d79a6de..8f1e368d99 100644
--- a/internal/entity/models/baichuan.go
+++ b/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
}
diff --git a/internal/entity/models/cohere.go b/internal/entity/models/cohere.go
index 82d4363def..9cc5e47706 100644
--- a/internal/entity/models/cohere.go
+++ b/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
}
diff --git a/internal/entity/models/cohere_test.go b/internal/entity/models/cohere_test.go
new file mode 100644
index 0000000000..81ea4dfd26
--- /dev/null
+++ b/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)
+ }
+}
diff --git a/internal/entity/models/novita.go b/internal/entity/models/novita.go
index 1a7d128d38..fb75cc6274 100644
--- a/internal/entity/models/novita.go
+++ b/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 = ""
novitaThinkClose = ""
@@ -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 ... 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 ... 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
}
diff --git a/internal/entity/models/orcarouter.go b/internal/entity/models/orcarouter.go
index 36b1fbd37a..3270d4622d 100644
--- a/internal/entity/models/orcarouter.go
+++ b/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 {