daimon_plugin_pgvector/
store.rs

1//! [`PgVectorStore`] — a pgvector-backed [`VectorStore`] implementation.
2
3use std::collections::HashMap;
4
5use daimon_core::vector_store::VectorStore;
6use daimon_core::{DaimonError, Document, Result, ScoredDocument};
7use deadpool_postgres::Pool;
8use pgvector::Vector;
9
10use crate::DistanceMetric;
11
12/// A [`VectorStore`] backed by PostgreSQL with the pgvector extension.
13///
14/// Use [`PgVectorStoreBuilder`](crate::PgVectorStoreBuilder) to construct.
15pub struct PgVectorStore {
16    pub(crate) pool: Pool,
17    pub(crate) table: String,
18    pub(crate) dimensions: usize,
19    pub(crate) distance_metric: DistanceMetric,
20}
21
22impl PgVectorStore {
23    /// Returns the distance operator used in SQL ORDER BY clauses.
24    fn distance_operator(&self) -> &'static str {
25        match self.distance_metric {
26            DistanceMetric::Cosine => "<=>",
27            DistanceMetric::L2 => "<->",
28            DistanceMetric::InnerProduct => "<#>",
29        }
30    }
31
32    /// Builds the SQL expression that converts pgvector's distance operator
33    /// (`$1` is the query vector) into a similarity score.
34    ///
35    /// See [`VectorStore::query`] for the rationale behind each transform.
36    /// L2 uses `1 / (1 + distance)` so the score stays in `(0, 1]` and
37    /// monotonic; a naive `1 - distance` would be unbounded and go negative.
38    fn score_expr(&self) -> String {
39        let op = self.distance_operator();
40        match self.distance_metric {
41            DistanceMetric::Cosine => format!("1.0 - (embedding {op} $1)"),
42            DistanceMetric::L2 => format!("1.0 / (1.0 + (embedding {op} $1))"),
43            DistanceMetric::InnerProduct => format!("-(embedding {op} $1)"),
44        }
45    }
46
47    /// Returns a reference to the underlying connection pool.
48    pub fn pool(&self) -> &Pool {
49        &self.pool
50    }
51
52    /// Returns the table name used by this store.
53    pub fn table(&self) -> &str {
54        &self.table
55    }
56
57    /// Returns the configured vector dimensions.
58    pub fn dimensions(&self) -> usize {
59        self.dimensions
60    }
61}
62
63impl VectorStore for PgVectorStore {
64    async fn upsert(&self, id: &str, embedding: Vec<f32>, document: Document) -> Result<()> {
65        if embedding.len() != self.dimensions {
66            return Err(DaimonError::Other(format!(
67                "embedding dimension mismatch: expected {}, got {}",
68                self.dimensions,
69                embedding.len()
70            )));
71        }
72
73        let client = self
74            .pool
75            .get()
76            .await
77            .map_err(|e| DaimonError::Other(format!("pgvector pool error: {e}")))?;
78
79        let vec = Vector::from(embedding);
80        let metadata = serde_json::to_value(&document.metadata)
81            .map_err(|e| DaimonError::Other(format!("metadata serialization error: {e}")))?;
82
83        let sql = format!(
84            "INSERT INTO {} (id, embedding, content, metadata) VALUES ($1, $2, $3, $4) \
85             ON CONFLICT (id) DO UPDATE SET embedding = EXCLUDED.embedding, \
86             content = EXCLUDED.content, metadata = EXCLUDED.metadata",
87            self.table
88        );
89
90        client
91            .execute(&sql as &str, &[&id, &vec, &document.content, &metadata])
92            .await
93            .map_err(|e| DaimonError::Other(format!("pgvector upsert error: {e}")))?;
94
95        Ok(())
96    }
97
98    async fn query(&self, embedding: Vec<f32>, top_k: usize) -> Result<Vec<ScoredDocument>> {
99        if embedding.len() != self.dimensions {
100            return Err(DaimonError::Other(format!(
101                "embedding dimension mismatch: expected {}, got {}",
102                self.dimensions,
103                embedding.len()
104            )));
105        }
106
107        let client = self
108            .pool
109            .get()
110            .await
111            .map_err(|e| DaimonError::Other(format!("pgvector pool error: {e}")))?;
112
113        let vec = Vector::from(embedding);
114        let op = self.distance_operator();
115        let score_expr = self.score_expr();
116
117        let sql = format!(
118            "SELECT id, content, metadata, {score_expr} AS score \
119             FROM {} ORDER BY embedding {op} $1 LIMIT $2",
120            self.table
121        );
122
123        let rows = client
124            .query(&sql as &str, &[&vec, &(top_k as i64)])
125            .await
126            .map_err(|e| DaimonError::Other(format!("pgvector query error: {e}")))?;
127
128        let mut results = Vec::with_capacity(rows.len());
129        for row in rows {
130            let content: String = row.get("content");
131            let metadata_val: serde_json::Value = row.get("metadata");
132            let score: f64 = row.get("score");
133
134            let metadata: HashMap<String, serde_json::Value> =
135                serde_json::from_value(metadata_val).unwrap_or_default();
136
137            let doc = Document {
138                content,
139                metadata,
140                score: Some(score),
141            };
142            results.push(ScoredDocument::new(doc, score));
143        }
144
145        Ok(results)
146    }
147
148    async fn delete(&self, id: &str) -> Result<bool> {
149        let client = self
150            .pool
151            .get()
152            .await
153            .map_err(|e| DaimonError::Other(format!("pgvector pool error: {e}")))?;
154
155        let sql = format!("DELETE FROM {} WHERE id = $1", self.table);
156        let deleted = client
157            .execute(&sql as &str, &[&id])
158            .await
159            .map_err(|e| DaimonError::Other(format!("pgvector delete error: {e}")))?;
160
161        Ok(deleted > 0)
162    }
163
164    async fn count(&self) -> Result<usize> {
165        let client = self
166            .pool
167            .get()
168            .await
169            .map_err(|e| DaimonError::Other(format!("pgvector pool error: {e}")))?;
170
171        let sql = format!("SELECT COUNT(*) AS cnt FROM {}", self.table);
172        let row = client
173            .query_one(&sql as &str, &[])
174            .await
175            .map_err(|e| DaimonError::Other(format!("pgvector count error: {e}")))?;
176
177        let count: i64 = row.get("cnt");
178        Ok(count as usize)
179    }
180}
181
182#[cfg(test)]
183mod tests {
184    use super::*;
185
186    #[test]
187    fn test_distance_operator() {
188        let store = PgVectorStore {
189            pool: create_dummy_pool(),
190            table: "t".into(),
191            dimensions: 3,
192            distance_metric: DistanceMetric::Cosine,
193        };
194        assert_eq!(store.distance_operator(), "<=>");
195
196        let store = PgVectorStore {
197            dimensions: 3,
198            distance_metric: DistanceMetric::L2,
199            ..store
200        };
201        assert_eq!(store.distance_operator(), "<->");
202
203        let store = PgVectorStore {
204            distance_metric: DistanceMetric::InnerProduct,
205            ..store
206        };
207        assert_eq!(store.distance_operator(), "<#>");
208    }
209
210    #[test]
211    fn test_score_expr_per_metric() {
212        let base = PgVectorStore {
213            pool: create_dummy_pool(),
214            table: "t".into(),
215            dimensions: 3,
216            distance_metric: DistanceMetric::Cosine,
217        };
218        assert_eq!(base.score_expr(), "1.0 - (embedding <=> $1)");
219
220        let l2 = PgVectorStore {
221            distance_metric: DistanceMetric::L2,
222            ..base
223        };
224        // L2 must use the bounded transform, not `1 - distance`.
225        assert_eq!(l2.score_expr(), "1.0 / (1.0 + (embedding <-> $1))");
226
227        let ip = PgVectorStore {
228            distance_metric: DistanceMetric::InnerProduct,
229            ..l2
230        };
231        assert_eq!(ip.score_expr(), "-(embedding <#> $1)");
232    }
233
234    fn create_dummy_pool() -> Pool {
235        let cfg = deadpool_postgres::Config {
236            host: Some("localhost".into()),
237            port: Some(5432),
238            dbname: Some("test".into()),
239            ..Default::default()
240        };
241        cfg.create_pool(None, tokio_postgres::NoTls).unwrap()
242    }
243}
daimon_plugin_pgvector/store.rs

daimon_plugin_pgvector/
store.rs
