klieo_memory_pgvector/
long_term.rs

1//! `PgvectorLongTerm` — `LongTermMemory` over one Postgres table with a
2//! `vector(dim)` column (HNSW, cosine) and `scope_kind` / `scope_value` filter
3//! columns. Metadata is stored as `jsonb`.
4//!
5//! The table is provisioned lazily on first `remember`; a read that races
6//! ahead of it maps the `undefined_table` error to an empty result.
7
8use crate::client::PgvectorHandle;
9use crate::embedder::Embedder;
10use crate::error::{is_undefined_table, store_err};
11use async_trait::async_trait;
12use klieo_core::error::MemoryError;
13use klieo_core::ids::FactId;
14use klieo_core::memory::{Fact, LongTermMemory, Scope};
15use klieo_memory_graph::FilterableLongTermMemory;
16use sqlx::Row;
17use std::sync::Arc;
18use uuid::Uuid;
19
20fn scope_columns(scope: &Scope) -> (&'static str, String) {
21    match scope {
22        Scope::Workspace(s) => ("workspace", s.clone()),
23        Scope::Agent(s) => ("agent", s.clone()),
24        Scope::Global => ("global", String::new()),
25    }
26}
27
28fn row_to_fact(text: String, metadata: serde_json::Value) -> Fact {
29    Fact::new(text).with_metadata(metadata)
30}
31
32/// Render an embedding as a pgvector text literal (`[0.1,0.2,...]`). Bound as a
33/// plain `text` parameter and cast with `::vector` in SQL, which avoids a
34/// dedicated pgvector client crate (and the sqlx-driver version pin it forces).
35fn vector_literal(embedding: &[f32]) -> String {
36    let mut out = String::with_capacity(embedding.len() * 8 + 2);
37    out.push('[');
38    for (i, value) in embedding.iter().enumerate() {
39        if i > 0 {
40            out.push(',');
41        }
42        out.push_str(&value.to_string());
43    }
44    out.push(']');
45    out
46}
47
48/// Parse a caller-facing [`FactId`] back into the `uuid` column type. Ids are
49/// minted as UUIDs by [`PgvectorLongTerm::remember`], so a parse failure means
50/// a foreign id and is surfaced as a typed error rather than silently dropped.
51fn parse_fact_id(id: &FactId) -> Result<Uuid, MemoryError> {
52    Uuid::parse_str(&id.to_string())
53        .map_err(|e| MemoryError::Store(format!("fact id is not a uuid: {e}")))
54}
55
56/// `LongTermMemory` + `FilterableLongTermMemory` over the single table
57/// described in the module docs; cosine recall via the HNSW index, scope
58/// enforced by the `scope_kind`/`scope_value` columns.
59pub struct PgvectorLongTerm {
60    handle: PgvectorHandle,
61    embedder: Arc<dyn Embedder>,
62    embedder_id: String,
63}
64
65impl PgvectorLongTerm {
66    pub(crate) fn new(
67        handle: PgvectorHandle,
68        embedder: Arc<dyn Embedder>,
69        embedder_id: String,
70    ) -> Self {
71        Self {
72            handle,
73            embedder,
74            embedder_id,
75        }
76    }
77
78    /// Embed one text and validate it against the embedder's declared
79    /// dimension before it reaches the `vector(dim)` column.
80    async fn embed_one(&self, text: &str) -> Result<Vec<f32>, MemoryError> {
81        let dim = self.embedder.dimension();
82        let vector = self
83            .embedder
84            .embed(&[text.to_string()])
85            .await?
86            .into_iter()
87            .next()
88            .ok_or_else(|| MemoryError::Embedding("embedder returned empty vec".into()))?;
89        if vector.len() != dim {
90            return Err(MemoryError::Embedding(format!(
91                "embedder produced {}-dim vector, expected {}",
92                vector.len(),
93                dim
94            )));
95        }
96        Ok(vector)
97    }
98}
99
100#[async_trait]
101impl LongTermMemory for PgvectorLongTerm {
102    async fn remember(&self, scope: Scope, fact: Fact) -> Result<FactId, MemoryError> {
103        self.handle
104            .ensure_table(self.embedder.dimension() as u64)
105            .await?;
106        let vector = self.embed_one(&fact.text).await?;
107        let id = Uuid::new_v4();
108        let (kind, value) = scope_columns(&scope);
109        let sql = format!(
110            "INSERT INTO {} (fact_id, text, metadata, embedding, scope_kind, scope_value) \
111             VALUES ($1, $2, $3, $4::vector, $5, $6)",
112            self.handle.table
113        );
114        sqlx::query(&sql)
115            .bind(id)
116            .bind(&fact.text)
117            .bind(&fact.metadata)
118            .bind(vector_literal(&vector))
119            .bind(kind)
120            .bind(value)
121            .execute(&self.handle.pool)
122            .await
123            .map_err(store_err)?;
124        Ok(FactId::new(id.to_string()))
125    }
126
127    async fn recall(&self, scope: Scope, query: &str, k: usize) -> Result<Vec<Fact>, MemoryError> {
128        if k == 0 {
129            return Ok(Vec::new());
130        }
131        let vector = self.embed_one(query).await?;
132        let (kind, value) = scope_columns(&scope);
133        let sql = format!(
134            "SELECT text, metadata FROM {} \
135             WHERE scope_kind = $1 AND scope_value = $2 \
136             ORDER BY embedding <=> $3::vector LIMIT $4",
137            self.handle.table
138        );
139        let rows = sqlx::query(&sql)
140            .bind(kind)
141            .bind(value)
142            .bind(vector_literal(&vector))
143            .bind(k as i64)
144            .fetch_all(&self.handle.pool)
145            .await;
146        rows_to_facts(rows)
147    }
148
149    async fn forget(&self, id: FactId) -> Result<(), MemoryError> {
150        let uuid = parse_fact_id(&id)?;
151        let sql = format!("DELETE FROM {} WHERE fact_id = $1", self.handle.table);
152        match sqlx::query(&sql)
153            .bind(uuid)
154            .execute(&self.handle.pool)
155            .await
156        {
157            Ok(_) => Ok(()),
158            Err(e) if is_undefined_table(&e) => Ok(()),
159            Err(e) => Err(store_err(e)),
160        }
161    }
162}
163
164#[async_trait]
165impl FilterableLongTermMemory for PgvectorLongTerm {
166    async fn recall_filtered(
167        &self,
168        scope: Scope,
169        query: &str,
170        k: usize,
171        candidate_ids: &[FactId],
172    ) -> Result<Vec<Fact>, MemoryError> {
173        if k == 0 || candidate_ids.is_empty() {
174            return Ok(Vec::new());
175        }
176        let candidates: Vec<Uuid> = candidate_ids
177            .iter()
178            .map(parse_fact_id)
179            .collect::<Result<_, _>>()?;
180        let vector = self.embed_one(query).await?;
181        let (kind, value) = scope_columns(&scope);
182        let sql = format!(
183            "SELECT text, metadata FROM {} \
184             WHERE scope_kind = $1 AND scope_value = $2 AND fact_id = ANY($3) \
185             ORDER BY embedding <=> $4::vector LIMIT $5",
186            self.handle.table
187        );
188        let rows = sqlx::query(&sql)
189            .bind(kind)
190            .bind(value)
191            .bind(candidates)
192            .bind(vector_literal(&vector))
193            .bind(k as i64)
194            .fetch_all(&self.handle.pool)
195            .await;
196        rows_to_facts(rows)
197    }
198
199    fn embedder_id(&self) -> &str {
200        &self.embedder_id
201    }
202}
203
204/// Decode a query result into `Vec<Fact>`, mapping `undefined_table` (a read
205/// that raced ahead of the first `remember`) to an empty result.
206fn rows_to_facts(
207    rows: Result<Vec<sqlx::postgres::PgRow>, sqlx::Error>,
208) -> Result<Vec<Fact>, MemoryError> {
209    let rows = match rows {
210        Ok(rows) => rows,
211        Err(e) if is_undefined_table(&e) => return Ok(Vec::new()),
212        Err(e) => return Err(store_err(e)),
213    };
214    rows.into_iter()
215        .map(|row| {
216            let text: String = row.try_get("text").map_err(store_err)?;
217            let metadata: serde_json::Value = row.try_get("metadata").map_err(store_err)?;
218            Ok(row_to_fact(text, metadata))
219        })
220        .collect()
221}
222
223#[cfg(test)]
224mod tests {
225    use super::*;
226
227    #[allow(dead_code)]
228    fn _trait_object_safe(_: &dyn LongTermMemory) {}
229
230    #[test]
231    fn scope_columns_maps_each_variant() {
232        assert_eq!(
233            scope_columns(&Scope::Workspace("ws".into())),
234            ("workspace", "ws".to_string())
235        );
236        assert_eq!(
237            scope_columns(&Scope::Agent("a".into())),
238            ("agent", "a".to_string())
239        );
240        assert_eq!(scope_columns(&Scope::Global), ("global", String::new()));
241    }
242
243    #[test]
244    fn row_to_fact_carries_text_and_metadata() {
245        let fact = row_to_fact("hi".to_string(), serde_json::json!({"k": 1}));
246        assert_eq!(fact.text, "hi");
247        assert_eq!(fact.metadata, serde_json::json!({"k": 1}));
248    }
249
250    #[test]
251    fn parse_fact_id_rejects_non_uuid() {
252        let err = parse_fact_id(&FactId::new("not-a-uuid".to_string())).unwrap_err();
253        assert!(matches!(err, MemoryError::Store(_)));
254    }
255
256    #[test]
257    fn vector_literal_renders_pgvector_text_form() {
258        assert_eq!(vector_literal(&[]), "[]");
259        assert_eq!(vector_literal(&[1.0]), "[1]");
260        assert_eq!(vector_literal(&[0.5, -1.25, 2.0]), "[0.5,-1.25,2]");
261    }
262
263    #[test]
264    fn parse_fact_id_accepts_minted_uuid() {
265        let id = uuid::Uuid::new_v4().to_string();
266        let parsed = parse_fact_id(&FactId::new(id.clone())).unwrap();
267        assert_eq!(parsed.to_string(), id);
268    }
269}
klieo_memory_pgvector/long_term.rs

klieo_memory_pgvector/
long_term.rs
