mermaid_cli/providers/model/
mod.rs

1//! Model adapters wrapped as `ModelProvider` implementations.
2//!
3//! Five providers today: Ollama, Anthropic, Gemini, Meta, and OpenAI-
4//! compat (covering OpenAI, OpenRouter, Groq, Cerebras, DeepInfra,
5//! Together, and user-defined endpoints). Each wraps the
6//! corresponding adapter in `mermaid_model::models::adapters`; the adapter
7//! owns the wire format and the wrapper owns the trait shape.
8
9pub mod anthropic;
10pub mod gemini;
11pub mod meta;
12pub mod ollama;
13pub mod openai_compat;
14
15use std::sync::Arc;
16
17use async_trait::async_trait;
18
19use mermaid_domain::{ChatRequest, TurnId};
20use mermaid_model::models::adapters::ModelLimits;
21use mermaid_model::models::adapters::ollama_sizing::NumCtxSource;
22use mermaid_model::models::{ModelError, Result, TokenUsage};
23use mermaid_runtime::NewProviderProbe;
24
25use super::ctx::{FinalResponse, StreamContext, StreamEvent};
26use mermaid_model::models::ModelCapabilities;
27
28/// Resolved context sizing for a turn. For most providers `model_max ==
29/// effective` (the static advertised window). For Ollama they differ:
30/// `model_max` is the probed architectural window, while `effective` is what we
31/// actually enforce as `num_ctx` (auto-fitted to memory, capped, or an
32/// override). Compaction and the status bar use `effective`; "model supports up
33/// to X" uses `model_max`.
34#[derive(Debug, Clone, Copy, Default)]
35pub struct ContextSizing {
36    pub model_max: Option<usize>,
37    pub effective: Option<usize>,
38    /// How `effective` was chosen (Ollama only). `None` for static/advertised.
39    pub source: Option<NumCtxSource>,
40    /// The model's per-response output ceiling, when the provider exposes one
41    /// (`/models` metadata, or a documented static table). Rides the same
42    /// resolve→reducer pipeline as the window so `provider_capabilities` can be
43    /// refreshed live.
44    pub max_output: Option<usize>,
45}
46
47/// Where a loaded model actually sits in memory, from a post-turn probe (Ollama
48/// `/api/ps`). `total_bytes` is weights + KV + buffers; `size_vram_bytes` is the
49/// part resident in VRAM. `size_vram_bytes < total_bytes` means the model spilled
50/// to CPU/RAM (partial offload → slow). Only Ollama reports this.
51#[derive(Debug, Clone, Copy)]
52pub struct ModelPlacement {
53    pub size_vram_bytes: u64,
54    pub total_bytes: u64,
55    /// Auto-converge target: when the model spilled, the largest `num_ctx` that
56    /// would fit instead — or `None` if it already fits or shrinking can't help
57    /// (weights-bound). Computed against the *measured* footprint.
58    pub suggested_num_ctx: Option<u32>,
59}
60
61/// Provider-facing interface. A `ModelProvider` impl owns whatever
62/// HTTP client / state it needs and exposes `chat()` — that's the
63/// whole surface.
64#[async_trait]
65pub trait ModelProvider: Send + Sync {
66    /// `ModelCapabilities` the provider advertises. The reducer reads this
67    /// when building the outgoing `ChatRequest` (e.g. whether to
68    /// attach reasoning controls).
69    fn capabilities(&self) -> &ModelCapabilities;
70
71    /// Resolve the *effective* context window for a turn (what the model will
72    /// actually enforce). The default returns the static advertised window;
73    /// Ollama overrides this to probe the model's real window and auto-fit
74    /// `num_ctx` to host memory, honoring the request's per-model
75    /// `ollama_num_ctx` override. `None` means "let the backend decide".
76    ///
77    /// Awaited only on the effect runtime (never the reducer), so a probe never
78    /// blocks the UI.
79    async fn resolve_context_window(&self, request: &ChatRequest) -> ContextSizing {
80        let _ = request;
81        let max = self.capabilities().max_context_tokens;
82        ContextSizing {
83            model_max: max,
84            effective: max,
85            source: None,
86            max_output: self.capabilities().max_output_tokens,
87        }
88    }
89
90    /// Best-effort: where the loaded model currently sits in memory. The default
91    /// returns `None` (unknown / not applicable); Ollama overrides it to probe
92    /// `/api/ps`. Awaited only on the effect runtime, *after* a turn (when the
93    /// model is resident), so it never blocks the UI.
94    async fn verify_placement(&self, current_num_ctx: Option<usize>) -> Option<ModelPlacement> {
95        let _ = current_num_ctx;
96        None
97    }
98
99    /// Best-effort: whether the active model can actually see images. `None`
100    /// means "unknown / not applicable" — the default for providers that don't
101    /// probe, and for cloud providers whose vision support is already known
102    /// good; `Some(false)` is what drives the no-vision-model warning. Awaited
103    /// only on the effect runtime, so a probe never blocks the UI.
104    async fn supports_vision(&self) -> Option<bool> {
105        None
106    }
107
108    /// Stream a chat turn. Typed events flow through
109    /// `ctx.sink`; the returned `FinalResponse` carries token usage
110    /// and the Anthropic thinking-signature (opaque blob required to
111    /// continue extended thinking across turns).
112    ///
113    /// Cancellation: the provider MUST select! on `ctx.token.
114    /// cancelled()` inside any await that could block for more than
115    /// a few hundred ms. This is the contract that replaces the old
116    /// `check_interrupt` polling pattern.
117    async fn chat(&self, request: ChatRequest, ctx: StreamContext) -> Result<FinalResponse>;
118}
119
120/// Run a one-shot, non-interactive model call and collect its streamed text
121/// into a `String`. For internal calls whose output is NOT shown to the user
122/// as it streams — context compaction and the Auto-mode safety classifier.
123/// Drains a private event channel (ignoring reasoning / tool-call events) and
124/// returns the collected text plus final token usage. The `token` lets the
125/// caller cancel the call (e.g. on Ctrl+C) like any other turn work.
126pub(crate) async fn collect_text(
127    provider: Arc<dyn ModelProvider>,
128    turn: TurnId,
129    request: ChatRequest,
130    token: tokio_util::sync::CancellationToken,
131) -> Result<(String, Option<TokenUsage>)> {
132    let (stream_tx, mut stream_rx) = tokio::sync::mpsc::channel::<StreamEvent>(128);
133    let ctx = StreamContext::new(token, stream_tx, turn);
134    let collector = tokio::task::spawn(async move {
135        let mut text = String::new();
136        let mut usage = None;
137        while let Some(event) = stream_rx.recv().await {
138            match event {
139                StreamEvent::Text(chunk) => text.push_str(&chunk),
140                StreamEvent::Done {
141                    usage: done_usage, ..
142                } => usage = done_usage,
143                // Status is a user-facing plumbing notice, not content —
144                // a text collector has nowhere to surface it.
145                StreamEvent::Reasoning(_) | StreamEvent::ToolCall(_) | StreamEvent::Status(_) => {},
146            }
147        }
148        (text, usage)
149    });
150
151    let response = provider.chat(request, ctx).await;
152    let (text, stream_usage) = collector
153        .await
154        .map_err(|err| ModelError::StreamError(format!("collect_text collector failed: {err}")))?;
155    match response {
156        Ok(final_response) => Ok((text, final_response.usage.or(stream_usage))),
157        Err(err) => Err(err),
158    }
159}
160
161pub use anthropic::AnthropicProvider;
162pub use gemini::GeminiProvider;
163pub use meta::MetaProvider;
164pub use ollama::OllamaProvider;
165pub use openai_compat::OpenAICompatProvider;
166
167/// True when a `provider_probes` row is older than the probe TTL (shared by
168/// the Ollama context probe and the per-provider limits probes). An
169/// unparseable timestamp is treated as stale so it re-probes.
170pub(crate) fn probe_is_stale(probed_at: &str) -> bool {
171    use chrono::{DateTime, Utc};
172    match DateTime::parse_from_rfc3339(probed_at) {
173        Ok(t) => {
174            Utc::now()
175                .signed_duration_since(t.with_timezone(&Utc))
176                .num_days()
177                >= mermaid_model::constants::PROVIDER_PROBE_TTL_DAYS
178        },
179        // Unparseable timestamp → treat as stale and re-probe.
180        Err(_) => true,
181    }
182}
183
184/// Limits learned from one live probe of a provider's models endpoint, cached
185/// per (provider, model) in `provider_probes`. A successful fetch WITHOUT
186/// limit metadata (or a definitive "model not listed") is cached too — as
187/// `None`s — so providers that don't expose limits aren't re-fetched every
188/// turn. Fetch *failures* are never cached.
189#[derive(serde::Serialize, serde::Deserialize)]
190pub(crate) struct CachedLimits {
191    pub(crate) max_context_tokens: Option<usize>,
192    pub(crate) max_output_tokens: Option<usize>,
193}
194
195pub(crate) const LIMITS_PROBE_KEY: &str = "limits_probe";
196
197/// Load fresh cached limits, off the async runtime. Best-effort → `None`.
198pub(crate) async fn load_limits_from_db(provider: String, model: String) -> Option<CachedLimits> {
199    tokio::task::spawn_blocking(move || {
200        let rec = mermaid_runtime::with_shared_store(|store| {
201            store
202                .provider_probes()
203                .get(&provider, &model, LIMITS_PROBE_KEY)
204        })
205        .ok()??;
206        if probe_is_stale(&rec.probed_at) {
207            return None;
208        }
209        serde_json::from_str::<CachedLimits>(&rec.capability_value).ok()
210    })
211    .await
212    .ok()
213    .flatten()
214}
215
216/// Persist probed limits for subsequent sessions. Best-effort.
217pub(crate) async fn save_limits_to_db(provider: String, model: String, limits: &CachedLimits) {
218    let value = match serde_json::to_string(limits) {
219        Ok(v) => v,
220        Err(_) => return,
221    };
222    let _ = tokio::task::spawn_blocking(move || -> Option<()> {
223        mermaid_runtime::with_shared_store(|store| {
224            store.provider_probes().upsert(NewProviderProbe {
225                provider,
226                model_id: model,
227                capability_key: LIMITS_PROBE_KEY.into(),
228                capability_value: value,
229                confidence: "probed".into(),
230                error: None,
231            })
232        })
233        .ok()?;
234        Some(())
235    })
236    .await;
237}
238
239/// Cache-first limit resolution: return fresh cached limits when present,
240/// otherwise run `fetch` against the provider's models endpoint. A successful
241/// fetch is cached even when all-`None` (definitive "provider exposes
242/// nothing"); a failed fetch is NOT cached and resolves to `None` so the next
243/// turn retries.
244pub(crate) async fn resolve_limits_cached<F, Fut>(
245    provider: &str,
246    model: &str,
247    fetch: F,
248) -> Option<CachedLimits>
249where
250    F: FnOnce() -> Fut,
251    Fut: std::future::Future<Output = Result<ModelLimits>>,
252{
253    if let Some(cached) = load_limits_from_db(provider.to_string(), model.to_string()).await {
254        return Some(cached);
255    }
256    match fetch().await {
257        Ok(limits) => {
258            let cached = CachedLimits {
259                max_context_tokens: limits.max_context_tokens,
260                max_output_tokens: limits.max_output_tokens,
261            };
262            save_limits_to_db(provider.to_string(), model.to_string(), &cached).await;
263            Some(cached)
264        },
265        // Network/parse failure: don't cache; caller falls back to static.
266        Err(_) => None,
267    }
268}
269
270/// Extract a model's real per-response output ceiling from a provider's 400
271/// rejection body. Fires ONLY on unambiguous output-cap wordings:
272///
273/// - Ollama Cloud / MiniMax: `max_tokens (521276) exceeds model's maximum
274///   output tokens (131072) for model …`
275/// - OpenAI-compat: `max_tokens is too large: … This model supports at most
276///   16384 completion tokens …`
277///
278/// Anything else — in particular context-limit wordings ("prompt is too
279/// long", "maximum context length") — returns `None`: learning a window as
280/// an output cap would poison the cache, and a missed match just means
281/// today's behavior (the error surfaces). Values outside a sanity range are
282/// rejected as parser noise.
283pub(crate) fn parse_output_cap_message(body: &str) -> Option<usize> {
284    let cap = if let Some(rest) = text_after(body, "exceeds model's maximum output tokens") {
285        leading_integer(rest)
286    } else if body.contains("max_tokens is too large") {
287        text_after(body, "supports at most").and_then(leading_integer)
288    } else {
289        None
290    }?;
291    (1_024..10_000_000).contains(&cap).then_some(cap)
292}
293
294/// The slice of `haystack` after the first occurrence of `marker`.
295fn text_after<'a>(haystack: &'a str, marker: &str) -> Option<&'a str> {
296    haystack.find(marker).map(|i| &haystack[i + marker.len()..])
297}
298
299/// The first integer in `s`, required to start within a few characters —
300/// both documented wordings put the number right after the marker (`" ("` /
301/// `" "`), and a distant number would belong to something else.
302fn leading_integer(s: &str) -> Option<usize> {
303    let start = s.find(|c: char| c.is_ascii_digit()).filter(|&i| i <= 8)?;
304    s[start..]
305        .chars()
306        .take_while(char::is_ascii_digit)
307        .collect::<String>()
308        .parse()
309        .ok()
310}
311
312/// Decide the output cap for a one-shot retry after learning `learned` from
313/// a 400. AUTO (`requested == 0`, including "field omitted") retries at the
314/// learned cap; an explicit ask above the cap retries clamped to it; an ask
315/// already within the cap returns `None` — the 400 was about something else,
316/// so retrying the same request would loop.
317pub(crate) fn retry_cap(requested: usize, learned: usize) -> Option<usize> {
318    (requested == 0 || requested > learned).then_some(learned)
319}
320
321/// `parse_output_cap_message` gated to actual HTTP 400s — the only status
322/// where the body names a rejected parameter rather than a transient fault.
323pub(crate) fn output_cap_from_error(err: &ModelError) -> Option<usize> {
324    match err {
325        ModelError::Backend(mermaid_model::models::BackendError::HttpError {
326            status: 400,
327            message,
328            ..
329        }) => parse_output_cap_message(message),
330        _ => None,
331    }
332}
333
334/// Persist an output cap learned from a provider's 400 rejection (the error
335/// body names the model's real ceiling). Merges into any existing cached
336/// row — reading it raw, ignoring the TTL, since a stale window is still
337/// better than dropping it — and upserts as "probed". Best-effort.
338pub(crate) async fn learn_output_cap(provider: String, model: String, cap: usize) {
339    let _ = tokio::task::spawn_blocking(move || -> Option<()> {
340        let existing = mermaid_runtime::with_shared_store(|store| {
341            store
342                .provider_probes()
343                .get(&provider, &model, LIMITS_PROBE_KEY)
344        })
345        .ok()
346        .flatten()
347        .and_then(|rec| serde_json::from_str::<CachedLimits>(&rec.capability_value).ok());
348        let merged = CachedLimits {
349            max_context_tokens: existing.and_then(|l| l.max_context_tokens),
350            max_output_tokens: Some(cap),
351        };
352        let value = serde_json::to_string(&merged).ok()?;
353        mermaid_runtime::with_shared_store(|store| {
354            store.provider_probes().upsert(NewProviderProbe {
355                provider,
356                model_id: model,
357                capability_key: LIMITS_PROBE_KEY.into(),
358                capability_value: value,
359                confidence: "probed".into(),
360                error: None,
361            })
362        })
363        .ok()?;
364        Some(())
365    })
366    .await;
367}
368
369#[cfg(test)]
370mod tests {
371    use super::*;
372
373    // The incident wording (Ollama Cloud, minimax-m3), raw and JSON-wrapped.
374    const MINIMAX_RAW: &str =
375        "max_tokens (521276) exceeds model's maximum output tokens (131072) for model minimax-m3";
376    const MINIMAX_JSON: &str = r#"{"error":"max_tokens (521276) exceeds model's maximum output tokens (131072) for model minimax-m3 (ref: a05c9ffb-168f)"}"#;
377    const OPENAI_STYLE: &str = r#"{"error":{"message":"max_tokens is too large: 200000. This model supports at most 16384 completion tokens, whereas you provided 200000.","type":"invalid_request_error"}}"#;
378
379    #[test]
380    fn parse_output_cap_matches_documented_wordings() {
381        assert_eq!(parse_output_cap_message(MINIMAX_RAW), Some(131_072));
382        assert_eq!(parse_output_cap_message(MINIMAX_JSON), Some(131_072));
383        assert_eq!(parse_output_cap_message(OPENAI_STYLE), Some(16_384));
384    }
385
386    #[test]
387    fn parse_output_cap_never_matches_context_limit_wordings() {
388        // Learning a context window as an output cap would poison the cache —
389        // these must all be None even though they mention token limits.
390        for body in [
391            "prompt is too long: 210000 tokens > 200000 maximum",
392            "This model's maximum context length is 128000 tokens",
393            "input length and max_tokens exceed context limit: 190000 + 20000 > 200000",
394            "the request exceeds the maximum context window of 131072 tokens",
395            "rate limit exceeded, try again in 20s",
396            "",
397        ] {
398            assert_eq!(parse_output_cap_message(body), None, "matched: {body}");
399        }
400    }
401
402    #[test]
403    fn parse_output_cap_rejects_nonsense_values() {
404        // Sub-1024 and absurd values are parser noise, not real ceilings.
405        assert_eq!(
406            parse_output_cap_message("exceeds model's maximum output tokens (512)"),
407            None
408        );
409        assert_eq!(
410            parse_output_cap_message("exceeds model's maximum output tokens (99999999999)"),
411            None
412        );
413        // A number too far from the marker belongs to something else.
414        assert_eq!(
415            parse_output_cap_message(
416                "exceeds model's maximum output tokens for this deployment tier which is 131072"
417            ),
418            None
419        );
420    }
421
422    #[test]
423    fn retry_cap_triple() {
424        // AUTO (0 / omitted) → retry at the learned cap.
425        assert_eq!(retry_cap(0, 131_072), Some(131_072));
426        // Explicit ask above the cap → clamp.
427        assert_eq!(retry_cap(521_276, 131_072), Some(131_072));
428        // Ask already within the cap → the 400 was about something else.
429        assert_eq!(retry_cap(4_096, 131_072), None);
430    }
431
432    #[test]
433    fn output_cap_from_error_gates_on_http_400() {
434        let err_400 = ModelError::Backend(mermaid_model::models::BackendError::HttpError {
435            status: 400,
436            message: MINIMAX_JSON.to_string(),
437            debug: Default::default(),
438        });
439        assert_eq!(output_cap_from_error(&err_400), Some(131_072));
440        // Same body on a 500 is a transient fault, not a learned limit.
441        let err_500 = ModelError::Backend(mermaid_model::models::BackendError::HttpError {
442            status: 500,
443            message: MINIMAX_JSON.to_string(),
444            debug: Default::default(),
445        });
446        assert_eq!(output_cap_from_error(&err_500), None);
447        assert_eq!(output_cap_from_error(&ModelError::Cancelled), None);
448    }
449}
mermaid_cli/providers/model/mod.rs

mermaid_cli/providers/model/
mod.rs
