wlw/ai-agent
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: enhance tool search functionality and streamline tool management in AI agent

Browse Source
This commit is contained in:
mlogclub
2026-04-10 15:40:21 +08:00
parent 4722caab37
commit 31337fa714
5 changed files with 57 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/ai/runtime/internal/impl/callbacks/agent_trace_handler.go
+23 -16
View File
@@ -59,7 +59,7 @@ func (h *RuntimeTraceHandler) WrapInvokableToolCall(_ context.Context, endpoint
}
h.collector.AddToolItem(item)
if metadata, ok := h.resolveToolMetadata(item.ToolName); ok && strings.TrimSpace(metadata.ToolCode) == toolx.BuiltinToolSearchToolCode {
h.collector.AddToolSearchItem(buildToolSearchTraceItem(argumentsInJSON, result, err))
h.collector.AddToolSearchItem(h.buildToolSearchTraceItem(argumentsInJSON, result, err))
}
return result, err
}, nil
@@ -73,6 +73,13 @@ func (h *RuntimeTraceHandler) resolveToolMetadata(modelToolName string) (ToolMet
if modelToolName == "" {
return ToolMetadata{}, false
}
if modelToolName == toolx.BuiltinToolSearchToolName {
return ToolMetadata{
ToolCode: toolx.BuiltinToolSearchToolCode,
ServerCode: toolx.BuiltinToolCatalogServerCode,
ToolName: toolx.BuiltinToolSearchToolName,
}, true
}
metadata, ok := h.toolMetadataBy[modelToolName]
return metadata, ok
}
@@ -101,19 +108,13 @@ func previewToolText(text string, limit int) string {
return string(runes[:limit]) + "..."
}
func buildToolSearchTraceItem(argumentsInJSON string, result string, runErr error) ToolSearchTraceItem {
func (h *RuntimeTraceHandler) buildToolSearchTraceItem(argumentsInJSON string, result string, runErr error) ToolSearchTraceItem {
item := ToolSearchTraceItem{
Action: "search",
Status: "ok",
}
args := parseToolArguments(argumentsInJSON)
item.Query = strings.TrimSpace(readToolSearchString(args, "query"))
item.TargetToolCode = strings.TrimSpace(readToolSearchString(args, "toolCode"))
if item.TargetToolCode != "" {
item.Action = "invoke"
item.TargetServerCode, item.TargetToolName = toolx.SplitMCPToolCode(item.TargetToolCode)
} else {
item.Action = "search"
}
item.Query = strings.TrimSpace(readToolSearchString(args, "regex_pattern"))
if runErr != nil {
item.Status = "error"
item.ErrorMessage = runErr.Error()
@@ -123,8 +124,8 @@ func buildToolSearchTraceItem(argumentsInJSON string, result string, runErr erro
if err := json.Unmarshal([]byte(strings.TrimSpace(result)), &payload); err != nil {
return item
}
candidateItems, _ := payload["candidates"].([]any)
item.CandidateToolCodes = extractCandidateToolCodes(candidateItems)
selectedTools, _ := payload["selectedTools"].([]any)
item.CandidateToolCodes = h.extractSelectedToolCodes(selectedTools)
return item
}
@@ -140,18 +141,24 @@ func readToolSearchString(data map[string]any, key string) string {
return text
}
func extractCandidateToolCodes(items []any) []string {
func (h *RuntimeTraceHandler) extractSelectedToolCodes(items []any) []string {
if len(items) == 0 {
return nil
}
ret := make([]string, 0, len(items))
for _, item := range items {
payload, ok := item.(map[string]any)
toolName, ok := item.(string)
if !ok {
continue
}
toolCode, _ := payload["toolCode"].(string)
toolCode = strings.TrimSpace(toolCode)
toolName = strings.TrimSpace(toolName)
if toolName == "" {
continue
}
toolCode := toolName
if metadata, ok := h.resolveToolMetadata(toolName); ok && strings.TrimSpace(metadata.ToolCode) != "" {
toolCode = strings.TrimSpace(metadata.ToolCode)
}
if toolCode == "" {
continue
}