agent-team/internal/agent/agent.go

package agent

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

	"github.com/sdaduanbilei/agent-team/internal/llm"
	"gopkg.in/yaml.v3"
)

type Config struct {
	Name        string   `yaml:"name"`
	Description string   `yaml:"description"`
	Provider    string   `yaml:"provider"`
	Model       string   `yaml:"model"`
	BaseURL     string   `yaml:"base_url"`
	APIKeyEnv   string   `yaml:"api_key_env"`
	Skills      []string `yaml:"skills"`
	CanChallenge bool    `yaml:"can_challenge"`
}

type Agent struct {
	Config       Config
	AgentDoc     string // AGENT.md 正文（frontmatter 之后的内容）
	Soul         string // SOUL.md 人设
	Dir          string // agents/<name>/
	KnowledgeDir string // agents/{team}/knowledge/ 目录路径
	client       *llm.Client
}

func Load(dir string) (*Agent, error) {
	agentMD, err := os.ReadFile(filepath.Join(dir, "AGENT.md"))
	if err != nil {
		return nil, err
	}
	cfg, body, err := parseFrontmatterWithBody(agentMD)
	if err != nil {
		return nil, fmt.Errorf("parse AGENT.md: %w", err)
	}
	soul, _ := os.ReadFile(filepath.Join(dir, "SOUL.md"))

	if cfg.Provider == "" {
		cfg.Provider = "deepseek"
	}
	if cfg.APIKeyEnv == "" {
		cfg.APIKeyEnv = "DEEPSEEK_API_KEY"
	}

	client, err := llm.New(cfg.Provider, cfg.Model, cfg.BaseURL, cfg.APIKeyEnv)
	if err != nil {
		return nil, err
	}
	return &Agent{Config: cfg, AgentDoc: body, Soul: string(soul), Dir: dir, client: client}, nil
}

func parseFrontmatterWithBody(data []byte) (Config, string, error) {
	var cfg Config
	if !bytes.HasPrefix(data, []byte("---")) {
		return cfg, "", fmt.Errorf("missing frontmatter")
	}
	parts := bytes.SplitN(data, []byte("---"), 3)
	if len(parts) < 3 {
		return cfg, "", fmt.Errorf("invalid frontmatter")
	}
	if err := yaml.Unmarshal(parts[1], &cfg); err != nil {
		return cfg, "", err
	}
	body := strings.TrimSpace(string(parts[2]))
	return cfg, body, nil
}

// memoryMaxChars 是注入 prompt 的记忆硬限制（字符数）
const memoryMaxChars = 4000

// Memory 智能加载记忆：优先 core.md，然后当月文件，硬限制 4000 字符。
func (a *Agent) Memory() string {
	memDir := filepath.Join(a.Dir, "memory")
	if _, err := os.Stat(memDir); err != nil {
		return ""
	}

	var sb strings.Builder

	// 1. 优先加载 core.md（核心经验）
	if data, err := os.ReadFile(filepath.Join(memDir, "core.md")); err == nil {
		sb.Write(data)
		sb.WriteString("\n")
	}

	// 2. 加载当月文件
	currentMonth := time.Now().Format("2006-01") + ".md"
	if data, err := os.ReadFile(filepath.Join(memDir, currentMonth)); err == nil {
		sb.Write(data)
		sb.WriteString("\n")
	}

	// 3. 硬限制截断
	result := sb.String()
	runes := []rune(result)
	if len(runes) > memoryMaxChars {
		result = string(runes[:memoryMaxChars]) + "\n...(记忆已截断)"
	}
	return result
}

// SaveMemory appends/updates a memory file.
func (a *Agent) SaveMemory(filename, content string) error {
	memDir := filepath.Join(a.Dir, "memory")
	os.MkdirAll(memDir, 0755)
	return os.WriteFile(filepath.Join(memDir, filename), []byte(content), 0644)
}

// Chat sends messages and streams response tokens via onToken callback.
func (a *Agent) Chat(ctx context.Context, msgs []llm.Message, onToken func(string)) (string, error) {
	result, err := a.client.Stream(ctx, msgs, onToken)
	if err != nil {
		return "", err
	}
	return result.Content, nil
}

// ChatWithUsage 同 Chat 但额外返回 token 用量。
func (a *Agent) ChatWithUsage(ctx context.Context, msgs []llm.Message, onToken func(string)) (string, llm.Usage, error) {
	result, err := a.client.Stream(ctx, msgs, onToken)
	if err != nil {
		return "", llm.Usage{}, err
	}
	return result.Content, result.Usage, nil
}

// ChatWithTools 支持 tool calling 的对话
func (a *Agent) ChatWithTools(ctx context.Context, msgs []llm.Message, tools []llm.Tool, onToken func(string)) (*llm.StreamResult, error) {
	return a.client.StreamWithTools(ctx, msgs, tools, onToken)
}

