mermaid_cli/providers/model/
gemini.rs

1//! Gemini provider — wraps `models::adapters::gemini::GeminiAdapter`.
2//!
3//! Google's Gemini family uses a different wire format from OpenAI-
4//! compat (`:streamGenerateContent?alt=sse` + protobuf-ish JSON
5//! shape). The adapter handles all of that; this wrapper just
6//! forwards.
7
8use async_trait::async_trait;
9
10use mermaid_domain::ChatRequest;
11use mermaid_model::models::adapters::gemini::GeminiAdapter;
12use mermaid_model::models::{Model, ModelConfig, ModelError, Result};
13
14use super::super::ctx::{FinalResponse, StreamContext, StreamEvent};
15use super::{ContextSizing, ModelProvider, RejectionCache, resolve_limits_cached};
16use mermaid_model::models::ModelCapabilities;
17
18/// Gemini's AI Studio root, and the env vars its key lives in. `LEGACY_API_KEY_ENV`
19/// predates Google's rename and is still accepted when `GOOGLE_API_KEY` is unset —
20/// but only when the user has not pointed at a specific var themselves.
21pub const DEFAULT_BASE_URL: &str = "https://generativelanguage.googleapis.com/v1beta";
22pub const DEFAULT_API_KEY_ENV: &str = "GOOGLE_API_KEY";
23pub const LEGACY_API_KEY_ENV: &str = "GEMINI_API_KEY";
24
25pub struct GeminiProvider {
26    adapter: GeminiAdapter,
27    capabilities: ModelCapabilities,
28    /// Links the adapter's learned rejections to the cross-session cache.
29    rejections: RejectionCache,
30}
31
32impl GeminiProvider {
33    /// Wrap a fresh [`GeminiAdapter`] as a `ModelProvider`.
34    ///
35    /// # Errors
36    ///
37    /// Only [`GeminiAdapter::new`]'s — the HTTP client build. The API is not
38    /// contacted here, so an invalid key or unreachable `base_url` still
39    /// constructs and fails on the first request.
40    pub fn new(api_key: String, model_name: String, base_url: String) -> Result<Self> {
41        let adapter = GeminiAdapter::new(api_key, model_name, base_url)?;
42        let capabilities = adapter.capabilities().clone();
43        Ok(Self {
44            adapter,
45            capabilities,
46            rejections: RejectionCache::default(),
47        })
48    }
49}
50
51#[async_trait]
52impl ModelProvider for GeminiProvider {
53    fn capabilities(&self) -> &ModelCapabilities {
54        &self.capabilities
55    }
56
57    /// Live limit discovery via Gemini's models endpoint (`GET {base}/models/
58    /// {id}` → `inputTokenLimit` window + `outputTokenLimit` output ceiling).
59    /// Cache-first via `provider_probes` (TTL-bounded), one live fetch on a
60    /// miss; a fetch failure resolves all-`None`.
61    async fn resolve_context_window(&self, request: &ChatRequest) -> ContextSizing {
62        let _ = request;
63        let model = Model::name(&self.adapter).to_string();
64        let limits =
65            resolve_limits_cached("gemini", &model, || self.adapter.fetch_model_limits()).await;
66        let window = limits.as_ref().and_then(|l| l.max_context_tokens);
67        ContextSizing {
68            model_max: window,
69            effective: window,
70            source: None,
71            max_output: limits.as_ref().and_then(|l| l.max_output_tokens),
72            compacts_natively: false,
73        }
74    }
75
76    async fn chat(&self, request: ChatRequest, ctx: StreamContext) -> Result<FinalResponse> {
77        let config = ModelConfig::from(&request);
78        let chat_fut = self
79            .adapter
80            .chat(&request.messages, &config, Some(ctx.sink.clone()));
81
82        // Seed what earlier sessions learned this model rejects, and persist
83        // whatever this turn learns (see `learning`).
84        let learned_model = Model::name(&self.adapter).to_string();
85        let memory = self.adapter.param_memory();
86        self.rejections.seed("gemini", &learned_model, memory).await;
87
88        let response = tokio::select! {
89            biased;
90            _ = ctx.token.cancelled() => {
91                return Err(ModelError::Cancelled);
92            },
93            r = chat_fut => {
94                self.rejections.persist("gemini", &learned_model, memory).await;
95                r?
96            },
97        };
98
99        let usage = response.usage.clone();
100        let stop_reason = response.stop_reason.clone();
101        // The terminal Done goes on the same sink the adapter just finished
102        // writing to, so it cannot overtake a still-queued ToolCall.
103        let _ = ctx
104            .sink
105            .send(StreamEvent::Done {
106                usage: usage.clone(),
107                provider_continuation: None,
108                stop_reason: stop_reason.clone(),
109            })
110            .await;
111
112        Ok(FinalResponse {
113            usage,
114            provider_continuation: None,
115            tool_calls: response.tool_calls.unwrap_or_default(),
116            stop_reason,
117        })
118    }
119}
mermaid_cli/providers/model/gemini.rs

mermaid_cli/providers/model/
gemini.rs
