uni_xervo/provider/llamacpp/
embed.rs

1//! [`ModelTask::Embed`](crate::api::ModelTask::Embed) on `POST /v1/embeddings`.
2//!
3//! `llama-server` exposes an OpenAI-compatible embeddings endpoint, but it
4//! differs from hosted embedding APIs in one way that matters for production
5//! ingestion: **it never truncates embedding input**. Encoder models such as
6//! BGE (BERT) are non-causal, so the whole prompt has to fit in a single
7//! physical batch (`--ubatch-size`) and inside the context window
8//! (`--ctx-size`); anything longer is rejected with an HTTP error instead of
9//! being cut to size.
10//!
11//! This module therefore runs a deterministic, tokenizer-aware pre-pass before
12//! every embedding call:
13//!
14//! 1. Texts that provably fit are sent as plain strings. Every WordPiece token
15//!    consumes at least one character, so a text with at most
16//!    `max_input_tokens - 2` characters cannot exceed the budget once the two
17//!    special tokens (`[CLS]`/`[SEP]`) are added.
18//! 2. Longer texts are sent to the server's native `POST /tokenize` endpoint
19//!    with `add_special: true`, which returns the exact id sequence the server
20//!    would embed, special tokens included.
21//! 3. If that sequence exceeds `max_input_tokens`, the pre-pass keeps the
22//!    first `max_input_tokens - 1` ids and re-appends the trailing special
23//!    token, then sends the **token array** (not text) to `/v1/embeddings`.
24//!    llama.cpp passes integer arrays through verbatim, so the server embeds
25//!    exactly the bounded sequence — no lossy detokenize round trip.
26//!
27//! Batches preserve input order and produce exactly one vector per input.
28//!
29//! # Options
30//!
31//! Beyond the shared keys documented on the [parent module](super):
32//!
33//! | key | type | required | meaning |
34//! |-----|------|----------|---------|
35//! | `tokenizer_base_url` | string | no | server root for `/tokenize`; defaults to `base_url` with a trailing `/v1` removed |
36//! | `max_input_tokens` | integer ≥ 3 | yes | total token budget **including** special tokens (512 for BGE) |
37//! | `embedding_dimensions` | integer > 0 | yes | expected vector width (384 for BGE Small) |
38//!
39//! `max_input_tokens` must not exceed the server's `--ubatch-size` or
40//! `--ctx-size`; if it does, the server's rejection surfaces as
41//! [`RuntimeError::InferenceError`] naming the option.
42
43use super::{
44    LlamaCppCommon, MIN_MAX_INPUT_TOKENS, PROVIDER_ID, post_json, required_option,
45    resolve_tokenizer_base_url,
46};
47use crate::error::{Result, RuntimeError};
48use crate::provider::remote_common::parse_openai_embeddings_response;
49use crate::traits::{EmbedResult, EmbeddingModel};
50use async_trait::async_trait;
51use futures::future::try_join_all;
52use reqwest::Client;
53use serde_json::{Value, json};
54
55// ---------------------------------------------------------------------------
56// Configuration
57// ---------------------------------------------------------------------------
58
59/// Resolved, validated embedding configuration for one alias.
60#[derive(Debug, Clone)]
61pub(crate) struct EmbedConfig {
62    /// Options shared with every other task.
63    pub common: LlamaCppCommon,
64    /// Server root for `/tokenize`, no trailing slash.
65    pub tokenizer_base_url: String,
66    /// Total token budget including special tokens.
67    pub max_input_tokens: usize,
68    /// Expected embedding width.
69    pub dimensions: u32,
70}
71
72impl EmbedConfig {
73    /// Parse the alias `options`. The catalog validator already enforces the
74    /// same rules, but `ModelProvider::load` is public, so this is defensive.
75    pub(crate) fn from_options(options: &Value) -> Result<Self> {
76        let common = LlamaCppCommon::from_options(options)?;
77
78        let tokenizer_base_url = match options.get("tokenizer_base_url").and_then(|v| v.as_str()) {
79            Some(raw) if !raw.trim().is_empty() => raw.trim().trim_end_matches('/').to_string(),
80            _ => resolve_tokenizer_base_url(&common.base_url),
81        };
82
83        let max_input_tokens = options
84            .get("max_input_tokens")
85            .and_then(|v| v.as_u64())
86            .ok_or_else(|| required_option("max_input_tokens"))?
87            as usize;
88        if max_input_tokens < MIN_MAX_INPUT_TOKENS {
89            return Err(RuntimeError::Config(format!(
90                "Option 'max_input_tokens' for provider '{}' must be at least {} \
91                 (two special tokens plus one content token)",
92                PROVIDER_ID, MIN_MAX_INPUT_TOKENS
93            )));
94        }
95
96        let dimensions = options
97            .get("embedding_dimensions")
98            .and_then(|v| v.as_u64())
99            .filter(|d| *d > 0 && *d <= u32::MAX as u64)
100            .ok_or_else(|| required_option("embedding_dimensions"))?
101            as u32;
102
103        Ok(Self {
104            common,
105            tokenizer_base_url,
106            max_input_tokens,
107            dimensions,
108        })
109    }
110}
111
112// ---------------------------------------------------------------------------
113// Input planning (pure)
114// ---------------------------------------------------------------------------
115
116/// Wire form of one element of the `/v1/embeddings` `input` array.
117#[derive(Debug, Clone, PartialEq, Eq)]
118pub(crate) enum EmbedInput {
119    /// Send the original text; the server adds special tokens itself.
120    Text(String),
121    /// Send a pre-bounded token id sequence; the server embeds it verbatim.
122    Tokens(Vec<u32>),
123}
124
125impl EmbedInput {
126    pub(crate) fn to_json(&self) -> Value {
127        match self {
128            EmbedInput::Text(t) => Value::String(t.clone()),
129            EmbedInput::Tokens(ids) => json!(ids),
130        }
131    }
132}
133
134/// `true` when `text` might exceed the budget and must be tokenized to know.
135///
136/// Every WordPiece token consumes at least one character, so a text with at
137/// most `max_input_tokens - 2` characters (two slots reserved for the special
138/// tokens) provably fits without a tokenizer round trip.
139pub(crate) fn needs_tokenize(text: &str, max_input_tokens: usize) -> bool {
140    text.chars().count() > max_input_tokens.saturating_sub(2)
141}
142
143/// Decide the wire form for `text` given its full `add_special: true` token
144/// sequence. Fits → original text. Too long → the first
145/// `max_input_tokens - 1` ids plus the original trailing special token.
146pub(crate) fn plan_input(text: &str, tokens: &[u32], max_input_tokens: usize) -> EmbedInput {
147    if tokens.is_empty() || tokens.len() <= max_input_tokens {
148        return EmbedInput::Text(text.to_string());
149    }
150    let keep = max_input_tokens.saturating_sub(1).max(1);
151    let mut bounded: Vec<u32> = tokens[..keep].to_vec();
152    bounded.push(*tokens.last().expect("non-empty"));
153    EmbedInput::Tokens(bounded)
154}
155
156/// Check every vector has the configured width.
157pub(crate) fn validate_dimensions(vectors: &[Vec<f32>], expected: u32) -> Result<()> {
158    for (i, v) in vectors.iter().enumerate() {
159        if v.len() != expected as usize {
160            return Err(RuntimeError::ApiError(format!(
161                "llama.cpp returned a {}-dimensional vector at index {}, expected {} \
162                 (check the embedding_dimensions option)",
163                v.len(),
164                i,
165                expected
166            )));
167        }
168    }
169    Ok(())
170}
171
172// ---------------------------------------------------------------------------
173// Embedding model
174// ---------------------------------------------------------------------------
175
176/// Embedding model backed by a llama.cpp server.
177pub struct LlamaCppEmbeddingModel {
178    client: Client,
179    cb: crate::reliability::CircuitBreakerWrapper,
180    model_id: String,
181    cfg: EmbedConfig,
182}
183
184impl LlamaCppEmbeddingModel {
185    pub(crate) fn new(
186        client: Client,
187        cb: crate::reliability::CircuitBreakerWrapper,
188        model_id: String,
189        cfg: EmbedConfig,
190    ) -> Self {
191        Self {
192            client,
193            cb,
194            model_id,
195            cfg,
196        }
197    }
198
199    /// POST JSON, naming `max_input_tokens` in any input-too-long error.
200    async fn post_json(&self, url: String, payload: &Value) -> Result<Value> {
201        post_json(
202            &self.client,
203            &self.cfg.common,
204            url,
205            payload,
206            Some(self.cfg.max_input_tokens),
207        )
208        .await
209    }
210
211    /// Call `/tokenize` with special tokens so the count matches what the
212    /// server would embed for the plain string.
213    async fn tokenize(&self, text: &str) -> Result<Vec<u32>> {
214        let body = self
215            .post_json(
216                format!("{}/tokenize", self.cfg.tokenizer_base_url),
217                &json!({
218                    // `model` is required in router mode so the request reaches
219                    // the tokenizer of the model that will do the embedding;
220                    // single-model servers ignore it.
221                    "model": self.model_id,
222                    "content": text,
223                    "add_special": true,
224                    "parse_special": false,
225                }),
226            )
227            .await?;
228        let tokens = body
229            .get("tokens")
230            .and_then(|t| t.as_array())
231            .ok_or_else(|| {
232                RuntimeError::ApiError(
233                    "llama.cpp /tokenize response malformed: missing 'tokens' array".to_string(),
234                )
235            })?;
236        tokens
237            .iter()
238            .enumerate()
239            .map(|(i, t)| {
240                t.as_u64()
241                    .filter(|v| *v <= u32::MAX as u64)
242                    .map(|v| v as u32)
243                    .ok_or_else(|| {
244                        RuntimeError::ApiError(format!(
245                            "llama.cpp /tokenize response malformed: token {} is not an integer id",
246                            i
247                        ))
248                    })
249            })
250            .collect()
251    }
252
253    /// Decide the wire form of every input, tokenizing only those that might
254    /// exceed the budget. `try_join_all` preserves input order.
255    async fn plan_inputs(&self, texts: &[String]) -> Result<Vec<EmbedInput>> {
256        let max = self.cfg.max_input_tokens;
257        try_join_all(texts.iter().map(|text| async move {
258            if !needs_tokenize(text, max) {
259                return Ok(EmbedInput::Text(text.clone()));
260            }
261            let tokens = self.tokenize(text).await?;
262            Ok(plan_input(text, &tokens, max))
263        }))
264        .await
265    }
266
267    async fn post_embeddings(&self, inputs: &[EmbedInput]) -> Result<Value> {
268        let input: Vec<Value> = inputs.iter().map(EmbedInput::to_json).collect();
269        self.post_json(
270            format!("{}/embeddings", self.cfg.common.base_url),
271            &json!({ "model": self.model_id, "input": input }),
272        )
273        .await
274    }
275
276    async fn embed_pipeline(&self, texts: &[String]) -> Result<EmbedResult> {
277        let inputs = self.plan_inputs(texts).await?;
278        let body = self.post_embeddings(&inputs).await?;
279        let result = parse_openai_embeddings_response("llama.cpp", &body, Some(texts.len()))?;
280        validate_dimensions(&result.vectors, self.cfg.dimensions)?;
281        Ok(result)
282    }
283}
284
285#[async_trait]
286impl EmbeddingModel for LlamaCppEmbeddingModel {
287    async fn embed(&self, texts: &[&str]) -> Result<EmbedResult> {
288        if texts.is_empty() {
289            return Ok(EmbedResult {
290                vectors: Vec::new(),
291                usage: None,
292            });
293        }
294        let texts: Vec<String> = texts.iter().map(|s| s.to_string()).collect();
295
296        // One breaker call covers tokenize + embed. Deterministic input-size
297        // rejections (`InferenceError`, produced only by `map_llamacpp_status`)
298        // are a configuration problem, not a server outage, so they are
299        // returned as `Ok(Err(..))` from the closure — the breaker sees a
300        // success and stays closed — and flattened here.
301        self.cb
302            .call(move || async move {
303                match self.embed_pipeline(&texts).await {
304                    Err(e @ RuntimeError::InferenceError(_)) => Ok(Err(e)),
305                    other => other.map(Ok),
306                }
307            })
308            .await?
309    }
310
311    fn dimensions(&self) -> u32 {
312        self.cfg.dimensions
313    }
314}
315
316impl crate::traits::ModelInfo for LlamaCppEmbeddingModel {
317    fn model_id(&self) -> &str {
318        &self.model_id
319    }
320}
321
322// ---------------------------------------------------------------------------
323// Tests (no HTTP)
324// ---------------------------------------------------------------------------
325
326#[cfg(test)]
327mod tests {
328    use super::*;
329    use crate::provider::llamacpp::DEFAULT_REQUEST_TIMEOUT_SECS;
330    use std::time::Duration;
331
332    fn options() -> Value {
333        json!({
334            "base_url": "http://127.0.0.1:8080/v1",
335            "max_input_tokens": 512,
336            "embedding_dimensions": 384
337        })
338    }
339
340    // --- config -------------------------------------------------------------
341
342    #[test]
343    fn config_happy_path_and_defaults() {
344        let cfg = EmbedConfig::from_options(&options()).unwrap();
345        assert_eq!(cfg.common.base_url, "http://127.0.0.1:8080/v1");
346        assert_eq!(cfg.tokenizer_base_url, "http://127.0.0.1:8080");
347        assert_eq!(cfg.max_input_tokens, 512);
348        assert_eq!(cfg.dimensions, 384);
349        assert!(cfg.common.api_key.is_none());
350        assert_eq!(
351            cfg.common.request_timeout,
352            Duration::from_secs(DEFAULT_REQUEST_TIMEOUT_SECS)
353        );
354    }
355
356    #[test]
357    fn config_strips_trailing_slash_and_honours_explicit_tokenizer_url_and_timeout() {
358        let cfg = EmbedConfig::from_options(&json!({
359            "base_url": "http://h:1/v1/",
360            "tokenizer_base_url": "http://tok:2/",
361            "max_input_tokens": 3,
362            "embedding_dimensions": 4,
363            "request_timeout_secs": 7
364        }))
365        .unwrap();
366        assert_eq!(cfg.common.base_url, "http://h:1/v1");
367        assert_eq!(cfg.tokenizer_base_url, "http://tok:2");
368        assert_eq!(cfg.common.request_timeout, Duration::from_secs(7));
369    }
370
371    #[test]
372    fn config_requires_base_url_max_tokens_and_dimensions() {
373        for missing in ["base_url", "max_input_tokens", "embedding_dimensions"] {
374            let mut o = options();
375            o.as_object_mut().unwrap().remove(missing);
376            let err = EmbedConfig::from_options(&o).expect_err("must fail");
377            match err {
378                RuntimeError::Config(msg) => assert!(msg.contains(missing), "{msg}"),
379                other => panic!("expected Config, got {other:?}"),
380            }
381        }
382        assert!(EmbedConfig::from_options(&Value::Null).is_err());
383    }
384
385    #[test]
386    fn config_rejects_tiny_max_input_tokens() {
387        let mut o = options();
388        o["max_input_tokens"] = json!(2);
389        let err = EmbedConfig::from_options(&o).err().unwrap();
390        assert!(err.to_string().contains("at least 3"), "{err}");
391    }
392
393    #[test]
394    fn config_rejects_zero_dimensions() {
395        let mut o = options();
396        o["embedding_dimensions"] = json!(0);
397        assert!(EmbedConfig::from_options(&o).is_err());
398    }
399
400    // --- input planning -----------------------------------------------------
401
402    #[test]
403    fn needs_tokenize_boundaries_use_char_count() {
404        assert!(!needs_tokenize("", 10));
405        assert!(!needs_tokenize("12345678", 10)); // == max - 2
406        assert!(needs_tokenize("123456789", 10)); // == max - 1
407        // 8 CJK chars are 24 bytes but only 8 chars.
408        assert!(!needs_tokenize("你好世界你好世界", 10));
409        assert!(needs_tokenize("你好世界你好世界你", 10));
410        // Degenerate budgets never underflow.
411        assert!(needs_tokenize("a", 1));
412        assert!(!needs_tokenize("", 0));
413    }
414
415    #[test]
416    fn plan_input_fits_sends_text() {
417        let tokens: Vec<u32> = (0..10).collect();
418        assert_eq!(
419            plan_input("t", &tokens, 10),
420            EmbedInput::Text("t".to_string())
421        );
422        assert_eq!(plan_input("t", &[], 10), EmbedInput::Text("t".to_string()));
423    }
424
425    #[test]
426    fn plan_input_truncates_keeping_leading_ids_and_trailing_special() {
427        // [CLS]=101, content 1..=9, [SEP]=102 → 11 tokens, budget 10.
428        let mut tokens = vec![101u32];
429        tokens.extend(1..=9);
430        tokens.push(102);
431        match plan_input("t", &tokens, 10) {
432            EmbedInput::Tokens(ids) => {
433                assert_eq!(ids.len(), 10);
434                assert_eq!(ids[0], 101);
435                assert_eq!(&ids[1..9], &[1, 2, 3, 4, 5, 6, 7, 8]);
436                assert_eq!(*ids.last().unwrap(), 102);
437            }
438            other => panic!("expected Tokens, got {other:?}"),
439        }
440    }
441
442    #[test]
443    fn embed_input_json_shapes() {
444        assert_eq!(
445            EmbedInput::Text("hi".into()).to_json(),
446            Value::String("hi".into())
447        );
448        assert_eq!(EmbedInput::Tokens(vec![1, 2]).to_json(), json!([1, 2]));
449    }
450
451    // --- dimension validation -----------------------------------------------
452
453    #[test]
454    fn dimension_validation() {
455        assert!(validate_dimensions(&[vec![0.0; 4], vec![1.0; 4]], 4).is_ok());
456        let err = validate_dimensions(&[vec![0.0; 4], vec![1.0; 3]], 4)
457            .err()
458            .unwrap();
459        let msg = err.to_string();
460        assert!(
461            msg.contains("3-dimensional") && msg.contains("index 1"),
462            "{msg}"
463        );
464        assert!(msg.contains("embedding_dimensions"), "{msg}");
465    }
466}
uni_xervo/provider/llamacpp/embed.rs

uni_xervo/provider/llamacpp/
embed.rs
