cel_memory_postgres/
provider.rs

1//! `PostgresMemoryProvider` — PostgreSQL + pgvector backing storage.
2
3use std::collections::HashMap;
4use std::sync::Arc;
5
6use async_trait::async_trait;
7use cel_memory::{
8    CallerScope, MemoryChunk, MemoryError, MemoryProvider, MemoryQuery, MemorySession, MemoryStats,
9    MemoryTier, NewMemoryChunk, NewMemorySession, Result as MemoryResult, SessionFilter,
10    SessionOutcome,
11};
12use chrono::Utc;
13use pgvector::Vector;
14use sqlx::postgres::PgRow;
15use sqlx::{PgPool, Postgres, QueryBuilder, Row};
16use uuid::Uuid;
17
18use crate::error::PostgresMemoryError;
19use crate::util::{
20    chunk_matches_query, kind_str, outcome_str, row_to_chunk, rrf, source_str, str_to_outcome,
21    tier_str, EMBEDDING_DIM,
22};
23use cel_memory::Embedder;
24
25
26const GET_CHUNK_SQL: &str =
27    "SELECT id, created_at, kind, tier, source, session_id, project_root, caller_id, content, \
28     metadata, importance, pinned, shareable, superseded_by, embedding_model, embedding_dim \
29     FROM memory_chunks WHERE id = $1";
30
31const CHUNKS_BY_IDS_SQL: &str =
32    "SELECT id, created_at, kind, tier, source, session_id, project_root, caller_id, content, \
33     metadata, importance, pinned, shareable, superseded_by, embedding_model, embedding_dim \
34     FROM memory_chunks WHERE id = ANY($1)";
35
36/// PostgreSQL-backed [`MemoryProvider`] using pgvector + tsvector retrieval.
37pub struct PostgresMemoryProvider {
38    pool: PgPool,
39    embedder: Arc<dyn Embedder>,
40    write_hook: Option<Arc<dyn cel_memory::MemoryWriteHook>>,
41    summarizer: Option<Arc<dyn cel_memory::Summarizer>>,
42}
43
44impl PostgresMemoryProvider {
45    /// Connect to PostgreSQL, run migrations, and return a ready provider.
46    ///
47    /// The embedder must declare dimension [`EMBEDDING_DIM`] (384) to match the
48    /// schema. Use [`cel_memory::MockEmbedder`] in tests.
49    pub async fn connect(
50        database_url: &str,
51        embedder: Arc<dyn Embedder>,
52    ) -> Result<Self, PostgresMemoryError> {
53        if embedder.dim() != EMBEDDING_DIM {
54            return Err(PostgresMemoryError::DimMismatch {
55                expected: EMBEDDING_DIM,
56                actual: embedder.dim(),
57            });
58        }
59        let pool = sqlx::postgres::PgPoolOptions::new()
60            .max_connections(8)
61            .connect(database_url)
62            .await?;
63        sqlx::migrate!("./migrations")
64            .run(&pool)
65            .await
66            .map_err(|e| PostgresMemoryError::Migration(e.to_string()))?;
67        Ok(Self {
68            pool,
69            embedder,
70            write_hook: None,
71            summarizer: None,
72        })
73    }
74
75    /// Attach a write hook consulted before every persist.
76    pub fn with_write_hook(mut self, hook: Arc<dyn cel_memory::MemoryWriteHook>) -> Self {
77        self.write_hook = Some(hook);
78        self
79    }
80
81    /// Attach a summarizer for session summaries and rollups (Phase 3).
82    pub fn with_summarizer(mut self, summarizer: Arc<dyn cel_memory::Summarizer>) -> Self {
83        self.summarizer = Some(summarizer);
84        self
85    }
86
87    /// Pool accessor for integration tests.
88    #[doc(hidden)]
89    pub fn pool(&self) -> &PgPool {
90        &self.pool
91    }
92}
93
94#[async_trait]
95impl MemoryProvider for PostgresMemoryProvider {
96    async fn retrieve(&self, query: MemoryQuery) -> MemoryResult<Vec<MemoryChunk>> {
97        if query.text.trim().is_empty() {
98            return Err(MemoryError::InvalidArgument(
99                "query.text must not be empty".into(),
100            ));
101        }
102        let k = query.k.max(1);
103        let candidate_k = (3 * k).max(16) as i64;
104        let q_embedding = self.embedder.embed(&query.text).await?;
105        let q_vec = Vector::from(q_embedding);
106        let vector_sql = match query.caller_scope {
107            CallerScope::Global => {
108                "SELECT c.id
109                 FROM memory_chunks c
110                 INNER JOIN memory_vectors v ON v.chunk_id = c.id
111                 ORDER BY v.embedding <=> $1
112                 LIMIT $2"
113            }
114            CallerScope::Own | CallerScope::OwnPlusShared => {
115                "SELECT c.id
116                 FROM memory_chunks c
117                 INNER JOIN memory_vectors v ON v.chunk_id = c.id
118                 WHERE c.caller_id = $2
119                 ORDER BY v.embedding <=> $1
120                 LIMIT $3"
121            }
122        };
123        let vector_rows = match query.caller_scope {
124            CallerScope::Global => {
125                sqlx::query(vector_sql)
126                    .bind(q_vec.clone())
127                    .bind(candidate_k)
128                    .fetch_all(&self.pool)
129                    .await
130            }
131            CallerScope::Own => {
132                sqlx::query(vector_sql)
133                    .bind(q_vec.clone())
134                    .bind(&query.caller_id)
135                    .bind(candidate_k)
136                    .fetch_all(&self.pool)
137                    .await
138            }
139            CallerScope::OwnPlusShared => {
140                sqlx::query(
141                    "SELECT c.id
142                     FROM memory_chunks c
143                     INNER JOIN memory_vectors v ON v.chunk_id = c.id
144                     WHERE (c.caller_id = $2 OR c.shareable = TRUE)
145                     ORDER BY v.embedding <=> $1
146                     LIMIT $3",
147                )
148                .bind(q_vec.clone())
149                .bind(&query.caller_id)
150                .bind(candidate_k)
151                .fetch_all(&self.pool)
152                .await
153            }
154        }
155        .map_err(|e| MemoryError::Storage(e.to_string()))?;
156
157        let fts_rows = match query.caller_scope {
158            CallerScope::Global => {
159                sqlx::query(
160                    "SELECT c.id
161                     FROM memory_chunks c
162                     WHERE to_tsvector('english', c.content) @@ plainto_tsquery('english', $1)
163                     ORDER BY ts_rank(
164                         to_tsvector('english', c.content),
165                         plainto_tsquery('english', $1)
166                     ) DESC
167                     LIMIT $2",
168                )
169                .bind(&query.text)
170                .bind(candidate_k)
171                .fetch_all(&self.pool)
172                .await
173            }
174            CallerScope::Own => {
175                sqlx::query(
176                    "SELECT c.id
177                     FROM memory_chunks c
178                     WHERE c.caller_id = $1
179                       AND to_tsvector('english', c.content) @@ plainto_tsquery('english', $2)
180                     ORDER BY ts_rank(
181                         to_tsvector('english', c.content),
182                         plainto_tsquery('english', $2)
183                     ) DESC
184                     LIMIT $3",
185                )
186                .bind(&query.caller_id)
187                .bind(&query.text)
188                .bind(candidate_k)
189                .fetch_all(&self.pool)
190                .await
191            }
192            CallerScope::OwnPlusShared => {
193                sqlx::query(
194                    "SELECT c.id
195                     FROM memory_chunks c
196                     WHERE (c.caller_id = $1 OR c.shareable = TRUE)
197                       AND to_tsvector('english', c.content) @@ plainto_tsquery('english', $2)
198                     ORDER BY ts_rank(
199                         to_tsvector('english', c.content),
200                         plainto_tsquery('english', $2)
201                     ) DESC
202                     LIMIT $3",
203                )
204                .bind(&query.caller_id)
205                .bind(&query.text)
206                .bind(candidate_k)
207                .fetch_all(&self.pool)
208                .await
209            }
210        }
211        .map_err(|e| MemoryError::Storage(e.to_string()))?;
212
213        let mut scores: HashMap<String, f32> = HashMap::new();
214        for (rank, row) in vector_rows.into_iter().enumerate() {
215            let id: String = row
216                .try_get("id")
217                .map_err(|e| MemoryError::Storage(e.to_string()))?;
218            *scores.entry(id).or_insert(0.0) += rrf(rank, 60.0);
219        }
220        for (rank, row) in fts_rows.into_iter().enumerate() {
221            let id: String = row
222                .try_get("id")
223                .map_err(|e| MemoryError::Storage(e.to_string()))?;
224            *scores.entry(id).or_insert(0.0) += rrf(rank, 60.0);
225        }
226
227        if scores.is_empty() {
228            return Ok(Vec::new());
229        }
230
231        let ids: Vec<String> = scores.keys().cloned().collect();
232        let rows = sqlx::query(CHUNKS_BY_IDS_SQL)
233            .bind(&ids)
234            .fetch_all(&self.pool)
235            .await
236            .map_err(|e| MemoryError::Storage(e.to_string()))?;
237
238        let mut filtered: Vec<MemoryChunk> = rows
239            .iter()
240            .filter_map(|row| row_to_chunk(row).ok())
241            .filter(|c| chunk_matches_query(c, &query))
242            .collect();
243
244        filtered.sort_by(|a, b| {
245            let sa = scores.get(&a.id).copied().unwrap_or(0.0);
246            let sb = scores.get(&b.id).copied().unwrap_or(0.0);
247            sb.partial_cmp(&sa).unwrap_or(std::cmp::Ordering::Equal)
248        });
249        filtered.truncate(k);
250        Ok(filtered)
251    }
252
253    async fn get(&self, chunk_id: &str) -> MemoryResult<Option<MemoryChunk>> {
254        let row = sqlx::query(GET_CHUNK_SQL)
255            .bind(chunk_id)
256            .fetch_optional(&self.pool)
257            .await
258            .map_err(|e| MemoryError::Storage(e.to_string()))?;
259        row.map(|r| row_to_chunk(&r)).transpose()
260    }
261
262    async fn get_session(&self, session_id: &str) -> MemoryResult<Option<MemorySession>> {
263        let row = sqlx::query(
264            "SELECT id, started_at, ended_at, caller_id, title, summary, outcome, metadata
265             FROM memory_sessions WHERE id = $1",
266        )
267        .bind(session_id)
268        .fetch_optional(&self.pool)
269        .await
270        .map_err(|e| MemoryError::Storage(e.to_string()))?;
271        row.map(|r| row_to_session(&r)).transpose()
272    }
273
274    async fn list_sessions(&self, filter: SessionFilter) -> MemoryResult<Vec<MemorySession>> {
275        let mut builder: QueryBuilder<Postgres> = QueryBuilder::new(
276            "SELECT id, started_at, ended_at, caller_id, title, summary, outcome, metadata \
277             FROM memory_sessions WHERE TRUE",
278        );
279        if let Some(caller) = &filter.caller_id {
280            builder.push(" AND caller_id = ");
281            builder.push_bind(caller);
282        }
283        if filter.open_only {
284            builder.push(" AND outcome = 'open'");
285        } else if let Some(outcome) = filter.outcome {
286            builder.push(" AND outcome = ");
287            builder.push_bind(outcome_str(outcome));
288        }
289        builder.push(" ORDER BY started_at DESC");
290        let rows = builder
291            .build()
292            .fetch_all(&self.pool)
293            .await
294            .map_err(|e| MemoryError::Storage(e.to_string()))?;
295        rows.iter().map(row_to_session).collect()
296    }
297
298    async fn write(&self, new_chunk: NewMemoryChunk) -> MemoryResult<MemoryChunk> {
299        if new_chunk.content.trim().is_empty() {
300            return Err(MemoryError::InvalidArgument(
301                "content must not be empty".into(),
302            ));
303        }
304        if let Some(hook) = &self.write_hook {
305            match hook.before_write(&new_chunk).await? {
306                cel_memory::WriteDecision::Allow => {}
307                cel_memory::WriteDecision::Redact { reason } => {
308                    return Ok(MemoryChunk {
309                        id: Uuid::now_v7().to_string(),
310                        created_at: Utc::now(),
311                        kind: new_chunk.kind,
312                        tier: MemoryTier::Session,
313                        source: new_chunk.source,
314                        session_id: new_chunk.session_id,
315                        project_root: new_chunk.project_root,
316                        caller_id: new_chunk.caller_id,
317                        content: format!("<redacted: {reason}>"),
318                        metadata: serde_json::json!({"redacted": true, "reason": reason}),
319                        importance: 0.0,
320                        pinned: false,
321                        shareable: false,
322                        superseded_by: None,
323                        embedding_model: "none".into(),
324                        embedding_dim: 0,
325                    });
326                }
327            }
328        }
329
330        let id = Uuid::now_v7().to_string();
331        let created_at = Utc::now();
332        let importance = cel_memory::score_importance(&new_chunk);
333        let embedding = self.embedder.embed(&new_chunk.content).await?;
334        if embedding.len() != EMBEDDING_DIM {
335            return Err(MemoryError::Internal(format!(
336                "embedder produced dim {}, expected {EMBEDDING_DIM}",
337                embedding.len()
338            )));
339        }
340
341        let chunk = MemoryChunk {
342            id: id.clone(),
343            created_at,
344            kind: new_chunk.kind,
345            tier: MemoryTier::Session,
346            source: new_chunk.source,
347            session_id: new_chunk.session_id.clone(),
348            project_root: new_chunk.project_root.clone(),
349            caller_id: new_chunk.caller_id.clone(),
350            content: new_chunk.content.clone(),
351            metadata: new_chunk.metadata.clone(),
352            importance,
353            pinned: new_chunk.pinned,
354            shareable: new_chunk.shareable,
355            superseded_by: None,
356            embedding_model: self.embedder.model_name().to_string(),
357            embedding_dim: EMBEDDING_DIM as u32,
358        };
359
360        let mut tx = self
361            .pool
362            .begin()
363            .await
364            .map_err(|e| MemoryError::Storage(e.to_string()))?;
365
366        sqlx::query(
367            "INSERT INTO memory_chunks(
368                id, created_at, kind, tier, source, session_id, project_root,
369                caller_id, content, metadata, importance, pinned, shareable,
370                superseded_by, embedding_model, embedding_dim
371            ) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16)",
372        )
373        .bind(&chunk.id)
374        .bind(chunk.created_at)
375        .bind(kind_str(chunk.kind))
376        .bind(tier_str(chunk.tier))
377        .bind(source_str(chunk.source))
378        .bind(&chunk.session_id)
379        .bind(&chunk.project_root)
380        .bind(&chunk.caller_id)
381        .bind(&chunk.content)
382        .bind(&chunk.metadata)
383        .bind(chunk.importance)
384        .bind(chunk.pinned)
385        .bind(chunk.shareable)
386        .bind::<Option<String>>(None)
387        .bind(&chunk.embedding_model)
388        .bind(chunk.embedding_dim as i32)
389        .execute(&mut *tx)
390        .await
391        .map_err(|e| MemoryError::Storage(e.to_string()))?;
392
393        sqlx::query("INSERT INTO memory_vectors (chunk_id, embedding) VALUES ($1, $2)")
394            .bind(&chunk.id)
395            .bind(Vector::from(embedding))
396            .execute(&mut *tx)
397            .await
398            .map_err(|e| MemoryError::Storage(e.to_string()))?;
399
400        tx.commit()
401            .await
402            .map_err(|e| MemoryError::Storage(e.to_string()))?;
403        Ok(chunk)
404    }
405
406    async fn write_batch(&self, chunks: Vec<NewMemoryChunk>) -> MemoryResult<Vec<MemoryChunk>> {
407        let mut out = Vec::with_capacity(chunks.len());
408        for chunk in chunks {
409            out.push(self.write(chunk).await?);
410        }
411        Ok(out)
412    }
413
414    async fn open_session(&self, init: NewMemorySession) -> MemoryResult<MemorySession> {
415        let session = MemorySession {
416            id: Uuid::now_v7().to_string(),
417            started_at: Utc::now(),
418            ended_at: None,
419            caller_id: init.caller_id,
420            title: init.title,
421            summary: None,
422            outcome: SessionOutcome::Open,
423            metadata: init.metadata,
424        };
425        sqlx::query(
426            "INSERT INTO memory_sessions (id, started_at, ended_at, caller_id, title, summary, outcome, metadata)
427             VALUES ($1,$2,$3,$4,$5,$6,$7,$8)",
428        )
429        .bind(&session.id)
430        .bind(session.started_at)
431        .bind(session.ended_at)
432        .bind(&session.caller_id)
433        .bind(&session.title)
434        .bind(&session.summary)
435        .bind(outcome_str(session.outcome))
436        .bind(&session.metadata)
437        .execute(&self.pool)
438        .await
439        .map_err(|e| MemoryError::Storage(e.to_string()))?;
440        Ok(session)
441    }
442
443    async fn close_session(&self, session_id: &str, outcome: SessionOutcome) -> MemoryResult<()> {
444        let ended_at = Utc::now();
445        let rows =
446            sqlx::query("UPDATE memory_sessions SET ended_at = $1, outcome = $2 WHERE id = $3")
447                .bind(ended_at)
448                .bind(outcome_str(outcome))
449                .bind(session_id)
450                .execute(&self.pool)
451                .await
452                .map_err(|e| MemoryError::Storage(e.to_string()))?;
453        if rows.rows_affected() == 0 {
454            return Err(MemoryError::NotFound(session_id.into()));
455        }
456        Ok(())
457    }
458
459    async fn rename_session(&self, session_id: &str, title: &str) -> MemoryResult<()> {
460        let rows = sqlx::query("UPDATE memory_sessions SET title = $1 WHERE id = $2")
461            .bind(title)
462            .bind(session_id)
463            .execute(&self.pool)
464            .await
465            .map_err(|e| MemoryError::Storage(e.to_string()))?;
466        if rows.rows_affected() == 0 {
467            return Err(MemoryError::NotFound(session_id.into()));
468        }
469        Ok(())
470    }
471
472    async fn stats(&self) -> MemoryResult<MemoryStats> {
473        let total_chunks: i64 = sqlx::query_scalar("SELECT COUNT(*)::bigint FROM memory_chunks")
474            .fetch_one(&self.pool)
475            .await
476            .map_err(|e| MemoryError::Storage(e.to_string()))?;
477        let total_sessions: i64 =
478            sqlx::query_scalar("SELECT COUNT(*)::bigint FROM memory_sessions")
479                .fetch_one(&self.pool)
480                .await
481                .map_err(|e| MemoryError::Storage(e.to_string()))?;
482        let embedding_model: Option<String> = sqlx::query_scalar(
483            "SELECT embedding_model FROM memory_chunks ORDER BY created_at DESC LIMIT 1",
484        )
485        .fetch_optional(&self.pool)
486        .await
487        .map_err(|e| MemoryError::Storage(e.to_string()))?;
488
489        Ok(MemoryStats {
490            total_chunks: total_chunks as usize,
491            total_sessions: total_sessions as usize,
492            embedding_model,
493            ..MemoryStats::default()
494        })
495    }
496
497    async fn summarize_session(&self, _session_id: &str) -> MemoryResult<MemoryChunk> {
498        if self.summarizer.is_none() {
499            return Err(MemoryError::NotImplemented(
500                "PostgresMemoryProvider::summarize_session — attach summarizer via with_summarizer",
501            ));
502        }
503        Err(MemoryError::NotImplemented(
504            "PostgresMemoryProvider::summarize_session — Phase 3",
505        ))
506    }
507
508    async fn rollup_day(&self, _date: chrono::NaiveDate) -> MemoryResult<Vec<MemoryChunk>> {
509        Err(MemoryError::NotImplemented(
510            "PostgresMemoryProvider::rollup_day — Phase 3",
511        ))
512    }
513
514    async fn rollup_rule_week(
515        &self,
516        _rule_id: &str,
517        _week_start: chrono::NaiveDate,
518    ) -> MemoryResult<MemoryChunk> {
519        Err(MemoryError::NotImplemented(
520            "PostgresMemoryProvider::rollup_rule_week — Phase 3",
521        ))
522    }
523
524    async fn run_aging_sweep(&self) -> MemoryResult<cel_memory::AgingReport> {
525        Err(MemoryError::NotImplemented(
526            "PostgresMemoryProvider::run_aging_sweep — Phase 2",
527        ))
528    }
529
530    async fn re_embed_all(&self, _target_model: &str) -> MemoryResult<cel_memory::ReEmbedReport> {
531        Err(MemoryError::NotImplemented(
532            "PostgresMemoryProvider::re_embed_all — Phase 4",
533        ))
534    }
535
536    async fn export(
537        &self,
538        _filter: cel_memory::ExportFilter,
539    ) -> MemoryResult<cel_memory::ExportBundle> {
540        Err(MemoryError::NotImplemented(
541            "PostgresMemoryProvider::export — Phase 2",
542        ))
543    }
544
545    async fn pin(&self, _chunk_id: &str, _pinned: bool) -> MemoryResult<()> {
546        Err(MemoryError::NotImplemented(
547            "PostgresMemoryProvider::pin — Phase 2",
548        ))
549    }
550
551    async fn update_importance(&self, _chunk_id: &str, _importance: f32) -> MemoryResult<()> {
552        Err(MemoryError::NotImplemented(
553            "PostgresMemoryProvider::update_importance — Phase 2",
554        ))
555    }
556
557    async fn supersede(&self, _old_id: &str, _new_id: &str) -> MemoryResult<()> {
558        Err(MemoryError::NotImplemented(
559            "PostgresMemoryProvider::supersede — Phase 2",
560        ))
561    }
562
563    async fn record_access(
564        &self,
565        _chunk_id: &str,
566        _retrieved_by: &str,
567        _used: bool,
568    ) -> MemoryResult<()> {
569        Err(MemoryError::NotImplemented(
570            "PostgresMemoryProvider::record_access — Phase 2",
571        ))
572    }
573
574    async fn delete(
575        &self,
576        _chunk_id: &str,
577        _reason: cel_memory::EvictionReason,
578    ) -> MemoryResult<()> {
579        Err(MemoryError::NotImplemented(
580            "PostgresMemoryProvider::delete — Phase 2",
581        ))
582    }
583
584    async fn delete_matching(
585        &self,
586        _predicate: cel_memory::MemoryPredicate,
587        _reason: cel_memory::EvictionReason,
588    ) -> MemoryResult<usize> {
589        Err(MemoryError::NotImplemented(
590            "PostgresMemoryProvider::delete_matching — Phase 2",
591        ))
592    }
593
594    async fn purge_all(&self) -> MemoryResult<cel_memory::PurgeReport> {
595        Err(MemoryError::NotImplemented(
596            "PostgresMemoryProvider::purge_all — Phase 2",
597        ))
598    }
599}
600
601fn row_to_session(row: &PgRow) -> MemoryResult<MemorySession> {
602    let metadata: serde_json::Value = row.try_get("metadata").unwrap_or(serde_json::json!({}));
603    Ok(MemorySession {
604        id: row
605            .try_get("id")
606            .map_err(|e| MemoryError::Storage(e.to_string()))?,
607        started_at: row
608            .try_get("started_at")
609            .map_err(|e| MemoryError::Storage(e.to_string()))?,
610        ended_at: row.try_get("ended_at").ok(),
611        caller_id: row
612            .try_get("caller_id")
613            .map_err(|e| MemoryError::Storage(e.to_string()))?,
614        title: row.try_get("title").ok(),
615        summary: row.try_get("summary").ok(),
616        outcome: str_to_outcome(
617            &row.try_get::<String, _>("outcome")
618                .map_err(|e| MemoryError::Storage(e.to_string()))?,
619        )?,
620        metadata,
621    })
622}
cel_memory_postgres/provider.rs

cel_memory_postgres/
provider.rs
