Merge pull request #1971 from docker/memory-improvements

Add search, update, categories, and default path to memory tool

Author: dgageot

docs/configuration/tools/index.md

@@ -101,17 +101,29 @@ toolsets:
 
 ### Memory
 
-Persistent key-value storage backed by SQLite. Data survives across sessions, letting agents remember context, user preferences, and past decisions.
+Persistent key-value storage backed by SQLite. Data survives across sessions, letting agents remember context, user preferences, and past decisions. Memories can be organized with categories and searched by keyword.
+
+Each agent gets its own database at `~/.cagent/memory/<agent-name>/memory.db` by default.
 
 ```yaml
 toolsets:
   - type: memory
-    path: ./agent_memory.db # optional: custom database path
+    path: ./agent_memory.db # optional: override the default location
 ```
 
-| Property | Type   | Default   | Description                                                            |
-| -------- | ------ | --------- | ---------------------------------------------------------------------- |
-| `path`   | string | automatic | Path to the SQLite database file. If omitted, uses a default location. |
+| Property | Type   | Default                                      | Description                          |
+| -------- | ------ | -------------------------------------------- | ------------------------------------ |
+| `path`   | string | `~/.cagent/memory/<agent-name>/memory.db`    | Path to the SQLite database file     |
+
+| Operation          | Description                                                         |
+| ------------------ | ------------------------------------------------------------------- |
+| `add_memory`       | Store a new memory with optional category                           |
+| `get_memories`     | Retrieve all stored memories                                        |
+| `delete_memory`    | Delete a specific memory by ID                                      |
+| `search_memories`  | Search memories by keywords and/or category (more efficient than get_all) |
+| `update_memory`    | Update an existing memory's content and/or category by ID           |
+
+Memories support an optional `category` field (e.g., `preference`, `fact`, `project`, `decision`) for organization and filtering.
 
 ### Fetch
 

pkg/acp/registry.go

