YunShu/tool.go

package main

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
	"time"
)

var registeredTools = make(map[string]*ToolDef)

func RegisterTool(td *ToolDef) {
	registeredTools[td.Name] = td
}

func ListRegisteredTools() []*ToolDef {
	list := make([]*ToolDef, 0, len(registeredTools))
	for _, td := range registeredTools {
		list = append(list, td)
	}
	return list
}

func GetToolDefs(names []string) []ToolDef {
	defs := make([]ToolDef, 0, len(names))
	for _, name := range names {
		if td, ok := registeredTools[name]; ok {
			defs = append(defs, *td)
		}
	}
	return defs
}

func ExecuteTool(tc ToolCall) (string, error) {
	td, ok := registeredTools[tc.Function.Name]
	if !ok {
		return "", fmt.Errorf("未知工具: %s", tc.Function.Name)
	}

	var args map[string]interface{}
	if err := json.Unmarshal([]byte(tc.Function.Arguments), &args); err != nil {
		return "", fmt.Errorf("解析工具参数失败: %w", err)
	}

	return td.Execute(args)
}

func init() {
	RegisterTool(&ToolDef{
		Name:        "http-get",
		Description: "发送 HTTP GET 请求获取数据",
		Parameters: ToolParameter{
			Type: "object",
			Properties: map[string]ToolProperty{
				"url":     {Type: "string", Description: "请求的完整 URL 地址"},
				"headers": {Type: "string", Description: "请求头，JSON 格式的键值对字符串，如 {\"User-Agent\": \"...\"}"},
			},
			Required: []string{"url"},
		},
		Execute: func(args map[string]interface{}) (string, error) {
			url, _ := args["url"].(string)
			if url == "" {
				return "", fmt.Errorf("缺少 url 参数")
			}

			req, err := http.NewRequest("GET", url, nil)
			if err != nil {
				return "", fmt.Errorf("创建请求失败: %w", err)
			}

			if headersStr, ok := args["headers"].(string); ok && headersStr != "" {
				var headers map[string]string
				if err := json.Unmarshal([]byte(headersStr), &headers); err == nil {
					for k, v := range headers {
						req.Header.Set(k, v)
					}
				}
			}

			client := &http.Client{Timeout: 15 * time.Second}
			resp, err := client.Do(req)
			if err != nil {
				return "", fmt.Errorf("请求失败: %w", err)
			}
			defer resp.Body.Close()

			body, err := io.ReadAll(resp.Body)
			if err != nil {
				return "", fmt.Errorf("读取响应失败: %w", err)
			}

			if resp.StatusCode != 200 {
				return fmt.Sprintf("HTTP %d: %s", resp.StatusCode, string(body)), nil
			}

			return string(body), nil
		},
	})

	RegisterTool(&ToolDef{
		Name:        "skill",
		Description: "加载指定名称的 Skill 知识内容到当前上下文，获取专业知识",
		Parameters: ToolParameter{
			Type: "object",
			Properties: map[string]ToolProperty{
				"name": {Type: "string", Description: "Skill 名称，如 msn-weather-api"},
			},
			Required: []string{"name"},
		},
		Execute: func(args map[string]interface{}) (string, error) {
			name, _ := args["name"].(string)
			if name == "" {
				return "", fmt.Errorf("缺少 name 参数")
			}
			return LoadSkill(name)
		},
	})

	RegisterTool(&ToolDef{
		Name:        "read-file",
		Description: "读取本地文件内容",
		Parameters: ToolParameter{
			Type: "object",
			Properties: map[string]ToolProperty{
				"path": {Type: "string", Description: "文件路径，相对于项目根目录"},
			},
			Required: []string{"path"},
		},
		Execute: func(args map[string]interface{}) (string, error) {
			path, _ := args["path"].(string)
			if path == "" {
				return "", fmt.Errorf("缺少 path 参数")
			}
			data, err := os.ReadFile(path)
			if err != nil {
				return "", fmt.Errorf("读取文件失败: %w", err)
			}
			return strings.TrimSpace(string(data)), nil
		},
	})

	RegisterTool(&ToolDef{
		Name:        "geocode",
		Description: "查询城市或地点的经纬度坐标，返回 lat/lon/name/country。支持中文城市名（如 北京、上海、成都）和英文名",
		Parameters: ToolParameter{
			Type: "object",
			Properties: map[string]ToolProperty{
				"city": {Type: "string", Description: "城市名称，支持中文（如 北京）或英文（如 Beijing）"},
			},
			Required: []string{"city"},
		},
		Execute: func(args map[string]interface{}) (string, error) {
			city, _ := args["city"].(string)
			if city == "" {
				return "", fmt.Errorf("缺少 city 参数")
			}

			url := fmt.Sprintf("https://wttr.in/%s?format=j1", city)
			req, err := http.NewRequest("GET", url, nil)
			if err != nil {
				return "", fmt.Errorf("创建请求失败: %w", err)
			}

			client := &http.Client{Timeout: 15 * time.Second}
			resp, err := client.Do(req)
			if err != nil {
				return "", fmt.Errorf("请求 wttr.in 失败: %w", err)
			}
			defer resp.Body.Close()

			body, err := io.ReadAll(resp.Body)
			if err != nil {
				return "", fmt.Errorf("读取响应失败: %w", err)
			}

			if resp.StatusCode != 200 {
				return "", fmt.Errorf("wttr.in 返回 HTTP %d: %s", resp.StatusCode, string(body))
			}

			var result struct {
				NearestArea []struct {
					AreaName []struct {
						Value string `json:"value"`
					} `json:"areaName"`
					Country []struct {
						Value string `json:"value"`
					} `json:"country"`
					Latitude  string `json:"latitude"`
					Longitude string `json:"longitude"`
					Population string `json:"population"`
				} `json:"nearest_area"`
			}

			if err := json.Unmarshal(body, &result); err != nil {
				return "", fmt.Errorf("解析 wttr.in 响应失败: %w", err)
			}

			if len(result.NearestArea) == 0 {
				return "", fmt.Errorf("未找到城市: %s", city)
			}

			area := result.NearestArea[0]
			name := ""
			if len(area.AreaName) > 0 {
				name = area.AreaName[0].Value
			}
			country := ""
			if len(area.Country) > 0 {
				country = area.Country[0].Value
			}

			out, _ := json.Marshal(map[string]interface{}{
				"lat":     area.Latitude,
				"lon":     area.Longitude,
				"name":    name,
				"country": country,
			})
			return string(out), nil
		},
	})
}