docling-rag 0.41.0

Pluggable RAG subsystem for docling.rs: chunking, embeddings, vector store, and semantic search.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

//! In-memory vector store. Never persisted; used by tests and quick evals.

use super::{top_k_by_cosine, VectorStore};
use crate::model::{Chunk, Document, Scored};
use crate::Result;
use async_trait::async_trait;
use std::sync::RwLock;

/// A store that keeps everything in process memory.
#[derive(Default)]
pub struct MemoryStore {
    docs: RwLock<Vec<Document>>,
    chunks: RwLock<Vec<Chunk>>,
}

impl MemoryStore {
    /// Create an empty in-memory store.
    pub fn new() -> Self {
        MemoryStore::default()
    }
}

#[async_trait]
impl VectorStore for MemoryStore {
    async fn migrate(&self) -> Result<()> {
        Ok(())
    }

    async fn upsert_document(&self, doc: &Document) -> Result<()> {
        let mut docs = self.docs.write().unwrap();
        if let Some(existing) = docs.iter_mut().find(|d| d.id == doc.id) {
            *existing = doc.clone();
        } else {
            docs.push(doc.clone());
        }
        Ok(())
    }

    async fn find_document_by_hash(&self, hash: &str) -> Result<Option<String>> {
        Ok(self
            .docs
            .read()
            .unwrap()
            .iter()
            .find(|d| d.hash == hash)
            .map(|d| d.id.clone()))
    }

    async fn insert_chunks(&self, chunks: &[Chunk]) -> Result<()> {
        self.chunks.write().unwrap().extend_from_slice(chunks);
        Ok(())
    }

    async fn vector_search(&self, query: &[f32], k: usize) -> Result<Vec<Scored>> {
        let chunks = self.chunks.read().unwrap();
        let candidates = chunks.iter().filter_map(|c| {
            c.embedding.as_ref().map(|e| {
                let mut bare = c.clone();
                bare.embedding = None;
                (bare, e.clone())
            })
        });
        Ok(top_k_by_cosine(query, candidates.collect::<Vec<_>>(), k))
    }

    async fn all_chunks(&self) -> Result<Vec<Chunk>> {
        Ok(self
            .chunks
            .read()
            .unwrap()
            .iter()
            .map(|c| {
                let mut bare = c.clone();
                bare.embedding = None;
                bare
            })
            .collect())
    }

    async fn count_chunks(&self) -> Result<usize> {
        Ok(self.chunks.read().unwrap().len())
    }

    async fn count_documents(&self) -> Result<usize> {
        Ok(self.docs.read().unwrap().len())
    }

    async fn list_documents(&self) -> Result<Vec<Document>> {
        Ok(self.docs.read().unwrap().clone())
    }

    async fn delete_document(&self, doc_id: &str) -> Result<()> {
        self.docs.write().unwrap().retain(|d| d.id != doc_id);
        self.chunks.write().unwrap().retain(|c| c.doc_id != doc_id);
        Ok(())
    }

    async fn delete_documents_by_source(&self, source_uri: &str) -> Result<()> {
        let ids: Vec<String> = {
            let docs = self.docs.read().unwrap();
            docs.iter()
                .filter(|d| d.source_uri == source_uri)
                .map(|d| d.id.clone())
                .collect()
        };
        self.docs
            .write()
            .unwrap()
            .retain(|d| d.source_uri != source_uri);
        self.chunks
            .write()
            .unwrap()
            .retain(|c| !ids.contains(&c.doc_id));
        Ok(())
    }

    async fn clear(&self) -> Result<()> {
        self.docs.write().unwrap().clear();
        self.chunks.write().unwrap().clear();
        Ok(())
    }
}