@@ -14,7 +14,7 @@ import (
 func createToolsetRegistry(agent *Agent) *teamloader.ToolsetRegistry {
 	registry := teamloader.NewDefaultToolsetRegistry()
 
-	registry.Register("filesystem", func(ctx context.Context, toolset latest.Toolset, parentDir string, runConfig *config.RuntimeConfig) (tools.ToolSet, error) {
+	registry.Register("filesystem", func(ctx context.Context, toolset latest.Toolset, parentDir string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
 		wd := runConfig.WorkingDir
 		if wd == "" {
 			var err error

pkg/config/latest/validate.go

@@ -116,9 +116,7 @@ func (t *Toolset) validate() error {
 	case "shell":
 		// no additional validation needed
 	case "memory":
-		if t.Path == "" {
-			return errors.New("memory toolset requires a path to be set")
-		}
+		// path is optional; defaults to ~/.cagent/memory/<agent-name>/memory.db
 	case "tasks":
 		// path defaults to ./tasks.json if not set
 	case "mcp":

pkg/creator/agent.go

@@ -120,7 +120,7 @@ func createToolsetRegistry(workingDir string) *teamloader.ToolsetRegistry {
 	}
 
 	registry := teamloader.NewDefaultToolsetRegistry()
-	registry.Register("filesystem", func(context.Context, latest.Toolset, string, *config.RuntimeConfig) (tools.ToolSet, error) {
+	registry.Register("filesystem", func(context.Context, latest.Toolset, string, *config.RuntimeConfig, string) (tools.ToolSet, error) {
 		return tracker, nil
 	})
 

pkg/creator/agent_test.go

@@ -158,7 +158,7 @@ func TestFileWriteTracker(t *testing.T) {
 	require.NotNil(t, registry)
 
 	// Create the toolset through the registry
-	toolset, err := registry.CreateTool(ctx, latest.Toolset{Type: "filesystem"}, runConfig.WorkingDir, runConfig)
+	toolset, err := registry.CreateTool(ctx, latest.Toolset{Type: "filesystem"}, runConfig.WorkingDir, runConfig, "test-agent")
 	require.NoError(t, err)
 	require.NotNil(t, toolset)
 

pkg/memory/database/database.go

@@ -5,16 +5,22 @@ import (
 	"errors"
 )
 
-var ErrEmptyID = errors.New("memory ID cannot be empty")
+var (
+	ErrEmptyID        = errors.New("memory ID cannot be empty")
+	ErrMemoryNotFound = errors.New("memory not found")
+)
 
 type UserMemory struct {
-	ID        string `description:"The ID of the memory"`
-	CreatedAt string `description:"The creation timestamp of the memory"`
-	Memory    string `description:"The content of the memory"`
+	ID        string `json:"id" description:"The ID of the memory"`
+	CreatedAt string `json:"created_at" description:"The creation timestamp of the memory"`
+	Memory    string `json:"memory" description:"The content of the memory"`
+	Category  string `json:"category,omitempty" description:"The category of the memory"`
 }
 
 type Database interface {
 	AddMemory(ctx context.Context, memory UserMemory) error
 	GetMemories(ctx context.Context) ([]UserMemory, error)
 	DeleteMemory(ctx context.Context, memory UserMemory) error
+	SearchMemories(ctx context.Context, query, category string) ([]UserMemory, error)
+	UpdateMemory(ctx context.Context, memory UserMemory) error
 }

pkg/memory/database/sqlite/sqlite.go

@@ -3,6 +3,8 @@ package sqlite
 import (
 	"context"
 	"database/sql"
+	"fmt"
+	"strings"
 
 	"github.com/docker/cagent/pkg/memory/database"
 	"github.com/docker/cagent/pkg/sqliteutil"
@@ -26,20 +28,28 @@ func NewMemoryDatabase(path string) (database.Database, error) {
 		return nil, err
 	}
 
+	// Add category column if it doesn't exist (transparent migration)
+	if _, err := db.ExecContext(context.Background(), "ALTER TABLE memories ADD COLUMN category TEXT DEFAULT ''"); err != nil {
+		if !strings.Contains(err.Error(), "duplicate column name") {
+			db.Close()
+			return nil, fmt.Errorf("memory database migration failed: %w", err)
+		}
+	}
+
 	return &MemoryDatabase{db: db}, nil
 }
 
 func (m *MemoryDatabase) AddMemory(ctx context.Context, memory database.UserMemory) error {
 	if memory.ID == "" {
 		return database.ErrEmptyID
 	}
-	_, err := m.db.ExecContext(ctx, "INSERT INTO memories (id, created_at, memory) VALUES (?, ?, ?)",
-		memory.ID, memory.CreatedAt, memory.Memory)
+	_, err := m.db.ExecContext(ctx, "INSERT INTO memories (id, created_at, memory, category) VALUES (?, ?, ?, ?)",
+		memory.ID, memory.CreatedAt, memory.Memory, memory.Category)
 	return err
 }
 
 func (m *MemoryDatabase) GetMemories(ctx context.Context) ([]database.UserMemory, error) {
-	rows, err := m.db.QueryContext(ctx, "SELECT id, created_at, memory FROM memories")
+	rows, err := m.db.QueryContext(ctx, "SELECT id, created_at, memory, COALESCE(category, '') FROM memories")
 	if err != nil {
 		return nil, err
 	}
@@ -48,7 +58,7 @@ func (m *MemoryDatabase) GetMemories(ctx context.Context) ([]database.UserMemory
 	var memories []database.UserMemory
 	for rows.Next() {
 		var memory database.UserMemory
-		err := rows.Scan(&memory.ID, &memory.CreatedAt, &memory.Memory)
+		err := rows.Scan(&memory.ID, &memory.CreatedAt, &memory.Memory, &memory.Category)
 		if err != nil {
 			return nil, err
 		}
@@ -66,3 +76,73 @@ func (m *MemoryDatabase) DeleteMemory(ctx context.Context, memory database.UserM
 	_, err := m.db.ExecContext(ctx, "DELETE FROM memories WHERE id = ?", memory.ID)
 	return err
 }
+
+func (m *MemoryDatabase) SearchMemories(ctx context.Context, query, category string) ([]database.UserMemory, error) {
+	var conditions []string
+	var args []any
+
+	if query != "" {
+		words := strings.Fields(query)
+		for _, word := range words {
+			conditions = append(conditions, "LOWER(memory) LIKE LOWER(?) ESCAPE '\\'")
+			escaped := strings.ReplaceAll(word, `\`, `\\`)
+			escaped = strings.ReplaceAll(escaped, `%`, `\%`)
+			escaped = strings.ReplaceAll(escaped, `_`, `\_`)
+			args = append(args, "%"+escaped+"%")
+		}
+	}
+
+	if category != "" {
+		conditions = append(conditions, "LOWER(category) = LOWER(?)")
+		args = append(args, category)
+	}
+
+	stmt := "SELECT id, created_at, memory, COALESCE(category, '') FROM memories"
+	if len(conditions) > 0 {
+		stmt += " WHERE " + strings.Join(conditions, " AND ")
+	}
+
+	rows, err := m.db.QueryContext(ctx, stmt, args...)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	var memories []database.UserMemory
+	for rows.Next() {
+		var memory database.UserMemory
+		err := rows.Scan(&memory.ID, &memory.CreatedAt, &memory.Memory, &memory.Category)
+		if err != nil {
+			return nil, err
+		}
+		memories = append(memories, memory)
+	}
+
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+
+	return memories, nil
+}
+
+func (m *MemoryDatabase) UpdateMemory(ctx context.Context, memory database.UserMemory) error {
+	if memory.ID == "" {
+		return database.ErrEmptyID
+	}
+
+	result, err := m.db.ExecContext(ctx, "UPDATE memories SET memory = ?, category = ? WHERE id = ?",
+		memory.Memory, memory.Category, memory.ID)
+	if err != nil {
+		return err
+	}
+
+	rows, err := result.RowsAffected()
+	if err != nil {
+		return err
+	}
+	if rows == 0 {
+		return fmt.Errorf("%w: %s", database.ErrMemoryNotFound, memory.ID)
+	}
+
+	return nil
+}