aichat/toolrouter/tools.go

package toolrouter

import (
	"context"
	"fmt"
	"strings"
	"time"

	searchagent "aichat/agents/search"
	sqlquery "aichat/agents/sql"
	timeagent "aichat/agents/time"
	"aichat/completion"
	"aichat/llm"
	"aichat/message"
	"aichat/stream"
	"aichat/utils"

	"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
)

type AgentTool struct {
	name       string
	definition *model.Tool
	execute    func(context.Context, string) (string, error)
}

func NewAgentTool(name string, definition *model.Tool, execute func(context.Context, string) (string, error)) AgentTool {
	return AgentTool{name: name, definition: definition, execute: execute}
}

func (t AgentTool) Name() string { return t.name }

func (t AgentTool) Definition() *model.Tool { return t.definition }

func AvailableAgentTools(state *State, profile *llm.Profile, searchState *searchagent.State, sqlState *sqlquery.State, emit stream.EmitFunc) []AgentTool {
	return availableAgentTools(state, profile, searchState, sqlState, emit)
}

func availableAgentTools(state *State, profile *llm.Profile, searchState *searchagent.State, sqlState *sqlquery.State, emit stream.EmitFunc) []AgentTool {
	if state == nil || state.cfg == nil || !state.cfg.Enabled {
		return nil
	}
	tools := make([]AgentTool, 0, len(state.cfg.Tools))
	for _, item := range state.cfg.Tools {
		if !item.Enabled {
			continue
		}
		description := strings.TrimSpace(item.Description)
		switch item.Name {
		case timeagent.ToolName:
			tools = append(tools, AgentTool{
				name:       timeagent.ToolName,
				definition: timeagent.ToolDefinition(description),
				execute: func(ctx context.Context, args string) (string, error) {
					result, err := timeagent.ExecuteTool(args, time.Now())
					if err == nil && emit != nil {
						emit(stream.Frame{Type: "trace", Tool: timeagent.ToolName, Stage: "resolve", Status: "success", Message: "已获取当前时间上下文"})
					}
					return result, err
				},
			})
		case searchagent.ToolName:
			if searchState == nil || !searchState.Enabled() {
				continue
			}
			tools = append(tools, AgentTool{
				name:       searchagent.ToolName,
				definition: searchState.ToolDefinition(description),
				execute: func(ctx context.Context, args string) (string, error) {
					if emit != nil {
						emit(stream.Frame{Type: "trace", Tool: searchagent.ToolName, Stage: "request", Status: "running", Message: "正在联网搜索"})
					}
					result, err := searchState.ExecuteTool(ctx, args)
					if emit != nil {
						status := "success"
						messageText := "联网搜索完成"
						if err != nil {
							status = "error"
							messageText = "联网搜索失败"
						}
						emit(stream.Frame{Type: "trace", Tool: searchagent.ToolName, Stage: "results", Status: status, Message: messageText})
					}
					return result, err
				},
			})
		case sqlquery.ToolName:
			if sqlState == nil || !sqlState.Enabled() {
				continue
			}
			tools = append(tools, AgentTool{
				name:       sqlquery.ToolName,
				definition: sqlState.ToolDefinition(description),
				execute: func(ctx context.Context, args string) (string, error) {
					if emit != nil {
						emit(stream.Frame{Type: "trace", Tool: sqlquery.ToolName, Stage: "execute", Status: "running", Message: "正在查询数据库"})
					}
					generator := func(ctx context.Context, prompt string, maxTokens int) (string, error) {
						return completion.CompleteText(ctx, profile, []message.ChatMessage{{Role: "system", Content: prompt}}, maxTokens)
					}
					result, err := sqlState.ExecuteTool(ctx, args, generator)
					if emit != nil {
						status := "success"
						messageText := "数据库查询完成"
						if err != nil {
							status = "error"
							messageText = "数据库查询失败"
						}
						emit(stream.Frame{Type: "trace", Tool: sqlquery.ToolName, Stage: "execute", Status: status, Message: messageText})
					}
					return result, err
				},
			})
		}
	}
	return tools
}

func ExecuteAgentToolCall(ctx context.Context, call *model.ToolCall, tools map[string]AgentTool, emit stream.EmitFunc) string {
	if call == nil || call.Type != model.ToolTypeFunction {
		result := "工具调用无效：仅支持 function 类型工具。"
		if emit != nil {
			emit(stream.Frame{Type: "trace", Tool: "agent_tools", Stage: "execute", Status: "error", Message: result})
		}
		return result
	}
	toolName := call.Function.Name
	if emit != nil {
		emit(stream.Frame{Type: "trace", Tool: toolName, Stage: "arguments", Status: "running", Message: "准备执行工具", Data: map[string]any{"tool_call_id": call.ID, "arguments": call.Function.Arguments}})
	}
	tool, ok := tools[toolName]
	if !ok {
		result := fmt.Sprintf("工具调用失败：未知工具 %s。", toolName)
		if emit != nil {
			emit(stream.Frame{Type: "trace", Tool: toolName, Stage: "execute", Status: "error", Message: result})
		}
		return result
	}
	started := time.Now()
	result, err := tool.execute(ctx, call.Function.Arguments)
	durationMs := time.Since(started).Milliseconds()
	if err != nil {
		messageText := fmt.Sprintf("工具 %s 执行失败：%v", tool.name, err)
		if emit != nil {
			emit(stream.Frame{Type: "trace", Tool: tool.name, Stage: "execute", Status: "error", Message: "工具执行失败", Data: map[string]any{"tool_call_id": call.ID, "duration_ms": durationMs, "error": err.Error()}})
		}
		return messageText
	}
	if strings.TrimSpace(result) == "" {
		result = fmt.Sprintf("工具 %s 执行完成，但没有返回内容。", tool.name)
	}
	if emit != nil {
		emit(stream.Frame{Type: "trace", Tool: tool.name, Stage: "result", Status: "success", Message: "工具执行完成", Data: map[string]any{"tool_call_id": call.ID, "duration_ms": durationMs, "result_preview": utils.TruncateString(result, 1200)}})
	}
	return result
}