agents_and_robots/shell/memory/sqlite.go

// Package shellmem implements persistent memory storage using SQLite.
package shellmem

import (
	"context"
	"database/sql"
	"embed"
	"fmt"
	"io/fs"
	"log/slog"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"github.com/enmanuel/agents/pkg/llm"
	"github.com/enmanuel/agents/pkg/memory"
)

//go:embed migrations/*.sql
var migrationsFS embed.FS

func applyMigrations(db *sql.DB) error {
	files, err := fs.Glob(migrationsFS, "migrations/*.sql")
	if err != nil {
		return err
	}
	sort.Strings(files)
	for _, f := range files {
		b, err := migrationsFS.ReadFile(f)
		if err != nil {
			return err
		}
		if _, err := db.Exec(string(b)); err != nil {
			if strings.Contains(err.Error(), "duplicate column") || strings.Contains(err.Error(), "already exists") {
				continue
			}
			return fmt.Errorf("migrate %s: %w", f, err)
		}
	}
	return nil
}

// SQLiteStore implements memory.Store using SQLite.
type SQLiteStore struct {
	db     *sql.DB
	logger *slog.Logger
}

// New opens (or creates) a SQLite database at dbPath and runs migrations.
func New(dbPath string, logger *slog.Logger) (*SQLiteStore, error) {
	log := logger.With("component", "memory", "db_path", dbPath)
	log.Info("memory_open")
	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 := applyMigrations(db); err != nil {
		db.Close()
		return nil, fmt.Errorf("migrate memory db: %w", err)
	}
	log.Info("memory_ready")
	return &SQLiteStore{db: db, logger: log}, nil
}

func (s *SQLiteStore) SaveFact(ctx context.Context, f memory.Fact) error {
	s.logger.Debug("memory_save_fact", "subject", f.Subject, "key", f.Key)
	_, 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(),
	)
	if err != nil {
		s.logger.Error("memory_save_fact_error", "subject", f.Subject, "key", f.Key, "err", err)
	}
	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)
	}
	s.logger.Debug("memory_recall", "subject", subject, "count", len(facts))
	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 {
	s.logger.Debug("memory_save_msg", "room", m.RoomID, "role", m.Role)
	_, 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(),
	)
	if err != nil {
		s.logger.Error("memory_save_msg_error", "room", m.RoomID, "err", err)
	}
	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]
	}
	s.logger.Debug("memory_load_msgs", "room", roomID, "count", len(msgs))
	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 {
	s.logger.Info("memory_closed")
	return s.db.Close()
}