YunShu/catalog.go

package main

import (
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"gopkg.in/yaml.v3"
)

// ============================================================
// YAML 目录结构
// ============================================================

type ToolsCatalog struct {
	Version string          `yaml:"version"`
	Auto    CatalogSection  `yaml:"auto"`
	Manual  CatalogSection  `yaml:"manual,omitempty"`
}

type CatalogSection struct {
	Tools  []CatalogTool  `yaml:"tools"`
	Skills []CatalogSkill `yaml:"skills"`
	Agents []CatalogAgent `yaml:"agents"`
}

type CatalogTool struct {
	Name        string                    `yaml:"name"`
	Type        string                    `yaml:"type"`
	Status      string                    `yaml:"status"`
	Description string                    `yaml:"description"`
	Parameters  map[string]ParameterField `yaml:"parameters"`
	Returns     string                    `yaml:"returns"`
	Source      string                    `yaml:"source"`
}

type ParameterField struct {
	Type        string `yaml:"type"`
	Required    bool   `yaml:"required"`
	Description string `yaml:"description"`
}

type CatalogSkill struct {
	Name        string `yaml:"name"`
	Path        string `yaml:"path"`
	Description string `yaml:"description"`
	Status      string `yaml:"status"`
}

type CatalogAgent struct {
	Name        string   `yaml:"name"`
	Type        string   `yaml:"type"`
	Path        string   `yaml:"path"`
	Description string   `yaml:"description"`
	Tools       []string `yaml:"tools"`
	Status      string   `yaml:"status"`
}

// ============================================================
// 构建目录
// ============================================================

func BuildCatalog() *ToolsCatalog {
	return &ToolsCatalog{
		Version: "1.0",
		Auto: CatalogSection{
			Tools:  buildToolList(),
			Skills: scanSkills(),
			Agents: scanAgents(),
		},
	}
}

func buildToolList() []CatalogTool {
	tools := ListRegisteredTools()
	list := make([]CatalogTool, 0, len(tools))
	for _, t := range tools {
		ct := CatalogTool{
			Name:        t.Name,
			Type:        "builtin",
			Status:      "active",
			Description: t.Description,
			Source:      "src/tool.go",
		}

		// 从 Schema (map[string]any) 提取参数
		ct.Parameters = make(map[string]ParameterField)
		if t.Parameters == nil {
			list = append(list, ct)
			continue
		}

		props, _ := t.Parameters["properties"].(map[string]any)
		required := make(map[string]bool)
		if reqList, ok := t.Parameters["required"].([]any); ok {
			for _, r := range reqList {
				if s, ok := r.(string); ok {
					required[s] = true
				}
			}
		}

		for name, propRaw := range props {
			prop, ok := propRaw.(map[string]any)
			if !ok {
				continue
			}
			typ, _ := prop["type"].(string)
			desc, _ := prop["description"].(string)
			ct.Parameters[name] = ParameterField{
				Type:        typ,
				Required:    required[name],
				Description: desc,
			}
		}

		list = append(list, ct)
	}
	return list
}

func scanSkills() []CatalogSkill {
	// 搜索路径：项目目录 → 用户配置目录
	dirs := []string{
		"skills",
		filepath.Join(ConfigDir(), "skills"),
	}

	seen := make(map[string]bool)
	var list []CatalogSkill

	for _, dir := range dirs {
		entries, err := os.ReadDir(dir)
		if err != nil {
			continue
		}
		for _, e := range entries {
			if !e.IsDir() || seen[e.Name()] {
				continue
			}
			seen[e.Name()] = true

			// 尝试读取 SKILL.md 获取描述
			skillPath := filepath.Join(dir, e.Name(), "SKILL.md")
			desc := fmt.Sprintf("skill: %s", e.Name())

			if data, err := os.ReadFile(skillPath); err == nil {
				content := string(data)
				if fm, _, err := parseFrontmatterSimple(content); err == nil {
					if d, ok := fm["description"]; ok {
						desc = fmt.Sprintf("%v", d)
					}
				}
			}

			list = append(list, CatalogSkill{
				Name:        e.Name(),
				Path:        fmt.Sprintf("skills/%s/SKILL.md", e.Name()),
				Description: desc,
				Status:      "active",
			})
		}
	}
	return list
}

