first commit

This commit is contained in:
2026-03-08 15:40:34 +07:00
commit 8dc496b626
159 changed files with 27932 additions and 0 deletions
+259
View File
@@ -0,0 +1,259 @@
package memory
import (
"encoding/json"
"fmt"
"os"
"path/filepath"
"sort"
"strings"
"sync"
"time"
)
type Memory struct {
ID int `json:"id"`
Content string `json:"content"`
Tags []string `json:"tags,omitempty"`
CreatedAt time.Time `json:"created_at"`
LastUsed time.Time `json:"last_used"`
}
type Store struct {
mu sync.Mutex
path string
memories []Memory
nextID int
}
func NewStore(path string) *Store {
if path == "" {
home, err := os.UserHomeDir()
if err != nil {
home = "."
}
path = filepath.Join(home, ".config", "ai-agent", "memories.json")
}
s := &Store{path: path}
s.load()
return s
}
func (s *Store) Save(content string, tags []string) (int, error) {
s.mu.Lock()
defer s.mu.Unlock()
s.nextID++
mem := Memory{
ID: s.nextID,
Content: content,
Tags: tags,
CreatedAt: time.Now(),
LastUsed: time.Now(),
}
s.memories = append(s.memories, mem)
if err := s.persist(); err != nil {
return 0, err
}
return mem.ID, nil
}
func (s *Store) Recall(query string, maxResults int) []Memory {
s.mu.Lock()
defer s.mu.Unlock()
if maxResults <= 0 {
maxResults = 5
}
queryLower := strings.ToLower(query)
words := strings.Fields(queryLower)
type scored struct {
mem Memory
score int
}
var results []scored
for i := range s.memories {
mem := s.memories[i]
score := 0
contentLower := strings.ToLower(mem.Content)
for _, w := range words {
if strings.Contains(contentLower, w) {
score += 2
}
}
for _, tag := range mem.Tags {
tagLower := strings.ToLower(tag)
for _, w := range words {
if strings.Contains(tagLower, w) {
score += 3
}
}
}
if score > 0 {
results = append(results, scored{mem: mem, score: score})
}
}
sort.Slice(results, func(i, j int) bool {
if results[i].score != results[j].score {
return results[i].score > results[j].score
}
return results[i].mem.LastUsed.After(results[j].mem.LastUsed)
})
if len(results) > maxResults {
results = results[:maxResults]
}
now := time.Now()
out := make([]Memory, len(results))
for i, r := range results {
out[i] = r.mem
for j := range s.memories {
if s.memories[j].ID == r.mem.ID {
s.memories[j].LastUsed = now
break
}
}
}
_ = s.persist()
return out
}
func (s *Store) Recent(n int) []Memory {
s.mu.Lock()
defer s.mu.Unlock()
if len(s.memories) == 0 {
return nil
}
sorted := make([]Memory, len(s.memories))
copy(sorted, s.memories)
sort.Slice(sorted, func(i, j int) bool {
return sorted[i].LastUsed.After(sorted[j].LastUsed)
})
if n > len(sorted) {
n = len(sorted)
}
return sorted[:n]
}
func (s *Store) Count() int {
s.mu.Lock()
defer s.mu.Unlock()
return len(s.memories)
}
func (s *Store) Delete(id int) (bool, error) {
s.mu.Lock()
defer s.mu.Unlock()
for i, mem := range s.memories {
if mem.ID == id {
s.memories = append(s.memories[:i], s.memories[i+1:]...)
return true, s.persist()
}
}
return false, nil
}
func (s *Store) DeleteByTag(tag string) (int, error) {
s.mu.Lock()
defer s.mu.Unlock()
tagLower := strings.ToLower(tag)
var remaining []Memory
deleted := 0
for _, mem := range s.memories {
found := false
for _, t := range mem.Tags {
if strings.ToLower(t) == tagLower {
found = true
break
}
}
if found {
deleted++
} else {
remaining = append(remaining, mem)
}
}
s.memories = remaining
if deleted > 0 {
return deleted, s.persist()
}
return deleted, nil
}
func (s *Store) Update(id int, content string, tags []string) (bool, error) {
s.mu.Lock()
defer s.mu.Unlock()
for i, mem := range s.memories {
if mem.ID == id {
if content != "" {
s.memories[i].Content = content
}
if tags != nil {
s.memories[i].Tags = tags
}
s.memories[i].LastUsed = time.Now()
return true, s.persist()
}
}
return false, nil
}
func (s *Store) Prune(olderThan time.Duration) (int, error) {
s.mu.Lock()
defer s.mu.Unlock()
cutoff := time.Now().Add(-olderThan)
var remaining []Memory
deleted := 0
for _, mem := range s.memories {
if mem.CreatedAt.Before(cutoff) {
deleted++
} else {
remaining = append(remaining, mem)
}
}
s.memories = remaining
if deleted > 0 {
return deleted, s.persist()
}
return deleted, nil
}
func (s *Store) Get(id int) (Memory, bool) {
s.mu.Lock()
defer s.mu.Unlock()
for _, mem := range s.memories {
if mem.ID == id {
return mem, true
}
}
return Memory{}, false
}
func (s *Store) load() {
data, err := os.ReadFile(s.path)
if err != nil {
return
}
var memories []Memory
if err := json.Unmarshal(data, &memories); err != nil {
return
}
s.memories = memories
for _, m := range s.memories {
if m.ID > s.nextID {
s.nextID = m.ID
}
}
}
func (s *Store) persist() error {
dir := filepath.Dir(s.path)
if err := os.MkdirAll(dir, 0o755); err != nil {
return fmt.Errorf("create memory dir: %w", err)
}
data, err := json.MarshalIndent(s.memories, "", " ")
if err != nil {
return fmt.Errorf("marshal memories: %w", err)
}
if err := os.WriteFile(s.path, data, 0o644); err != nil {
return fmt.Errorf("write memories: %w", err)
}
return nil
}
+381
View File
@@ -0,0 +1,381 @@
package memory
import (
"path/filepath"
"testing"
"time"
)
func TestStore_Save_And_Count(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "memories.json")
s := NewStore(path)
if s.Count() != 0 {
t.Fatalf("new store Count = %d, want 0", s.Count())
}
id1, err := s.Save("first memory", []string{"tag1"})
if err != nil {
t.Fatalf("Save returned error: %v", err)
}
if id1 != 1 {
t.Errorf("first Save id = %d, want 1", id1)
}
if s.Count() != 1 {
t.Errorf("Count after first Save = %d, want 1", s.Count())
}
id2, err := s.Save("second memory", []string{"tag2"})
if err != nil {
t.Fatalf("Save returned error: %v", err)
}
if id2 != 2 {
t.Errorf("second Save id = %d, want 2", id2)
}
if s.Count() != 2 {
t.Errorf("Count after second Save = %d, want 2", s.Count())
}
}
func TestStore_Recall(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "memories.json")
s := NewStore(path)
s.Save("the user prefers Go language", []string{"preference", "golang"})
s.Save("project uses PostgreSQL database", []string{"tech", "database"})
s.Save("user name is Alice", []string{"name"})
t.Run("content match", func(t *testing.T) {
results := s.Recall("Go", 10)
if len(results) == 0 {
t.Fatal("expected results for 'Go' query")
}
found := false
for _, r := range results {
if r.Content == "the user prefers Go language" {
found = true
}
}
if !found {
t.Error("expected to find 'the user prefers Go language'")
}
})
t.Run("tag match", func(t *testing.T) {
results := s.Recall("golang", 10)
if len(results) == 0 {
t.Fatal("expected results for 'golang' tag query")
}
if results[0].Content != "the user prefers Go language" {
t.Errorf("top result = %q, want 'the user prefers Go language'", results[0].Content)
}
})
t.Run("combined scoring", func(t *testing.T) {
// "database" matches both content and tag for PostgreSQL entry.
results := s.Recall("database", 10)
if len(results) == 0 {
t.Fatal("expected results for 'database' query")
}
if results[0].Content != "project uses PostgreSQL database" {
t.Errorf("top result = %q, want 'project uses PostgreSQL database'",
results[0].Content)
}
})
t.Run("maxResults limit", func(t *testing.T) {
results := s.Recall("user", 1)
if len(results) > 1 {
t.Errorf("maxResults=1 but got %d results", len(results))
}
})
t.Run("default maxResults 5 when 0", func(t *testing.T) {
// With 3 memories, should return all 3 (default limit is 5).
results := s.Recall("user", 0)
if len(results) > 5 {
t.Errorf("default maxResults should be 5, got %d results", len(results))
}
})
t.Run("case insensitive", func(t *testing.T) {
results := s.Recall("ALICE", 10)
if len(results) == 0 {
t.Fatal("expected case-insensitive match for 'ALICE'")
}
if results[0].Content != "user name is Alice" {
t.Errorf("result = %q, want 'user name is Alice'", results[0].Content)
}
})
t.Run("no matches", func(t *testing.T) {
results := s.Recall("xyzzyzxyz", 10)
if len(results) != 0 {
t.Errorf("expected no results for nonsense query, got %d", len(results))
}
})
}
func TestStore_Recall_TieBreakByRecency(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "memories.json")
s := NewStore(path)
// Save two memories with the same scoring potential.
s.Save("alpha topic info", []string{"info"})
// Small delay so LastUsed differs.
time.Sleep(10 * time.Millisecond)
s.Save("beta topic info", []string{"info"})
results := s.Recall("info", 10)
if len(results) < 2 {
t.Fatalf("expected at least 2 results, got %d", len(results))
}
// Both match tag "info" equally (+3), so more recent (beta) should come first.
if results[0].Content != "beta topic info" {
t.Errorf("expected more recent 'beta topic info' first, got %q", results[0].Content)
}
}
func TestStore_Recent(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "memories.json")
s := NewStore(path)
s.Save("old memory", nil)
time.Sleep(10 * time.Millisecond)
s.Save("new memory", nil)
t.Run("ordering by LastUsed", func(t *testing.T) {
recent := s.Recent(2)
if len(recent) != 2 {
t.Fatalf("Recent(2) returned %d, want 2", len(recent))
}
if recent[0].Content != "new memory" {
t.Errorf("first recent = %q, want 'new memory'", recent[0].Content)
}
if recent[1].Content != "old memory" {
t.Errorf("second recent = %q, want 'old memory'", recent[1].Content)
}
})
t.Run("limit exceeds count returns all", func(t *testing.T) {
recent := s.Recent(100)
if len(recent) != 2 {
t.Errorf("Recent(100) returned %d, want 2", len(recent))
}
})
t.Run("empty store", func(t *testing.T) {
emptyPath := filepath.Join(dir, "empty.json")
empty := NewStore(emptyPath)
recent := empty.Recent(5)
if recent != nil {
t.Errorf("empty Recent should return nil, got %v", recent)
}
})
}
func TestStore_Persistence_RoundTrip(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "memories.json")
s1 := NewStore(path)
s1.Save("persistent memory", []string{"test"})
s1.Save("another memory", []string{"test2"})
// Create new store from same path.
s2 := NewStore(path)
if s2.Count() != 2 {
t.Errorf("reloaded Count = %d, want 2", s2.Count())
}
// Verify data is intact.
recent := s2.Recent(2)
contents := map[string]bool{}
for _, m := range recent {
contents[m.Content] = true
}
if !contents["persistent memory"] {
t.Error("missing 'persistent memory' after reload")
}
if !contents["another memory"] {
t.Error("missing 'another memory' after reload")
}
// Verify IDs continue.
id, err := s2.Save("third", nil)
if err != nil {
t.Fatalf("Save after reload: %v", err)
}
if id != 3 {
t.Errorf("continued id = %d, want 3", id)
}
}
func TestStore_Delete(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "memories.json")
s := NewStore(path)
id, _ := s.Save("to be deleted", []string{"temp"})
if s.Count() != 1 {
t.Fatalf("expected 1 memory, got %d", s.Count())
}
deleted, err := s.Delete(id)
if err != nil {
t.Fatalf("Delete returned error: %v", err)
}
if !deleted {
t.Error("Delete returned false for existing memory")
}
if s.Count() != 0 {
t.Errorf("Count after delete = %d, want 0", s.Count())
}
// Try deleting non-existent.
deleted, err = s.Delete(999)
if err != nil {
t.Fatalf("Delete returned error: %v", err)
}
if deleted {
t.Error("Delete should return false for non-existent memory")
}
}
func TestStore_Update(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "memories.json")
s := NewStore(path)
id, _ := s.Save("original content", []string{"original"})
updated, err := s.Update(id, "updated content", []string{"updated"})
if err != nil {
t.Fatalf("Update returned error: %v", err)
}
if !updated {
t.Error("Update returned false for existing memory")
}
// Verify update.
mem, found := s.Get(id)
if !found {
t.Fatal("memory not found after update")
}
if mem.Content != "updated content" {
t.Errorf("Content = %q, want 'updated content'", mem.Content)
}
if len(mem.Tags) != 1 || mem.Tags[0] != "updated" {
t.Errorf("Tags = %v, want ['updated']", mem.Tags)
}
// Try updating non-existent.
updated, err = s.Update(999, "test", nil)
if err != nil {
t.Fatalf("Update returned error: %v", err)
}
if updated {
t.Error("Update should return false for non-existent memory")
}
}
func TestStore_DeleteByTag(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "memories.json")
s := NewStore(path)
s.Save("keep this 1", []string{"keep"})
s.Save("delete this", []string{"temp"})
s.Save("keep this 2", []string{"keep"})
s.Save("delete this too", []string{"temp"})
s.Save("also keep", []string{"permanent"})
deleted, err := s.DeleteByTag("temp")
if err != nil {
t.Fatalf("DeleteByTag returned error: %v", err)
}
if deleted != 2 {
t.Errorf("DeleteByTag deleted = %d, want 2", deleted)
}
if s.Count() != 3 {
t.Errorf("Count after delete = %d, want 3", s.Count())
}
// Verify only temp memories are gone.
results := s.Recall("keep", 10)
if len(results) != 3 {
t.Errorf("Recall returned %d, want 3", len(results))
}
}
func TestStore_Get(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "memories.json")
s := NewStore(path)
id, _ := s.Save("test memory", []string{"tag"})
mem, found := s.Get(id)
if !found {
t.Fatal("Get returned false for existing memory")
}
if mem.Content != "test memory" {
t.Errorf("Content = %q, want 'test memory'", mem.Content)
}
if len(mem.Tags) != 1 || mem.Tags[0] != "tag" {
t.Errorf("Tags = %v, want ['tag']", mem.Tags)
}
// Try getting non-existent.
_, found = s.Get(999)
if found {
t.Error("Get should return false for non-existent memory")
}
}
func TestStore_UpdatePartial(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "memories.json")
s := NewStore(path)
id, _ := s.Save("original content", []string{"original", "tags"})
// Update only content, keep tags.
updated, err := s.Update(id, "new content", nil)
if err != nil {
t.Fatalf("Update returned error: %v", err)
}
if !updated {
t.Error("Update returned false")
}
mem, _ := s.Get(id)
if mem.Content != "new content" {
t.Errorf("Content = %q, want 'new content'", mem.Content)
}
// Tags should remain unchanged when nil is passed.
if len(mem.Tags) != 2 {
t.Errorf("Tags = %v, want 2 tags", mem.Tags)
}
// Update only tags, keep content.
updated, err = s.Update(id, "", []string{"only", "tags"})
if err != nil {
t.Fatalf("Update returned error: %v", err)
}
if !updated {
t.Error("Update returned false")
}
mem, _ = s.Get(id)
if mem.Content != "new content" {
t.Errorf("Content changed unexpectedly to %q", mem.Content)
}
if len(mem.Tags) != 2 || mem.Tags[0] != "only" || mem.Tags[1] != "tags" {
t.Errorf("Tags = %v, want ['only', 'tags']", mem.Tags)
}
}
+102
View File
@@ -0,0 +1,102 @@
package memory
import (
"ai-agent/internal/llm"
)
func BuiltinToolDefs() []llm.ToolDef {
return []llm.ToolDef{
{
Name: "memory_save",
Description: "Save an important fact, user preference, or piece of context to persistent memory. Use this proactively when the user shares information worth remembering across sessions.",
Parameters: map[string]any{
"type": "object",
"properties": map[string]any{
"content": map[string]any{
"type": "string",
"description": "The fact or information to remember.",
},
"tags": map[string]any{
"type": "array",
"items": map[string]any{"type": "string"},
"description": "Optional tags for categorization (e.g., 'preference', 'project', 'name').",
},
},
"required": []string{"content"},
},
},
{
Name: "memory_recall",
Description: "Search persistent memory for previously saved facts. Use this when you need to recall user preferences, project details, or other saved context.",
Parameters: map[string]any{
"type": "object",
"properties": map[string]any{
"query": map[string]any{
"type": "string",
"description": "Search query to find relevant memories.",
},
},
"required": []string{"query"},
},
},
{
Name: "memory_delete",
Description: "Delete a memory by its ID. Use memory_recall or memory_list first to find the ID of the memory you want to delete.",
Parameters: map[string]any{
"type": "object",
"properties": map[string]any{
"id": map[string]any{
"type": "integer",
"description": "The ID of the memory to delete (use memory_list or memory_recall to find IDs).",
},
},
"required": []string{"id"},
},
},
{
Name: "memory_update",
Description: "Update an existing memory's content or tags. Use memory_recall or memory_list first to find the ID.",
Parameters: map[string]any{
"type": "object",
"properties": map[string]any{
"id": map[string]any{
"type": "integer",
"description": "The ID of the memory to update (use memory_list or memory_recall to find IDs).",
},
"content": map[string]any{
"type": "string",
"description": "New content for the memory.",
},
"tags": map[string]any{
"type": "array",
"items": map[string]any{"type": "string"},
"description": "New tags for the memory.",
},
},
"required": []string{"id"},
},
},
{
Name: "memory_list",
Description: "List all stored memories with their IDs, content, and tags. Use this to see what has been saved.",
Parameters: map[string]any{
"type": "object",
"properties": map[string]any{
"limit": map[string]any{
"type": "integer",
"description": "Maximum number of memories to return (default: 20).",
},
},
},
},
}
}
func IsBuiltinTool(name string) bool {
switch name {
case "memory_save", "memory_recall", "memory_delete", "memory_update", "memory_list":
return true
default:
return false
}
}
+48
View File
@@ -0,0 +1,48 @@
package memory
import "testing"
func TestBuiltinToolDefs(t *testing.T) {
defs := BuiltinToolDefs()
if len(defs) != 5 {
t.Fatalf("BuiltinToolDefs() returned %d defs, want 5", len(defs))
}
names := map[string]bool{}
for _, d := range defs {
names[d.Name] = true
}
expected := []string{"memory_save", "memory_recall", "memory_delete", "memory_update", "memory_list"}
for _, name := range expected {
if !names[name] {
t.Errorf("missing %s tool definition", name)
}
}
}
func TestIsBuiltinTool(t *testing.T) {
tests := []struct {
name string
tool string
want bool
}{
{name: "memory_save", tool: "memory_save", want: true},
{name: "memory_recall", tool: "memory_recall", want: true},
{name: "memory_delete", tool: "memory_delete", want: true},
{name: "memory_update", tool: "memory_update", want: true},
{name: "memory_list", tool: "memory_list", want: true},
{name: "unknown tool", tool: "unknown", want: false},
{name: "empty string", tool: "", want: false},
{name: "partial match", tool: "memory_", want: false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := IsBuiltinTool(tt.tool)
if got != tt.want {
t.Errorf("IsBuiltinTool(%q) = %v, want %v", tt.tool, got, tt.want)
}
})
}
}