rig_core/providers/ollama/
embedding.rs

1//! Ollama's embedding wire, `POST /api/embed`, and its embedding models.
2//!
3//! ```
4//! use rig_core::providers::ollama::{ALL_MINILM, OllamaConfig};
5//! let model = OllamaConfig::new().client().embedding(ALL_MINILM, None);
6//! assert_eq!(model.wire.ndims, 384);
7//! ```
8
9use crate::error::EncodeError;
10use crate::operation::Embedding;
11use crate::wire::{Body, Capabilities, Descriptor, DocumentDecoder, Encoded, Framing, Mode, Wire};
12use serde::{Deserialize, Serialize};
13
14use super::{OllamaConfig, PROVIDER_NAME};
15
16/// The `all-minilm` embedding model.
17pub const ALL_MINILM: &str = "all-minilm";
18/// The `nomic-embed-text` embedding model.
19pub const NOMIC_EMBED_TEXT: &str = "nomic-embed-text";
20/// The `mxbai-embed-large` embedding model.
21pub const MXBAI_EMBED_LARGE: &str = "mxbai-embed-large";
22/// The `bge-m3` multilingual embedding model.
23pub const BGE_M3: &str = "bge-m3";
24/// The `embeddinggemma` embedding model.
25pub const EMBEDDINGGEMMA: &str = "embeddinggemma";
26/// The `qwen3-embedding` embedding model family; dimensions vary by size, so pass them explicitly.
27pub const QWEN3_EMBEDDING: &str = "qwen3-embedding";
28
29fn model_dimensions_from_identifier(identifier: &str) -> Option<usize> {
30    match identifier {
31        ALL_MINILM => Some(384),
32        NOMIC_EMBED_TEXT => Some(768),
33        MXBAI_EMBED_LARGE => Some(1024),
34        BGE_M3 => Some(1024),
35        EMBEDDINGGEMMA => Some(768),
36        _ => None,
37    }
38}
39
40#[derive(Debug, Clone, Serialize, Deserialize)]
41pub struct EmbeddingResponse {
42    pub model: String,
43    pub embeddings: Vec<Vec<f64>>,
44    #[serde(default)]
45    pub total_duration: Option<u64>,
46    #[serde(default)]
47    pub load_duration: Option<u64>,
48    #[serde(default)]
49    pub prompt_eval_count: Option<u64>,
50}
51
52/// The most texts `POST /api/embed` accepts in one call.
53const MAX_DOCUMENTS: usize = 1024;
54
55impl OllamaConfig {
56    /// Build an embedding wire reporting the supplied width, known model width,
57    /// or zero if unknown. The width is metadata and is not sent to the daemon.
58    pub(crate) fn embedding(&self, model: impl Into<String>, ndims: Option<usize>) -> Embeddings {
59        let model = model.into();
60        let ndims = ndims
61            .or_else(|| model_dimensions_from_identifier(&model))
62            .unwrap_or_default();
63        Embeddings {
64            provider: self.clone(),
65            model,
66            ndims,
67        }
68    }
69}
70
71/// The embedding wire: `POST /api/embed`.
72#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
73pub struct Embeddings {
74    /// The daemon this wire speaks to.
75    pub provider: OllamaConfig,
76    /// The model to address.
77    pub model: String,
78    /// The width this wire reports, from the caller or the model's published
79    /// dimensions. `0` means neither named one.
80    pub ndims: usize,
81}
82
83impl Wire for Embeddings {
84    type Op = Embedding;
85    type Payload = crate::wire::Encoded;
86    type Frame = crate::wire::WireFrame;
87    type Event = EmbeddingResponse;
88    type Decoder<'id> = DocumentDecoder<Embedding, Self::Event>;
89    type Reassembler = crate::wire::document::Unreassembled;
90
91    fn describe(&self) -> Descriptor<'_> {
92        Descriptor::new(PROVIDER_NAME)
93            .model(self.model.as_str())
94            .capabilities(Capabilities::embedding(MAX_DOCUMENTS, self.ndims))
95    }
96
97    fn encode(&self, texts: Vec<String>, _mode: Mode) -> Result<Encoded, EncodeError> {
98        let body = serde_json::json!({ "model": self.model, "input": texts });
99        let request = self
100            .provider
101            .request(http::Method::POST, "/api/embed")
102            .body(Body::Bytes(serde_json::to_vec(&body)?))?;
103        Ok(Encoded::new(request, Framing::Whole))
104    }
105
106    /// Decodes one `/api/embed` reply.
107    fn decoder<'id>(&self) -> Self::Decoder<'id> {
108        DocumentDecoder::new(
109            |body| {
110                crate::providers::internal::wire::classify_marker_keyed_frame(body, &["embeddings"])
111            },
112            |reply| {
113                // Ollama counts the prompt it embedded and nothing else: every token
114                // of an embedding is input.
115                let usage = crate::completion::Usage {
116                    input_tokens: reply.prompt_eval_count,
117                    total_tokens: reply.prompt_eval_count,
118                    ..Default::default()
119                };
120                Ok(crate::embeddings::EmbeddingResponse {
121                    model: Some(reply.model),
122                    usage,
123                    ..crate::embeddings::EmbeddingResponse::from_vectors(reply.embeddings)
124                })
125            },
126        )
127    }
128}
129
130#[cfg(test)]
131mod tests;
rig_core/providers/ollama/embedding.rs

rig_core/providers/ollama/
embedding.rs