func scanAgents() []CatalogAgent {
	dirs := []string{
		"agents",
		filepath.Join(ConfigDir(), "agents"),
	}

	seen := make(map[string]bool)
	var list []CatalogAgent

	for _, dir := range dirs {
		entries, err := os.ReadDir(dir)
		if err != nil {
			continue
		}
		for _, e := range entries {
			if e.IsDir() || !strings.HasSuffix(e.Name(), ".md") || seen[e.Name()] {
				continue
			}
			seen[e.Name()] = true

			agentPath := filepath.Join(dir, e.Name())
			data, err := os.ReadFile(agentPath)
			if err != nil {
				continue
			}

			fm, _, err := parseFrontmatterSimple(string(data))
			if err != nil {
				continue
			}

			agentType := "main"
			if t, ok := fm["type"]; ok {
				agentType = fmt.Sprintf("%v", t)
			}

			cat := CatalogAgent{
				Name:   fmt.Sprintf("%v", fm["name"]),
				Type:   agentType,
				Path:   fmt.Sprintf("agents/%s", e.Name()),
				Status: "active",
			}

			if d, ok := fm["description"]; ok {
				cat.Description = fmt.Sprintf("%v", d)
			}
			if tools, ok := fm["tools"]; ok {
				if list, ok := tools.([]interface{}); ok {
					for _, t := range list {
						cat.Tools = append(cat.Tools, fmt.Sprintf("%v", t))
					}
				}
			}

			list = append(list, cat)
		}
	}
	return list
}

// ============================================================
// 写入 tools.yml
// ============================================================

func GenerateToolsYAML() {
	catalog := BuildCatalog()

	// 尝试读取已有的 manual 节，保留用户编辑
	existingPath := filepath.Join(ConfigDir(), "tools.yml")
	if data, err := os.ReadFile(existingPath); err == nil {
		var existing ToolsCatalog
		if err := yaml.Unmarshal(data, &existing); err == nil {
			catalog.Manual = existing.Manual
		}
	}

	dir := ConfigDir()
	if err := os.MkdirAll(dir, 0755); err != nil {
		return
	}

	// 用 2 空格缩进编码
	var buf bytes.Buffer
	encoder := yaml.NewEncoder(&buf)
	encoder.SetIndent(2)
	if err := encoder.Encode(catalog); err != nil {
		return
	}
	encoder.Close()

	// 在文件头加注释
	header := "# ⚠️ auto 节由云枢 Agent 自动生成，每次启动覆写；manual 节保留用户自定义\n"
	yamlPath := filepath.Join(dir, "tools.yml")
	os.WriteFile(yamlPath, []byte(header+buf.String()), 0644)
}

// ============================================================
// 注入目录到 system prompt
// ============================================================

// BuildSubAgentPrompt 生成可用子 Agent 列表，动态注入到主 Agent 的 system prompt
func BuildSubAgentPrompt(subs []*AgentDef) string {
	if len(subs) == 0 {
		return ""
	}
	var b strings.Builder
	b.WriteString("\n\n## 可用子 Agent\n")
	b.WriteString("\n以下子 Agent 可通过 task 工具调度：\n")
	for _, s := range subs {
		b.WriteString(fmt.Sprintf("- **%s**: %s\n", s.Name, s.Description))
	}
	b.WriteString("\n用 `task(\"agent_name\", {args})` 调度。不自己回答领域问题。\n")
	return b.String()
}

// BuildInjectPrompt 生成能力边界目录，追加到 system prompt 末尾
func BuildInjectPrompt(toolNames []string) string {
	var b strings.Builder
	b.WriteString("\n\n## 能力边界\n\n")

	// 工具
	allTools := ListRegisteredTools()
	b.WriteString("### 可用工具\n")
	for _, t := range allTools {
		available := false
		for _, name := range toolNames {
			if t.Name == name {
				available = true
				break
			}
		}
		if available {
			b.WriteString(fmt.Sprintf("- %s: %s\n", t.Name, t.Description))
		}
	}

	// skill
	skills := scanSkills()
	if len(skills) > 0 {
		b.WriteString("\n### 可用技能\n")
		for _, s := range skills {
			b.WriteString(fmt.Sprintf("- %s: %s\n", s.Name, s.Description))
		}
	}

	// 数据文件
	b.WriteString("\n### 数据文件\n")
	b.WriteString("- data/tools.yml: 完整工具/技能/Agent 目录（可读此文件获取详细参数）\n")

	return b.String()
}

// ============================================================
// 辅助：简易 frontmatter 解析（只读 YAML map）
// ============================================================

func parseFrontmatterSimple(content string) (map[string]interface{}, string, error) {
	lines := strings.Split(content, "\n")
	if len(lines) < 2 || strings.TrimSpace(lines[0]) != "---" {
		return nil, content, nil
	}
	endIdx := -1
	for i := 1; i < len(lines); i++ {
		if strings.TrimSpace(lines[i]) == "---" {
			endIdx = i
			break
		}
	}
	if endIdx == -1 {
		return nil, content, nil
	}
	frontmatter := strings.Join(lines[1:endIdx], "\n")
	body := strings.Join(lines[endIdx+1:], "\n")

	var fm map[string]interface{}
	if err := yaml.Unmarshal([]byte(frontmatter), &fm); err != nil {
		return nil, body, nil
	}
	return fm, body, nil
}