uni_xervo/provider/llamacpp/
mod.rs

1//! Remote provider for a llama.cpp [`llama-server`](https://github.com/ggml-org/llama.cpp/tree/master/tools/server).
2//!
3//! `llama-server` speaks a near-complete OpenAI-compatible API over one loaded
4//! model, plus a handful of native endpoints of its own. This module is the
5//! provider shell: the option plumbing every task shares, the HTTP seam, and
6//! the error classifier. Each task lives in its own sibling module:
7//!
8//! | task | module | endpoint |
9//! |------|--------|----------|
10//! | [`ModelTask::Embed`] | [`embed`] | `POST /v1/embeddings` + native `/tokenize` |
11//! | [`ModelTask::Generate`] | [`chat`] | `POST /v1/chat/completions` |
12//! | [`ModelTask::Rerank`] | [`rerank`] | `POST /v1/rerank` |
13//! | [`ModelTask::DocumentExtract`] | [`docs`] | `POST /v1/chat/completions` |
14//!
15//! # Shared options
16//!
17//! | key | type | required | meaning |
18//! |-----|------|----------|---------|
19//! | `base_url` | string | yes | OpenAI-compatible root including `/v1`, e.g. `http://127.0.0.1:8080/v1` |
20//! | `api_key_env` | string | no | env var holding a bearer token; omit when the server runs without `--api-key` |
21//! | `request_timeout_secs` | integer > 0 | no | per-HTTP-request timeout, default 60 |
22//!
23//! Per-task options are documented on the module that reads them.
24//!
25//! # Input length
26//!
27//! `llama-server` never silently truncates. A prompt that does not fit the
28//! physical batch (`--ubatch-size`) or the context window (`--ctx-size`) is
29//! rejected with an HTTP error, which [`map_llamacpp_status`] turns into a
30//! non-retryable [`RuntimeError::InferenceError`] naming the remedy. That is a
31//! deterministic configuration problem rather than a server outage, so these
32//! rejections are also kept off the circuit breaker — see
33//! [`embed::LlamaCppEmbeddingModel::embed`].
34
35mod chat;
36mod docs;
37mod embed;
38mod rerank;
39
40pub use chat::LlamaCppGeneratorModel;
41pub use docs::LlamaCppDocumentExtractor;
42pub use embed::LlamaCppEmbeddingModel;
43pub use rerank::LlamaCppRerankerModel;
44
45use crate::api::{ModelAliasSpec, ModelTask};
46use crate::error::{Result, RuntimeError};
47use crate::provider::remote_common::RemoteProviderBase;
48use crate::traits::{
49    DocumentExtractionModel, EmbeddingModel, GeneratorModel, LoadedModelHandle, ModelProvider,
50    ProviderCapabilities, ProviderHealth, RerankerModel,
51};
52use async_trait::async_trait;
53use reqwest::Client;
54use serde_json::Value;
55use std::sync::Arc;
56use std::time::Duration;
57
58/// Provider id used in [`ModelAliasSpec::provider_id`].
59pub const PROVIDER_ID: &str = "remote/llamacpp";
60
61/// Default per-request HTTP timeout when `request_timeout_secs` is unset.
62pub const DEFAULT_REQUEST_TIMEOUT_SECS: u64 = 60;
63
64/// Smallest meaningful `max_input_tokens`: two special tokens plus one
65/// content token.
66pub const MIN_MAX_INPUT_TOKENS: usize = 3;
67
68/// Remote provider for a llama.cpp `llama-server`.
69pub struct RemoteLlamaCppProvider {
70    base: RemoteProviderBase,
71}
72
73impl Default for RemoteLlamaCppProvider {
74    fn default() -> Self {
75        Self {
76            base: RemoteProviderBase::new(),
77        }
78    }
79}
80
81impl RemoteLlamaCppProvider {
82    pub fn new() -> Self {
83        Self::default()
84    }
85
86    #[cfg(test)]
87    fn insert_test_breaker(&self, key: crate::api::ModelRuntimeKey, age: Duration) {
88        self.base.insert_test_breaker(key, age);
89    }
90
91    #[cfg(test)]
92    fn breaker_count(&self) -> usize {
93        self.base.breaker_count()
94    }
95
96    #[cfg(test)]
97    fn force_cleanup_now_for_test(&self) {
98        self.base.force_cleanup_now_for_test();
99    }
100}
101
102#[async_trait]
103impl ModelProvider for RemoteLlamaCppProvider {
104    fn provider_id(&self) -> &'static str {
105        PROVIDER_ID
106    }
107
108    fn capabilities(&self) -> ProviderCapabilities {
109        ProviderCapabilities {
110            supported_tasks: vec![
111                ModelTask::Embed,
112                ModelTask::Generate,
113                ModelTask::Rerank,
114                ModelTask::DocumentExtract,
115            ],
116        }
117    }
118
119    async fn load(&self, spec: &ModelAliasSpec) -> Result<LoadedModelHandle> {
120        match spec.task {
121            ModelTask::Embed => {
122                let model = embed::LlamaCppEmbeddingModel::new(
123                    self.base.client.clone(),
124                    self.base.circuit_breaker_for(spec),
125                    spec.model_id.clone(),
126                    embed::EmbedConfig::from_options(&spec.options)?,
127                );
128                let handle: Arc<dyn EmbeddingModel> = Arc::new(model);
129                Ok(Arc::new(handle) as LoadedModelHandle)
130            }
131            ModelTask::Generate => {
132                let model = chat::LlamaCppGeneratorModel::new(
133                    self.base.client.clone(),
134                    self.base.circuit_breaker_for(spec),
135                    spec.model_id.clone(),
136                    chat::ChatConfig::from_options(&spec.options)?,
137                );
138                let handle: Arc<dyn GeneratorModel> = Arc::new(model);
139                Ok(Arc::new(handle) as LoadedModelHandle)
140            }
141            ModelTask::Rerank => {
142                let model = rerank::LlamaCppRerankerModel::new(
143                    self.base.client.clone(),
144                    self.base.circuit_breaker_for(spec),
145                    spec.model_id.clone(),
146                    rerank::RerankConfig::from_options(&spec.options)?,
147                );
148                let handle: Arc<dyn RerankerModel> = Arc::new(model);
149                Ok(Arc::new(handle) as LoadedModelHandle)
150            }
151            ModelTask::DocumentExtract => {
152                let model = docs::LlamaCppDocumentExtractor::new(
153                    self.base.client.clone(),
154                    self.base.circuit_breaker_for(spec),
155                    spec.model_id.clone(),
156                    docs::DocsConfig::from_options(&spec.options)?,
157                );
158                let handle: Arc<dyn DocumentExtractionModel> = Arc::new(model);
159                Ok(Arc::new(handle) as LoadedModelHandle)
160            }
161            other => Err(RuntimeError::CapabilityMismatch(format!(
162                "llama.cpp provider does not support task {:?}",
163                other
164            ))),
165        }
166    }
167
168    async fn health(&self) -> ProviderHealth {
169        ProviderHealth::Healthy
170    }
171}
172
173// ---------------------------------------------------------------------------
174// Shared configuration
175// ---------------------------------------------------------------------------
176
177/// The part of an alias's `options` that every task reads.
178#[derive(Debug, Clone)]
179pub(crate) struct LlamaCppCommon {
180    /// OpenAI-compatible root including `/v1`, no trailing slash.
181    pub base_url: String,
182    /// Bearer token, if `api_key_env` was configured.
183    pub api_key: Option<String>,
184    /// Per-HTTP-request timeout.
185    pub request_timeout: Duration,
186}
187
188/// Build the "option is required" error for `key`.
189pub(crate) fn required_option(key: &str) -> RuntimeError {
190    RuntimeError::Config(format!(
191        "Option '{}' for provider '{}' is required",
192        key, PROVIDER_ID
193    ))
194}
195
196impl LlamaCppCommon {
197    /// Parse the options every task shares. The catalog validator already
198    /// enforces the same rules, but `ModelProvider::load` is public, so this
199    /// is defensive.
200    pub(crate) fn from_options(options: &Value) -> Result<Self> {
201        let base_url = options
202            .get("base_url")
203            .and_then(|v| v.as_str())
204            .map(|s| s.trim().trim_end_matches('/').to_string())
205            .filter(|s| !s.is_empty())
206            .ok_or_else(|| required_option("base_url"))?;
207
208        let api_key = match options.get("api_key_env").and_then(|v| v.as_str()) {
209            None => None,
210            Some(env_name) => Some(
211                std::env::var(env_name)
212                    .map_err(|_| RuntimeError::Config(format!("{} env var not set", env_name)))?,
213            ),
214        };
215
216        let request_timeout = Duration::from_secs(
217            options
218                .get("request_timeout_secs")
219                .and_then(|v| v.as_u64())
220                .filter(|s| *s > 0)
221                .unwrap_or(DEFAULT_REQUEST_TIMEOUT_SECS),
222        );
223
224        Ok(Self {
225            base_url,
226            api_key,
227            request_timeout,
228        })
229    }
230}
231
232/// Derive the server root for native (non-`/v1`) endpoints from an
233/// OpenAI-style `base_url`: strip trailing slashes, then exactly one trailing
234/// `/v1` path segment.
235///
236/// ```text
237/// http://h:8080/v1      -> http://h:8080
238/// http://h:8080/v1/     -> http://h:8080
239/// http://h:8080/api/v1  -> http://h:8080/api
240/// http://h:8080         -> http://h:8080   (unchanged)
241/// http://h:8080/v10     -> http://h:8080/v10 (unchanged)
242/// ```
243pub(crate) fn resolve_tokenizer_base_url(base_url: &str) -> String {
244    let trimmed = base_url.trim_end_matches('/');
245    match trimmed.strip_suffix("/v1") {
246        Some(root) if !root.is_empty() => root.to_string(),
247        _ => trimmed.to_string(),
248    }
249}
250
251// ---------------------------------------------------------------------------
252// HTTP seam
253// ---------------------------------------------------------------------------
254
255/// Attach the bearer token when one is configured.
256fn apply_auth(common: &LlamaCppCommon, req: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
257    match &common.api_key {
258        Some(key) => req.header("Authorization", format!("Bearer {}", key)),
259        None => req,
260    }
261}
262
263/// POST JSON to `url`, apply timeout and auth, and map non-2xx responses.
264///
265/// `token_budget` is the alias's configured input budget when the task has one
266/// (embedding's `max_input_tokens`); it only shapes the advice in an
267/// input-too-long error. Tasks with no such option pass `None`.
268pub(crate) async fn post_json(
269    client: &Client,
270    common: &LlamaCppCommon,
271    url: String,
272    payload: &Value,
273    token_budget: Option<usize>,
274) -> Result<Value> {
275    let response = apply_auth(common, client.post(url))
276        .timeout(common.request_timeout)
277        .json(payload)
278        .send()
279        .await
280        .map_err(map_reqwest_error)?;
281
282    let status = response.status().as_u16();
283    let text = response.text().await.map_err(map_reqwest_error)?;
284    if !(200..300).contains(&status) {
285        return Err(map_llamacpp_status(status, &text, token_budget));
286    }
287    serde_json::from_str(&text)
288        .map_err(|e| RuntimeError::ApiError(format!("llama.cpp returned malformed JSON: {}", e)))
289}
290
291// ---------------------------------------------------------------------------
292// Error mapping (pure)
293// ---------------------------------------------------------------------------
294
295/// Classification of a llama.cpp error body.
296#[derive(Debug, Clone, Copy, PartialEq, Eq)]
297pub(crate) enum LlamaCppErrorKind {
298    /// Prompt longer than the physical batch (`--ubatch-size`); HTTP 500 in
299    /// current builds.
300    InputTooLarge,
301    /// Prompt longer than the context window (`--ctx-size`).
302    ExceedsContext,
303    /// The endpoint exists but this server process was not started with the
304    /// flag that enables it — `/rerank` without `--reranking`, for instance.
305    /// HTTP 501 in current builds.
306    NotSupported,
307    /// Anything else.
308    Other,
309}
310
311const ERROR_SNIPPET_MAX_CHARS: usize = 200;
312
313/// Parse `{"error":{"code","message","type"}}` leniently and classify it.
314/// Returns the kind plus a human-readable message (or a truncated raw snippet
315/// when the body is not that shape).
316pub(crate) fn classify_error_body(body: &str) -> (LlamaCppErrorKind, String) {
317    let parsed: Option<Value> = serde_json::from_str(body).ok();
318    let err_obj = parsed.as_ref().and_then(|v| v.get("error"));
319    let message = err_obj
320        .and_then(|e| e.get("message"))
321        .and_then(|m| m.as_str())
322        .map(|s| s.to_string())
323        .unwrap_or_else(|| body.chars().take(ERROR_SNIPPET_MAX_CHARS).collect());
324    let err_type = err_obj
325        .and_then(|e| e.get("type"))
326        .and_then(|t| t.as_str())
327        .unwrap_or("");
328
329    let lower = message.to_ascii_lowercase();
330    let kind = if lower.contains("too large to process") {
331        LlamaCppErrorKind::InputTooLarge
332    } else if err_type == "exceed_context_size"
333        || lower.contains("exceeds the available context size")
334    {
335        LlamaCppErrorKind::ExceedsContext
336    } else if err_type == "not_supported_error" || lower.contains("does not support") {
337        LlamaCppErrorKind::NotSupported
338    } else {
339        LlamaCppErrorKind::Other
340    };
341    (kind, message)
342}
343
344/// Map a non-2xx llama.cpp response to a [`RuntimeError`].
345///
346/// Two families of rejection are deterministic configuration problems rather
347/// than outages, so both become [`RuntimeError::InferenceError`] — not
348/// retryable and not counted against the circuit breaker — regardless of the
349/// HTTP status the server chose:
350///
351/// - **Input size**: the server's `--ubatch-size`/`--ctx-size` is smaller than
352///   the request. `token_budget` names the offending option when the task has
353///   one; without it the advice can only point at the server's own flags.
354/// - **Endpoint not enabled**: the route exists but the process was started
355///   without the flag that activates it, which current builds report as HTTP
356///   501. Retrying cannot help — the server has to be restarted — and the
357///   message already names the flag, so it is passed through verbatim.
358pub(crate) fn map_llamacpp_status(
359    status: u16,
360    body: &str,
361    token_budget: Option<usize>,
362) -> RuntimeError {
363    let (kind, message) = classify_error_body(body);
364    match kind {
365        LlamaCppErrorKind::InputTooLarge | LlamaCppErrorKind::ExceedsContext => {
366            match token_budget {
367                Some(max_input_tokens) => RuntimeError::InferenceError(format!(
368                    "llama.cpp server rejected the input length ({}); max_input_tokens={} is \
369                     larger than the server allows — lower max_input_tokens or start \
370                     llama-server with --ubatch-size and --ctx-size of at least that many tokens",
371                    message, max_input_tokens
372                )),
373                None => RuntimeError::InferenceError(format!(
374                    "llama.cpp server rejected the input length ({}); the request does not fit \
375                     the server's limits — send less input, or start llama-server with a larger \
376                     --ctx-size and --ubatch-size",
377                    message
378                )),
379            }
380        }
381        LlamaCppErrorKind::NotSupported => RuntimeError::InferenceError(format!(
382            "llama.cpp server rejected the request: {}",
383            message
384        )),
385        LlamaCppErrorKind::Other => match status {
386            429 => RuntimeError::RateLimited,
387            401 | 403 => RuntimeError::Unauthorized,
388            500..=599 => RuntimeError::Unavailable,
389            _ => RuntimeError::ApiError(format!("llama.cpp API error: {}: {}", status, message)),
390        },
391    }
392}
393
394/// Map a reqwest transport error.
395fn map_reqwest_error(e: reqwest::Error) -> RuntimeError {
396    if e.is_timeout() {
397        RuntimeError::Timeout
398    } else {
399        RuntimeError::ApiError(e.to_string())
400    }
401}
402
403// ---------------------------------------------------------------------------
404// Tests (no HTTP)
405// ---------------------------------------------------------------------------
406
407#[cfg(test)]
408mod tests {
409    use super::*;
410    use crate::api::ModelRuntimeKey;
411    use serde_json::json;
412
413    static ENV_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
414
415    fn options() -> Value {
416        json!({
417            "base_url": "http://127.0.0.1:8080/v1",
418            "max_input_tokens": 512,
419            "embedding_dimensions": 384
420        })
421    }
422
423    fn spec(alias: &str, task: ModelTask, model_id: &str, options: Value) -> ModelAliasSpec {
424        ModelAliasSpec {
425            alias: alias.to_string(),
426            task,
427            provider_id: PROVIDER_ID.to_string(),
428            model_id: model_id.to_string(),
429            revision: None,
430            warmup: crate::api::WarmupPolicy::Lazy,
431            required: false,
432            timeout: None,
433            load_timeout: None,
434            retry: None,
435            options,
436        }
437    }
438
439    // --- provider / breaker -------------------------------------------------
440
441    #[tokio::test]
442    async fn provider_id_and_capabilities() {
443        let p = RemoteLlamaCppProvider::new();
444        assert_eq!(p.provider_id(), "remote/llamacpp");
445        assert_eq!(
446            p.capabilities().supported_tasks,
447            vec![
448                ModelTask::Embed,
449                ModelTask::Generate,
450                ModelTask::Rerank,
451                ModelTask::DocumentExtract
452            ]
453        );
454        assert!(matches!(p.health().await, ProviderHealth::Healthy));
455    }
456
457    #[tokio::test]
458    async fn load_rejects_unsupported_tasks() {
459        let p = RemoteLlamaCppProvider::new();
460        for task in [ModelTask::Raw, ModelTask::EmbedImage] {
461            let err = p
462                .load(&spec("x/y", task, "bge", options()))
463                .await
464                .expect_err("must fail");
465            assert!(
466                matches!(err, RuntimeError::CapabilityMismatch(_)),
467                "{err:?}"
468            );
469        }
470    }
471
472    #[tokio::test]
473    async fn load_exposes_configured_dimensions_and_model_id() {
474        let p = RemoteLlamaCppProvider::new();
475        let handle = p
476            .load(&spec("embed/a", ModelTask::Embed, "bge-small", options()))
477            .await
478            .unwrap();
479        let model = handle.downcast_ref::<Arc<dyn EmbeddingModel>>().unwrap();
480        assert_eq!(model.dimensions(), 384);
481        assert_eq!(model.model_id(), "bge-small");
482    }
483
484    #[tokio::test]
485    async fn load_generate_yields_a_generator_handle() {
486        let p = RemoteLlamaCppProvider::new();
487        // Generation needs no embed-only options.
488        let handle = p
489            .load(&spec(
490                "gen/a",
491                ModelTask::Generate,
492                "smolvlm",
493                json!({ "base_url": "http://127.0.0.1:8080/v1" }),
494            ))
495            .await
496            .unwrap();
497        let model = handle.downcast_ref::<Arc<dyn GeneratorModel>>().unwrap();
498        assert_eq!(model.model_id(), "smolvlm");
499    }
500
501    #[tokio::test]
502    async fn load_rerank_yields_a_reranker_handle() {
503        let p = RemoteLlamaCppProvider::new();
504        // Reranking needs no embed-only options either.
505        let handle = p
506            .load(&spec(
507                "rerank/a",
508                ModelTask::Rerank,
509                "jina-reranker-v1-tiny-en",
510                json!({ "base_url": "http://127.0.0.1:8080/v1" }),
511            ))
512            .await
513            .unwrap();
514        let model = handle.downcast_ref::<Arc<dyn RerankerModel>>().unwrap();
515        assert_eq!(model.model_id(), "jina-reranker-v1-tiny-en");
516    }
517
518    #[tokio::test]
519    async fn load_document_extract_yields_an_extractor_handle() {
520        let p = RemoteLlamaCppProvider::new();
521        let handle = p
522            .load(&spec(
523                "docs/a",
524                ModelTask::DocumentExtract,
525                "olmocr-2",
526                json!({ "base_url": "http://127.0.0.1:8080/v1", "style": "olmocr" }),
527            ))
528            .await
529            .unwrap();
530        let model = handle
531            .downcast_ref::<Arc<dyn DocumentExtractionModel>>()
532            .unwrap();
533        assert_eq!(model.model_id(), "olmocr-2");
534    }
535
536    #[tokio::test]
537    async fn breaker_reused_for_same_runtime_key() {
538        let p = RemoteLlamaCppProvider::new();
539        let _ = p
540            .load(&spec("embed/a", ModelTask::Embed, "bge", options()))
541            .await
542            .unwrap();
543        let _ = p
544            .load(&spec("embed/b", ModelTask::Embed, "bge", options()))
545            .await
546            .unwrap();
547        assert_eq!(p.breaker_count(), 1);
548    }
549
550    #[tokio::test]
551    async fn breaker_isolated_by_model() {
552        let p = RemoteLlamaCppProvider::new();
553        let _ = p
554            .load(&spec("embed/a", ModelTask::Embed, "bge", options()))
555            .await
556            .unwrap();
557        let _ = p
558            .load(&spec("embed/b", ModelTask::Embed, "nomic", options()))
559            .await
560            .unwrap();
561        assert_eq!(p.breaker_count(), 2);
562    }
563
564    #[tokio::test]
565    async fn breaker_cleanup_evicts_stale_entries() {
566        let p = RemoteLlamaCppProvider::new();
567        let stale = spec("embed/stale", ModelTask::Embed, "old", options());
568        let fresh = spec("embed/fresh", ModelTask::Embed, "new", options());
569        p.insert_test_breaker(
570            ModelRuntimeKey::new(&stale),
571            RemoteProviderBase::BREAKER_TTL + Duration::from_secs(5),
572        );
573        p.insert_test_breaker(ModelRuntimeKey::new(&fresh), Duration::from_secs(1));
574        assert_eq!(p.breaker_count(), 2);
575        p.force_cleanup_now_for_test();
576        let _ = p.load(&fresh).await.unwrap();
577        assert_eq!(p.breaker_count(), 1);
578    }
579
580    // --- shared config ------------------------------------------------------
581
582    #[test]
583    fn common_config_happy_path_and_defaults() {
584        let cfg = LlamaCppCommon::from_options(&options()).unwrap();
585        assert_eq!(cfg.base_url, "http://127.0.0.1:8080/v1");
586        assert!(cfg.api_key.is_none());
587        assert_eq!(
588            cfg.request_timeout,
589            Duration::from_secs(DEFAULT_REQUEST_TIMEOUT_SECS)
590        );
591    }
592
593    #[test]
594    fn common_config_strips_trailing_slash_and_honours_timeout() {
595        let cfg = LlamaCppCommon::from_options(&json!({
596            "base_url": "http://h:1/v1/",
597            "request_timeout_secs": 7
598        }))
599        .unwrap();
600        assert_eq!(cfg.base_url, "http://h:1/v1");
601        assert_eq!(cfg.request_timeout, Duration::from_secs(7));
602    }
603
604    #[test]
605    fn common_config_requires_base_url() {
606        let err = LlamaCppCommon::from_options(&json!({})).expect_err("must fail");
607        match err {
608            RuntimeError::Config(msg) => assert!(msg.contains("base_url"), "{msg}"),
609            other => panic!("expected Config, got {other:?}"),
610        }
611        assert!(LlamaCppCommon::from_options(&Value::Null).is_err());
612
613        // A blank or whitespace-only value is as absent as a missing key.
614        for blank in ["", "   "] {
615            assert!(
616                LlamaCppCommon::from_options(&json!({ "base_url": blank })).is_err(),
617                "{blank:?} should be rejected"
618            );
619        }
620    }
621
622    #[test]
623    fn common_config_ignores_a_zero_timeout_and_falls_back_to_the_default() {
624        let cfg = LlamaCppCommon::from_options(&json!({
625            "base_url": "http://h:1/v1",
626            "request_timeout_secs": 0
627        }))
628        .unwrap();
629        assert_eq!(
630            cfg.request_timeout,
631            Duration::from_secs(DEFAULT_REQUEST_TIMEOUT_SECS)
632        );
633    }
634
635    #[tokio::test]
636    async fn common_config_api_key_env_present_and_missing() {
637        let _lock = ENV_LOCK.lock().await;
638        let mut o = options();
639        o["api_key_env"] = json!("UNI_XERVO_LLAMACPP_TEST_KEY");
640
641        // SAFETY: protected by ENV_LOCK
642        unsafe { std::env::remove_var("UNI_XERVO_LLAMACPP_TEST_KEY") };
643        let err = LlamaCppCommon::from_options(&o).err().unwrap();
644        assert!(
645            err.to_string().contains("UNI_XERVO_LLAMACPP_TEST_KEY"),
646            "{err}"
647        );
648
649        // SAFETY: protected by ENV_LOCK
650        unsafe { std::env::set_var("UNI_XERVO_LLAMACPP_TEST_KEY", "sekrit") };
651        let cfg = LlamaCppCommon::from_options(&o).unwrap();
652        assert_eq!(cfg.api_key.as_deref(), Some("sekrit"));
653        // SAFETY: protected by ENV_LOCK
654        unsafe { std::env::remove_var("UNI_XERVO_LLAMACPP_TEST_KEY") };
655    }
656
657    #[test]
658    fn tokenizer_base_url_derivation() {
659        assert_eq!(
660            resolve_tokenizer_base_url("http://h:8080/v1"),
661            "http://h:8080"
662        );
663        assert_eq!(
664            resolve_tokenizer_base_url("http://h:8080/v1/"),
665            "http://h:8080"
666        );
667        assert_eq!(
668            resolve_tokenizer_base_url("http://h:8080/api/v1"),
669            "http://h:8080/api"
670        );
671        assert_eq!(resolve_tokenizer_base_url("http://h:8080"), "http://h:8080");
672        assert_eq!(
673            resolve_tokenizer_base_url("http://h:8080/v10"),
674            "http://h:8080/v10"
675        );
676        assert_eq!(resolve_tokenizer_base_url("/v1"), "/v1");
677    }
678
679    // --- error mapping ------------------------------------------------------
680
681    #[test]
682    fn classify_error_bodies() {
683        let too_large = r#"{"error":{"code":500,"message":"input is too large to process. increase the physical batch size","type":"server_error"}}"#;
684        let (k, m) = classify_error_body(too_large);
685        assert_eq!(k, LlamaCppErrorKind::InputTooLarge);
686        assert!(m.starts_with("input is too large"));
687
688        let ctx_type =
689            r#"{"error":{"code":400,"message":"whatever","type":"exceed_context_size"}}"#;
690        assert_eq!(
691            classify_error_body(ctx_type).0,
692            LlamaCppErrorKind::ExceedsContext
693        );
694        let ctx_msg = r#"{"error":{"message":"prompt exceeds the available context size. increase context size"}}"#;
695        assert_eq!(
696            classify_error_body(ctx_msg).0,
697            LlamaCppErrorKind::ExceedsContext
698        );
699
700        let (k, m) = classify_error_body("<html>502 bad gateway</html>");
701        assert_eq!(k, LlamaCppErrorKind::Other);
702        assert_eq!(m, "<html>502 bad gateway</html>");
703
704        let long = "x".repeat(1000);
705        assert_eq!(
706            classify_error_body(&long).1.chars().count(),
707            ERROR_SNIPPET_MAX_CHARS
708        );
709    }
710
711    #[test]
712    fn status_mapping() {
713        let too_large = r#"{"error":{"message":"input is too large to process. increase the physical batch size"}}"#;
714        let e = map_llamacpp_status(500, too_large, Some(512));
715        assert!(matches!(e, RuntimeError::InferenceError(_)), "{e:?}");
716        assert!(!e.is_retryable());
717        assert!(e.to_string().contains("max_input_tokens=512"), "{e}");
718
719        let e = map_llamacpp_status(
720            400,
721            r#"{"error":{"type":"exceed_context_size","message":"x"}}"#,
722            Some(512),
723        );
724        assert!(matches!(e, RuntimeError::InferenceError(_)), "{e:?}");
725
726        assert!(matches!(
727            map_llamacpp_status(500, "boom", Some(512)),
728            RuntimeError::Unavailable
729        ));
730        assert!(matches!(
731            map_llamacpp_status(503, "", Some(512)),
732            RuntimeError::Unavailable
733        ));
734        assert!(matches!(
735            map_llamacpp_status(429, "", Some(512)),
736            RuntimeError::RateLimited
737        ));
738        assert!(matches!(
739            map_llamacpp_status(401, "", Some(512)),
740            RuntimeError::Unauthorized
741        ));
742        assert!(matches!(
743            map_llamacpp_status(403, "", Some(512)),
744            RuntimeError::Unauthorized
745        ));
746        match map_llamacpp_status(404, r#"{"error":{"message":"no route"}}"#, Some(512)) {
747            RuntimeError::ApiError(m) => {
748                assert!(m.contains("404") && m.contains("no route"), "{m}")
749            }
750            other => panic!("{other:?}"),
751        }
752    }
753
754    /// A 501 from an endpoint the server was not started for is permanent:
755    /// `Unavailable` would have it retried and would trip the breaker.
756    #[test]
757    fn an_endpoint_that_is_not_enabled_is_not_retryable() {
758        let body = r#"{"error":{"code":501,"message":"This server does not support reranking. Start it with `--reranking`","type":"not_supported_error"}}"#;
759        assert_eq!(classify_error_body(body).0, LlamaCppErrorKind::NotSupported);
760
761        let e = map_llamacpp_status(501, body, None);
762        assert!(matches!(e, RuntimeError::InferenceError(_)), "{e:?}");
763        assert!(!e.is_retryable(), "{e:?}");
764        // The server's own advice is the useful part, so keep it.
765        assert!(e.to_string().contains("--reranking"), "{e}");
766    }
767
768    #[test]
769    fn status_mapping_without_a_token_budget_advises_server_flags_instead() {
770        let too_large = r#"{"error":{"message":"input is too large to process. increase the physical batch size"}}"#;
771        let e = map_llamacpp_status(500, too_large, None);
772        assert!(matches!(e, RuntimeError::InferenceError(_)), "{e:?}");
773        assert!(!e.is_retryable());
774
775        let msg = e.to_string();
776        assert!(!msg.contains("max_input_tokens"), "{msg}");
777        assert!(msg.contains("--ctx-size"), "{msg}");
778
779        // Classification of everything else is unaffected by the budget.
780        assert!(matches!(
781            map_llamacpp_status(429, "", None),
782            RuntimeError::RateLimited
783        ));
784    }
785}
uni_xervo/provider/llamacpp/mod.rs

uni_xervo/provider/llamacpp/
mod.rs
