Skip to content

refactor: enhance token estimation and tool result handling in conversation service #490

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
chenyme merged 2 commits into from
Jul 15, 2026
Merged

refactor: enhance token estimation and tool result handling in conversation service #490

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0483421
refactor: enhance token estimation and tool result handling in conver…
chenyme Jul 15, 2026
71ce58e
Merge branch 'dev' into mcp
chenyme Jul 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 17 additions & 2 deletions backend/internal/application/conversation/service_message_context.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,9 +74,24 @@ func estimateMessageTokens(message llm.Message) int64 {
for _, part := range message.Parts {
tokens += estimateContentPartTokens(part)
}
return tokens + estimateTokens(message.ReasoningContent)
} else {
tokens += estimateTokens(message.Content)
}
tokens += estimateTokens(message.ReasoningContent)
for _, call := range message.ToolCalls {
tokens += estimateTokens(call.ToolCallID)
tokens += estimateTokens(call.ToolName)
tokens += estimateTokens(call.ArgumentsJSON)
tokens += 8
}
for _, result := range message.ToolResults {
tokens += estimateTokens(result.ToolCallID)
tokens += estimateTokens(result.ToolName)
tokens += estimateTokens(result.OutputJSON)
tokens += estimateTokens(result.Error)
tokens += 8
}
return tokens + estimateTokens(message.Content) + estimateTokens(message.ReasoningContent)
return tokens
}

