aither_gemini/
provider.rs

1use crate::{GEMINI_API_BASE_URL, Gemini, config::AuthMode, error::GeminiError};
2use aither_core::llm::{
3    LanguageModelProvider, model::Profile as ModelProfile, provider::Profile as ProviderProfile,
4};
5use aither_models::lookup as lookup_model_info;
6use serde::Deserialize;
7use std::{future::Future, sync::Arc};
8use zenwave::{Client, client};
9
10/// Provider capable of listing and instantiating `Gemini` models.
11#[derive(Clone, Debug)]
12pub struct GeminiProvider {
13    inner: Arc<ProviderConfig>,
14}
15
16impl GeminiProvider {
17    /// Create a new provider with the given API key.
18    pub fn new(api_key: impl Into<String>) -> Self {
19        Self {
20            inner: Arc::new(ProviderConfig {
21                api_key: api_key.into(),
22                base_url: GEMINI_API_BASE_URL.to_string(),
23                auth: AuthMode::Query,
24            }),
25        }
26    }
27
28    /// Override the REST base URL.
29    #[must_use]
30    pub fn base_url(mut self, url: impl Into<String>) -> Self {
31        Arc::make_mut(&mut self.inner).base_url = url.into();
32        self
33    }
34
35    /// Select header-based authentication.
36    #[must_use]
37    pub fn with_auth_mode(mut self, mode: AuthMode) -> Self {
38        Arc::make_mut(&mut self.inner).auth = mode;
39        self
40    }
41}
42
43impl LanguageModelProvider for GeminiProvider {
44    type Model = Gemini;
45    type Error = GeminiError;
46
47    fn list_models(&self) -> impl Future<Output = Result<Vec<ModelProfile>, Self::Error>> + Send {
48        let cfg = self.inner.clone();
49        async move {
50            let mut endpoint = format!("{}/models", cfg.base_url.trim_end_matches('/'));
51            if cfg.auth == AuthMode::Query {
52                let separator = if endpoint.contains('?') { '&' } else { '?' };
53                endpoint.push(separator);
54                endpoint.push_str("key=");
55                endpoint.push_str(&cfg.api_key);
56            }
57
58            let mut backend = client();
59            let mut builder = backend.get(endpoint).map_err(GeminiError::Http)?;
60
61            if cfg.auth == AuthMode::Header {
62                builder = builder
63                    .header("x-goog-api-key", cfg.api_key.clone())
64                    .map_err(GeminiError::Http)?;
65            }
66
67            let response: ModelListResponse = builder.json().await.map_err(GeminiError::Http)?;
68
69            Ok(response
70                .models
71                .into_iter()
72                .filter_map(|model| {
73                    let model_id = model.name.trim_start_matches("models/").to_string();
74                    let context_length = model
75                        .input_token_limit
76                        .and_then(|value| u32::try_from(value).ok())
77                        .or_else(|| lookup_model_info(&model_id).and_then(aither_models::ModelEntry::max_input_tokens))
78                        .or_else(|| lookup_model_info(&model.name).and_then(aither_models::ModelEntry::max_input_tokens));
79                    let Some(context_length) = context_length else {
80                        tracing::warn!(
81                            model = %model.name,
82                            "skip model from Gemini list: missing input_token_limit and no aither-models fallback"
83                        );
84                        return None;
85                    };
86                    Some(ModelProfile::new(
87                        model_id,
88                        "google",
89                        model.name,
90                        model.description.unwrap_or_default(),
91                        context_length,
92                    ))
93                })
94                .collect())
95        }
96    }
97
98    fn get_model(
99        &self,
100        name: &str,
101    ) -> impl Future<Output = Result<Self::Model, Self::Error>> + Send {
102        let cfg = self.inner.clone();
103        let name = name.to_string();
104        async move {
105            let backend = Gemini::new(cfg.api_key.clone())
106                .with_base_url(cfg.base_url.clone())
107                .with_auth_mode(cfg.auth)
108                .with_text_model(name);
109            Ok(backend)
110        }
111    }
112
113    fn profile() -> ProviderProfile {
114        ProviderProfile::new("gemini", "Google Gemini API provider")
115    }
116}
117
118#[derive(Debug, Clone)]
119struct ProviderConfig {
120    api_key: String,
121    base_url: String,
122    auth: AuthMode,
123}
124
125#[derive(Debug, Deserialize)]
126struct ModelListResponse {
127    models: Vec<ModelDescriptor>,
128}
129
130/// A model entry as the list-models endpoint returns it.
131///
132/// Fields this crate does not read are still declared so the payload
133/// deserializes faithfully.
134#[allow(dead_code)]
135#[derive(Debug, Deserialize)]
136#[serde(rename_all = "camelCase")]
137struct ModelDescriptor {
138    name: String,
139    display_name: Option<String>,
140    description: Option<String>,
141    input_token_limit: Option<i32>,
142    output_token_limit: Option<i32>,
143    supported_generation_methods: Option<Vec<String>>,
144}
aither_gemini/provider.rs

aither_gemini/
provider.rs
