feat: implement memory management system with SQLite persistence, including conversation windows and episodic facts

2026-03-06 00:39:22 +00:00
parent d26be78c46
commit cb9489e633
7 changed files with 642 additions and 8 deletions
@@ -0,0 +1,175 @@
+// Package shellmem implements persistent memory storage using SQLite.
+package shellmem
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+	"os"
+	"path/filepath"
+	"time"
+
+	"github.com/enmanuel/agents/pkg/llm"
+	"github.com/enmanuel/agents/pkg/memory"
+)
+
+const schema = `
+CREATE TABLE IF NOT EXISTS facts (
+    agent_id   TEXT NOT NULL,
+    subject    TEXT NOT NULL,
+    key        TEXT NOT NULL,
+    value      TEXT NOT NULL,
+    updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    PRIMARY KEY (agent_id, subject, key)
+);
+
+CREATE TABLE IF NOT EXISTS messages (
+    id         INTEGER PRIMARY KEY AUTOINCREMENT,
+    agent_id   TEXT NOT NULL,
+    room_id    TEXT NOT NULL,
+    role       TEXT NOT NULL,
+    content    TEXT NOT NULL,
+    created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
+);
+
+CREATE INDEX IF NOT EXISTS idx_messages_room ON messages(agent_id, room_id, created_at DESC);
+CREATE INDEX IF NOT EXISTS idx_facts_subject ON facts(agent_id, subject);
+`
+
+// SQLiteStore implements memory.Store using SQLite.
+type SQLiteStore struct {
+	db *sql.DB
+}
+
+// New opens (or creates) a SQLite database at dbPath and runs migrations.
+func New(dbPath string) (*SQLiteStore, error) {
+	if err := os.MkdirAll(filepath.Dir(dbPath), 0o755); err != nil {
+		return nil, fmt.Errorf("create memory db dir: %w", err)
+	}
+	db, err := sql.Open("sqlite3", dbPath)
+	if err != nil {
+		return nil, fmt.Errorf("open memory db: %w", err)
+	}
+	if _, err := db.Exec(schema); err != nil {
+		db.Close()
+		return nil, fmt.Errorf("migrate memory db: %w", err)
+	}
+	return &SQLiteStore{db: db}, nil
+}
+
+func (s *SQLiteStore) SaveFact(ctx context.Context, f memory.Fact) error {
+	_, err := s.db.ExecContext(ctx,
+		`INSERT OR REPLACE INTO facts (agent_id, subject, key, value, updated_at)
+		 VALUES (?, ?, ?, ?, ?)`,
+		f.AgentID, f.Subject, f.Key, f.Value, time.Now().UTC(),
+	)
+	return err
+}
+
+func (s *SQLiteStore) RecallFacts(ctx context.Context, agentID, subject string, key *string) ([]memory.Fact, error) {
+	var rows *sql.Rows
+	var err error
+	if key != nil {
+		rows, err = s.db.QueryContext(ctx,
+			`SELECT agent_id, subject, key, value, updated_at FROM facts
+			 WHERE agent_id = ? AND subject = ? AND key = ?`,
+			agentID, subject, *key,
+		)
+	} else {
+		rows, err = s.db.QueryContext(ctx,
+			`SELECT agent_id, subject, key, value, updated_at FROM facts
+			 WHERE agent_id = ? AND subject = ?`,
+			agentID, subject,
+		)
+	}
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	var facts []memory.Fact
+	for rows.Next() {
+		var f memory.Fact
+		if err := rows.Scan(&f.AgentID, &f.Subject, &f.Key, &f.Value, &f.UpdatedAt); err != nil {
+			return nil, err
+		}
+		facts = append(facts, f)
+	}
+	return facts, rows.Err()
+}
+
+func (s *SQLiteStore) DeleteFacts(ctx context.Context, agentID, subject string, key *string) error {
+	if key != nil {
+		_, err := s.db.ExecContext(ctx,
+			`DELETE FROM facts WHERE agent_id = ? AND subject = ? AND key = ?`,
+			agentID, subject, *key,
+		)
+		return err
+	}
+	_, err := s.db.ExecContext(ctx,
+		`DELETE FROM facts WHERE agent_id = ? AND subject = ?`,
+		agentID, subject,
+	)
+	return err
+}
+
+func (s *SQLiteStore) SaveMessage(ctx context.Context, m memory.HistoryMessage) error {
+	_, err := s.db.ExecContext(ctx,
+		`INSERT INTO messages (agent_id, room_id, role, content, created_at)
+		 VALUES (?, ?, ?, ?, ?)`,
+		m.AgentID, m.RoomID, string(m.Role), m.Content, time.Now().UTC(),
+	)
+	return err
+}
+
+func (s *SQLiteStore) LoadMessages(ctx context.Context, agentID, roomID string, limit int) ([]memory.HistoryMessage, error) {
+	rows, err := s.db.QueryContext(ctx,
+		`SELECT agent_id, room_id, role, content, created_at FROM messages
+		 WHERE agent_id = ? AND room_id = ?
+		 ORDER BY created_at DESC LIMIT ?`,
+		agentID, roomID, limit,
+	)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	var msgs []memory.HistoryMessage
+	for rows.Next() {
+		var m memory.HistoryMessage
+		var role string
+		if err := rows.Scan(&m.AgentID, &m.RoomID, &role, &m.Content, &m.CreatedAt); err != nil {
+			return nil, err
+		}
+		m.Role = llm.Role(role)
+		msgs = append(msgs, m)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+
+	// Reverse to chronological order
+	for i, j := 0, len(msgs)-1; i < j; i, j = i+1, j-1 {
+		msgs[i], msgs[j] = msgs[j], msgs[i]
+	}
+	return msgs, nil
+}
+
+func (s *SQLiteStore) DeleteMessages(ctx context.Context, agentID string, roomID *string) error {
+	if roomID != nil {
+		_, err := s.db.ExecContext(ctx,
+			`DELETE FROM messages WHERE agent_id = ? AND room_id = ?`,
+			agentID, *roomID,
+		)
+		return err
+	}
+	_, err := s.db.ExecContext(ctx,
+		`DELETE FROM messages WHERE agent_id = ?`,
+		agentID,
+	)
+	return err
+}
+
+func (s *SQLiteStore) Close() error {
+	return s.db.Close()
+}