func estimatePromptTokens(messages []llm.Message) int64 {
Expand Down
91 changes: 67 additions & 24 deletions backend/internal/application/conversation/service_message_send.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1121,29 +1121,53 @@ func (s *Service) sendMessageInternal(
}
llmCallCount := 1
toolLedger := newToolExecutionLedger()
toolHistoryTrimmedForRun := false

for len(upstreamOutput.ToolCalls) > 0 && llmCallCount < maxLLMCalls && remainingToolCalls > 0 {
pendingToolCalls := upstreamOutput.ToolCalls
if len(pendingToolCalls) > remainingToolCalls {
pendingToolCalls = pendingToolCalls[:remainingToolCalls]
}
reasoningContent := ""
if reasoningContentPassback {
reasoningContent = outputReasoningContent(upstreamOutput)
}
assistantToolMessage := llm.Message{
Role: "assistant",
Content: assistantText,
ReasoningContent: reasoningContent,
ToolCalls: pendingToolCalls,
}
toolResultTokenBudget := resolveToolResultTokenBudget(
generateInput,
llmMessages,
assistantToolMessage,
route.UpstreamModel,
route.ModelCapabilitiesJSON,
)
toolCtx, toolSpan := platformtracing.Start(ctx, "conversation.tool.execute",
trace.WithAttributes(
attribute.Int64("conversation.id", int64(input.ConversationID)),
attribute.Int64("user.id", int64(input.UserID)),
attribute.Int("conversation.tool.request_count", len(upstreamOutput.ToolCalls)),
attribute.Int("conversation.tool.remaining_count", remainingToolCalls),
attribute.Int64("conversation.tool.result_token_budget", toolResultTokenBudget),
),
)
toolResult := s.executeAssistantToolCalls(toolCtx, executeAssistantToolCallsInput{
UserID: input.UserID,
ConversationID: input.ConversationID,
MessageID: assistantMessage.ID,
RequestID: input.RequestID,
RunID: runID,
ToolCalls: upstreamOutput.ToolCalls,
ToolCallLimit: remainingToolCalls,
TraceRecorder: traceRecorder,
ToolNameMap: toolRuntime.nameMap,
MCPConfigs: toolRuntime.mcpConfigs,
ToolSchemas: toolRuntime.schemas,
Ledger: toolLedger,
UserID: input.UserID,
ConversationID: input.ConversationID,
MessageID: assistantMessage.ID,
RequestID: input.RequestID,
RunID: runID,
ToolCalls: pendingToolCalls,
ToolCallLimit: remainingToolCalls,
TraceRecorder: traceRecorder,
ToolNameMap: toolRuntime.nameMap,
MCPConfigs: toolRuntime.mcpConfigs,
ToolSchemas: toolRuntime.schemas,
Ledger: toolLedger,
ResultTokenBudget: toolResultTokenBudget,
})
toolSpan.SetAttributes(
attribute.Int("conversation.tool.executed_count", len(toolResult.Rows)),
Expand All @@ -1163,22 +1187,35 @@ func (s *Service) sendMessageInternal(
if len(toolResult.ToolResults) == 0 {
break
}
reasoningContent := ""
if reasoningContentPassback {
reasoningContent = outputReasoningContent(upstreamOutput)
}
assistantToolMessage.ToolCalls = toolResult.ExecutedToolCalls
llmMessages = append(llmMessages,
llm.Message{
Role: "assistant",
Content: assistantText,
ReasoningContent: reasoningContent,
ToolCalls: toolResult.ExecutedToolCalls,
},
assistantToolMessage,
llm.Message{
Role: "tool",
ToolResults: toolResult.ToolResults,
},
)
var toolHistoryTrimmed bool
llmMessages, toolHistoryTrimmed = trimToolFollowUpHistory(
generateInput,
llmMessages,
route.UpstreamModel,
route.ModelCapabilitiesJSON,
)
if toolHistoryTrimmed {
toolHistoryTrimmedForRun = true
sendSpan.SetAttributes(attribute.Bool("conversation.tool.history_trimmed", true))
}
var toolResultsRebalanced bool
llmMessages, toolResultsRebalanced = rebalanceToolFollowUpResults(
generateInput,
llmMessages,
route.UpstreamModel,
route.ModelCapabilitiesJSON,
)
if toolResultsRebalanced {
sendSpan.SetAttributes(attribute.Bool("conversation.tool.results_rebalanced", true))
}

followUpInput := generateInput
if llmCallCount+1 >= maxLLMCalls {
Expand All @@ -1187,7 +1224,7 @@ func (s *Service) sendMessageInternal(
followUpInput.DisableTools = true
followUpInput.PreviousResponseID = ""
applyOpenAIResponsesInstructions(route, routeConfig.Endpoint, &followUpInput)
} else if routeConfig.Endpoint == llm.EndpointResponses && supportsPreviousResponseIDRoute(route) && strings.TrimSpace(upstreamOutput.ResponseID) != "" {
} else if !toolHistoryTrimmed && !toolResultsRebalanced && routeConfig.Endpoint == llm.EndpointResponses && supportsPreviousResponseIDRoute(route) && strings.TrimSpace(upstreamOutput.ResponseID) != "" {
followUpInput.PreviousResponseID = strings.TrimSpace(upstreamOutput.ResponseID)
followUpInput.Messages = []llm.Message{{Role: "tool", ToolResults: toolResult.ToolResults}}
} else {
Expand Down Expand Up @@ -1284,6 +1321,12 @@ func (s *Service) sendMessageInternal(
Tools: toolRuntime.definitions,
Options: filteredOptions,
})
responseIDForPersistence := upstreamOutput.ResponseID
// 历史裁剪后的上游 response 不再代表数据库可重建的完整历史，禁止跨轮复用。
if toolHistoryTrimmedForRun {
responseIDForPersistence = ""
statefulPromptFingerprint = ""
}

run.InputTokens = effectiveInputTokens
run.OutputTokens = effectiveOutputTokens
Expand Down Expand Up @@ -1341,7 +1384,7 @@ func (s *Service) sendMessageInternal(
OutputTokens: effectiveOutputTokens,
ReasoningTokens: totalUsage.ReasoningTokens,
AssistantLatency: assistantLatencyMS,
ResponseID: upstreamOutput.ResponseID,
ResponseID: responseIDForPersistence,
StatefulPromptFingerprint: statefulPromptFingerprint,
ToolCallRows: toolCallRows,
PersistedToolCallKeys: persistedToolCallKeys,
Expand Down
182 changes: 182 additions & 0 deletions backend/internal/application/conversation/service_message_usage.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,3 +117,185 @@ func estimateToolDefinitionTokens(tools []llm.ToolDefinition) int64 {
}
return tokens
}

// resolveToolResultTokenBudget 计算当前用户轮次的全部工具结果可使用的模型输入预算。
// 新批次先使用该上限，回灌前再对同轮全部结果统一分配，不额外透支有效上下文。
func resolveToolResultTokenBudget(
generateInput llm.GenerateInput,
messages []llm.Message,
pendingAssistant llm.Message,
modelName string,
capabilitiesJSON string,
) int64 {
budgetMessages := toolResultPayloadPlaceholders(prioritizeCurrentToolMessages(messages))
placeholderResults := make([]llm.ToolResult, 0, len(pendingAssistant.ToolCalls))
for _, call := range pendingAssistant.ToolCalls {
placeholderResults = append(placeholderResults, llm.ToolResult{
ToolCallID: call.ToolCallID,
ToolName: call.ToolName,
OutputJSON: "{}",
})
}
budgetMessages = append(
budgetMessages,
pendingAssistant,
llm.Message{Role: "tool", ToolResults: placeholderResults},
)
available := int64(llm.EffectiveContextBudgetFromCapabilities(modelName, capabilitiesJSON)) -
estimateToolFollowUpInputTokens(generateInput, budgetMessages)
if available < 0 {
return 0
}
return available
}

// rebalanceToolFollowUpResults 在完整工具回灌请求超预算时，统一压缩当前轮的全部工具结果。
func rebalanceToolFollowUpResults(
generateInput llm.GenerateInput,
messages []llm.Message,
modelName string,
capabilitiesJSON string,
) ([]llm.Message, bool) {
effectiveBudget := int64(llm.EffectiveContextBudgetFromCapabilities(modelName, capabilitiesJSON))
if estimateToolFollowUpInputTokens(generateInput, messages) <= effectiveBudget {
return messages, false
}

_, currentUserIndex := toolHistoryBounds(messages)
if currentUserIndex < 0 {
return messages, false
}
fixedMessages := toolResultPayloadPlaceholders(messages)
resultBudget := effectiveBudget - estimateToolFollowUpInputTokens(generateInput, fixedMessages)
if resultBudget < 0 {
resultBudget = 0
}

type resultRef struct {
messageIndex int
resultIndex int
}
result := append([]llm.Message(nil), messages...)
refs := make([]resultRef, 0)
slots := make([]toolExecutionSlot, 0)
for messageIndex := currentUserIndex + 1; messageIndex < len(result); messageIndex++ {
if len(result[messageIndex].ToolResults) == 0 {
continue
}
result[messageIndex].ToolResults = append([]llm.ToolResult(nil), result[messageIndex].ToolResults...)
for resultIndex, toolResult := range result[messageIndex].ToolResults {
refs = append(refs, resultRef{messageIndex: messageIndex, resultIndex: resultIndex})
slots = append(slots, toolExecutionSlot{result: toolResult})
}
}
if len(slots) == 0 {
return messages, false
}

enforceToolResultAggregateBudget(slots, resultBudget)
changed := false
for index, ref := range refs {
if result[ref.messageIndex].ToolResults[ref.resultIndex] != slots[index].result {
changed = true
result[ref.messageIndex].ToolResults[ref.resultIndex] = slots[index].result
}
}
if !changed {
return messages, false
}
return result, true
}

// toolResultPayloadPlaceholders 保留工具结果的协议结构，但移除可变正文以计算固定上下文开销。
func toolResultPayloadPlaceholders(messages []llm.Message) []llm.Message {
result := append([]llm.Message(nil), messages...)
for messageIndex := range result {
if len(result[messageIndex].ToolResults) == 0 {
continue
}
placeholders := make([]llm.ToolResult, 0, len(result[messageIndex].ToolResults))
for _, toolResult := range result[messageIndex].ToolResults {
placeholders = append(placeholders, llm.ToolResult{
ToolCallID: toolResult.ToolCallID,
ToolName: toolResult.ToolName,
OutputJSON: "{}",
Status: toolResult.Status,
})
}
result[messageIndex].ToolResults = placeholders
}
return result
}

// trimToolFollowUpHistory 仅在工具回灌请求超预算时删除最老的完整历史轮次。
func trimToolFollowUpHistory(
generateInput llm.GenerateInput,
messages []llm.Message,
modelName string,
capabilitiesJSON string,
) ([]llm.Message, bool) {
effectiveBudget := int64(llm.EffectiveContextBudgetFromCapabilities(modelName, capabilitiesJSON))
estimatedTokens := estimateToolFollowUpInputTokens(generateInput, messages)
if estimatedTokens <= effectiveBudget {
return messages, false
}

systemEnd, currentUserIndex := toolHistoryBounds(messages)
if currentUserIndex <= systemEnd {
return messages, false
}
for cutFrom := systemEnd; cutFrom < currentUserIndex; cutFrom++ {
estimatedTokens -= estimateMessageTokens(messages[cutFrom])
nextIndex := cutFrom + 1
if nextIndex < currentUserIndex && messages[nextIndex].Role != "user" {
continue
}
if estimatedTokens <= effectiveBudget || nextIndex == currentUserIndex {
trimmed := make([]llm.Message, 0, systemEnd+len(messages)-nextIndex)
trimmed = append(trimmed, messages[:systemEnd]...)
trimmed = append(trimmed, messages[nextIndex:]...)
return trimmed, true
}
}
return messages, false
}

// prioritizeCurrentToolMessages 返回系统指令和当前用户轮次，供工具结果计算最大可用预算。
func prioritizeCurrentToolMessages(messages []llm.Message) []llm.Message {
systemEnd, currentUserIndex := toolHistoryBounds(messages)
if currentUserIndex <= systemEnd {
return append([]llm.Message(nil), messages...)
}
result := make([]llm.Message, 0, systemEnd+len(messages)-currentUserIndex)
result = append(result, messages[:systemEnd]...)
result = append(result, messages[currentUserIndex:]...)
return result
}

// toolHistoryBounds 定位系统前缀结束位置和当前轮用户消息。
func toolHistoryBounds(messages []llm.Message) (int, int) {
systemEnd := 0
for systemEnd < len(messages) && messages[systemEnd].Role == "system" {
systemEnd++
}
currentUserIndex := -1
for index := len(messages) - 1; index >= systemEnd; index-- {
if messages[index].Role == "user" {
currentUserIndex = index
break
}
}
return systemEnd, currentUserIndex
}

// estimateToolFollowUpInputTokens 按全量请求形状估算工具回灌输入。
func estimateToolFollowUpInputTokens(generateInput llm.GenerateInput, messages []llm.Message) int64 {
budgetMessages := messages
if strings.TrimSpace(generateInput.Instructions) != "" {
_, budgetMessages = extractOpenAIResponsesInstructions(messages)
}
budgetInput := generateInput
budgetInput.Messages = budgetMessages
budgetInput.PreviousResponseID = ""
return estimateGenerateInputTokens(budgetInput)
}
Loading
Loading