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

@@ -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

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 != "" {

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]"