docling-rag 0.41.2

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

//! Pluggable embedding providers.
//!
//! All providers implement [`Embedder`]. The default is [`ollama`] (`bge-m3`,
//! 1024-dim); [`gemini`] and the feature-gated `onnx` provider are alternatives,
//! and [`hash`] is a deterministic offline embedder for tests and evaluation.

mod hash;

#[cfg(feature = "onnx-embed")]
mod onnx;

pub mod gemini;
pub mod ollama;

pub use hash::HashEmbedder;

use crate::config::EmbedProvider;
use crate::{RagConfig, RagError, Result};
use async_trait::async_trait;
use std::sync::Arc;

/// A source of embedding vectors. Implementations must be cheap to `clone` via
/// `Arc` and safe to share across tasks.
#[async_trait]
pub trait Embedder: Send + Sync {
    /// Embed a batch of texts, returning one vector per input in order.
    async fn embed(&self, texts: &[String]) -> Result<Vec<Vec<f32>>>;

    /// The dimensionality of every returned vector.
    fn dim(&self) -> usize;

    /// A short identifier (`"ollama:bge-m3"`, `"hash"`, …) for logs and eval reports.
    fn id(&self) -> &str;

    /// Convenience: embed a single text.
    async fn embed_one(&self, text: &str) -> Result<Vec<f32>> {
        let mut v = self.embed(std::slice::from_ref(&text.to_string())).await?;
        v.pop()
            .ok_or_else(|| RagError::Embedding("provider returned no vector".into()))
    }
}

/// Build the embedder selected by `cfg.embed_provider`.
pub fn from_config(cfg: &RagConfig) -> Result<Arc<dyn Embedder>> {
    match cfg.embed_provider {
        EmbedProvider::Ollama => Ok(Arc::new(ollama::OllamaEmbedder::from_config(cfg))),
        EmbedProvider::Gemini => Ok(Arc::new(gemini::GeminiEmbedder::from_config(cfg)?)),
        EmbedProvider::Hash => Ok(Arc::new(HashEmbedder::new(cfg.embed_dim))),
        EmbedProvider::Onnx => {
            #[cfg(feature = "onnx-embed")]
            {
                Ok(Arc::new(onnx::OnnxEmbedder::from_config(cfg)?))
            }
            #[cfg(not(feature = "onnx-embed"))]
            {
                Err(RagError::FeatureDisabled(
                    "onnx".into(),
                    "onnx-embed".into(),
                ))
            }
        }
    }
}