From 944d72628439abde4c664e5e8ad2843e72174f40 Mon Sep 17 00:00:00 2001 From: Hz_ Date: Mon, 27 Jul 2026 17:16:50 +0800 Subject: [PATCH] fix(go-models): record provider token usage (#17423) ## Summary - Parse chat, embedding, and rerank usage from provider responses - Record usage with the correct model type even when no usage sink is provided - Cover SiliconFlow, Aliyun, Huawei Cloud, Qiniu, and VolcEngine response formats --- internal/entity/models/aliyun.go | 157 ++++++++++------ internal/entity/models/base_model.go | 22 ++- internal/entity/models/huaweicloud.go | 137 ++++++++++---- internal/entity/models/non_chat_usage_test.go | 124 +++++++++++++ .../models/openai_compatible_usage_test.go | 110 ++++++++++++ internal/entity/models/qiniu.go | 118 +++++++----- internal/entity/models/siliconflow.go | 168 +++++++++++------- internal/entity/models/siliconflow_test.go | 108 +++++++++++ internal/entity/models/volcengine.go | 147 +++++++++------ internal/entity/models/volcengine_test.go | 65 +++++++ 10 files changed, 901 insertions(+), 255 deletions(-) create mode 100644 internal/entity/models/non_chat_usage_test.go create mode 100644 internal/entity/models/openai_compatible_usage_test.go diff --git a/internal/entity/models/aliyun.go b/internal/entity/models/aliyun.go index 5e4d4935ed..fc2cdc208e 100644 --- a/internal/entity/models/aliyun.go +++ b/internal/entity/models/aliyun.go @@ -52,6 +52,37 @@ func (a *AliyunModel) Name() string { return "Tongyi-Qianwen" } +// AliyunChatResponse mirrors DashScope's OpenAI-compatible chat response. +type AliyunChatResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Created int64 `json:"created"` + Model string `json:"model"` + Choices []struct { + FinishReason string `json:"finish_reason"` + Index int `json:"index"` + Logprobs any `json:"logprobs"` + Message struct { + Content string `json:"content"` + ReasoningContent string `json:"reasoning_content"` + Role string `json:"role"` + ToolCalls []map[string]any `json:"tool_calls"` + } `json:"message"` + } `json:"choices"` + SystemFingerprint string `json:"system_fingerprint"` + Usage struct { + CompletionTokens int `json:"completion_tokens"` + PromptTokens int `json:"prompt_tokens"` + TotalTokens int `json:"total_tokens"` + CompletionTokensDetails struct { + ReasoningTokens int `json:"reasoning_tokens"` + } `json:"completion_tokens_details"` + PromptTokensDetails struct { + CachedTokens int `json:"cached_tokens"` + } `json:"prompt_tokens_details"` + } `json:"usage"` +} + func (a *AliyunModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) { if err := a.baseModel.APIConfigCheck(apiConfig); err != nil { return nil, err @@ -121,52 +152,47 @@ func (a *AliyunModel) ChatWithMessages(ctx context.Context, modelName string, me return nil, fmt.Errorf("API request failed with status %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 parse response: %w", err) - } - - choices, ok := result["choices"].([]interface{}) - if !ok || len(choices) == 0 { - return nil, fmt.Errorf("no choices in response") - } - - 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") - } - - answer, hasAnswer := messageMap["content"].(string) - toolCalls := extractToolCalls(messageMap) - if !hasAnswer && len(toolCalls) == 0 { - return nil, fmt.Errorf("response contains neither content nor tool calls") - } - - var reasonContent string - if chatModelConfig != nil && chatModelConfig.Thinking != nil && *chatModelConfig.Thinking { - reasonContent, ok = messageMap["reasoning_content"].(string) - if !ok { - return nil, fmt.Errorf("invalid content format") + return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, chatConfig *ChatConfig) (chatResponseParts, error) { + var result AliyunChatResponse + if err := json.Unmarshal(body, &result); err != nil { + return chatResponseParts{}, fmt.Errorf("failed to parse response: %w", err) } - // if first char of reasonContent is \n remove the '\n' - if reasonContent != "" && reasonContent[0] == '\n' { - reasonContent = reasonContent[1:] + if len(result.Choices) == 0 { + return chatResponseParts{}, fmt.Errorf("no choices in response") } - } - chatResponse := &ChatResponse{ - Answer: &answer, - ReasonContent: &reasonContent, - ToolCalls: toolCalls, - } + choice := &result.Choices[0] + content := choice.Message.Content + if content == "" && len(choice.Message.ToolCalls) == 0 { + return chatResponseParts{}, fmt.Errorf("response contains neither content nor tool calls") + } + reasonContent := "" + if chatConfig != nil && chatConfig.Thinking != nil && *chatConfig.Thinking { + reasonContent = choice.Message.ReasoningContent + if reasonContent == "" { + reasoning, answer := GetThinkingAndAnswer(chatConfig.ModelClass, &content) + if reasoning != nil { + reasonContent = *reasoning + content = *answer + } + } + if reasonContent != "" && reasonContent[0] == '\n' { + reasonContent = reasonContent[1:] + } + } - return chatResponse, nil + return chatResponseParts{ + RequestID: result.ID, + Content: &content, + ReasonContent: &reasonContent, + ToolCalls: choice.Message.ToolCalls, + Usage: &TokenUsage{ + PromptTokens: result.Usage.PromptTokens, + CompletionTokens: result.Usage.CompletionTokens, + TotalTokens: result.Usage.TotalTokens, + }, + }, nil + }) } // ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel) @@ -187,6 +213,7 @@ func (a *AliyunModel) ChatStreamlyWithSender(ctx context.Context, modelName stri // Build request body with streaming enabled reqBody := buildRequestBody(chatModelConfig, modelName, messages, true) + reqBody["stream_options"] = map[string]interface{}{"include_usage": true} if chatModelConfig != nil { if chatModelConfig.Stream != nil && !*chatModelConfig.Stream { @@ -240,6 +267,14 @@ func (a *AliyunModel) ChatStreamlyWithSender(ctx context.Context, modelName stri 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 @@ -332,11 +367,11 @@ func aliyunToolChoice(modelName string, messages []Message, configured *string) } type aliyunEmbeddingResponse struct { - Data []EmbeddingData `json:"data"` - Model string `json:"model"` - Object string `json:"object"` - Usage aliyunUsage `json:"usage"` - ID string `json:"id"` + Data []aliyunEmbeddingData `json:"data"` + Model string `json:"model"` + Object string `json:"object"` + Usage aliyunUsage `json:"usage"` + ID string `json:"id"` } type aliyunEmbeddingData struct { @@ -420,6 +455,10 @@ func (a *AliyunModel) Embed(ctx context.Context, modelName *string, texts []stri embeddingData.Index = dataElem.Index embeddings = append(embeddings, embeddingData) } + recordResponseUsage(modelUsage, parsed.ID, &TokenUsage{ + PromptTokens: parsed.Usage.PromptTokens, + TotalTokens: parsed.Usage.TotalTokens, + }, "embedding") return embeddings, nil } @@ -433,10 +472,14 @@ type aliyunRerankRequest struct { } type aliyunRerankResponse struct { + ID string `json:"id"` Results []struct { Index int `json:"index"` RelevanceScore float64 `json:"relevance_score"` } `json:"results"` + Usage struct { + TotalTokens int `json:"total_tokens"` + } `json:"usage"` } func (a *AliyunModel) Rerank(ctx context.Context, modelName *string, query string, documents []string, apiConfig *APIConfig, rerankConfig *RerankConfig, modelUsage *common.ModelUsage) (*RerankResponse, error) { @@ -459,8 +502,11 @@ func (a *AliyunModel) Rerank(ctx context.Context, modelName *string, query strin url := fmt.Sprintf("%s/%s", strings.TrimSuffix(baseURL, "/"), a.baseModel.URLSuffix.Rerank) - var topN = rerankConfig.TopN - if rerankConfig.TopN == 0 { + topN := len(documents) + if rerankConfig != nil && rerankConfig.TopN > 0 { + topN = rerankConfig.TopN + } + if topN == 0 { topN = len(documents) } @@ -503,12 +549,19 @@ func (a *AliyunModel) Rerank(ctx context.Context, modelName *string, query strin return nil, fmt.Errorf("Aliyun rerank API error: %s, body: %s", resp.Status, string(body)) } - var rerankResponse RerankResponse - if err = json.Unmarshal(body, &rerankResponse); err != nil { + var parsed aliyunRerankResponse + if err = json.Unmarshal(body, &parsed); err != nil { return nil, fmt.Errorf("failed to parse response: %w", err) } + rerankResponse := &RerankResponse{Data: make([]RerankResult, 0, len(parsed.Results))} + for _, item := range parsed.Results { + rerankResponse.Data = append(rerankResponse.Data, RerankResult{Index: item.Index, RelevanceScore: item.RelevanceScore}) + } + recordResponseUsage(modelUsage, parsed.ID, &TokenUsage{ + TotalTokens: parsed.Usage.TotalTokens, + }, "rerank") - return &rerankResponse, nil + return rerankResponse, nil } // TranscribeAudio transcribe audio diff --git a/internal/entity/models/base_model.go b/internal/entity/models/base_model.go index 6169e86b83..28c9cd27b5 100644 --- a/internal/entity/models/base_model.go +++ b/internal/entity/models/base_model.go @@ -63,9 +63,7 @@ func parseChatCompletionResponse(body []byte, chatConfig *ChatConfig, modelUsage return nil, err } - if err := collectChatModelUsage(modelUsage, parts.RequestID, parts.Usage); err != nil { - common.Error("Failed to collect model usage", err) - } + recordResponseUsage(modelUsage, parts.RequestID, parts.Usage, "chat") return &ChatResponse{ Answer: parts.Content, @@ -75,14 +73,19 @@ func parseChatCompletionResponse(body []byte, chatConfig *ChatConfig, modelUsage }, nil } -// collectChatModelUsage records one completed chat response when the caller -// supplied a usage sink. -func collectChatModelUsage(modelUsage *common.ModelUsage, requestID string, usage *TokenUsage) error { +// recordResponseUsage records the request ID and token usage returned by a +// completed model response. +func recordResponseUsage(modelUsage *common.ModelUsage, requestID string, usage *TokenUsage, modelType string) { if modelUsage == nil { - return nil + return + } + if modelUsage.Type == "" { + modelUsage.Type = modelType } modelUsage.RequestID = requestID - return collectModelUsage(modelUsage, usage) + if err := collectModelUsage(modelUsage, usage); err != nil { + common.Error("Failed to collect model usage", err) + } } // collectModelUsage records token usage and response time for one model call. @@ -125,6 +128,9 @@ func applyStreamUsage(chatConfig *ChatConfig, modelUsage *common.ModelUsage, usa if chatConfig != nil { chatConfig.UsageResult = usage } + if modelUsage == nil { + return + } if err := collectModelUsage(modelUsage, usage); err != nil { common.Error("Failed to collect model usage", err) } diff --git a/internal/entity/models/huaweicloud.go b/internal/entity/models/huaweicloud.go index 7cb9f8d7d3..2c23a1270b 100644 --- a/internal/entity/models/huaweicloud.go +++ b/internal/entity/models/huaweicloud.go @@ -50,6 +50,37 @@ func (h *HuaweiCloudModel) Name() string { return "huaweicloud" } +// HuaweiCloudChatResponse captures the OpenAI-compatible fields consumed by RAGFlow. +type HuaweiCloudChatResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Created int64 `json:"created"` + Model string `json:"model"` + Choices []struct { + FinishReason string `json:"finish_reason"` + Index int `json:"index"` + Logprobs any `json:"logprobs"` + Message struct { + Content string `json:"content"` + ReasoningContent string `json:"reasoning_content"` + Role string `json:"role"` + ToolCalls []map[string]any `json:"tool_calls"` + } `json:"message"` + } `json:"choices"` + SystemFingerprint string `json:"system_fingerprint"` + Usage struct { + CompletionTokens int `json:"completion_tokens"` + PromptTokens int `json:"prompt_tokens"` + TotalTokens int `json:"total_tokens"` + CompletionTokensDetails struct { + ReasoningTokens int `json:"reasoning_tokens"` + } `json:"completion_tokens_details"` + PromptTokensDetails struct { + CachedTokens int `json:"cached_tokens"` + } `json:"prompt_tokens_details"` + } `json:"usage"` +} + func huaweiCloudRegion(api *APIConfig) string { region := "default" if api != nil && api.Region != nil && *api.Region != "" { @@ -220,44 +251,44 @@ func (h *HuaweiCloudModel) ChatWithMessages(ctx context.Context, modelName strin return nil, fmt.Errorf("Huawei Cloud chat API error: status %d, body: %s", rep.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 HuaweiCloudChatResponse + 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") - } - - 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") - } - toolCalls := extractToolCalls(messageMap) - content, hasContent := messageMap["content"].(string) - if !hasContent && len(toolCalls) == 0 { - return nil, fmt.Errorf("invalid content format") - } - - reasonContent := "" - if r, ok := messageMap["reasoning_content"].(string); ok { - reasonContent = r + choice := &result.Choices[0] + content := choice.Message.Content + if content == "" && len(choice.Message.ToolCalls) == 0 { + return chatResponseParts{}, fmt.Errorf("invalid content format") + } + reasonContent := choice.Message.ReasoningContent + if chatConfig != nil && chatConfig.Thinking != nil && *chatConfig.Thinking && reasonContent == "" { + reasoning, answer := GetThinkingAndAnswer(chatConfig.ModelClass, &content) + if reasoning != nil { + reasonContent = *reasoning + content = *answer + } + } if reasonContent != "" && reasonContent[0] == '\n' { reasonContent = reasonContent[1:] } - } - return &ChatResponse{ - Answer: &content, - ReasonContent: &reasonContent, - ToolCalls: toolCalls, - }, nil + return chatResponseParts{ + RequestID: result.ID, + Content: &content, + ReasonContent: &reasonContent, + ToolCalls: choice.Message.ToolCalls, + Usage: &TokenUsage{ + PromptTokens: result.Usage.PromptTokens, + CompletionTokens: result.Usage.CompletionTokens, + TotalTokens: result.Usage.TotalTokens, + }, + }, nil + }) } func (h *HuaweiCloudModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { @@ -329,6 +360,14 @@ func (h *HuaweiCloudModel) ChatStreamlyWithSender(ctx context.Context, modelName return fmt.Errorf("huaweicloud: upstream stream error: %v", apiErr) } + 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 @@ -376,10 +415,27 @@ func (h *HuaweiCloudModel) ChatStreamlyWithSender(ctx context.Context, modelName } type huaweiCloudEmbeddingResponse 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"` + } `json:"usage"` +} + +type huaweiCloudRerankResponse struct { + ID string `json:"id"` + Results []struct { + Index int `json:"index"` + RelevanceScore float64 `json:"relevance_score"` + } `json:"results"` + Usage struct { + PromptTokens int `json:"prompt_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage"` } func (h *HuaweiCloudModel) Embed(ctx context.Context, modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig, modelUsage *common.ModelUsage) ([]EmbeddingData, error) { @@ -476,6 +532,10 @@ func (h *HuaweiCloudModel) Embed(ctx context.Context, modelName *string, texts [ return nil, fmt.Errorf("missing embedding index %d", i) } } + recordResponseUsage(modelUsage, parsed.ID, &TokenUsage{ + PromptTokens: parsed.Usage.PromptTokens, + TotalTokens: parsed.Usage.TotalTokens, + }, "embedding") return embeddings, nil } @@ -541,12 +601,7 @@ func (h *HuaweiCloudModel) Rerank(ctx context.Context, modelName *string, query return nil, fmt.Errorf("Huawei Cloud rerank API error: status %d, body: %s", resp.StatusCode, string(body)) } - var parsed struct { - Results []struct { - Index int `json:"index"` - RelevanceScore float64 `json:"relevance_score"` - } `json:"results"` - } + var parsed huaweiCloudRerankResponse if err = json.Unmarshal(body, &parsed); err != nil { return nil, fmt.Errorf("failed to decode response: %w", err) } @@ -569,6 +624,10 @@ func (h *HuaweiCloudModel) Rerank(ctx context.Context, modelName *string, query RelevanceScore: item.RelevanceScore, }) } + recordResponseUsage(modelUsage, parsed.ID, &TokenUsage{ + PromptTokens: parsed.Usage.PromptTokens, + TotalTokens: parsed.Usage.TotalTokens, + }, "rerank") return result, nil } diff --git a/internal/entity/models/non_chat_usage_test.go b/internal/entity/models/non_chat_usage_test.go new file mode 100644 index 0000000000..3eb88c6ff9 --- /dev/null +++ b/internal/entity/models/non_chat_usage_test.go @@ -0,0 +1,124 @@ +package models + +import ( + "net/http" + "net/http/httptest" + "ragflow/internal/common" + "testing" +) + +func TestProviderEmbeddingAndRerankUsage(t *testing.T) { + type providerCase struct { + name string + newModel func(string) ModelDriver + embeddingBody string + rerankBody string + embedInput int + embedTotal int + rerankInput int + rerankOutput int + rerankTotal int + supportsRerank bool + } + + cases := []providerCase{ + { + name: "siliconflow", + newModel: func(baseURL string) ModelDriver { + return NewSiliconflowModel(map[string]string{"default": baseURL}, URLSuffix{Embedding: "embedding", Rerank: "rerank"}) + }, + embeddingBody: `{"id":"embed-sf","data":[{"embedding":[0.1,0.2],"index":0}],"usage":{"prompt_tokens":7,"total_tokens":7}}`, + rerankBody: `{"id":"rerank-sf","results":[{"index":0,"relevance_score":0.9}],"meta":{"tokens":{"input_tokens":7,"output_tokens":2}}}`, + embedInput: 7, embedTotal: 7, rerankInput: 7, rerankOutput: 2, rerankTotal: 9, + supportsRerank: true, + }, + { + name: "aliyun", + newModel: func(baseURL string) ModelDriver { + return NewAliyunModel(map[string]string{"default": baseURL}, URLSuffix{Embedding: "embedding", Rerank: "rerank"}) + }, + embeddingBody: `{"id":"embed-aliyun","data":[{"embedding":[0.1,0.2],"index":0}],"usage":{"prompt_tokens":7,"total_tokens":7}}`, + rerankBody: `{"id":"rerank-aliyun","results":[{"index":0,"relevance_score":0.9}],"usage":{"total_tokens":9}}`, + embedInput: 7, embedTotal: 7, rerankTotal: 9, + supportsRerank: true, + }, + { + name: "huaweicloud", + newModel: func(baseURL string) ModelDriver { + return NewHuaweiCloudModel(map[string]string{"default": baseURL}, URLSuffix{Embedding: "embedding", Rerank: "rerank"}) + }, + embeddingBody: `{"id":"embed-huawei","data":[{"embedding":[0.1,0.2],"index":0}],"usage":{"prompt_tokens":7,"total_tokens":7}}`, + rerankBody: `{"id":"rerank-huawei","results":[{"index":0,"relevance_score":0.9}],"usage":{"prompt_tokens":9,"total_tokens":9}}`, + embedInput: 7, embedTotal: 7, rerankInput: 9, rerankTotal: 9, + supportsRerank: true, + }, + { + name: "volcengine", + newModel: func(baseURL string) ModelDriver { + return NewVolcEngine(map[string]string{"default": baseURL}, URLSuffix{Embedding: "embedding"}) + }, + embeddingBody: `{"id":"embed-volc","data":{"embedding":[0.1,0.2]},"usage":{"prompt_tokens":7,"total_tokens":7}}`, + embedInput: 7, embedTotal: 7, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch r.URL.Path { + case "/embedding": + _, _ = w.Write([]byte(tc.embeddingBody)) + case "/rerank": + _, _ = w.Write([]byte(tc.rerankBody)) + default: + http.NotFound(w, r) + } + })) + defer server.Close() + + apiKey := "test-key" + modelName := "test-model" + model := tc.newModel(server.URL) + + embeddingUsage := &common.ModelUsage{} + embeddings, err := model.Embed(t.Context(), &modelName, []string{"document"}, &APIConfig{ApiKey: &apiKey}, &EmbeddingConfig{}, embeddingUsage) + if err != nil { + t.Fatalf("Embed: %v", err) + } + if len(embeddings) != 1 || len(embeddings[0].Embedding) != 2 { + t.Fatalf("embeddings=%#v", embeddings) + } + assertModelUsage(t, embeddingUsage, tc.embedInput, 0, tc.embedTotal) + if embeddingUsage.Type != "embedding" { + t.Fatalf("embedding usage type=%q, want embedding", embeddingUsage.Type) + } + if !tc.supportsRerank { + return + } + + rerankUsage := &common.ModelUsage{} + reranked, err := model.Rerank(t.Context(), &modelName, "query", []string{"document"}, &APIConfig{ApiKey: &apiKey}, &RerankConfig{TopN: 1}, rerankUsage) + if err != nil { + t.Fatalf("Rerank: %v", err) + } + if len(reranked.Data) != 1 || reranked.Data[0].Index != 0 { + t.Fatalf("reranked=%#v", reranked) + } + assertModelUsage(t, rerankUsage, tc.rerankInput, tc.rerankOutput, tc.rerankTotal) + if rerankUsage.Type != "rerank" { + t.Fatalf("rerank usage type=%q, want rerank", rerankUsage.Type) + } + }) + } +} + +func assertModelUsage(t *testing.T, usage *common.ModelUsage, input, output, total int) { + t.Helper() + if usage.RequestID == "" { + t.Fatal("usage request ID is empty") + } + if usage.InputTokens != input || usage.OutputTokens != output || usage.TotalTokens != total { + t.Fatalf("usage=(%d,%d,%d), want (%d,%d,%d)", usage.InputTokens, usage.OutputTokens, usage.TotalTokens, input, output, total) + } +} diff --git a/internal/entity/models/openai_compatible_usage_test.go b/internal/entity/models/openai_compatible_usage_test.go new file mode 100644 index 0000000000..9b75555984 --- /dev/null +++ b/internal/entity/models/openai_compatible_usage_test.go @@ -0,0 +1,110 @@ +// +// 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" + "io" + "net/http" + "net/http/httptest" + "testing" +) + +func TestOpenAICompatibleProvidersExtractStreamingUsage(t *testing.T) { + providers := []struct { + name string + new func(string) ModelDriver + path string + }{ + { + name: "siliconflow", + new: func(baseURL string) ModelDriver { + return NewSiliconflowModel(map[string]string{"default": baseURL}, URLSuffix{Chat: "chat/completions"}) + }, + path: "/chat/completions", + }, + { + name: "aliyun", + new: func(baseURL string) ModelDriver { + return NewAliyunModel(map[string]string{"default": baseURL}, URLSuffix{Chat: "chat/completions"}) + }, + path: "/chat/completions", + }, + { + name: "qiniu", + new: func(baseURL string) ModelDriver { + return NewQiniuModel(map[string]string{"default": baseURL}, URLSuffix{Chat: "chat/completions"}) + }, + path: "/chat/completions", + }, + { + name: "volcengine", + new: func(baseURL string) ModelDriver { + return NewVolcEngine(map[string]string{"default": baseURL}, URLSuffix{Chat: "chat/completions"}) + }, + path: "/chat/completions", + }, + } + + for _, provider := range providers { + provider := provider + t.Run(provider.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + t.Errorf("method=%s, want POST", r.Method) + } + if got := r.Header.Get("Authorization"); got != "Bearer test-key" { + t.Errorf("Authorization=%q, want Bearer test-key", got) + } + if r.URL.Path != provider.path { + t.Errorf("path=%q, want %q", r.URL.Path, provider.path) + } + var requestBody map[string]any + if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil { + t.Fatalf("decode request: %v", err) + } + streamOptions, ok := requestBody["stream_options"].(map[string]any) + if !ok || streamOptions["include_usage"] != true { + t.Errorf("stream_options=%#v, want include_usage=true", requestBody["stream_options"]) + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"answer\"},\"finish_reason\":null}]}\n\n") + _, _ = io.WriteString(w, "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":5,\"total_tokens\":8}}\n\n") + _, _ = io.WriteString(w, "data: [DONE]\n\n") + })) + defer server.Close() + + apiKey := "test-key" + config := &ChatConfig{} + err := provider.new(server.URL).ChatStreamlyWithSender( + t.Context(), + "model", + []Message{{Role: "user", Content: "hello"}}, + &APIConfig{ApiKey: &apiKey}, + config, + nil, + func(*string, *string) error { return nil }, + ) + if err != nil { + t.Fatalf("ChatStreamlyWithSender: %v", err) + } + if config.UsageResult == nil || config.UsageResult.PromptTokens != 3 || config.UsageResult.CompletionTokens != 5 || config.UsageResult.TotalTokens != 8 { + t.Fatalf("UsageResult=%#v, want prompt=3 completion=5 total=8", config.UsageResult) + } + }) + } +} diff --git a/internal/entity/models/qiniu.go b/internal/entity/models/qiniu.go index 2f5aa28793..11e70e299e 100644 --- a/internal/entity/models/qiniu.go +++ b/internal/entity/models/qiniu.go @@ -50,6 +50,37 @@ func (q *QiniuModel) Name() string { return "qiniu" } +// QiniuChatResponse captures the OpenAI-compatible fields consumed by RAGFlow. +type QiniuChatResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Created int64 `json:"created"` + Model string `json:"model"` + Choices []struct { + FinishReason string `json:"finish_reason"` + Index int `json:"index"` + Logprobs any `json:"logprobs"` + Message struct { + Content string `json:"content"` + ReasoningContent string `json:"reasoning_content"` + Role string `json:"role"` + ToolCalls []map[string]any `json:"tool_calls"` + } `json:"message"` + } `json:"choices"` + SystemFingerprint string `json:"system_fingerprint"` + Usage struct { + CompletionTokens int `json:"completion_tokens"` + PromptTokens int `json:"prompt_tokens"` + TotalTokens int `json:"total_tokens"` + CompletionTokensDetails struct { + ReasoningTokens int `json:"reasoning_tokens"` + } `json:"completion_tokens_details"` + PromptTokensDetails struct { + CachedTokens int `json:"cached_tokens"` + } `json:"prompt_tokens_details"` + } `json:"usage"` +} + var qiniuQwenThinkingModels = map[string]struct{}{ "qwen3-next-80b-a3b-thinking": {}, "qwen3-235b-a22b-thinking-2507": {}, @@ -191,51 +222,47 @@ func (q *QiniuModel) ChatWithMessages(ctx context.Context, modelName string, mes 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) - } - - choices, ok := result["choices"].([]interface{}) - if !ok || len(choices) == 0 { - return nil, fmt.Errorf("no choices in response") - } - - 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") - } - - toolCalls := extractToolCalls(messageMap) - content, hasContent := messageMap["content"].(string) - if !hasContent && len(toolCalls) == 0 { - return nil, fmt.Errorf("invalid content format") - } - - var reasonContent string - if chatModelConfig != nil && chatModelConfig.Thinking != nil && *chatModelConfig.Thinking { - reasonContent, ok = messageMap["reasoning_content"].(string) - if !ok && len(toolCalls) == 0 { - return nil, fmt.Errorf("invalid reasoning content format") + return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, chatConfig *ChatConfig) (chatResponseParts, error) { + var result QiniuChatResponse + 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") } - if reasonContent != "" && reasonContent[0] == '\n' { - reasonContent = reasonContent[1:] + choice := &result.Choices[0] + content := choice.Message.Content + if content == "" && len(choice.Message.ToolCalls) == 0 { + return chatResponseParts{}, fmt.Errorf("invalid content format") + } + reasonContent := "" + if chatConfig != nil && chatConfig.Thinking != nil && *chatConfig.Thinking { + reasonContent = choice.Message.ReasoningContent + if reasonContent == "" { + reasoning, answer := GetThinkingAndAnswer(chatConfig.ModelClass, &content) + if reasoning != nil { + reasonContent = *reasoning + content = *answer + } + } + if reasonContent != "" && reasonContent[0] == '\n' { + reasonContent = reasonContent[1:] + } } - } - chatResponse := &ChatResponse{ - Answer: &content, - ReasonContent: &reasonContent, - ToolCalls: toolCalls, - } - - return chatResponse, nil + return chatResponseParts{ + RequestID: result.ID, + Content: &content, + ReasonContent: &reasonContent, + ToolCalls: choice.Message.ToolCalls, + Usage: &TokenUsage{ + PromptTokens: result.Usage.PromptTokens, + CompletionTokens: result.Usage.CompletionTokens, + TotalTokens: result.Usage.TotalTokens, + }, + }, nil + }) } func (q *QiniuModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { @@ -254,6 +281,7 @@ func (q *QiniuModel) ChatStreamlyWithSender(ctx context.Context, modelName strin url := fmt.Sprintf("%s/%s", baseURL, q.baseModel.URLSuffix.Chat) reqBody := buildRequestBody(chatModelConfig, modelName, messages, true) + reqBody["stream_options"] = map[string]interface{}{"include_usage": true} if chatModelConfig != nil { applyQiniuThinkingConfig(reqBody, modelName, chatModelConfig) @@ -291,6 +319,14 @@ func (q *QiniuModel) ChatStreamlyWithSender(ctx context.Context, modelName strin 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 diff --git a/internal/entity/models/siliconflow.go b/internal/entity/models/siliconflow.go index 3d0f1d4501..ff85cabc83 100644 --- a/internal/entity/models/siliconflow.go +++ b/internal/entity/models/siliconflow.go @@ -55,6 +55,38 @@ func (s *SiliconflowModel) Name() string { return "SILICONFLOW" } +// SiliconflowChatResponse mirrors the response returned by SiliconFlow's +// POST /chat/completions endpoint. SiliconFlow uses the OpenAI-compatible +// schema and includes reasoning and cache token details when available. +type SiliconflowChatResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Created int64 `json:"created"` + Model string `json:"model"` + 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"` + } `json:"choices"` + SystemFingerprint string `json:"system_fingerprint"` + Usage struct { + CompletionTokens int `json:"completion_tokens"` + PromptTokens int `json:"prompt_tokens"` + TotalTokens int `json:"total_tokens"` + CompletionTokensDetails struct { + ReasoningTokens int `json:"reasoning_tokens"` + } `json:"completion_tokens_details"` + PromptTokensDetails struct { + CachedTokens int `json:"cached_tokens"` + } `json:"prompt_tokens_details"` + } `json:"usage"` +} + // ChatWithMessages sends multiple messages with roles and returns response func (s *SiliconflowModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) { if err := s.baseModel.APIConfigCheck(apiConfig); err != nil { @@ -114,58 +146,48 @@ func (s *SiliconflowModel) ChatWithMessages(ctx context.Context, modelName strin return nil, fmt.Errorf("API request failed with status %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 parse response: %w", err) - } + return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, chatConfig *ChatConfig) (chatResponseParts, error) { + var result SiliconflowChatResponse + 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] + content := choice.Message.Content + if 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") - } - - toolCalls := extractToolCalls(messageMap) - content, hasContent := messageMap["content"].(string) - if !hasContent && len(toolCalls) == 0 { - return nil, fmt.Errorf("invalid content format") - } - - var reasonContent string - if chatModelConfig != nil && chatModelConfig.Thinking != nil && *chatModelConfig.Thinking { - reasonContent, ok = messageMap["reasoning_content"].(string) - if !ok { - // If reasoning_content not in response, try parsing from content tags - reasoning, answer := GetThinkingAndAnswer(chatModelConfig.ModelClass, &content) - if reasoning != nil { - reasonContent = *reasoning - content = *answer + reasonContent := "" + if chatConfig != nil && chatConfig.Thinking != nil && *chatConfig.Thinking { + reasonContent = choice.Message.ReasoningContent + if reasonContent == "" { + reasoning, answer := GetThinkingAndAnswer(chatConfig.ModelClass, &content) + if reasoning != nil { + reasonContent = *reasoning + content = *answer + } } - } else { - // if first char of reasonContent is \n remove the '\n' if reasonContent != "" && reasonContent[0] == '\n' { reasonContent = reasonContent[1:] } } - } - chatResponse := &ChatResponse{ - Answer: &content, - ReasonContent: &reasonContent, - ToolCalls: toolCalls, - } - - return chatResponse, nil + return chatResponseParts{ + RequestID: result.ID, + Content: &content, + ReasonContent: &reasonContent, + ToolCalls: choice.Message.ToolCalls, + Usage: &TokenUsage{ + PromptTokens: result.Usage.PromptTokens, + CompletionTokens: result.Usage.CompletionTokens, + TotalTokens: result.Usage.TotalTokens, + }, + }, nil + }) } // ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel) @@ -186,8 +208,12 @@ func (s *SiliconflowModel) ChatStreamlyWithSender(ctx context.Context, modelName // Build request body with streaming enabled reqBody := buildRequestBody(chatModelConfig, modelName, messages, true) + reqBody["stream_options"] = map[string]interface{}{"include_usage": true} if chatModelConfig != nil { chatModelConfig.ToolCallsResult = nil + if chatModelConfig.Thinking != nil { + reqBody["enable_thinking"] = *chatModelConfig.Thinking + } } if strings.Contains(strings.ToLower(modelName), "qwen3") && (chatModelConfig == nil || chatModelConfig.Thinking == nil) { reqBody["enable_thinking"] = false @@ -226,6 +252,14 @@ func (s *SiliconflowModel) 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(chatModelConfig, modelUsage, tokenUsage) + } + choices, ok := event["choices"].([]interface{}) if !ok || len(choices) == 0 { return nil @@ -281,22 +315,18 @@ func (s *SiliconflowModel) ChatStreamlyWithSender(ctx context.Context, modelName } type siliconflowEmbeddingResponse struct { - Object string `json:"object"` - Model string `json:"model"` - Data []siliconflowEmbeddingData `json:"data"` - Usage siliconflowUsage `json:"usage"` -} - -type siliconflowEmbeddingData struct { - Object string `json:"object"` - Embedding []float64 `json:"embedding"` - Index int `json:"index"` -} - -type siliconflowUsage struct { - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - TotalTokens int `json:"total_tokens"` + ID string `json:"id"` + Object string `json:"object"` + Model string `json:"model"` + Data []struct { + Object string `json:"object"` + Embedding []float64 `json:"embedding"` + Index int `json:"index"` + } `json:"data"` + Usage struct { + PromptTokens int `json:"prompt_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage"` } // siliconflowMaxBatchSize is the per-request input limit documented at @@ -381,6 +411,10 @@ func (s *SiliconflowModel) Embed(ctx context.Context, modelName *string, texts [ embeddingData.Index = dataElem.Index embeddings = append(embeddings, embeddingData) } + recordResponseUsage(modelUsage, parsed.ID, &TokenUsage{ + PromptTokens: parsed.Usage.PromptTokens, + TotalTokens: parsed.Usage.TotalTokens, + }, "embedding") return embeddings, nil } @@ -572,12 +606,15 @@ func (s *SiliconflowModel) Rerank(ctx context.Context, modelName *string, query if len(documents) == 0 { return &RerankResponse{}, nil } + if modelName == nil || *modelName == "" { + return nil, fmt.Errorf("model name is required") + } apiKey := *apiConfig.ApiKey - var topN = rerankConfig.TopN - if rerankConfig.TopN == 0 { - topN = len(documents) + topN := len(documents) + if rerankConfig != nil && rerankConfig.TopN > 0 { + topN = rerankConfig.TopN } reqBody := SiliconflowRerankRequest{ @@ -639,6 +676,11 @@ func (s *SiliconflowModel) Rerank(ctx context.Context, modelName *string, query RelevanceScore: result.RelevanceScore, }) } + recordResponseUsage(modelUsage, siliconflowRerankResp.ID, &TokenUsage{ + PromptTokens: siliconflowRerankResp.Meta.Tokens.InputTokens, + CompletionTokens: siliconflowRerankResp.Meta.Tokens.OutputTokens, + TotalTokens: siliconflowRerankResp.Meta.Tokens.InputTokens + siliconflowRerankResp.Meta.Tokens.OutputTokens, + }, "rerank") return &rerankResponse, nil } diff --git a/internal/entity/models/siliconflow_test.go b/internal/entity/models/siliconflow_test.go index d40863f1b6..3623e61ca3 100644 --- a/internal/entity/models/siliconflow_test.go +++ b/internal/entity/models/siliconflow_test.go @@ -17,6 +17,8 @@ package models import ( + "encoding/json" + "io" "net/http" "net/http/httptest" "testing" @@ -47,3 +49,109 @@ func TestSiliconflowChatRejectsMissingContent(t *testing.T) { t.Fatal("expected missing content error") } } + +func TestSiliconflowChatWithMessagesExtractsResponseAndUsage(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _ = json.NewEncoder(w).Encode(map[string]any{ + "id": "chatcmpl-siliconflow", + "object": "chat.completion", + "created": 1, + "model": "Qwen/Qwen3-8B", + "choices": []map[string]any{{ + "index": 0, + "message": map[string]any{ + "role": "assistant", + "content": "answer", + "reasoning_content": "\nreasoning", + }, + "finish_reason": "stop", + }}, + "usage": map[string]any{ + "prompt_tokens": 3, + "completion_tokens": 5, + "total_tokens": 8, + "completion_tokens_details": map[string]any{ + "reasoning_tokens": 2, + }, + "prompt_tokens_details": map[string]any{ + "cached_tokens": 1, + }, + }, + }) + })) + defer server.Close() + + apiKey := "test-key" + thinking := true + response, err := NewSiliconflowModel( + map[string]string{"default": server.URL}, + URLSuffix{Chat: "chat/completions"}, + ).ChatWithMessages( + t.Context(), + "Qwen/Qwen3-8B", + []Message{{Role: "user", Content: "hi"}}, + &APIConfig{ApiKey: &apiKey}, + &ChatConfig{Thinking: &thinking}, + nil, + ) + if err != nil { + t.Fatalf("ChatWithMessages: %v", err) + } + if response.Answer == nil || *response.Answer != "answer" { + t.Fatalf("Answer=%#v, want answer", response.Answer) + } + if response.ReasonContent == nil || *response.ReasonContent != "reasoning" { + t.Fatalf("ReasonContent=%#v, want reasoning", response.ReasonContent) + } + if response.Usage == nil || response.Usage.PromptTokens != 3 || response.Usage.CompletionTokens != 5 || response.Usage.TotalTokens != 8 { + t.Fatalf("Usage=%#v, want prompt=3 completion=5 total=8", response.Usage) + } +} + +func TestSiliconflowChatStreamlyWithSenderCollectsUsage(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var requestBody map[string]any + if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil { + t.Errorf("decode request: %v", err) + return + } + if requestBody["enable_thinking"] != true { + t.Errorf("enable_thinking=%#v, want true", requestBody["enable_thinking"]) + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"answer\"}}]}\n\n") + _, _ = io.WriteString(w, "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":5,\"total_tokens\":8}}\n\n") + _, _ = io.WriteString(w, "data: [DONE]\n\n") + })) + defer server.Close() + + apiKey := "test-key" + thinking := true + config := &ChatConfig{Thinking: &thinking} + err := NewSiliconflowModel( + map[string]string{"default": server.URL}, + URLSuffix{Chat: "chat/completions"}, + ).ChatStreamlyWithSender( + t.Context(), + "Qwen/Qwen3-8B", + []Message{{Role: "user", Content: "hi"}}, + &APIConfig{ApiKey: &apiKey}, + config, + nil, + func(content *string, reasoning *string) error { + if reasoning != nil { + t.Errorf("reasoning=%q", *reasoning) + } + if content != nil && *content != "answer" && *content != "[DONE]" { + t.Errorf("content=%q", *content) + } + return nil + }, + ) + if err != nil { + t.Fatalf("ChatStreamlyWithSender: %v", err) + } + if config.UsageResult == nil || config.UsageResult.PromptTokens != 3 || config.UsageResult.CompletionTokens != 5 || config.UsageResult.TotalTokens != 8 { + t.Fatalf("UsageResult=%#v, want prompt=3 completion=5 total=8", config.UsageResult) + } +} diff --git a/internal/entity/models/volcengine.go b/internal/entity/models/volcengine.go index 3a5cccc6cd..013c1eed1b 100644 --- a/internal/entity/models/volcengine.go +++ b/internal/entity/models/volcengine.go @@ -52,6 +52,37 @@ func (v *VolcEngine) Name() string { return "volcengine" } +// VolcEngineChatResponse mirrors Ark's OpenAI-compatible chat response. +type VolcEngineChatResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Created int64 `json:"created"` + Model string `json:"model"` + Choices []struct { + FinishReason string `json:"finish_reason"` + Index int `json:"index"` + Logprobs any `json:"logprobs"` + Message struct { + Content string `json:"content"` + ReasoningContent string `json:"reasoning_content"` + Role string `json:"role"` + ToolCalls []map[string]any `json:"tool_calls"` + } `json:"message"` + } `json:"choices"` + SystemFingerprint string `json:"system_fingerprint"` + Usage struct { + CompletionTokens int `json:"completion_tokens"` + PromptTokens int `json:"prompt_tokens"` + TotalTokens int `json:"total_tokens"` + CompletionTokensDetails struct { + ReasoningTokens int `json:"reasoning_tokens"` + } `json:"completion_tokens_details"` + PromptTokensDetails struct { + CachedTokens int `json:"cached_tokens"` + } `json:"prompt_tokens_details"` + } `json:"usage"` +} + // getAPIKey extracts the actual API key from VolcEngine's stored format. // VolcEngine stores the api_key as a JSON string like: // @@ -165,61 +196,44 @@ func (v *VolcEngine) ChatWithMessages(ctx context.Context, modelName string, mes return nil, fmt.Errorf("API request failed with status %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) - } + return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, chatConfig *ChatConfig) (chatResponseParts, error) { + var result VolcEngineChatResponse + 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 || len(choices) == 0 { - return nil, fmt.Errorf("no choices in response") - } - - 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") - } - - content, _ := messageMap["content"].(string) - - var reasonContent string - if chatModelConfig != nil && chatModelConfig.Thinking != nil && *chatModelConfig.Thinking { - reasonContent, ok = messageMap["reasoning_content"].(string) - if !ok { - return nil, fmt.Errorf("invalid reasonContent format") + choice := &result.Choices[0] + content := choice.Message.Content + if content == "" && len(choice.Message.ToolCalls) == 0 { + return chatResponseParts{}, fmt.Errorf("invalid content format") + } + reasonContent := choice.Message.ReasoningContent + if chatConfig != nil && chatConfig.Thinking != nil && *chatConfig.Thinking && reasonContent == "" { + reasoning, answer := GetThinkingAndAnswer(chatConfig.ModelClass, &content) + if reasoning != nil { + reasonContent = *reasoning + content = *answer + } } if reasonContent != "" && reasonContent[0] == '\n' { reasonContent = reasonContent[1:] } - } - var toolCalls []map[string]any - if tcs, ok := messageMap["tool_calls"].([]any); ok { - for _, tc := range tcs { - if toolCall, ok := tc.(map[string]any); ok { - toolCalls = append(toolCalls, toolCall) - } - } - } - - chatResponse := &ChatResponse{ - Answer: &content, - ReasonContent: &reasonContent, - ToolCalls: toolCalls, - } - if pt, ct, tt := extractUsageFromMap(result); tt > 0 { - chatResponse.Usage = &TokenUsage{ - PromptTokens: pt, CompletionTokens: ct, TotalTokens: tt, - } - } - - return chatResponse, nil + return chatResponseParts{ + RequestID: result.ID, + Content: &content, + ReasonContent: &reasonContent, + ToolCalls: choice.Message.ToolCalls, + Usage: &TokenUsage{ + PromptTokens: result.Usage.PromptTokens, + CompletionTokens: result.Usage.CompletionTokens, + TotalTokens: result.Usage.TotalTokens, + }, + }, nil + }) } // ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel) @@ -240,6 +254,7 @@ func (v *VolcEngine) ChatStreamlyWithSender(ctx context.Context, modelName strin // Build request body with streaming enabled reqBody := buildRequestBody(modelConfig, modelName, messages, true) + reqBody["stream_options"] = map[string]interface{}{"include_usage": true} if modelConfig != nil { // TODO VolcEngine has `auto` mode @@ -275,6 +290,10 @@ func (v *VolcEngine) ChatStreamlyWithSender(ctx context.Context, modelName strin thinkingFlag = "enabled" reqBody["reasoning_effort"] = "xhigh" break + case "max": + thinkingFlag = "enabled" + reqBody["reasoning_effort"] = "max" + break default: return fmt.Errorf("invalid effort level") } @@ -318,9 +337,18 @@ func (v *VolcEngine) ChatStreamlyWithSender(ctx context.Context, modelName strin } accumulatedToolCalls := make(map[int]map[string]any) - if _, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error { + sawTerminal := false + 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 @@ -330,6 +358,9 @@ func (v *VolcEngine) ChatStreamlyWithSender(ctx context.Context, modelName strin if !ok { return nil } + if finishReason, ok := firstChoice["finish_reason"].(string); ok && finishReason != "" { + sawTerminal = true + } delta, ok := firstChoice["delta"].(map[string]interface{}) if !ok { @@ -401,9 +432,13 @@ func (v *VolcEngine) ChatStreamlyWithSender(ctx context.Context, modelName strin } return nil - }); err != nil { + }) + if err != nil { return fmt.Errorf("failed to scan response body: %w", err) } + if !done && !sawTerminal { + return fmt.Errorf("volcengine: stream ended before [DONE] or finish_reason") + } if len(accumulatedToolCalls) > 0 && modelConfig != nil { indices := make([]int, 0, len(accumulatedToolCalls)) @@ -461,6 +496,9 @@ func (v *VolcEngine) Embed(ctx context.Context, modelName *string, texts []strin if len(texts) == 0 { return []EmbeddingData{}, nil } + if modelName == nil || *modelName == "" { + return nil, fmt.Errorf("model name is required") + } resolvedBaseURL, err := v.baseModel.GetBaseURL(apiConfig) if err != nil { @@ -469,6 +507,8 @@ func (v *VolcEngine) Embed(ctx context.Context, modelName *string, texts []strin url := fmt.Sprintf("%s/%s", resolvedBaseURL, v.baseModel.URLSuffix.Embedding) var embeddings []EmbeddingData + var totalUsage TokenUsage + var requestID string for i, text := range texts { @@ -493,7 +533,6 @@ func (v *VolcEngine) Embed(ctx context.Context, modelName *string, texts []strin err, ) } - // Run each per-text request in its own scope so the context's // deadline is cancelled at the end of every iteration instead of // piling up deferred cancels until the whole batch finishes. @@ -534,12 +573,16 @@ func (v *VolcEngine) Embed(ctx context.Context, modelName *string, texts []strin if err != nil { return nil, err } + requestID = parsed.ID + totalUsage.PromptTokens += parsed.Usage.PromptTokens + totalUsage.TotalTokens += parsed.Usage.TotalTokens var embeddingData EmbeddingData embeddingData.Index = i embeddingData.Embedding = parsed.Data.Embedding embeddings = append(embeddings, embeddingData) } + recordResponseUsage(modelUsage, requestID, &totalUsage, "embedding") return embeddings, nil } diff --git a/internal/entity/models/volcengine_test.go b/internal/entity/models/volcengine_test.go index dbde616be5..d2e88a75f0 100644 --- a/internal/entity/models/volcengine_test.go +++ b/internal/entity/models/volcengine_test.go @@ -2,6 +2,7 @@ package models import ( "encoding/json" + "io" "net/http" "net/http/httptest" "os" @@ -111,3 +112,67 @@ func TestVolcEngineListModelsRequiresModelsSuffix(t *testing.T) { t.Fatalf("expected missing models suffix error, got %v", err) } } + +func TestVolcEngineChatStreamSupportsMaxEffortAndUsage(t *testing.T) { + ctx := t.Context() + srv := newVolcEngineServer(t, func(t *testing.T, r *http.Request, w http.ResponseWriter) { + var body map[string]interface{} + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Fatalf("decode request: %v", err) + } + if body["reasoning_effort"] != "max" { + t.Errorf("reasoning_effort=%v, want max", body["reasoning_effort"]) + } + streamOptions, ok := body["stream_options"].(map[string]interface{}) + if !ok || streamOptions["include_usage"] != true { + t.Errorf("stream_options=%#v, want include_usage=true", body["stream_options"]) + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"answer\"},\"finish_reason\":\"stop\"}]}\n\n") + _, _ = io.WriteString(w, "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":5,\"total_tokens\":8}}\n\n") + _, _ = io.WriteString(w, "data: [DONE]\n\n") + }) + defer srv.Close() + + apiKey := "test-key" + thinking := true + effort := "max" + config := &ChatConfig{Thinking: &thinking, Effort: &effort} + if err := newVolcEngineForTest(srv.URL).ChatStreamlyWithSender( + ctx, + "doubao-seed-2-0-pro-260215", + []Message{{Role: "user", Content: "hello"}}, + &APIConfig{ApiKey: &apiKey}, + config, + nil, + func(*string, *string) error { return nil }, + ); err != nil { + t.Fatalf("ChatStreamlyWithSender: %v", err) + } + if config.UsageResult == nil || config.UsageResult.TotalTokens != 8 { + t.Fatalf("UsageResult=%#v, want total tokens 8", config.UsageResult) + } +} + +func TestVolcEngineChatStreamRejectsTruncatedResponse(t *testing.T) { + ctx := t.Context() + srv := newVolcEngineServer(t, func(t *testing.T, _ *http.Request, w http.ResponseWriter) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"partial\"}}]}\n\n") + }) + defer srv.Close() + + apiKey := "test-key" + err := newVolcEngineForTest(srv.URL).ChatStreamlyWithSender( + ctx, + "doubao-seed-2-0-pro-260215", + []Message{{Role: "user", Content: "hello"}}, + &APIConfig{ApiKey: &apiKey}, + &ChatConfig{}, + nil, + func(*string, *string) error { return nil }, + ) + if err == nil || !strings.Contains(err.Error(), "stream ended before [DONE] or finish_reason") { + t.Fatalf("error=%v, want truncated stream error", err) + } +}