From 65daa8d1f98fa4e1373a86e5affcf08ef1650ff8 Mon Sep 17 00:00:00 2001 From: Siddhant Khare Date: Tue, 24 Feb 2026 18:37:19 +0000 Subject: [PATCH 1/2] feat: add Postgres/Supabase backend for memory store Adds PostgresStore alongside the existing SQLiteStore, both implementing the memory.Store interface. This allows persistent memory storage in Supabase instead of ephemeral SQLite on Render. Changes: - pkg/memory/postgres.go: full Store implementation using pgx driver - pkg/memory/postgres_migration.sql: schema for Supabase SQL editor - cmd/api.go, cmd/mcp.go: new --memory-backend and --memory-dsn flags - cmd/api_memory.go: MemoryAPI now uses Store interface, not SQLiteStore - cmd/memory.go: memoryStoreFromConfig supports 'sqlite' and 'postgres' Usage: distill api --memory --memory-backend postgres \ --memory-dsn 'postgres://user:pass@host:5432/db?sslmode=require' Co-authored-by: Ona --- cmd/api.go | 18 +- cmd/api_memory.go | 2 +- cmd/mcp.go | 16 +- cmd/memory.go | 21 +- go.mod | 5 + go.sum | 11 + pkg/memory/postgres.go | 448 ++++++++++++++++++++++++++++++ pkg/memory/postgres_migration.sql | 26 ++ 8 files changed, 531 insertions(+), 16 deletions(-) create mode 100644 pkg/memory/postgres.go create mode 100644 pkg/memory/postgres_migration.sql diff --git a/cmd/api.go b/cmd/api.go index 6e221e5..60689d1 100644 --- a/cmd/api.go +++ b/cmd/api.go @@ -46,6 +46,8 @@ func init() { apiCmd.Flags().String("embedding-model", "text-embedding-3-small", "OpenAI embedding model") apiCmd.Flags().String("api-keys", "", "Comma-separated list of valid API keys (or use DISTILL_API_KEYS)") apiCmd.Flags().Bool("memory", false, "Enable persistent memory store") + apiCmd.Flags().String("memory-backend", "sqlite", "Memory store backend: sqlite or postgres") + apiCmd.Flags().String("memory-dsn", "", "Memory store DSN (file path for sqlite, connection string for postgres)") apiCmd.Flags().Bool("session", false, "Enable session management") apiCmd.Flags().String("session-db", "distill-sessions.db", "SQLite database path for session store") @@ -184,15 +186,23 @@ func runAPI(cmd *cobra.Command, args []string) error { // Setup memory store (opt-in) enableMemory, _ := cmd.Flags().GetBool("memory") if enableMemory { - memDBPath := viper.GetString("memory.db_path") - if memDBPath == "" { - memDBPath = "distill-memory.db" + memBackend, _ := cmd.Flags().GetString("memory-backend") + memDSN, _ := cmd.Flags().GetString("memory-dsn") + // Fall back to config file values + if memDSN == "" { + memDSN = viper.GetString("memory.dsn") + } + if memDSN == "" { + memDSN = viper.GetString("memory.db_path") + } + if memBackend == "sqlite" && memDSN == "" { + memDSN = "distill-memory.db" } memThreshold := viper.GetFloat64("memory.dedup_threshold") if memThreshold == 0 { memThreshold = 0.15 } - memStore, err := memoryStoreFromConfig(memDBPath, memThreshold) + memStore, err := memoryStoreFromConfig(memBackend, memDSN, memThreshold) if err != nil { return fmt.Errorf("failed to create memory store: %w", err) } diff --git a/cmd/api_memory.go b/cmd/api_memory.go index df95f5d..3e00f3d 100644 --- a/cmd/api_memory.go +++ b/cmd/api_memory.go @@ -13,7 +13,7 @@ import ( // MemoryAPI handles memory-related HTTP endpoints. type MemoryAPI struct { - store *memory.SQLiteStore + store memory.Store embedder retriever.EmbeddingProvider } diff --git a/cmd/mcp.go b/cmd/mcp.go index 659f4ac..fb94dcc 100644 --- a/cmd/mcp.go +++ b/cmd/mcp.go @@ -92,7 +92,9 @@ func init() { // Memory store mcpCmd.Flags().Bool("memory", false, "Enable persistent memory store") - mcpCmd.Flags().String("memory-db", "distill-memory.db", "SQLite database path for memory store") + mcpCmd.Flags().String("memory-backend", "sqlite", "Memory store backend: sqlite or postgres") + mcpCmd.Flags().String("memory-dsn", "", "Memory store DSN (file path for sqlite, connection string for postgres") + mcpCmd.Flags().String("memory-db", "distill-memory.db", "SQLite database path for memory store (deprecated, use --memory-dsn)") mcpCmd.Flags().Bool("session", false, "Enable session management") mcpCmd.Flags().String("session-db", "distill-sessions.db", "SQLite database path for session store") @@ -108,7 +110,7 @@ type MCPServer struct { broker *contextlab.Broker embedder retriever.EmbeddingProvider cfg contextlab.BrokerConfig - memStore *memory.SQLiteStore + memStore memory.Store sessStore *session.SQLiteStore } @@ -159,10 +161,12 @@ func runMCP(cmd *cobra.Command, args []string) error { // Create memory store (opt-in) enableMemory, _ := cmd.Flags().GetBool("memory") if enableMemory { - memDBPath, _ := cmd.Flags().GetString("memory-db") - memCfg := memory.DefaultConfig() - memCfg.DedupThreshold = threshold - memStore, err := memory.NewSQLiteStore(memDBPath, memCfg) + memBackend, _ := cmd.Flags().GetString("memory-backend") + memDSN, _ := cmd.Flags().GetString("memory-dsn") + if memDSN == "" { + memDSN, _ = cmd.Flags().GetString("memory-db") + } + memStore, err := memoryStoreFromConfig(memBackend, memDSN, threshold) if err != nil { return fmt.Errorf("failed to create memory store: %w", err) } diff --git a/cmd/memory.go b/cmd/memory.go index 5120c42..640720b 100644 --- a/cmd/memory.go +++ b/cmd/memory.go @@ -250,11 +250,22 @@ func runMemoryStats(cmd *cobra.Command, args []string) error { // memoryStoreFromConfig creates a memory store from the API server config. // Used by the API server and MCP server. -func memoryStoreFromConfig(dbPath string, threshold float64) (*memory.SQLiteStore, error) { - if dbPath == "" { - dbPath = "distill-memory.db" - } +// backend: "sqlite" (default) or "postgres" +// dsn: database path (sqlite) or connection string (postgres) +func memoryStoreFromConfig(backend, dsn string, threshold float64) (memory.Store, error) { cfg := memory.DefaultConfig() cfg.DedupThreshold = threshold - return memory.NewSQLiteStore(dbPath, cfg) + + switch backend { + case "postgres": + if dsn == "" { + return nil, fmt.Errorf("--memory-dsn is required for postgres backend") + } + return memory.NewPostgresStore(dsn, cfg) + default: + if dsn == "" { + dsn = "distill-memory.db" + } + return memory.NewSQLiteStore(dsn, cfg) + } } diff --git a/go.mod b/go.mod index a2fe129..e20985d 100644 --- a/go.mod +++ b/go.mod @@ -39,6 +39,10 @@ require ( github.com/hashicorp/hcl v1.0.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect github.com/invopop/jsonschema v0.13.0 // indirect + github.com/jackc/pgpassfile v1.0.0 // indirect + github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect + github.com/jackc/pgx/v5 v5.8.0 // indirect + github.com/jackc/puddle/v2 v2.2.2 // indirect github.com/magiconair/properties v1.8.10 // indirect github.com/mailru/easyjson v0.7.7 // indirect github.com/mattn/go-isatty v0.0.20 // indirect @@ -72,6 +76,7 @@ require ( go.yaml.in/yaml/v2 v2.4.2 // indirect golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546 // indirect golang.org/x/net v0.49.0 // indirect + golang.org/x/sync v0.19.0 // indirect golang.org/x/sys v0.40.0 // indirect golang.org/x/term v0.39.0 // indirect golang.org/x/text v0.33.0 // indirect diff --git a/go.sum b/go.sum index 35bac8e..4ab98f3 100644 --- a/go.sum +++ b/go.sum @@ -42,6 +42,14 @@ github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2 github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= github.com/invopop/jsonschema v0.13.0 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E= github.com/invopop/jsonschema v0.13.0/go.mod h1:ffZ5Km5SWWRAIN6wbDXItl95euhFz2uON45H2qjYt+0= +github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= +github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM= +github.com/jackc/pgx/v5 v5.8.0 h1:TYPDoleBBme0xGSAX3/+NujXXtpZn9HBONkQC7IEZSo= +github.com/jackc/pgx/v5 v5.8.0/go.mod h1:QVeDInX2m9VyzvNeiCJVjCkNFqzsNb43204HshNSZKw= +github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= +github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y= github.com/juju/gnuflag v0.0.0-20171113085948-2ce1bb71843d/go.mod h1:2PavIy+JPciBPrBUjwbNvtwB6RQlve+hkpll6QSNmOE= github.com/k0kubun/go-ansi v0.0.0-20180517002512-3bf9e2903213/go.mod h1:vNUNkEQ1e29fT/6vq2aBdFsgNPmy8qMdSay1npru+Sw= @@ -119,6 +127,7 @@ github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSS github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= @@ -165,6 +174,8 @@ golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546 h1:mgKeJMpvi0yx/sU5GsxQ7p6s2 golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546/go.mod h1:j/pmGrbnkbPtQfxEe5D0VQhZC6qKbfKifgD0oM7sR70= golang.org/x/net v0.49.0 h1:eeHFmOGUTtaaPSGNmjBKpbng9MulQsJURQUAfUwY++o= golang.org/x/net v0.49.0/go.mod h1:/ysNB2EvaqvesRkuLAyjI1ycPZlQHM3q01F02UY/MV8= +golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4= +golang.org/x/sync v0.19.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.22.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/sys v0.40.0 h1:DBZZqJ2Rkml6QMQsZywtnjnnGvHza6BTfYFWY9kjEWQ= diff --git a/pkg/memory/postgres.go b/pkg/memory/postgres.go new file mode 100644 index 0000000..3b9c7cb --- /dev/null +++ b/pkg/memory/postgres.go @@ -0,0 +1,448 @@ +package memory + +import ( + "context" + "database/sql" + "encoding/json" + "fmt" + "strings" + "time" + + distillmath "github.com/Siddhant-K-code/distill/pkg/math" + _ "github.com/jackc/pgx/v5/stdlib" +) + +// PostgresStore implements Store using PostgreSQL (Supabase) for persistent storage. +type PostgresStore struct { + db *sql.DB + cfg Config +} + +// NewPostgresStore creates a new Postgres-backed memory store. +// dsn should be a Postgres connection string, e.g.: +// +// "postgres://user:pass@host:5432/dbname?sslmode=require" +func NewPostgresStore(dsn string, cfg Config) (*PostgresStore, error) { + if dsn == "" { + return nil, fmt.Errorf("postgres DSN is required") + } + + db, err := sql.Open("pgx", dsn) + if err != nil { + return nil, fmt.Errorf("open postgres: %w", err) + } + + db.SetMaxOpenConns(10) + db.SetMaxIdleConns(5) + db.SetConnMaxLifetime(5 * time.Minute) + + // Verify connection + if err := db.Ping(); err != nil { + _ = db.Close() + return nil, fmt.Errorf("ping postgres: %w", err) + } + + s := &PostgresStore{db: db, cfg: cfg} + if err := s.migrate(); err != nil { + _ = db.Close() + return nil, fmt.Errorf("migrate: %w", err) + } + + return s, nil +} + +func (s *PostgresStore) migrate() error { + schema := ` + CREATE TABLE IF NOT EXISTS memories ( + id TEXT PRIMARY KEY, + text TEXT NOT NULL, + embedding BYTEA, + source TEXT DEFAULT '', + session_id TEXT DEFAULT '', + metadata JSONB DEFAULT '{}', + decay_level INTEGER DEFAULT 0, + created_at TIMESTAMPTZ NOT NULL, + last_referenced TIMESTAMPTZ NOT NULL, + access_count INTEGER DEFAULT 0 + ); + CREATE TABLE IF NOT EXISTS memory_tags ( + memory_id TEXT NOT NULL REFERENCES memories(id) ON DELETE CASCADE, + tag TEXT NOT NULL, + PRIMARY KEY (memory_id, tag) + ); + CREATE INDEX IF NOT EXISTS idx_memory_tags_tag ON memory_tags(tag); + CREATE INDEX IF NOT EXISTS idx_memories_decay ON memories(decay_level); + CREATE INDEX IF NOT EXISTS idx_memories_created ON memories(created_at); + CREATE INDEX IF NOT EXISTS idx_memories_referenced ON memories(last_referenced); + ` + _, err := s.db.Exec(schema) + return err +} + +// Store adds entries with write-time deduplication. +func (s *PostgresStore) Store(ctx context.Context, req StoreRequest) (*StoreResult, error) { + result := &StoreResult{} + + for _, entry := range req.Entries { + if entry.Text == "" { + continue + } + + // Check for semantic duplicates if embedding is provided + if len(entry.Embedding) > 0 { + dupID, err := s.findDuplicate(ctx, entry.Embedding) + if err != nil { + return nil, fmt.Errorf("find duplicate: %w", err) + } + if dupID != "" { + _, err := s.db.ExecContext(ctx, + `UPDATE memories SET last_referenced = $1, access_count = access_count + 1 WHERE id = $2`, + time.Now().UTC(), dupID, + ) + if err != nil { + return nil, fmt.Errorf("update duplicate: %w", err) + } + result.Deduplicated++ + continue + } + } + + id := generateID() + now := time.Now().UTC() + + metaJSON, _ := json.Marshal(entry.Metadata) + embBlob := encodeEmbedding(entry.Embedding) + + sessionID := req.SessionID + + _, err := s.db.ExecContext(ctx, + `INSERT INTO memories (id, text, embedding, source, session_id, metadata, decay_level, created_at, last_referenced, access_count) + VALUES ($1, $2, $3, $4, $5, $6, 0, $7, $8, 0)`, + id, entry.Text, embBlob, entry.Source, sessionID, string(metaJSON), now, now, + ) + if err != nil { + return nil, fmt.Errorf("insert memory: %w", err) + } + + for _, tag := range entry.Tags { + _, err := s.db.ExecContext(ctx, + "INSERT INTO memory_tags (memory_id, tag) VALUES ($1, $2) ON CONFLICT DO NOTHING", + id, tag, + ) + if err != nil { + return nil, fmt.Errorf("insert tag: %w", err) + } + } + + result.Stored++ + } + + var total int + if err := s.db.QueryRowContext(ctx, "SELECT COUNT(*) FROM memories").Scan(&total); err != nil { + return nil, err + } + result.TotalMemories = total + + return result, nil +} + +func (s *PostgresStore) findDuplicate(ctx context.Context, embedding []float32) (string, error) { + rows, err := s.db.QueryContext(ctx, "SELECT id, embedding FROM memories WHERE embedding IS NOT NULL") + if err != nil { + return "", err + } + defer func() { _ = rows.Close() }() + + for rows.Next() { + var id string + var embBlob []byte + if err := rows.Scan(&id, &embBlob); err != nil { + return "", err + } + + existing := decodeEmbedding(embBlob) + if len(existing) == 0 { + continue + } + + dist := distillmath.CosineDistance(embedding, existing) + if dist < s.cfg.DedupThreshold { + return id, nil + } + } + + return "", rows.Err() +} + +// Recall retrieves memories matching a query, ranked by relevance and recency. +func (s *PostgresStore) Recall(ctx context.Context, req RecallRequest) (*RecallResult, error) { + if req.Query == "" && len(req.QueryEmbedding) == 0 { + return nil, ErrInvalidQuery + } + + maxResults := req.MaxResults + if maxResults <= 0 { + maxResults = 10 + } + + recencyWeight := req.RecencyWeight + if recencyWeight < 0 { + recencyWeight = 0 + } + if recencyWeight > 1 { + recencyWeight = 1 + } + + // Build query with optional tag filter + query := "SELECT m.id, m.text, m.embedding, m.source, m.decay_level, m.last_referenced FROM memories m" + var args []interface{} + + if len(req.Tags) > 0 { + placeholders := make([]string, len(req.Tags)) + for i, tag := range req.Tags { + placeholders[i] = fmt.Sprintf("$%d", i+1) + args = append(args, tag) + } + query += " WHERE m.id IN (SELECT memory_id FROM memory_tags WHERE tag IN (" + strings.Join(placeholders, ",") + "))" + } + + rows, err := s.db.QueryContext(ctx, query, args...) + if err != nil { + return nil, fmt.Errorf("query memories: %w", err) + } + defer func() { _ = rows.Close() }() + + type rawRow struct { + id, text, source string + embBlob []byte + decayLevel int + lastRef time.Time + } + var rawRows []rawRow + for rows.Next() { + var r rawRow + if err := rows.Scan(&r.id, &r.text, &r.embBlob, &r.source, &r.decayLevel, &r.lastRef); err != nil { + return nil, err + } + rawRows = append(rawRows, r) + } + if err := rows.Err(); err != nil { + return nil, err + } + + var candidates []scored + now := time.Now() + + for _, r := range rawRows { + tags, _ := s.loadTags(ctx, r.id) + + var similarity float64 + if len(req.QueryEmbedding) > 0 { + existing := decodeEmbedding(r.embBlob) + if len(existing) > 0 { + dist := distillmath.CosineDistance(req.QueryEmbedding, existing) + similarity = 1.0 - dist + } + } + + age := now.Sub(r.lastRef).Hours() + recency := 1.0 + if age > 0 { + recency = 1.0 / (1.0 + age/24.0) + } + + relevance := (1.0-recencyWeight)*similarity + recencyWeight*recency + + candidates = append(candidates, scored{ + memory: RecalledMemory{ + ID: r.id, + Text: r.text, + Source: r.source, + Tags: tags, + Relevance: relevance, + DecayLevel: DecayLevel(r.decayLevel), + LastReferenced: r.lastRef, + }, + relevance: relevance, + }) + } + + sortByRelevance(candidates) + + var results []RecalledMemory + tokenCount := 0 + for _, c := range candidates { + if len(results) >= maxResults { + break + } + tokens := estimateTokens(c.memory.Text) + if req.MaxTokens > 0 && tokenCount+tokens > req.MaxTokens { + break + } + results = append(results, c.memory) + tokenCount += tokens + } + + if len(results) > 0 { + ids := make([]string, len(results)) + for i, m := range results { + ids[i] = m.ID + } + s.touchMemories(ctx, ids) + } + + return &RecallResult{ + Memories: results, + Stats: RecallStats{ + Candidates: len(candidates), + Deduplicated: len(candidates) - len(results), + Returned: len(results), + TokenCount: tokenCount, + }, + }, nil +} + +// Forget removes memories matching the given criteria. +func (s *PostgresStore) Forget(ctx context.Context, req ForgetRequest) (*ForgetResult, error) { + var conditions []string + var args []interface{} + argIdx := 1 + + if len(req.IDs) > 0 { + placeholders := make([]string, len(req.IDs)) + for i, id := range req.IDs { + placeholders[i] = fmt.Sprintf("$%d", argIdx) + args = append(args, id) + argIdx++ + } + conditions = append(conditions, "id IN ("+strings.Join(placeholders, ",")+")") + } + + if len(req.Tags) > 0 { + placeholders := make([]string, len(req.Tags)) + for i, tag := range req.Tags { + placeholders[i] = fmt.Sprintf("$%d", argIdx) + args = append(args, tag) + argIdx++ + } + conditions = append(conditions, "id IN (SELECT memory_id FROM memory_tags WHERE tag IN ("+strings.Join(placeholders, ",")+"))") + } + + if !req.OlderThan.IsZero() { + conditions = append(conditions, fmt.Sprintf("created_at < $%d", argIdx)) + args = append(args, req.OlderThan.UTC()) + argIdx++ + } + + if len(conditions) == 0 { + return &ForgetResult{}, nil + } + + query := "DELETE FROM memories WHERE " + strings.Join(conditions, " AND ") + res, err := s.db.ExecContext(ctx, query, args...) + if err != nil { + return nil, fmt.Errorf("delete memories: %w", err) + } + + removed, _ := res.RowsAffected() + + var total int + if err := s.db.QueryRowContext(ctx, "SELECT COUNT(*) FROM memories").Scan(&total); err != nil { + return nil, err + } + + return &ForgetResult{ + Removed: int(removed), + TotalMemories: total, + }, nil +} + +// Stats returns memory store statistics. +func (s *PostgresStore) Stats(ctx context.Context) (*Stats, error) { + stats := &Stats{ + ByDecayLevel: make(map[int]int), + BySource: make(map[string]int), + } + + if err := s.db.QueryRowContext(ctx, "SELECT COUNT(*) FROM memories").Scan(&stats.TotalMemories); err != nil { + return nil, err + } + + rows, err := s.db.QueryContext(ctx, "SELECT decay_level, COUNT(*) FROM memories GROUP BY decay_level") + if err != nil { + return nil, err + } + for rows.Next() { + var level, count int + if err := rows.Scan(&level, &count); err != nil { + _ = rows.Close() + return nil, err + } + stats.ByDecayLevel[level] = count + } + _ = rows.Close() + + rows, err = s.db.QueryContext(ctx, "SELECT source, COUNT(*) FROM memories WHERE source != '' GROUP BY source") + if err != nil { + return nil, err + } + for rows.Next() { + var source string + var count int + if err := rows.Scan(&source, &count); err != nil { + _ = rows.Close() + return nil, err + } + stats.BySource[source] = count + } + _ = rows.Close() + + var oldest, newest sql.NullTime + _ = s.db.QueryRowContext(ctx, "SELECT MIN(created_at) FROM memories").Scan(&oldest) + _ = s.db.QueryRowContext(ctx, "SELECT MAX(created_at) FROM memories").Scan(&newest) + if oldest.Valid { + stats.OldestMemory = oldest.Time + } + if newest.Valid { + stats.NewestMemory = newest.Time + } + + return stats, nil +} + +// Close closes the database connection pool. +func (s *PostgresStore) Close() error { + return s.db.Close() +} + +func (s *PostgresStore) loadTags(ctx context.Context, memoryID string) ([]string, error) { + rows, err := s.db.QueryContext(ctx, "SELECT tag FROM memory_tags WHERE memory_id = $1", memoryID) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + + var tags []string + for rows.Next() { + var tag string + if err := rows.Scan(&tag); err != nil { + return nil, err + } + tags = append(tags, tag) + } + return tags, rows.Err() +} + +func (s *PostgresStore) touchMemories(ctx context.Context, ids []string) { + if len(ids) == 0 { + return + } + placeholders := make([]string, len(ids)) + args := []interface{}{time.Now().UTC()} + for i, id := range ids { + placeholders[i] = fmt.Sprintf("$%d", i+2) + args = append(args, id) + } + query := "UPDATE memories SET last_referenced = $1, access_count = access_count + 1 WHERE id IN (" + strings.Join(placeholders, ",") + ")" + _, _ = s.db.ExecContext(ctx, query, args...) +} diff --git a/pkg/memory/postgres_migration.sql b/pkg/memory/postgres_migration.sql new file mode 100644 index 0000000..60126b9 --- /dev/null +++ b/pkg/memory/postgres_migration.sql @@ -0,0 +1,26 @@ +-- Memory store tables for Postgres/Supabase backend. +-- Run this in Supabase SQL Editor or via migration. + +CREATE TABLE IF NOT EXISTS memories ( + id TEXT PRIMARY KEY, + text TEXT NOT NULL, + embedding BYTEA, + source TEXT DEFAULT '', + session_id TEXT DEFAULT '', + metadata JSONB DEFAULT '{}', + decay_level INTEGER DEFAULT 0, + created_at TIMESTAMPTZ NOT NULL, + last_referenced TIMESTAMPTZ NOT NULL, + access_count INTEGER DEFAULT 0 +); + +CREATE TABLE IF NOT EXISTS memory_tags ( + memory_id TEXT NOT NULL REFERENCES memories(id) ON DELETE CASCADE, + tag TEXT NOT NULL, + PRIMARY KEY (memory_id, tag) +); + +CREATE INDEX IF NOT EXISTS idx_memory_tags_tag ON memory_tags(tag); +CREATE INDEX IF NOT EXISTS idx_memories_decay ON memories(decay_level); +CREATE INDEX IF NOT EXISTS idx_memories_created ON memories(created_at); +CREATE INDEX IF NOT EXISTS idx_memories_referenced ON memories(last_referenced); From cf36938986fc72a10ae215157c8d0b9d15e060e2 Mon Sep 17 00:00:00 2001 From: Siddhant Khare Date: Tue, 24 Feb 2026 18:43:46 +0000 Subject: [PATCH 2/2] fix: lint error + docs for postgres memory backend - Remove dead argIdx increment (ineffectual assignment) - README: add postgres backend usage, config examples, DATABASE_URL env var - CHANGELOG: add unreleased section for postgres backend - FAQ: mention postgres/supabase option in memory description Co-authored-by: Ona --- CHANGELOG.md | 8 ++++++++ FAQ.md | 2 +- README.md | 21 +++++++++++++++++++-- pkg/memory/postgres.go | 1 - 4 files changed, 28 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0051264..26c57b0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,14 @@ All notable changes to Distill are documented here. +## [Unreleased] + +### Added + +- **PostgreSQL/Supabase memory backend** — Persistent memory storage using Postgres instead of ephemeral SQLite. New `--memory-backend postgres` and `--memory-dsn` flags on `api` and `mcp` commands. Implements the same `Store` interface as SQLiteStore. ([#40](https://github.com/Siddhant-K-code/distill/pull/40)) + +--- + ## [v0.4.0] - 2026-02-24 ### Added diff --git a/FAQ.md b/FAQ.md index e753cd3..808bbbb 100644 --- a/FAQ.md +++ b/FAQ.md @@ -22,7 +22,7 @@ LLMs are non-deterministic. The same input can produce different compressed outp ### What is Context Memory? -Persistent memory that accumulates knowledge across agent sessions. Store context once, recall it later by semantic similarity + recency. Memories are deduplicated on write and compressed over time through hierarchical decay (full text → summary → keywords → evicted). Enable with `--memory` on the `api` or `mcp` commands. +Persistent memory that accumulates knowledge across agent sessions. Store context once, recall it later by semantic similarity + recency. Memories are deduplicated on write and compressed over time through hierarchical decay (full text → summary → keywords → evicted). Enable with `--memory` on the `api` or `mcp` commands. Supports SQLite (default, local) and PostgreSQL/Supabase (`--memory-backend postgres`) for persistent storage. ### What are Sessions? diff --git a/README.md b/README.md index 59586fa..3b005ea 100644 --- a/README.md +++ b/README.md @@ -212,7 +212,16 @@ See [mcp/README.md](mcp/README.md) for more configuration options. Persistent memory that accumulates knowledge across agent sessions. Memories are deduplicated on write, ranked by relevance + recency on recall, and compressed over time through hierarchical decay. -Enable with the `--memory` flag on `api` or `mcp` commands. +Enable with the `--memory` flag on `api` or `mcp` commands. Supports SQLite (default, local) and PostgreSQL (persistent, recommended for production). + +```bash +# SQLite (default - local file, good for development) +distill api --memory + +# PostgreSQL (persistent - recommended for production/Supabase) +distill api --memory --memory-backend postgres \ + --memory-dsn 'postgres://user:pass@host:5432/db?sslmode=require' +``` ### CLI @@ -273,8 +282,13 @@ Accessing a memory resets its decay clock. Configure ages via `distill.yaml`: ```yaml memory: + # SQLite (default) db_path: distill-memory.db dedup_threshold: 0.15 + + # Or PostgreSQL/Supabase + # backend: postgres + # dsn: postgres://user:pass@host:5432/db?sslmode=require ``` ## Session Management @@ -431,7 +445,9 @@ auth: - ${DISTILL_API_KEY} memory: - db_path: distill-memory.db + db_path: distill-memory.db # SQLite (default) + # backend: postgres # Use 'postgres' for Supabase/PostgreSQL + # dsn: ${DATABASE_URL} # Postgres connection string dedup_threshold: 0.15 session: @@ -448,6 +464,7 @@ Environment variables can be referenced using `${VAR}` or `${VAR:-default}` synt OPENAI_API_KEY # For text → embedding conversion (see note below) PINECONE_API_KEY # For Pinecone backend QDRANT_URL # For Qdrant backend (default: localhost:6334) +DATABASE_URL # For Postgres memory backend (Supabase connection string) DISTILL_API_KEYS # Optional: protect your self-hosted instance (see below) ``` diff --git a/pkg/memory/postgres.go b/pkg/memory/postgres.go index 3b9c7cb..d55eb8a 100644 --- a/pkg/memory/postgres.go +++ b/pkg/memory/postgres.go @@ -331,7 +331,6 @@ func (s *PostgresStore) Forget(ctx context.Context, req ForgetRequest) (*ForgetR if !req.OlderThan.IsZero() { conditions = append(conditions, fmt.Sprintf("created_at < $%d", argIdx)) args = append(args, req.OlderThan.UTC()) - argIdx++ } if len(conditions) == 0 {