From be178071529f4e66591ff65c4ff51a148187bac4 Mon Sep 17 00:00:00 2001 From: mlogclub Date: Sat, 11 Apr 2026 22:50:35 +0800 Subject: [PATCH] feat: enhance tooling management with static and dynamic tool code handling --- .../ai/runtime/internal/engine/service.go | 33 ++++++++++ .../impl/callbacks/agent_trace_handler.go | 62 ++++++++++++++++--- .../impl/callbacks/runlog_callback.go | 11 ++++ .../internal/impl/callbacks/trace_callback.go | 3 + 4 files changed, 102 insertions(+), 7 deletions(-) diff --git a/internal/ai/runtime/internal/engine/service.go b/internal/ai/runtime/internal/engine/service.go index fa44a19..02fe0f8 100644 --- a/internal/ai/runtime/internal/engine/service.go +++ b/internal/ai/runtime/internal/engine/service.go @@ -97,6 +97,7 @@ func (s *Service) Run(ctx context.Context, req Request) (*Summary, error) { toolDefsByModelName[modelName] = toolCode } collector.Data.Input.ToolCodes = append(collector.Data.Input.ToolCodes, summary.ToolCodes...) + collector.SetTooling(staticToolCodeList(req.ToolSet), definitionToolCodes(filteredToolDefs), len(filteredToolDefs) > 0) collector.Data.Model.Provider = string(req.AIConfig.Provider) collector.Data.Model.Name = req.AIConfig.ModelName @@ -250,6 +251,7 @@ func (s *Service) Resume(ctx context.Context, req ResumeRequest) (*Summary, erro toolDefsByModelName[modelName] = toolCode } collector.Data.Input.ToolCodes = append(collector.Data.Input.ToolCodes, summary.ToolCodes...) + collector.SetTooling(staticToolCodeList(req.ToolSet), definitionToolCodes(toolDefs), len(toolDefs) > 0) collector.Data.Model.Provider = string(req.AIConfig.Provider) collector.Data.Model.Name = req.AIConfig.ModelName agent, err := s.agentFactory.BuildCustomerServiceAgent(ctx, factory.BuildCustomerServiceAgentInput{ @@ -559,3 +561,34 @@ func toolSetStaticToolCodes(toolSet *registry.ToolSet) map[string]string { } return toolSet.StaticToolCodes } + +func definitionToolCodes(definitions []adapter.MCPToolDefinition) []string { + if len(definitions) == 0 { + return nil + } + ret := make([]string, 0, len(definitions)) + for _, item := range definitions { + toolCode := strings.TrimSpace(item.ToolCode) + if toolCode == "" { + continue + } + ret = append(ret, toolCode) + } + return ret +} + +func staticToolCodeList(toolSet *registry.ToolSet) []string { + toolCodes := toolSetStaticToolCodes(toolSet) + if len(toolCodes) == 0 { + return nil + } + ret := make([]string, 0, len(toolCodes)) + for _, toolCode := range toolCodes { + toolCode = strings.TrimSpace(toolCode) + if toolCode == "" { + continue + } + ret = append(ret, toolCode) + } + return ret +} diff --git a/internal/ai/runtime/internal/impl/callbacks/agent_trace_handler.go b/internal/ai/runtime/internal/impl/callbacks/agent_trace_handler.go index 5f397d3..a9d179a 100644 --- a/internal/ai/runtime/internal/impl/callbacks/agent_trace_handler.go +++ b/internal/ai/runtime/internal/impl/callbacks/agent_trace_handler.go @@ -122,12 +122,19 @@ func previewToolText(text string, limit int) string { } func (h *RuntimeTraceHandler) buildToolSearchTraceItem(argumentsInJSON string, result string, runErr error) ToolSearchTraceItem { - item := ToolSearchTraceItem{ - Action: "search", - Status: "ok", - } + item := ToolSearchTraceItem{Status: "ok"} args := parseToolArguments(argumentsInJSON) - item.Query = strings.TrimSpace(readToolSearchString(args, "regex_pattern")) + item.Query = strings.TrimSpace(firstNonBlank( + readToolSearchString(args, "query"), + readToolSearchString(args, "regex_pattern"), + )) + item.TargetToolCode = strings.TrimSpace(readToolSearchString(args, "toolCode")) + item.TargetServerCode, item.TargetToolName = toolx.SplitMCPToolCode(item.TargetToolCode) + if item.TargetToolCode != "" { + item.Action = "invoke" + } else { + item.Action = "search" + } if runErr != nil { item.Status = "error" item.ErrorMessage = runErr.Error() @@ -137,8 +144,7 @@ func (h *RuntimeTraceHandler) buildToolSearchTraceItem(argumentsInJSON string, r if err := json.Unmarshal([]byte(strings.TrimSpace(result)), &payload); err != nil { return item } - selectedTools, _ := payload["selectedTools"].([]any) - item.CandidateToolCodes = h.extractSelectedToolCodes(selectedTools) + item.CandidateToolCodes = h.extractCandidateToolCodes(payload) return item } @@ -154,6 +160,29 @@ func readToolSearchString(data map[string]any, key string) string { return text } +func firstNonBlank(values ...string) string { + for _, value := range values { + value = strings.TrimSpace(value) + if value != "" { + return value + } + } + return "" +} + +func (h *RuntimeTraceHandler) extractCandidateToolCodes(payload map[string]any) []string { + if len(payload) == 0 { + return nil + } + if items, ok := payload["selectedTools"].([]any); ok { + return h.extractSelectedToolCodes(items) + } + if items, ok := payload["candidates"].([]any); ok { + return h.extractCandidateObjects(items) + } + return nil +} + func (h *RuntimeTraceHandler) extractSelectedToolCodes(items []any) []string { if len(items) == 0 { return nil @@ -179,3 +208,22 @@ func (h *RuntimeTraceHandler) extractSelectedToolCodes(items []any) []string { } return ret } + +func (h *RuntimeTraceHandler) extractCandidateObjects(items []any) []string { + if len(items) == 0 { + return nil + } + ret := make([]string, 0, len(items)) + for _, item := range items { + obj, ok := item.(map[string]any) + if !ok { + continue + } + toolCode := strings.TrimSpace(readToolSearchString(obj, "toolCode")) + if toolCode == "" { + continue + } + ret = append(ret, toolCode) + } + return ret +} diff --git a/internal/ai/runtime/internal/impl/callbacks/runlog_callback.go b/internal/ai/runtime/internal/impl/callbacks/runlog_callback.go index d0ac9f4..dabaa99 100644 --- a/internal/ai/runtime/internal/impl/callbacks/runlog_callback.go +++ b/internal/ai/runtime/internal/impl/callbacks/runlog_callback.go @@ -30,6 +30,17 @@ func (c *RuntimeTraceCollector) Marshal() string { return string(buf) } +func (c *RuntimeTraceCollector) SetTooling(staticToolCodes []string, dynamicToolCodes []string, toolSearchEnabled bool) { + if c == nil { + return + } + c.mu.Lock() + defer c.mu.Unlock() + c.Data.Input.StaticToolCodes = append([]string(nil), staticToolCodes...) + c.Data.Input.DynamicToolCodes = append([]string(nil), dynamicToolCodes...) + c.Data.Input.ToolSearchEnabled = toolSearchEnabled +} + func (c *RuntimeTraceCollector) AddToolItem(item ToolTraceItem) { if c == nil { return diff --git a/internal/ai/runtime/internal/impl/callbacks/trace_callback.go b/internal/ai/runtime/internal/impl/callbacks/trace_callback.go index 1727e0a..7833543 100644 --- a/internal/ai/runtime/internal/impl/callbacks/trace_callback.go +++ b/internal/ai/runtime/internal/impl/callbacks/trace_callback.go @@ -58,6 +58,9 @@ type RuntimeTraceData struct { HistoryMessageCount int `json:"historyMessageCount,omitempty"` KnowledgeBaseIDs []int64 `json:"knowledgeBaseIds,omitempty"` ToolCodes []string `json:"toolCodes,omitempty"` + StaticToolCodes []string `json:"staticToolCodes,omitempty"` + DynamicToolCodes []string `json:"dynamicToolCodes,omitempty"` + ToolSearchEnabled bool `json:"toolSearchEnabled,omitempty"` CurrentUserMessagePreview string `json:"currentUserMessagePreview,omitempty"` } `json:"input"` Retriever struct {