From e5c038a411638e2b72b341821c0468f497690895 Mon Sep 17 00:00:00 2001 From: Haruko386 Date: Tue, 28 Jul 2026 19:14:05 +0800 Subject: [PATCH] Go: add token usage for orcarouter, baichuan, cohere and novita (#17455) As title #17284 --- internal/entity/models/baichuan.go | 99 ++++++++++------ internal/entity/models/cohere.go | 161 ++++++++++++++++++++------ internal/entity/models/cohere_test.go | 127 ++++++++++++++++++++ internal/entity/models/novita.go | 133 ++++++++++++++------- internal/entity/models/orcarouter.go | 88 +++++++++----- 5 files changed, 464 insertions(+), 144 deletions(-) create mode 100644 internal/entity/models/cohere_test.go 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 {