From 339bba879373f5b1f1c0db18cf08582502a18c82 Mon Sep 17 00:00:00 2001 From: Hz_ Date: Fri, 24 Jul 2026 10:51:37 +0800 Subject: [PATCH] refactor(go-models): share chat response handling (#17289) ## Summary - Add provider-specific DeepSeek chat response handling aligned with its documented schema. - Share Zhipu and DeepSeek usage accounting and chat response helpers. ## Testing - `go test -c -o /tmp/ragflow-models.test ./internal/entity/models` #17284 --- internal/entity/models/base_model.go | 94 +++++++++++++++ internal/entity/models/deepseek.go | 119 ++++++++++--------- internal/entity/models/zhipu-ai.go | 163 ++++++--------------------- 3 files changed, 193 insertions(+), 183 deletions(-) diff --git a/internal/entity/models/base_model.go b/internal/entity/models/base_model.go index ae936d15b7..6169e86b83 100644 --- a/internal/entity/models/base_model.go +++ b/internal/entity/models/base_model.go @@ -24,9 +24,13 @@ import ( "fmt" "io" "net/http" + "ragflow/internal/common" + "ragflow/internal/engine/clickhouse" "sort" "strings" "time" + + "github.com/mitchellh/mapstructure" ) type BaseModel struct { @@ -36,6 +40,96 @@ type BaseModel struct { AllowEmptyAPIKey bool } +// chatResponseParts is the provider-normalized result of a non-streaming chat +// completion. Provider response structs remain provider-specific; only the +// common ChatResponse and model-usage handling is shared. +type chatResponseParts struct { + RequestID string + Content *string + ReasonContent *string + ToolCalls []map[string]interface{} + Usage *TokenUsage +} + +// chatResponseExtractor parses one provider-specific response and maps it +// into the common result used by RAGFlow's chat model abstraction. +type chatResponseExtractor func([]byte, *ChatConfig) (chatResponseParts, error) + +// parseChatCompletionResponse applies the shared ChatResponse and +// usage-accounting flow after the provider-specific response is parsed. +func parseChatCompletionResponse(body []byte, chatConfig *ChatConfig, modelUsage *common.ModelUsage, extract chatResponseExtractor) (*ChatResponse, error) { + parts, err := extract(body, chatConfig) + if err != nil { + return nil, err + } + + if err := collectChatModelUsage(modelUsage, parts.RequestID, parts.Usage); err != nil { + common.Error("Failed to collect model usage", err) + } + + return &ChatResponse{ + Answer: parts.Content, + ReasonContent: parts.ReasonContent, + ToolCalls: parts.ToolCalls, + Usage: parts.Usage, + }, nil +} + +// collectChatModelUsage records one completed chat response when the caller +// supplied a usage sink. +func collectChatModelUsage(modelUsage *common.ModelUsage, requestID string, usage *TokenUsage) error { + if modelUsage == nil { + return nil + } + modelUsage.RequestID = requestID + return collectModelUsage(modelUsage, usage) +} + +// collectModelUsage records token usage and response time for one model call. +// The caller owns setting RequestID because streaming providers can receive it +// in a different event from usage. +func collectModelUsage(modelUsage *common.ModelUsage, usage *TokenUsage) error { + if modelUsage == nil { + return nil + } + if usage != nil { + modelUsage.InputTokens = usage.PromptTokens + modelUsage.OutputTokens = usage.CompletionTokens + modelUsage.TotalTokens = usage.TotalTokens + } + modelUsage.ResponseTimeMS = time.Since(modelUsage.StartAt).Milliseconds() + return clickhouse.GetDriver().CollectModelUsage(modelUsage) +} + +// decodeOpenAICompatibleStreamUsage extracts aggregate token usage from one +// OpenAI-compatible streaming event. A missing usage field is not an error. +func decodeOpenAICompatibleStreamUsage(event map[string]any) (*TokenUsage, bool, error) { + rawUsage, ok := event["usage"].(map[string]any) + if !ok { + return nil, false, nil + } + usage := &TokenUsage{} + if err := mapstructure.Decode(rawUsage, usage); err != nil { + return nil, false, err + } + return usage, true, nil +} + +// applyStreamUsage exposes streamed token usage to the caller and records it +// for model-usage analytics when a usage event is received. Analytics failures +// are logged but do not interrupt the stream. +func applyStreamUsage(chatConfig *ChatConfig, modelUsage *common.ModelUsage, usage *TokenUsage) { + if usage == nil { + return + } + if chatConfig != nil { + chatConfig.UsageResult = usage + } + if err := collectModelUsage(modelUsage, usage); err != nil { + common.Error("Failed to collect model usage", err) + } +} + func (b *BaseModel) APIConfigCheck(apiConfig *APIConfig) error { if b.AllowEmptyAPIKey { return nil diff --git a/internal/entity/models/deepseek.go b/internal/entity/models/deepseek.go index 4c4558c749..2ab0212bfe 100644 --- a/internal/entity/models/deepseek.go +++ b/internal/entity/models/deepseek.go @@ -52,6 +52,39 @@ func (d *DeepSeekModel) Name() string { return "deepseek" } +// DeepSeekChatResponse mirrors the response returned by DeepSeek's +// POST /chat/completions endpoint. The shape follows the official OpenAI- +// compatible response schema, including reasoning content, tool calls, and +// the cache/reasoning token usage details. +type DeepSeekChatResponse 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"` + SystemFingerprint string `json:"system_fingerprint"` + Object string `json:"object"` + Usage struct { + CompletionTokens int `json:"completion_tokens"` + PromptTokens int `json:"prompt_tokens"` + TotalTokens int `json:"total_tokens"` + PromptCacheHitTokens int `json:"prompt_cache_hit_tokens"` + PromptCacheMissTokens int `json:"prompt_cache_miss_tokens"` + CompletionTokensDetails struct { + ReasoningTokens int `json:"reasoning_tokens"` + } `json:"completion_tokens_details"` + } `json:"usage"` +} + func (d *DeepSeekModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) { if err := d.baseModel.APIConfigCheck(apiConfig); err != nil { return nil, err @@ -138,62 +171,36 @@ func (d *DeepSeekModel) ChatWithMessages(ctx context.Context, modelName string, 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) (chatResponseParts, error) { + var result DeepSeekChatResponse + if err := json.Unmarshal(body, &result); err != nil { + return chatResponseParts{}, 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") - } + if len(result.Choices) == 0 { + return chatResponseParts{}, 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") - } - - var content string - if c, ok := messageMap["content"].(string); ok { - content = c - } - - var reasonContent string - if rc, ok := messageMap["reasoning_content"].(string); ok { - reasonContent = rc - // if first char of reasonContent is \n remove the '\n' + choice := &result.Choices[0] + reasonContent := choice.Message.ReasoningContent + // DeepSeek may prefix reasoning content with a newline in non-streaming + // responses; keep the existing provider behavior of removing that prefix. if reasonContent != "" && reasonContent[0] == '\n' { reasonContent = reasonContent[1:] } - } - var toolCalls []map[string]interface{} - if tcs, ok := messageMap["tool_calls"].([]interface{}); ok { - for _, tc := range tcs { - if tcMap, ok := tc.(map[string]interface{}); ok { - toolCalls = append(toolCalls, tcMap) - } - } - } - - 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: &choice.Message.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) @@ -280,6 +287,14 @@ func (d *DeepSeekModel) 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(chatModelConfig, modelUsage, tokenUsage) + } + choices, ok := event["choices"].([]interface{}) if !ok || len(choices) == 0 { return nil @@ -295,9 +310,7 @@ func (d *DeepSeekModel) ChatStreamlyWithSender(ctx context.Context, modelName st return nil } - if accumulateToolCallDeltas(delta, accumulatedToolCalls) { - return nil - } + accumulateToolCallDeltas(delta, accumulatedToolCalls) content, ok := delta["content"].(string) if ok && content != "" { diff --git a/internal/entity/models/zhipu-ai.go b/internal/entity/models/zhipu-ai.go index b6efb156ad..aa3de79fca 100644 --- a/internal/entity/models/zhipu-ai.go +++ b/internal/entity/models/zhipu-ai.go @@ -28,12 +28,7 @@ import ( "os" "path/filepath" "ragflow/internal/common" - "ragflow/internal/engine/clickhouse" - "sort" "strings" - "time" - - "github.com/mitchellh/mapstructure" ) // ZhipuAIModel implements ModelDriver for Zhipu AI @@ -150,53 +145,34 @@ func (z *ZhipuAIModel) ChatWithMessages(ctx context.Context, modelName string, m return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) } - // Parse response - var result ZhipuChatResponse - if err = json.Unmarshal(body, &result); err != nil { - return nil, fmt.Errorf("failed to parse response: %w", err) - } - - if result.Choices == nil || len(result.Choices) == 0 { - return nil, fmt.Errorf("empty response") - } - - content := &result.Choices[0].Message.Content - - var reasonContent *string - if chatModelConfig != nil && chatModelConfig.Thinking != nil && *chatModelConfig.Thinking { - reasonContent = &result.Choices[0].Message.ReasoningContent - } - - var toolCalls = result.Choices[0].Message.ToolCalls - - usage := &TokenUsage{ - PromptTokens: result.Usage.PromptTokens, - CompletionTokens: result.Usage.CompletionTokens, - TotalTokens: result.Usage.TotalTokens, - } - - if modelUsage != nil { - modelUsage.RequestID = result.RequestId - modelUsage.InputTokens = result.Usage.PromptTokens - modelUsage.OutputTokens = result.Usage.CompletionTokens - modelUsage.TotalTokens = result.Usage.TotalTokens - modelUsage.ResponseTimeMS = time.Since(modelUsage.StartAt).Milliseconds() - - clickhouseDriver := clickhouse.GetDriver() - err = clickhouseDriver.CollectModelUsage(modelUsage) - if err != nil { - return nil, fmt.Errorf("failed to collect model usage: %w", err) + return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, chatConfig *ChatConfig) (chatResponseParts, error) { + var result ZhipuChatResponse + if err := json.Unmarshal(body, &result); err != nil { + return chatResponseParts{}, fmt.Errorf("failed to parse response: %w", err) } - } - chatResponse := &ChatResponse{ - Answer: content, - ReasonContent: reasonContent, - Usage: usage, - ToolCalls: toolCalls, - } + if len(result.Choices) == 0 { + return chatResponseParts{}, fmt.Errorf("empty response") + } - return chatResponse, nil + choice := &result.Choices[0] + var reasonContent *string + if chatConfig != nil && chatConfig.Thinking != nil && *chatConfig.Thinking { + reasonContent = &choice.Message.ReasoningContent + } + + return chatResponseParts{ + RequestID: result.RequestId, + Content: &choice.Message.Content, + ReasonContent: reasonContent, + ToolCalls: choice.Message.ToolCalls, + Usage: &TokenUsage{ + PromptTokens: result.Usage.PromptTokens, + CompletionTokens: result.Usage.CompletionTokens, + TotalTokens: result.Usage.TotalTokens, + }, + }, nil + }) } // ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel) @@ -263,25 +239,12 @@ func (z *ZhipuAIModel) ChatStreamlyWithSender(ctx context.Context, modelName str if _, err = ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error { common.Info(fmt.Sprintf("%v", event)) - tokenUsageMap, ok := event["usage"].(map[string]interface{}) - if ok { - tokenUsage := TokenUsage{} - if err = mapstructure.Decode(tokenUsageMap, &tokenUsage); err != nil { - return err - } - if chatModelConfig != nil { - chatModelConfig.UsageResult = &tokenUsage - } - if modelUsage != nil { - modelUsage.InputTokens = tokenUsage.PromptTokens - modelUsage.OutputTokens = tokenUsage.CompletionTokens - modelUsage.TotalTokens = tokenUsage.TotalTokens - modelUsage.ResponseTimeMS = time.Since(modelUsage.StartAt).Milliseconds() - clickhouseDriver := clickhouse.GetDriver() - if err = clickhouseDriver.CollectModelUsage(modelUsage); err != nil { - return err - } - } + tokenUsage, found, usageErr := decodeOpenAICompatibleStreamUsage(event) + if usageErr != nil { + return usageErr + } + if found { + applyStreamUsage(chatModelConfig, modelUsage, tokenUsage) } choices, ok := event["choices"].([]interface{}) @@ -298,6 +261,7 @@ func (z *ZhipuAIModel) ChatStreamlyWithSender(ctx context.Context, modelName str if !ok { return nil } + accumulateToolCallDeltas(delta, accumulatedToolCalls) reasoningContent, ok := delta["reasoning_content"].(string) if ok && reasoningContent != "" { @@ -306,56 +270,6 @@ func (z *ZhipuAIModel) ChatStreamlyWithSender(ctx context.Context, modelName str } } - if tcs, ok := delta["tool_calls"].([]interface{}); ok { - for _, tc := range tcs { - tcMap, ok := tc.(map[string]interface{}) - if !ok { - continue - } - idxF, ok := tcMap["index"].(float64) - if !ok { - continue - } - idx := int(idxF) - existing, hasExisting := accumulatedToolCalls[idx] - if !hasExisting { - accumulatedToolCalls[idx] = cloneMap(tcMap) - continue - } - if id, ok := tcMap["id"].(string); ok && id != "" { - if eid, ok := existing["id"].(string); ok { - existing["id"] = eid + id - } else { - existing["id"] = id - } - } - if typ, ok := tcMap["type"].(string); ok && typ != "" { - existing["type"] = typ - } - if fn, ok := tcMap["function"].(map[string]interface{}); ok { - ef, ok := existing["function"].(map[string]interface{}) - if !ok { - ef = make(map[string]interface{}) - existing["function"] = ef - } - if name, ok := fn["name"].(string); ok && name != "" { - if en, ok := ef["name"].(string); ok { - ef["name"] = en + name - } else { - ef["name"] = name - } - } - if args, ok := fn["arguments"].(string); ok && args != "" { - if ea, ok := ef["arguments"].(string); ok { - ef["arguments"] = ea + args - } else { - ef["arguments"] = args - } - } - } - } - } - content, ok := delta["content"].(string) if ok && content != "" { if err = sender(&content, nil); err != nil { @@ -368,18 +282,7 @@ func (z *ZhipuAIModel) ChatStreamlyWithSender(ctx context.Context, modelName str return fmt.Errorf("failed to scan response body: %w", err) } - if len(accumulatedToolCalls) > 0 && chatModelConfig != nil { - indices := make([]int, 0, len(accumulatedToolCalls)) - for idx := range accumulatedToolCalls { - indices = append(indices, idx) - } - sort.Ints(indices) - toolCalls := make([]map[string]interface{}, 0, len(accumulatedToolCalls)) - for _, idx := range indices { - toolCalls = append(toolCalls, accumulatedToolCalls[idx]) - } - chatModelConfig.ToolCallsResult = &toolCalls - } + setSortedToolCallsResult(chatModelConfig, accumulatedToolCalls) // Send [DONE] marker for OpenAI compatibility endOfStream := "[DONE]"