Skip to main content

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::{ModelError, Result, TokenUsage};
21
22use super::capabilities::Capabilities;
23use super::ctx::{FinalResponse, StreamContext, StreamEvent};
24
25/// Provider-facing interface. A `ModelProvider` impl owns whatever
26/// HTTP client / state it needs and exposes `chat()` — that's the
27/// whole surface.
28#[async_trait]
29pub trait ModelProvider: Send + Sync {
30    /// Capabilities the provider advertises. The reducer reads this
31    /// when building the outgoing `ChatRequest` (e.g. whether to
32    /// attach reasoning controls).
33    fn capabilities(&self) -> &Capabilities;
34
35    /// Stream a chat turn. Typed events flow through
36    /// `ctx.sink`; the returned `FinalResponse` carries token usage
37    /// and the Anthropic thinking-signature (opaque blob required to
38    /// continue extended thinking across turns).
39    ///
40    /// Cancellation: the provider MUST select! on `ctx.token.
41    /// cancelled()` inside any await that could block for more than
42    /// a few hundred ms. This is the contract that replaces the old
43    /// `check_interrupt` polling pattern.
44    async fn chat(&self, request: ChatRequest, ctx: StreamContext) -> Result<FinalResponse>;
45}
46
47/// Run a one-shot, non-interactive model call and collect its streamed text
48/// into a `String`. For internal calls whose output is NOT shown to the user
49/// as it streams — context compaction and the Auto-mode safety classifier.
50/// Drains a private event channel (ignoring reasoning / tool-call events) and
51/// returns the collected text plus final token usage. The `token` lets the
52/// caller cancel the call (e.g. on Ctrl+C) like any other turn work.
53pub(crate) async fn collect_text(
54    provider: Arc<dyn ModelProvider>,
55    turn: TurnId,
56    request: ChatRequest,
57    token: tokio_util::sync::CancellationToken,
58) -> Result<(String, Option<TokenUsage>)> {
59    let (stream_tx, mut stream_rx) = tokio::sync::mpsc::channel::<StreamEvent>(128);
60    let ctx = StreamContext::new(token, stream_tx, turn);
61    let collector = tokio::task::spawn(async move {
62        let mut text = String::new();
63        let mut usage = None;
64        while let Some(event) = stream_rx.recv().await {
65            match event {
66                StreamEvent::Text(chunk) => text.push_str(&chunk),
67                StreamEvent::Done {
68                    usage: done_usage, ..
69                } => usage = done_usage,
70                StreamEvent::Reasoning(_)
71                | StreamEvent::ToolCall(_)
72                | StreamEvent::ThinkingSignature(_) => {},
73            }
74        }
75        (text, usage)
76    });
77
78    let response = provider.chat(request, ctx).await;
79    let (text, stream_usage) = collector.await.map_err(|err| {
80        ModelError::StreamError(format!("collect_text collector failed: {}", err))
81    })?;
82    match response {
83        Ok(final_response) => Ok((text, final_response.usage.or(stream_usage))),
84        Err(err) => Err(err),
85    }
86}
87
88pub use anthropic::AnthropicProvider;
89pub use gemini::GeminiProvider;
90pub use ollama::OllamaProvider;
91pub use openai_compat::OpenAICompatProvider;