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;