llm_kernel/embedding/
pgvector.rs

1//! pgvector `AsyncVectorIndex` — PostgreSQL + the `pgvector` extension.
2//!
3//! `PgVectorIndex` implements [`AsyncVectorIndex`] over a `sqlx::PgPool`,
4//! mirroring `qdrant`/`elastic` (async remote vector backend). Vectors live in
5//! a `{table}(id BIGINT PK, vec vector)` relation; search uses the cosine
6//! distance operator `<=>`. Requires `CREATE EXTENSION vector` in the target DB
7//! (the table + HNSW index are created automatically by [`PgVectorIndex::new`]).
8
9use async_trait::async_trait;
10use sqlx::PgPool;
11use sqlx::QueryBuilder;
12use sqlx::postgres::PgPoolOptions;
13
14use crate::embedding::vector_index::SearchHit;
15use crate::error::{KernelError, Result};
16
17/// 검색 결과 행(id + cosine 유사도) — 튜플 대신 구조체로 sqlx `FromRow` 안정 매핑.
18#[derive(sqlx::FromRow)]
19struct ScoreRow {
20    id: i64,
21    score: f64,
22}
23
24/// f32 슬라이스 → pgvector 문자열 리터럴 `[1,2,3]` (text → vector 입력 캐스트).
25/// `pgvector::Vector`의 sqlx `Type` 바인드가 의존 환경에 따라 충돌해 문자열로 회피.
26fn vec_literal(v: &[f32]) -> String {
27    let mut s = String::from("[");
28    for (i, f) in v.iter().enumerate() {
29        if i > 0 {
30            s.push(',');
31        }
32        s.push_str(&f.to_string());
33    }
34    s.push(']');
35    s
36}
37
38/// PostgreSQL vector index backed by the `pgvector` extension.
39///
40/// All operations are async over a shared `PgPool` (MVCC, connection-pooled).
41/// The relation named `table` holds `(id BIGINT PK, vec vector(N))` rows.
42pub struct PgVectorIndex {
43    pool: PgPool,
44    table: String,
45    dim: usize,
46}
47
48impl PgVectorIndex {
49    /// Connect to `url` (libpq connstring / `postgresql://…`), create the
50    /// vector table + HNSW cosine index if missing, and return a ready index.
51    ///
52    /// `dim` is enforced by the fixed `vector(dim)` column; vectors whose
53    /// length differs from `dim` are rejected by pgvector on insert.
54    pub async fn new(url: &str, table: &str, dim: usize) -> Result<Self> {
55        validate_table_name(table)?;
56        let pool = PgPoolOptions::new()
57            .max_connections(8)
58            .connect(url)
59            .await
60            .map_err(|e| KernelError::Embedding(format!("pgvector connect: {e}")))?;
61        let idx = Self {
62            pool,
63            table: table.to_string(),
64            dim,
65        };
66        idx.init_schema().await?;
67        Ok(idx)
68    }
69
70    /// `CREATE TABLE IF NOT EXISTS {table} (id BIGINT PK, vec vector)` + HNSW
71    /// cosine index. Idempotent.
72    async fn init_schema(&self) -> Result<()> {
73        // Identifier is caller-controlled (not user input at runtime) and
74        // validated in `new`, so `format!` is acceptable here. PG cannot bind
75        // identifiers. Callers pass a fixed, validated table name.
76        sqlx::query(&format!(
77            "CREATE TABLE IF NOT EXISTS {} (id BIGINT PRIMARY KEY, vec vector({}) NOT NULL)",
78            self.table, self.dim
79        ))
80        .execute(&self.pool)
81        .await
82        .map_err(|e| KernelError::Embedding(format!("pgvector create table: {e}")))?;
83        sqlx::query(&format!(
84            "CREATE INDEX IF NOT EXISTS idx_{}_vec ON {} USING hnsw (vec vector_cosine_ops)",
85            self.table, self.table
86        ))
87        .execute(&self.pool)
88        .await
89        .map_err(|e| KernelError::Embedding(format!("pgvector hnsw index: {e}")))?;
90        Ok(())
91    }
92
93    /// Underlying connection pool.
94    ///
95    /// Callers that need **cross-table transactional consistency** — e.g. pruning
96    /// a law's chunks plus their vectors atomically in one transaction — can
97    /// `pool().begin()` and run their own DML on the same pool the index uses.
98    /// This preserves the table-name encapsulation while letting the caller
99    /// coordinate a multi-statement transaction.
100    pub fn pool(&self) -> &PgPool {
101        &self.pool
102    }
103
104    /// Remove vectors by external IDs **within a caller-provided transaction**.
105    ///
106    /// Enables atomic cross-table deletes: begin a tx on [`pool`](Self::pool),
107    /// delete related rows in other relations (chunks, edges, …), call this with
108    /// the same `&mut PgConnection`, then `commit()`. A failure anywhere rolls
109    /// back the whole set — no orphaned vectors or chunks. The table name stays
110    /// encapsulated (`format!` over the validated identifier).
111    pub async fn remove_in_tx(&self, tx: &mut sqlx::PgConnection, ids: &[u64]) -> Result<()> {
112        if ids.is_empty() {
113            return Ok(());
114        }
115        let ids: Vec<i64> = ids.iter().map(|&i| to_pg_id(i)).collect::<Result<_>>()?;
116        sqlx::query(&format!("DELETE FROM {} WHERE id = ANY($1)", self.table))
117            .bind(&ids)
118            .execute(tx)
119            .await
120            .map_err(|e| KernelError::Embedding(format!("pgvector remove_in_tx: {e}")))?;
121        Ok(())
122    }
123}
124
125#[async_trait]
126impl crate::embedding::AsyncVectorIndex for PgVectorIndex {
127    async fn add(&self, vectors: &[Vec<f32>], ids: &[u64]) -> Result<()> {
128        if vectors.len() != ids.len() {
129            return Err(KernelError::Embedding(format!(
130                "vectors.len() ({}) must equal ids.len() ({})",
131                vectors.len(),
132                ids.len()
133            )));
134        }
135        if vectors.is_empty() {
136            return Ok(());
137        }
138        // Map u64 → i64 up front: pgvector stores BIGINT (i64), so values
139        // above i64::MAX cannot be represented and must be rejected rather
140        // than silently wrapped. Single batched INSERT → one round trip.
141        let pg_ids: Vec<i64> = ids.iter().map(|&id| to_pg_id(id)).collect::<Result<_>>()?;
142        let mut q = QueryBuilder::new("INSERT INTO ");
143        q.push(self.table.as_str());
144        q.push(" (id, vec) VALUES ");
145        // pgvector `vec` 컬럼은 text 리터럴 입력 시 `::vector` 캐스트가 필수다.
146        // `push_values` 는 값별 캐스트를 붙일 수 없어 수동으로 VALUES 튜플을 조립한다
147        // (캐스트 누락 시 "column vec is of type vector but expression is of type text").
148        for (i, (v, &id)) in vectors.iter().zip(pg_ids.iter()).enumerate() {
149            if i > 0 {
150                q.push(", ");
151            }
152            q.push("(");
153            q.push_bind(id);
154            q.push(", ");
155            q.push_bind(vec_literal(v));
156            q.push("::vector)");
157        }
158        q.push(" ON CONFLICT (id) DO UPDATE SET vec = EXCLUDED.vec");
159        q.build()
160            .execute(&self.pool)
161            .await
162            .map_err(|e| KernelError::Embedding(format!("pgvector add: {e}")))?;
163        Ok(())
164    }
165
166    async fn remove(&self, ids: &[u64]) -> Result<()> {
167        if ids.is_empty() {
168            return Ok(());
169        }
170        let ids: Vec<i64> = ids.iter().map(|&i| to_pg_id(i)).collect::<Result<_>>()?;
171        sqlx::query(&format!("DELETE FROM {} WHERE id = ANY($1)", self.table))
172            .bind(&ids)
173            .execute(&self.pool)
174            .await
175            .map_err(|e| KernelError::Embedding(format!("pgvector remove: {e}")))?;
176        Ok(())
177    }
178
179    async fn search(&self, query: &[f32], k: usize) -> Result<Vec<SearchHit>> {
180        let q = vec_literal(query);
181        // cosine distance <=> : 0 (동일) .. 2 (반대). score = 1 - distance.
182        let rows: Vec<ScoreRow> = sqlx::query_as(&format!(
183            "SELECT id, 1 - (vec <=> $1::vector) AS score FROM {} ORDER BY vec <=> $1::vector LIMIT $2",
184            self.table
185        ))
186        .bind(q)
187        .bind(k as i64)
188        .fetch_all(&self.pool)
189        .await
190        .map_err(|e| KernelError::Embedding(format!("pgvector search: {e}")))?;
191        Ok(rows
192            .into_iter()
193            .map(|r| SearchHit {
194                id: r.id as u64,
195                score: r.score as f32,
196            })
197            .collect())
198    }
199
200    async fn search_filtered(
201        &self,
202        query: &[f32],
203        k: usize,
204        allowlist: &[u64],
205    ) -> Result<Vec<SearchHit>> {
206        if allowlist.is_empty() {
207            return Ok(Vec::new());
208        }
209        let q = vec_literal(query);
210        let allow: Vec<i64> = allowlist
211            .iter()
212            .map(|&i| to_pg_id(i))
213            .collect::<Result<_>>()?;
214        let rows: Vec<ScoreRow> = sqlx::query_as(&format!(
215            "SELECT id, 1 - (vec <=> $1::vector) AS score FROM {} WHERE id = ANY($2) \
216             ORDER BY vec <=> $1::vector LIMIT $3",
217            self.table
218        ))
219        .bind(q)
220        .bind(&allow)
221        .bind(k as i64)
222        .fetch_all(&self.pool)
223        .await
224        .map_err(|e| KernelError::Embedding(format!("pgvector search_filtered: {e}")))?;
225        Ok(rows
226            .into_iter()
227            .map(|r| SearchHit {
228                id: r.id as u64,
229                score: r.score as f32,
230            })
231            .collect())
232    }
233
234    async fn len(&self) -> Result<usize> {
235        let n: i64 = sqlx::query_scalar(&format!("SELECT count(*) FROM {}", self.table))
236            .fetch_one(&self.pool)
237            .await
238            .map_err(|e| KernelError::Embedding(format!("pgvector len: {e}")))?;
239        Ok(n as usize)
240    }
241
242    fn dim(&self) -> usize {
243        self.dim
244    }
245}
246
247/// `u64` external ID → PG `BIGINT` (`i64`). IDs exceeding `i64::MAX` cannot
248/// be stored in a BIGINT column — reject rather than silently wrap.
249fn to_pg_id(id: u64) -> Result<i64> {
250    i64::try_from(id).map_err(|_| KernelError::Embedding(format!("id {id} exceeds BIGINT range")))
251}
252
253/// Validate that `table` is a plain, safe SQL identifier (ASCII alphanumeric +
254/// `_`, starting with a letter or `_`). It is interpolated into DDL/DML via
255/// `format!` (PG cannot bind identifiers), so reject anything that could break
256/// out of the identifier context.
257fn validate_table_name(table: &str) -> Result<()> {
258    let mut chars = table.chars();
259    let first_ok = matches!(chars.next(), Some(c) if c.is_ascii_alphabetic() || c == '_');
260    let valid = first_ok && chars.all(|c| c.is_ascii_alphanumeric() || c == '_');
261    if !valid {
262        return Err(KernelError::Embedding(format!(
263            "invalid table identifier: {table:?}"
264        )));
265    }
266    Ok(())
267}
268
269#[cfg(test)]
270mod tests {
271    use super::*;
272    use crate::embedding::AsyncVectorIndex;
273
274    /// `LLMKERNEL_PG_URL` 미설정 시 자동 skip (graph-pg pg.rs 패턴).
275    fn pg_url() -> Option<String> {
276        std::env::var("LLMKERNEL_PG_URL").ok()
277    }
278
279    /// add → search → search_filtered → remove 라운드트립. HNSW는 대량 데이터에서
280    /// 정확도가 보장되지만 소규모 테스트에선 정확 매칭이 간헐적일 수 있어
281    /// 여기선 id/회수 위주로 검증.
282    #[tokio::test]
283    async fn roundtrip_add_search_remove() {
284        let Some(url) = pg_url() else {
285            eprintln!("skip pgvector test: LLMKERNEL_PG_URL unset");
286            return;
287        };
288        let table = format!("lk_test_{}", line!());
289        let idx = PgVectorIndex::new(&url, &table, 3).await.expect("new");
290
291        let vecs = vec![
292            vec![1.0, 0.0, 0.0],
293            vec![0.0, 1.0, 0.0],
294            vec![0.0, 0.0, 1.0],
295        ];
296        let ids = vec![10u64, 20, 30];
297        idx.add(&vecs, &ids).await.expect("add");
298        assert_eq!(idx.len().await.unwrap(), 3);
299
300        // nearest to [1,0,0] → id 10 우선
301        let hits = idx.search(&[1.0, 0.0, 0.0], 1).await.unwrap();
302        assert_eq!(hits.len(), 1);
303        assert_eq!(hits[0].id, 10);
304
305        // filtered: allow [20,30] → 10 제외
306        let hits = idx
307            .search_filtered(&[1.0, 0.0, 0.0], 1, &[20, 30])
308            .await
309            .unwrap();
310        assert_eq!(hits.len(), 1);
311        assert_ne!(hits[0].id, 10);
312
313        idx.remove(&[10]).await.unwrap();
314        assert_eq!(idx.len().await.unwrap(), 2);
315
316        // cleanup
317        sqlx::query(&format!("DROP TABLE IF EXISTS {}", table))
318            .execute(&idx.pool)
319            .await
320            .ok();
321    }
322
323    /// `remove_in_tx`: 같은 풀에서 시작한 트랜잭션 내 삭제가 원자 반영되는지 검증.
324    /// klr prune(chunks + vectors 단일 tx)의 전제 — 커밋 전 롤백 시 벡터 보존 확인.
325    #[tokio::test]
326    async fn remove_in_tx_atomic_delete() {
327        let Some(url) = pg_url() else {
328            eprintln!("skip pgvector test: LLMKERNEL_PG_URL unset");
329            return;
330        };
331        let table = format!("lk_txtx_{}", line!());
332        let idx = PgVectorIndex::new(&url, &table, 3).await.expect("new");
333        idx.add(&[vec![1.0, 0.0, 0.0], vec![0.0, 1.0, 0.0]], &[11, 12])
334            .await
335            .expect("add");
336        assert_eq!(idx.len().await.unwrap(), 2);
337
338        // 커밋 경로: tx 내 remove → commit → 삭제 반영
339        let mut tx = idx.pool().begin().await.expect("begin tx");
340        idx.remove_in_tx(&mut tx, &[11])
341            .await
342            .expect("remove_in_tx");
343        tx.commit().await.expect("commit");
344        assert_eq!(idx.len().await.unwrap(), 1);
345
346        // 롤백 경로: tx 내 remove → rollback → 벡터 보존(원자성)
347        let mut tx2 = idx.pool().begin().await.expect("begin tx2");
348        idx.remove_in_tx(&mut tx2, &[12])
349            .await
350            .expect("remove_in_tx2");
351        tx2.rollback().await.expect("rollback");
352        assert_eq!(idx.len().await.unwrap(), 1);
353
354        // cleanup
355        sqlx::query(&format!("DROP TABLE IF EXISTS {}", table))
356            .execute(idx.pool())
357            .await
358            .ok();
359    }
360
361    #[test]
362    fn rejects_invalid_table_name() {
363        // valid identifiers accepted
364        assert!(validate_table_name("lk_test_1").is_ok());
365        assert!(validate_table_name("_vec").is_ok());
366        // rejected — would break out of identifier context
367        assert!(validate_table_name("").is_err());
368        assert!(validate_table_name("1bad").is_err());
369        assert!(validate_table_name("rm; DROP").is_err());
370        assert!(validate_table_name("weird\"name").is_err());
371        assert!(validate_table_name("sch.tbl").is_err());
372    }
373
374    #[test]
375    fn rejects_overflowing_id() {
376        assert_eq!(to_pg_id(0).unwrap(), 0);
377        assert_eq!(to_pg_id(42).unwrap(), 42);
378        assert_eq!(to_pg_id(i64::MAX as u64).unwrap(), i64::MAX);
379        assert!(to_pg_id((i64::MAX as u64) + 1).is_err());
380        assert!(to_pg_id(u64::MAX).is_err());
381    }
382}
llm_kernel/embedding/pgvector.rs

llm_kernel/embedding/
pgvector.rs
