docling_rag/embed/
mod.rs

1//! Pluggable embedding providers.
2//!
3//! All providers implement [`Embedder`]. The default is [`ollama`] (`bge-m3`,
4//! 1024-dim); [`gemini`] and the feature-gated `onnx` provider are alternatives,
5//! and [`hash`] is a deterministic offline embedder for tests and evaluation.
6
7mod hash;
8
9#[cfg(feature = "onnx-embed")]
10mod onnx;
11
12pub mod gemini;
13pub mod ollama;
14
15pub use hash::HashEmbedder;
16
17use crate::config::EmbedProvider;
18use crate::{RagConfig, RagError, Result};
19use async_trait::async_trait;
20use std::sync::Arc;
21
22/// A source of embedding vectors. Implementations must be cheap to `clone` via
23/// `Arc` and safe to share across tasks.
24#[async_trait]
25pub trait Embedder: Send + Sync {
26    /// Embed a batch of texts, returning one vector per input in order.
27    async fn embed(&self, texts: &[String]) -> Result<Vec<Vec<f32>>>;
28
29    /// The dimensionality of every returned vector.
30    fn dim(&self) -> usize;
31
32    /// A short identifier (`"ollama:bge-m3"`, `"hash"`, …) for logs and eval reports.
33    fn id(&self) -> &str;
34
35    /// Convenience: embed a single text.
36    async fn embed_one(&self, text: &str) -> Result<Vec<f32>> {
37        let mut v = self.embed(std::slice::from_ref(&text.to_string())).await?;
38        v.pop()
39            .ok_or_else(|| RagError::Embedding("provider returned no vector".into()))
40    }
41}
42
43/// Build the embedder selected by `cfg.embed_provider`.
44pub fn from_config(cfg: &RagConfig) -> Result<Arc<dyn Embedder>> {
45    match cfg.embed_provider {
46        EmbedProvider::Ollama => Ok(Arc::new(ollama::OllamaEmbedder::from_config(cfg))),
47        EmbedProvider::Gemini => Ok(Arc::new(gemini::GeminiEmbedder::from_config(cfg)?)),
48        EmbedProvider::Hash => Ok(Arc::new(HashEmbedder::new(cfg.embed_dim))),
49        EmbedProvider::Onnx => {
50            #[cfg(feature = "onnx-embed")]
51            {
52                Ok(Arc::new(onnx::OnnxEmbedder::from_config(cfg)?))
53            }
54            #[cfg(not(feature = "onnx-embed"))]
55            {
56                Err(RagError::FeatureDisabled(
57                    "onnx".into(),
58                    "onnx-embed".into(),
59                ))
60            }
61        }
62    }
63}
docling_rag/embed/mod.rs

docling_rag/embed/
mod.rs
