115 lines
2.9 KiB
Go
115 lines
2.9 KiB
Go
package agent
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
|
|
"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
|
|
Soul string // system prompt from SOUL.md
|
|
Dir string // agents/<name>/
|
|
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, err := parseFrontmatter(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, Soul: string(soul), Dir: dir, client: client}, nil
|
|
}
|
|
|
|
func parseFrontmatter(data []byte) (Config, 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")
|
|
}
|
|
return cfg, yaml.Unmarshal(parts[1], &cfg)
|
|
}
|
|
|
|
// Memory returns concatenated memory files content.
|
|
func (a *Agent) Memory() string {
|
|
memDir := filepath.Join(a.Dir, "memory")
|
|
entries, err := os.ReadDir(memDir)
|
|
if err != nil {
|
|
return ""
|
|
}
|
|
var sb strings.Builder
|
|
for _, e := range entries {
|
|
if strings.HasSuffix(e.Name(), ".md") {
|
|
data, _ := os.ReadFile(filepath.Join(memDir, e.Name()))
|
|
sb.Write(data)
|
|
sb.WriteString("\n")
|
|
}
|
|
}
|
|
return sb.String()
|
|
}
|
|
|
|
// 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) {
|
|
return a.client.Stream(ctx, msgs, onToken)
|
|
}
|
|
|
|
// BuildSystemPrompt constructs the full system prompt with soul + memory + injected context.
|
|
func (a *Agent) BuildSystemPrompt(extraContext string) string {
|
|
var sb strings.Builder
|
|
sb.WriteString(a.Soul)
|
|
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()
|
|
}
|