286 lines
8.0 KiB
Go
286 lines
8.0 KiB
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/ 目录路径
|
||
UseDBConfig bool // 为 true 时 BuildSystemPrompt 不从文件系统热重载
|
||
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 正文 = 功能描述/工作流程。
|
||
// UseDBConfig=true 时使用内存中的值（来自 SQLite），否则从文件系统热重载。
|
||
func (a *Agent) BuildSystemPrompt(extraContext string) string {
|
||
soul := a.Soul
|
||
agentDoc := a.AgentDoc
|
||
|
||
if !a.UseDBConfig {
|
||
// 文件系统热重载（仅在未使用 DB 配置时）
|
||
if data, err := os.ReadFile(filepath.Join(a.Dir, "SOUL.md")); err == nil {
|
||
soul = string(data)
|
||
}
|
||
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()
|
||
}
|