Port agent PRs to GO - 2 (#16565)

### Summary Port the following PRs to GO in this PR https://github.com/infiniflow/ragflow/pull/16420 https://github.com/infiniflow/ragflow/pull/13295
2026-07-28 03:38:11 +08:00 · 2026-07-02 20:20:11 +08:00
parent 24118ac0d1
commit dcbd0d260c
8 changed files with 524 additions and 11 deletions
--- a/internal/entity/models/chat_tools.go
+++ b/internal/entity/models/chat_tools.go
@@ -25,6 +25,20 @@ import (
 	"ragflow/internal/tokenizer"
 )

+// recordUsageFromResponse records one chat call's token usage to both
+// cm.LastUsage and the context-level run sink (if installed). Callers
+// should invoke this after each ChatWithMessages / ChatStreamlyWithSender
+// call so the canvas-layer aggregator and Langfuse both see the split.
+func recordUsageFromResponse(ctx context.Context, cm *ChatModel) {
+	if cm == nil {
+		return
+	}
+	if cm.LastUsage == nil {
+		return
+	}
+	tokenizer.RecordRunTokenUsage(ctx, cm.LastUsage.PromptTokens, cm.LastUsage.CompletionTokens, cm.LastUsage.TotalTokens)
+}
+
 const (
 	defaultMaxRetries = 3
 	defaultMaxRounds  = 5
@@ -77,7 +91,32 @@ func (cm *ChatModel) ChatWithTools(ctx context.Context, system string, history [
 }

 func runToolLoop(ctx context.Context, cm *ChatModel, history []Message, toolsList interface{}, chatCfg *ChatConfig, maxRounds int) (string, int, error) {
+	// Aggregate prompt/completion/total across all tool-calling rounds.
+	// Mirrors Python PR #16420 fix: previously the total was overwritten
+	// each round; now we accumulate so multi-round tool conversations
+	// report the correct grand total.
+	// Reset stale per-call usage from a previous call so a response
+	// without usage doesn't leak the prior call's data.
+	cm.LastUsage = nil
 	var totalTokens int
+	aggUsage := &ChatUsage{}
+
+	addRoundUsage := func(resp *ChatResponse) {
+		u := resp.Usage
+		if u == nil {
+			return
+		}
+		aggUsage.PromptTokens += u.PromptTokens
+		aggUsage.CompletionTokens += u.CompletionTokens
+		aggUsage.TotalTokens += u.TotalTokens
+		totalTokens = aggUsage.TotalTokens
+		// Store per-round delta (not cumulative) so RecordRunTokenUsage
+		// records each round's contribution exactly once.
+		cm.LastUsage = &ChatUsage{
+			PromptTokens: u.PromptTokens, CompletionTokens: u.CompletionTokens, TotalTokens: u.TotalTokens,
+		}
+		recordUsageFromResponse(ctx, cm)
+	}

 	for round := 0; round <= maxRounds; round++ {
 		select {
@@ -97,6 +136,7 @@ func runToolLoop(ctx context.Context, cm *ChatModel, history []Message, toolsLis
 		if resp == nil {
 			return "", totalTokens, fmt.Errorf("round %d: nil response", round)
 		}
+		addRoundUsage(resp)

 		if len(resp.ToolCalls) == 0 {
 			answer := ""
@@ -106,7 +146,11 @@ func runToolLoop(ctx context.Context, cm *ChatModel, history []Message, toolsLis
 			if resp.ReasonContent != nil && *resp.ReasonContent != "" {
 				answer = "<think>" + *resp.ReasonContent + "</think>" + answer
 			}
-			totalTokens += tokenizer.NumTokensFromString(answer)
+			// Fallback: if the provider didn't return usage info,
+			// estimate from the answer text using tiktoken.
+			if resp.Usage == nil || resp.Usage.TotalTokens == 0 {
+				totalTokens += tokenizer.NumTokensFromString(answer)
+			}
 			return answer, totalTokens, nil
 		}

@@ -126,7 +170,11 @@ func runToolLoop(ctx context.Context, cm *ChatModel, history []Message, toolsLis
 	if resp == nil || resp.Answer == nil {
 		return "", totalTokens, fmt.Errorf("final call: no answer")
 	}
-	totalTokens += tokenizer.NumTokensFromString(*resp.Answer)
+	addRoundUsage(resp)
+	// Fallback: use text-based estimation if no authoritative usage.
+	if resp.Usage == nil || resp.Usage.TotalTokens == 0 {
+		totalTokens += tokenizer.NumTokensFromString(*resp.Answer)
+	}
 	return *resp.Answer, totalTokens, nil
 }

@@ -177,7 +225,37 @@ func (cm *ChatModel) ChatStreamlyWithTools(ctx context.Context, system string, h
 }

 func runStreamToolLoop(ctx context.Context, cm *ChatModel, history []Message, toolsList interface{}, chatCfg *ChatConfig, maxRounds int, sender func(*string, *string) error) (int, error) {
+	// Aggregate token counts across every tool-calling round (each round is a
+	// separate provider request). Committing per round avoids the previous
+	// bug where a later round's total overwrote earlier rounds.
+	// Reset stale per-call usage from a previous call.
+	cm.LastUsage = nil
 	var totalTokens int
+	aggUsage := &ChatUsage{}
+
+	commitRound := func(cfg *ChatConfig, roundTokens int) {
+		// Prefer the authoritative usage from the API (extracted via
+		// stream_options.include_usage=true) over text-based token
+		// counting. Mirrors Python's usage_from_response accumulation
+		// in chat_model.py streaming handlers.
+		// Track per-round delta so RecordRunTokenUsage records each
+		// round's contribution exactly once (the sink does Add, not Set).
+		var deltaPrompt, deltaCompletion, deltaTotal int
+		if u := cfg.UsageResult; u != nil && u.TotalTokens > 0 {
+			deltaPrompt, deltaCompletion, deltaTotal = u.PromptTokens, u.CompletionTokens, u.TotalTokens
+			aggUsage.PromptTokens += deltaPrompt
+			aggUsage.CompletionTokens += deltaCompletion
+			aggUsage.TotalTokens += deltaTotal
+		} else {
+			deltaTotal = roundTokens
+			aggUsage.TotalTokens += deltaTotal
+		}
+		totalTokens = aggUsage.TotalTokens
+		cm.LastUsage = &ChatUsage{
+			PromptTokens: deltaPrompt, CompletionTokens: deltaCompletion, TotalTokens: deltaTotal,
+		}
+		recordUsageFromResponse(ctx, cm)
+	}

 	for round := 0; round <= maxRounds; round++ {
 		select {
@@ -192,10 +270,13 @@ func runStreamToolLoop(ctx context.Context, cm *ChatModel, history []Message, to
 		cfg.Stream = boolPtr(true)
 		var tcs []map[string]interface{}
 		cfg.ToolCallsResult = &tcs
+		var roundUsage ChatUsage
+		cfg.UsageResult = &roundUsage

 		reasoningStarted := false
 		var answer string
 		var pendingThinkClose bool
+		var roundTokens int

 		err := cm.ModelDriver.ChatStreamlyWithSender(*cm.ModelName, history, cm.APIConfig, &cfg, func(delta *string, reason *string) error {
 			if reason != nil && *reason != "" {
@@ -207,6 +288,7 @@ func runStreamToolLoop(ctx context.Context, cm *ChatModel, history []Message, to
 					}
 				}
 				pendingThinkClose = true
+				roundTokens += tokenizer.NumTokensFromString(*reason)
 				return sender(reason, nil)
 			}
 			// Reasoning ended, close the think block if open
@@ -221,7 +303,7 @@ func runStreamToolLoop(ctx context.Context, cm *ChatModel, history []Message, to
 				if *delta == "[DONE]" {
 					return nil
 				}
-				totalTokens += tokenizer.NumTokensFromString(*delta)
+				roundTokens += tokenizer.NumTokensFromString(*delta)
 				answer += *delta
 				if e := sender(delta, nil); e != nil {
 					return e
@@ -237,6 +319,9 @@ func runStreamToolLoop(ctx context.Context, cm *ChatModel, history []Message, to
 				return totalTokens, e
 			}
 		}
+		// Commit this round's token count to the running aggregate.
+		// Prefer authoritative API usage over text-based estimation.
+		commitRound(&cfg, roundTokens)
 		if err != nil {
 			return totalTokens, fmt.Errorf("round %d: %w", round, err)
 		}
@@ -263,7 +348,17 @@ func runStreamToolLoop(ctx context.Context, cm *ChatModel, history []Message, to
 	})
 	cfg := *chatCfg
 	cfg.Stream = boolPtr(true)
-	return totalTokens, cm.ModelDriver.ChatStreamlyWithSender(*cm.ModelName, history, cm.APIConfig, &cfg, sender)
+	var exceedUsage ChatUsage
+	cfg.UsageResult = &exceedUsage
+	var exceedTokens int
+	err := cm.ModelDriver.ChatStreamlyWithSender(*cm.ModelName, history, cm.APIConfig, &cfg, func(delta *string, reason *string) error {
+		if delta != nil && *delta != "" && *delta != "[DONE]" {
+			exceedTokens += tokenizer.NumTokensFromString(*delta)
+		}
+		return nil
+	})
+	commitRound(&cfg, exceedTokens)
+	return totalTokens, err
 }

 // appendToolResults executes tool calls concurrently, appends the assistant
--- a/internal/entity/models/llm.go
+++ b/internal/entity/models/llm.go
@@ -124,10 +124,23 @@ func (m *EinoChatModel) Generate(ctx context.Context, msgs []*schema.Message, op
 	if err := ctx.Err(); err != nil {
 		return nil, err
 	}
+	// Reset stale per-call usage before the call so that a response
+	// without a usage block doesn't leak the previous call's data.
+	// Mirrors Python's LLMBundle._reset_last_usage().
+	m.inner.LastUsage = nil
 	resp, err := m.inner.ModelDriver.ChatWithMessages(*m.inner.ModelName, internal, m.inner.APIConfig, m.chatCfg)
 	if err != nil {
 		return nil, fmt.Errorf("models: EinoChatModel.Generate(%s): %w", *m.inner.ModelName, err)
 	}
+	// Record the per-call token usage so the canvas-level aggregator (and
+	// Langfuse) can compute the run total. Mirrors Python's
+	// LLMBundle._report_usage() / self.mdl.last_usage pattern.
+	if resp != nil && resp.Usage != nil {
+		m.inner.LastUsage = &ChatUsage{
+			PromptTokens: resp.Usage.PromptTokens, CompletionTokens: resp.Usage.CompletionTokens, TotalTokens: resp.Usage.TotalTokens,
+		}
+		recordUsageFromResponse(ctx, m.inner)
+	}
 	return fromInternalResponse(resp), nil
 }

--- a/internal/entity/models/openai.go
+++ b/internal/entity/models/openai.go
@@ -211,6 +211,16 @@ func (o *OpenAIModel) ChatWithMessages(modelName string, messages []Message, api
 		ToolCalls:     toolCalls,
 	}

+	// Extract usage split (prompt/completion/total) from the raw API
+	// response for accurate per-call token accounting. Non-OpenAI
+	// providers that implement the OpenAI-compat API surface (DeepSeek,
+	// Moonshot, etc.) also return a "usage" key with the same shape.
+	if pt, ct, tt := extractUsageFromMap(result); tt > 0 {
+		chatResponse.Usage = &ChatUsage{
+			PromptTokens: pt, CompletionTokens: ct, TotalTokens: tt,
+		}
+	}
+
 	return chatResponse, nil
 }

@@ -247,11 +257,17 @@ func (o *OpenAIModel) ChatStreamlyWithSender(modelName string, messages []Messag
 		apiMessages[i] = apiMsg
 	}

-	// Build request body with streaming on by default
+	// Build request body with streaming on by default.
+	// stream_options.include_usage asks the provider to attach a
+	// usage block to the final streaming chunk (mirrors Python's
+	// chat_model.py _stream_options / stream_options.include_usage).
 	reqBody := map[string]interface{}{
 		"model":    modelName,
 		"messages": apiMessages,
 		"stream":   true,
+		"stream_options": map[string]interface{}{
+			"include_usage": true,
+		},
 	}

 	if chatModelConfig != nil {
@@ -316,6 +332,12 @@ func (o *OpenAIModel) ChatStreamlyWithSender(modelName string, messages []Messag

 	sawTerminal := false
 	accumulatedToolCalls := make(map[int]map[string]interface{})
+	// Capture the authoritative usage block from the final streaming
+	// chunk (when provider honours stream_options.include_usage=true).
+	// The last chunk in the stream carries the "usage" key alongside
+	// empty choices; we overwrite on every chunk so the final frame
+	// wins, matching Python's chat_model.py usage_from_response loop.
+	var streamUsage *ChatUsage
 	scanner := bufio.NewScanner(resp.Body)
 	for scanner.Scan() {
 		line := scanner.Text()
@@ -340,6 +362,13 @@ func (o *OpenAIModel) ChatStreamlyWithSender(modelName string, messages []Messag
 			continue
 		}

+		// Extract usage from this chunk. When stream_options.include_usage
+		// is true, the final chunk carries the full usage breakdown at the
+		// top level of the event alongside (possibly empty) choices.
+		if pt, ct, tt := extractUsageFromMap(event); tt > 0 {
+			streamUsage = &ChatUsage{PromptTokens: pt, CompletionTokens: ct, TotalTokens: tt}
+		}
+
 		choices, ok := event["choices"].([]interface{})
 		if !ok || len(choices) == 0 {
 			continue
@@ -421,6 +450,11 @@ func (o *OpenAIModel) ChatStreamlyWithSender(modelName string, messages []Messag
 		chatModelConfig.ToolCallsResult = &tcs
 	}

+	// Populate UsageResult with the authoritative usage from the stream.
+	if streamUsage != nil && chatModelConfig != nil {
+		chatModelConfig.UsageResult = streamUsage
+	}
+
 	// Send the [DONE] marker for OpenAI compatibility
 	endOfStream := "[DONE]"
 	if err := sender(&endOfStream, nil); err != nil {
@@ -1024,6 +1058,50 @@ func (o *OpenAIModel) ShowTask(taskID string, apiConfig *APIConfig) (*TaskRespon
 	return nil, fmt.Errorf("%s, no such method", o.Name())
 }

+// extractUsageFromMap reads the "usage" key from an OpenAI-style API
+// response and returns (prompt_tokens, completion_tokens, total_tokens).
+// All return values are zero when the response carries no usage block.
+func extractUsageFromMap(raw map[string]interface{}) (int, int, int) {
+	if raw == nil {
+		return 0, 0, 0
+	}
+	ru, ok := raw["usage"]
+	if !ok {
+		return 0, 0, 0
+	}
+	usage, ok := ru.(map[string]interface{})
+	if !ok {
+		return 0, 0, 0
+	}
+	get := func(keys ...string) int {
+		for _, k := range keys {
+			v, ok := usage[k]
+			if !ok {
+				continue
+			}
+			switch val := v.(type) {
+			case float64:
+				return int(val)
+			case int:
+				return val
+			case json.Number:
+				n, err := val.Int64()
+				if err == nil {
+					return int(n)
+				}
+			}
+		}
+		return 0
+	}
+	pt := get("prompt_tokens", "input_tokens")
+	ct := get("completion_tokens", "output_tokens")
+	tt := get("total_tokens")
+	if tt == 0 {
+		tt = pt + ct
+	}
+	return pt, ct, tt
+}
+
 func cloneMap(m map[string]interface{}) map[string]interface{} {
 	cp := make(map[string]interface{}, len(m))
 	for k, v := range m {
--- a/internal/entity/models/types.go
+++ b/internal/entity/models/types.go
@@ -60,6 +60,16 @@ type ChatResponse struct {
 	Answer        *string                  `json:"answer"`
 	ReasonContent *string                  `json:"reason_content"`
 	ToolCalls     []map[string]interface{} `json:"tool_calls,omitempty"`
+	Usage         *ChatUsage               `json:"usage,omitempty"`
+}
+
+// ChatUsage holds token usage split for one LLM call. Consumed by
+// LLMBundle for accurate Langfuse reporting and run aggregation.
+// Mirrors Python's common.token_utils.usage_from_response() split.
+type ChatUsage struct {
+	PromptTokens     int `json:"prompt_tokens"`
+	CompletionTokens int `json:"completion_tokens"`
+	TotalTokens      int `json:"total_tokens"`
 }

 type EmbeddingData struct {
@@ -154,6 +164,12 @@ type ChatConfig struct {
 	Tools           interface{}               `json:"tools,omitempty"`
 	ToolChoice      *string                   `json:"tool_choice,omitempty"`
 	ToolCallsResult *[]map[string]interface{} `json:"-"`
+	// UsageResult receives the token usage extracted from the final
+	// streaming chunk when stream_options.include_usage is true.
+	// The ChatStreamlyWithSender driver writes to this pointer (if
+	// non-nil) after the stream completes; callers read it the same
+	// way they read ToolCallsResult.
+	UsageResult *ChatUsage `json:"-"`
 }

 type APIConfig struct {
@@ -238,6 +254,10 @@ type ChatModel struct {
 	ModelName   *string
 	APIConfig   *APIConfig
 	ToolConfig  *ToolConfig
+	// LastUsage holds the token usage (prompt/completion/total) of the most
+	// recent chat call. Consumed by callers for accurate Langfuse reporting
+	// and per-run token aggregation. Reset before each call.
+	LastUsage *ChatUsage
 }

 // NewChatModel creates a new ChatModel