mermaid_cli/providers/model/
mod.rs

1//! Model adapters wrapped as `ModelProvider` implementations.
2//!
3//! Four providers today: Ollama, Anthropic, Gemini, 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 mod ollama;
12pub mod openai_compat;
13pub(crate) mod stream_bridge;
14
15use std::sync::Arc;
16
17use async_trait::async_trait;
18
19use crate::domain::{ChatRequest, TurnId};
20use crate::models::adapters::ollama_sizing::NumCtxSource;
21use crate::models::{ModelError, Result, TokenUsage};
22
23use super::capabilities::Capabilities;
24use super::ctx::{FinalResponse, StreamContext, StreamEvent};
25
26/// Resolved context sizing for a turn. For most providers `model_max ==
27/// effective` (the static advertised window). For Ollama they differ:
28/// `model_max` is the probed architectural window, while `effective` is what we
29/// actually enforce as `num_ctx` (auto-fitted to memory, capped, or an
30/// override). Compaction and the status bar use `effective`; "model supports up
31/// to X" uses `model_max`.
32#[derive(Debug, Clone, Copy, Default)]
33pub struct ContextSizing {
34    pub model_max: Option<usize>,
35    pub effective: Option<usize>,
36    /// How `effective` was chosen (Ollama only). `None` for static/advertised.
37    pub source: Option<NumCtxSource>,
38}
39
40/// Where a loaded model actually sits in memory, from a post-turn probe (Ollama
41/// `/api/ps`). `total_bytes` is weights + KV + buffers; `size_vram_bytes` is the
42/// part resident in VRAM. `size_vram_bytes < total_bytes` means the model spilled
43/// to CPU/RAM (partial offload → slow). Only Ollama reports this.
44#[derive(Debug, Clone, Copy)]
45pub struct ModelPlacement {
46    pub size_vram_bytes: u64,
47    pub total_bytes: u64,
48    /// Auto-converge target: when the model spilled, the largest `num_ctx` that
49    /// would fit instead — or `None` if it already fits or shrinking can't help
50    /// (weights-bound). Computed against the *measured* footprint.
51    pub suggested_num_ctx: Option<u32>,
52}
53
54/// Provider-facing interface. A `ModelProvider` impl owns whatever
55/// HTTP client / state it needs and exposes `chat()` — that's the
56/// whole surface.
57#[async_trait]
58pub trait ModelProvider: Send + Sync {
59    /// Capabilities the provider advertises. The reducer reads this
60    /// when building the outgoing `ChatRequest` (e.g. whether to
61    /// attach reasoning controls).
62    fn capabilities(&self) -> &Capabilities;
63
64    /// Resolve the *effective* context window for a turn (what the model will
65    /// actually enforce). The default returns the static advertised window;
66    /// Ollama overrides this to probe the model's real window and auto-fit
67    /// `num_ctx` to host memory, honoring the request's per-model
68    /// `ollama_num_ctx` override. `None` means "let the backend decide".
69    ///
70    /// Awaited only on the effect runtime (never the reducer), so a probe never
71    /// blocks the UI.
72    async fn resolve_context_window(&self, request: &ChatRequest) -> ContextSizing {
73        let _ = request;
74        let max = self.capabilities().max_context_tokens;
75        ContextSizing {
76            model_max: max,
77            effective: max,
78            source: None,
79        }
80    }
81
82    /// Best-effort: where the loaded model currently sits in memory. The default
83    /// returns `None` (unknown / not applicable); Ollama overrides it to probe
84    /// `/api/ps`. Awaited only on the effect runtime, *after* a turn (when the
85    /// model is resident), so it never blocks the UI.
86    async fn verify_placement(&self, current_num_ctx: Option<usize>) -> Option<ModelPlacement> {
87        let _ = current_num_ctx;
88        None
89    }
90
91    /// Stream a chat turn. Typed events flow through
92    /// `ctx.sink`; the returned `FinalResponse` carries token usage
93    /// and the Anthropic thinking-signature (opaque blob required to
94    /// continue extended thinking across turns).
95    ///
96    /// Cancellation: the provider MUST select! on `ctx.token.
97    /// cancelled()` inside any await that could block for more than
98    /// a few hundred ms. This is the contract that replaces the old
99    /// `check_interrupt` polling pattern.
100    async fn chat(&self, request: ChatRequest, ctx: StreamContext) -> Result<FinalResponse>;
101}
102
103/// Run a one-shot, non-interactive model call and collect its streamed text
104/// into a `String`. For internal calls whose output is NOT shown to the user
105/// as it streams — context compaction and the Auto-mode safety classifier.
106/// Drains a private event channel (ignoring reasoning / tool-call events) and
107/// returns the collected text plus final token usage. The `token` lets the
108/// caller cancel the call (e.g. on Ctrl+C) like any other turn work.
109pub(crate) async fn collect_text(
110    provider: Arc<dyn ModelProvider>,
111    turn: TurnId,
112    request: ChatRequest,
113    token: tokio_util::sync::CancellationToken,
114) -> Result<(String, Option<TokenUsage>)> {
115    let (stream_tx, mut stream_rx) = tokio::sync::mpsc::channel::<StreamEvent>(128);
116    let ctx = StreamContext::new(token, stream_tx, turn);
117    let collector = tokio::task::spawn(async move {
118        let mut text = String::new();
119        let mut usage = None;
120        while let Some(event) = stream_rx.recv().await {
121            match event {
122                StreamEvent::Text(chunk) => text.push_str(&chunk),
123                StreamEvent::Done {
124                    usage: done_usage, ..
125                } => usage = done_usage,
126                StreamEvent::Reasoning(_)
127                | StreamEvent::ToolCall(_)
128                | StreamEvent::ThinkingSignature(_) => {},
129            }
130        }
131        (text, usage)
132    });
133
134    let response = provider.chat(request, ctx).await;
135    let (text, stream_usage) = collector.await.map_err(|err| {
136        ModelError::StreamError(format!("collect_text collector failed: {}", err))
137    })?;
138    match response {
139        Ok(final_response) => Ok((text, final_response.usage.or(stream_usage))),
140        Err(err) => Err(err),
141    }
142}
143
144pub use anthropic::AnthropicProvider;
145pub use gemini::GeminiProvider;
146pub use ollama::OllamaProvider;
147pub use openai_compat::OpenAICompatProvider;
mermaid_cli/providers/model/mod.rs

mermaid_cli/providers/model/
mod.rs
