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    /// Number of chunks stored for one document — cheap enough to poll (the
49    /// API's ingest-progress endpoint hits it every couple of seconds).
50    async fn count_chunks_for(&self, doc_id: &str) -> Result<usize>;
51
52    /// The chunks at `ordinal - 1 ..= ordinal + 1` of one document, in
53    /// ordinal order — the search API's "extend context" neighborhood.
54    async fn chunk_neighborhood(&self, doc_id: &str, ordinal: i64) -> Result<Vec<Chunk>>;
55
56    /// Total number of stored documents.
57    async fn count_documents(&self) -> Result<usize>;
58
59    /// Every stored document with its metadata (including processing metrics).
60    async fn list_documents(&self) -> Result<Vec<Document>>;
61
62    /// Delete one document and all of its chunks (used to roll back a failed
63    /// streaming ingest). Deleting an unknown id is not an error.
64    async fn delete_document(&self, doc_id: &str) -> Result<()>;
65
66    /// Delete every document (and chunks) originating from `source_uri` — stale
67    /// rows from interrupted runs or previous versions of a changed file.
68    async fn delete_documents_by_source(&self, source_uri: &str) -> Result<()>;
69
70    /// Remove all documents and chunks.
71    async fn clear(&self) -> Result<()>;
72}
73
74/// Build and connect the store selected by `cfg.db_backend`, running migrations.
75pub async fn from_config(cfg: &RagConfig) -> Result<Arc<dyn VectorStore>> {
76    let store: Arc<dyn VectorStore> = match cfg.db_backend {
77        DbBackend::Memory => Arc::new(memory::MemoryStore::new()),
78        DbBackend::Sqlite => {
79            #[cfg(feature = "sqlite")]
80            {
81                Arc::new(sqlite::SqliteStore::connect(&cfg.database_url, cfg.embed_dim).await?)
82            }
83            #[cfg(not(feature = "sqlite"))]
84            {
85                return Err(crate::RagError::FeatureDisabled(
86                    "sqlite".into(),
87                    "sqlite".into(),
88                ));
89            }
90        }
91        DbBackend::Postgres => {
92            #[cfg(feature = "postgres")]
93            {
94                Arc::new(postgres::PostgresStore::connect(&cfg.database_url, cfg.embed_dim).await?)
95            }
96            #[cfg(not(feature = "postgres"))]
97            {
98                return Err(crate::RagError::FeatureDisabled(
99                    "postgres".into(),
100                    "postgres".into(),
101                ));
102            }
103        }
104    };
105    store.migrate().await?;
106    Ok(store)
107}
108
109/// Rank `(chunk, embedding)` candidates against a query by cosine similarity and
110/// keep the top `k`. Shared by the memory and SQLite brute-force backends.
111pub(crate) fn top_k_by_cosine(
112    query: &[f32],
113    candidates: impl IntoIterator<Item = (Chunk, Vec<f32>)>,
114    k: usize,
115) -> Vec<Scored> {
116    let mut scored: Vec<Scored> = candidates
117        .into_iter()
118        .map(|(chunk, emb)| Scored::new(chunk, crate::math::cosine(query, &emb)))
119        .collect();
120    scored.sort_by(|a, b| {
121        b.score
122            .partial_cmp(&a.score)
123            .unwrap_or(std::cmp::Ordering::Equal)
124    });
125    scored.truncate(k);
126    scored
127}
docling_rag/store/mod.rs

docling_rag/store/
mod.rs