// knowledgeMaxChars 是注入 prompt 的知识库硬限制（字符数）
const knowledgeMaxChars = 8000

// Knowledge 读取团队知识库目录下的所有 .md 文件，拼接返回，硬限制 8000 字符。
func (a *Agent) Knowledge() string {
	if a.KnowledgeDir == "" {
		return ""
	}
	entries, err := os.ReadDir(a.KnowledgeDir)
	if err != nil {
		return ""
	}
	var sb strings.Builder
	for _, e := range entries {
		if strings.HasSuffix(e.Name(), ".md") {
			data, _ := os.ReadFile(filepath.Join(a.KnowledgeDir, e.Name()))
			if len(data) > 0 {
				sb.WriteString("### " + e.Name() + "\n")
				sb.Write(data)
				sb.WriteString("\n\n")
			}
		}
	}
	result := sb.String()
	runes := []rune(result)
	if len(runes) > knowledgeMaxChars {
		result = string(runes[:knowledgeMaxChars]) + "\n...(知识库已截断)"
	}
	return result
}

// CompressMemory 压缩非当月的记忆文件：LLM 摘要后追加到 core.md，原文件移到 archive/。
func (a *Agent) CompressMemory(ctx context.Context) error {
	memDir := filepath.Join(a.Dir, "memory")
	entries, err := os.ReadDir(memDir)
	if err != nil {
		return err
	}

	currentMonth := time.Now().Format("2006-01") + ".md"
	var toCompress []string
	for _, e := range entries {
		name := e.Name()
		// 只处理月份文件（YYYY-MM.md），跳过 core.md、当月文件、archive 目录
		if !strings.HasSuffix(name, ".md") || name == "core.md" || name == currentMonth {
			continue
		}
		// 检查是否是月份格式
		if len(name) == 10 && name[4] == '-' {
			toCompress = append(toCompress, name)
		}
	}

	if len(toCompress) == 0 {
		return nil
	}

	archiveDir := filepath.Join(memDir, "archive")
	os.MkdirAll(archiveDir, 0755)

	for _, filename := range toCompress {
		data, err := os.ReadFile(filepath.Join(memDir, filename))
		if err != nil || len(data) == 0 {
			continue
		}

		// LLM 压缩
		content := string(data)
		if len([]rune(content)) > 2000 {
			content = string([]rune(content)[:2000])
		}
		compressPrompt := fmt.Sprintf(`将以下记忆压缩为 3-5 条最核心的经验教训，每条一行 bullet point，只保留最有价值的内容：

%s`, content)
		msgs := []llm.Message{
			llm.NewMsg("system", "你是一个记忆压缩助手。将详细记录压缩为核心经验。"),
			llm.NewMsg("user", compressPrompt),
		}
		summary, err := a.Chat(ctx, msgs, nil)
		if err != nil || summary == "" {
			continue
		}

		// 追加到 core.md
		coreFile := filepath.Join(memDir, "core.md")
		existing, _ := os.ReadFile(coreFile)
		coreContent := string(existing)
		if coreContent == "" {
			coreContent = "# 核心经验\n"
		}
		coreContent += fmt.Sprintf("\n## %s 经验摘要\n\n%s\n", strings.TrimSuffix(filename, ".md"), summary)
		os.WriteFile(coreFile, []byte(coreContent), 0644)

		// 原文件移到 archive/
		os.Rename(filepath.Join(memDir, filename), filepath.Join(archiveDir, filename))
	}

	return nil
}

// BuildSystemPrompt 构建完整的 system prompt。
// SOUL.md = 人设性格，AGENT.md 正文 = 功能描述/工作流程，均实时读取支持热更新。
func (a *Agent) BuildSystemPrompt(extraContext string) string {
	// 实时读取 SOUL.md
	soul := a.Soul
	if data, err := os.ReadFile(filepath.Join(a.Dir, "SOUL.md")); err == nil {
		soul = string(data)
	}

	// 实时读取 AGENT.md 正文
	agentDoc := a.AgentDoc
	if data, err := os.ReadFile(filepath.Join(a.Dir, "AGENT.md")); err == nil {
		if _, body, err := parseFrontmatterWithBody(data); err == nil && body != "" {
			agentDoc = body
		}
	}

	var sb strings.Builder
	sb.WriteString(soul)
	if agentDoc != "" {
		sb.WriteString("\n\n")
		sb.WriteString(agentDoc)
	}

	// 知识库注入（在 memory 之前）
	if knowledge := a.Knowledge(); knowledge != "" {
		sb.WriteString("\n\n<knowledge>\n")
		sb.WriteString(knowledge)
		sb.WriteString("</knowledge>")
	}

	if mem := a.Memory(); mem != "" {
		sb.WriteString("\n\n<memory>\n")
		sb.WriteString(mem)
		sb.WriteString("</memory>")
	}
	if extraContext != "" {
		sb.WriteString("\n\n")
		sb.WriteString(extraContext)
	}
	return sb.String()
}