agents_and_robots/shell/memory/sqlite.go

// 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()
}