docling_rag/store/
mod.rs

1//! Pluggable vector store.
2//!
3//! A store persists [`Document`]s (metadata) and their [`Chunk`]s (text +
4//! embedding) in two tables, and answers dense vector search. Keyword (BM25)
5//! search is layered on top in [`crate::retrieve`] using [`VectorStore::all_chunks`],
6//! so it works identically across every backend.
7
8pub mod memory;
9
10#[cfg(feature = "postgres")]
11pub mod postgres;
12#[cfg(feature = "sqlite")]
13pub mod sqlite;
14
15use crate::config::DbBackend;
16use crate::model::{Chunk, Document, Scored};
17use crate::{RagConfig, Result};
18use async_trait::async_trait;
19use std::sync::Arc;
20
21/// A document + chunk store with dense vector search.
22#[async_trait]
23pub trait VectorStore: Send + Sync {
24    /// Create tables/indexes if they do not exist. Safe to call repeatedly.
25    async fn migrate(&self) -> Result<()>;
26
27    /// Insert or replace a document row (keyed by `id`).
28    async fn upsert_document(&self, doc: &Document) -> Result<()>;
29
30    /// Return the id of an existing document with this content hash, if any.
31    /// Used to skip re-ingesting unchanged documents.
32    async fn find_document_by_hash(&self, hash: &str) -> Result<Option<String>>;
33
34    /// Bulk-insert chunks (each must carry a populated `embedding`).
35    async fn insert_chunks(&self, chunks: &[Chunk]) -> Result<()>;
36
37    /// Dense search: the `k` chunks whose embeddings are most cosine-similar to
38    /// `query`. Returned chunks omit their embedding vector to keep payloads small.
39    async fn vector_search(&self, query: &[f32], k: usize) -> Result<Vec<Scored>>;
40
41    /// Every chunk (id, doc_id, ordinal, text, metadata) with `embedding == None`.
42    /// Feeds the BM25 keyword index; keep it modest for large corpora.
43    async fn all_chunks(&self) -> Result<Vec<Chunk>>;
44
45    /// Total number of stored chunks.
46    async fn count_chunks(&self) -> Result<usize>;
47
48    /// Total number of stored documents.
49    async fn count_documents(&self) -> Result<usize>;
50
51    /// Every stored document with its metadata (including processing metrics).
52    async fn list_documents(&self) -> Result<Vec<Document>>;
53
54    /// Delete one document and all of its chunks (used to roll back a failed
55    /// streaming ingest). Deleting an unknown id is not an error.
56    async fn delete_document(&self, doc_id: &str) -> Result<()>;
57
58    /// Delete every document (and chunks) originating from `source_uri` — stale
59    /// rows from interrupted runs or previous versions of a changed file.
60    async fn delete_documents_by_source(&self, source_uri: &str) -> Result<()>;
61
62    /// Remove all documents and chunks.
63    async fn clear(&self) -> Result<()>;
64}
65
66/// Build and connect the store selected by `cfg.db_backend`, running migrations.
67pub async fn from_config(cfg: &RagConfig) -> Result<Arc<dyn VectorStore>> {
68    let store: Arc<dyn VectorStore> = match cfg.db_backend {
69        DbBackend::Memory => Arc::new(memory::MemoryStore::new()),
70        DbBackend::Sqlite => {
71            #[cfg(feature = "sqlite")]
72            {
73                Arc::new(sqlite::SqliteStore::connect(&cfg.database_url, cfg.embed_dim).await?)
74            }
75            #[cfg(not(feature = "sqlite"))]
76            {
77                return Err(crate::RagError::FeatureDisabled(
78                    "sqlite".into(),
79                    "sqlite".into(),
80                ));
81            }
82        }
83        DbBackend::Postgres => {
84            #[cfg(feature = "postgres")]
85            {
86                Arc::new(postgres::PostgresStore::connect(&cfg.database_url, cfg.embed_dim).await?)
87            }
88            #[cfg(not(feature = "postgres"))]
89            {
90                return Err(crate::RagError::FeatureDisabled(
91                    "postgres".into(),
92                    "postgres".into(),
93                ));
94            }
95        }
96    };
97    store.migrate().await?;
98    Ok(store)
99}
100
101/// Rank `(chunk, embedding)` candidates against a query by cosine similarity and
102/// keep the top `k`. Shared by the memory and SQLite brute-force backends.
103pub(crate) fn top_k_by_cosine(
104    query: &[f32],
105    candidates: impl IntoIterator<Item = (Chunk, Vec<f32>)>,
106    k: usize,
107) -> Vec<Scored> {
108    let mut scored: Vec<Scored> = candidates
109        .into_iter()
110        .map(|(chunk, emb)| Scored::new(chunk, crate::math::cosine(query, &emb)))
111        .collect();
112    scored.sort_by(|a, b| {
113        b.score
114            .partial_cmp(&a.score)
115            .unwrap_or(std::cmp::Ordering::Equal)
116    });
117    scored.truncate(k);
118    scored
119}
docling_rag/store/mod.rs

docling_rag/store/
mod.rs
