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