Skip to main content

aither_gemini/
config.rs

1use aither_core::llm::model::Ability;
2
3/// Gemini REST base URL used by the Developer API.
4pub const GEMINI_API_BASE_URL: &str = "https://generativelanguage.googleapis.com/v1beta";
5
6pub const USER_AGENT: &str = "aither-gemini/0.1";
7/// Google's latest and most capable Flash model, built for coding and agentic
8/// workflows. Thinking depth is steered with `thinkingLevel`, not a token budget.
9pub const DEFAULT_MODEL: &str = "gemini-3.7-flash";
10pub const DEFAULT_EMBEDDING_MODEL: &str = "gemini-embedding-001";
11/// Nano Banana 2.
12pub const DEFAULT_IMAGE_MODEL: &str = "gemini-3.1-flash-image";
13pub const DEFAULT_TTS_MODEL: &str = "gemini-3.1-flash-tts-preview";
14pub const DEFAULT_TTS_VOICE: &str = "Kore";
15
16/// Authentication strategy supported by the Gemini backend.
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub enum AuthMode {
19    /// Attach `?key=API_KEY` to every request (default).
20    Query,
21    /// Send the API key via `x-goog-api-key` header.
22    Header,
23}
24
25/// Native Gemini backend wired up to the `aither` traits.
26#[derive(Clone, Debug)]
27pub struct Gemini {
28    inner: GeminiConfig,
29}
30
31impl Gemini {
32    /// Create a backend using the default chat/embedding models.
33    #[must_use]
34    pub fn new(api_key: impl Into<String>) -> Self {
35        Self {
36            inner: GeminiConfig {
37                api_key: api_key.into(),
38                base_url: GEMINI_API_BASE_URL.to_string(),
39                auth: AuthMode::Query,
40                text_model: sanitize_model(DEFAULT_MODEL),
41                embedding_model: sanitize_model(DEFAULT_EMBEDDING_MODEL),
42                embedding_dimensions: 3072,
43                image_model: Some(sanitize_model(DEFAULT_IMAGE_MODEL)),
44                tts_model: Some(sanitize_model(DEFAULT_TTS_MODEL)),
45                tts_voice: DEFAULT_TTS_VOICE.to_string(),
46                native_abilities: vec![Ability::Pdf],
47            },
48        }
49    }
50
51    /// Override the REST base URL (useful for sandboxes or proxies).
52    #[must_use]
53    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
54        self.inner.base_url = base_url.into();
55        self
56    }
57
58    /// Select header-based authentication.
59    #[must_use]
60    pub const fn with_auth_mode(mut self, mode: AuthMode) -> Self {
61        self.inner.auth = mode;
62        self
63    }
64
65    /// Override the default chat model.
66    #[must_use]
67    pub fn with_text_model(mut self, model: impl Into<String>) -> Self {
68        self.inner.text_model = sanitize_model(model);
69        self
70    }
71
72    /// Override the embedding model and (optionally) its dimensionality.
73    #[must_use]
74    pub fn with_embedding_model(mut self, model: impl Into<String>, dim: usize) -> Self {
75        self.inner.embedding_model = sanitize_model(model);
76        self.inner.embedding_dimensions = dim;
77        self
78    }
79
80    /// Override the optional image model.
81    #[must_use]
82    pub fn with_image_model(mut self, model: impl Into<String>) -> Self {
83        self.inner.image_model = Some(sanitize_model(model));
84        self
85    }
86
87    /// Disable image generation support.
88    #[must_use]
89    pub fn without_image_model(mut self) -> Self {
90        self.inner.image_model = None;
91        self
92    }
93
94    /// Override the optional text-to-speech model + voice.
95    #[must_use]
96    pub fn with_tts(mut self, model: impl Into<String>, voice: impl Into<String>) -> Self {
97        self.inner.tts_model = Some(sanitize_model(model));
98        self.inner.tts_voice = voice.into();
99        self
100    }
101
102    pub(crate) const fn config(&self) -> &GeminiConfig {
103        &self.inner
104    }
105
106    /// Declare native capabilities exposed by the selected Gemini model.
107    #[must_use]
108    pub fn with_native_capabilities(
109        mut self,
110        abilities: impl IntoIterator<Item = Ability>,
111    ) -> Self {
112        for ability in abilities {
113            if !self.inner.native_abilities.contains(&ability) {
114                self.inner.native_abilities.push(ability);
115            }
116        }
117        self
118    }
119
120    /// Declare a single native capability.
121    #[must_use]
122    pub fn with_native_capability(self, ability: Ability) -> Self {
123        self.with_native_capabilities([ability])
124    }
125}
126
127#[derive(Debug, Clone)]
128pub struct GeminiConfig {
129    pub(crate) api_key: String,
130    pub(crate) base_url: String,
131    pub(crate) auth: AuthMode,
132    pub(crate) text_model: String,
133    pub(crate) embedding_model: String,
134    pub(crate) embedding_dimensions: usize,
135    pub(crate) image_model: Option<String>,
136    pub(crate) tts_model: Option<String>,
137    pub(crate) tts_voice: String,
138    pub(crate) native_abilities: Vec<Ability>,
139}
140
141impl GeminiConfig {
142    pub(crate) fn endpoint(&self, suffix: &str) -> String {
143        let mut url = format!(
144            "{}/{}",
145            self.base_url.trim_end_matches('/'),
146            suffix.trim_start_matches('/')
147        );
148        if self.auth == AuthMode::Query {
149            let separator = if url.contains('?') { '&' } else { '?' };
150            url.push(separator);
151            url.push_str("key=");
152            url.push_str(&self.api_key);
153        }
154        url
155    }
156
157    pub(crate) fn model_endpoint(&self, model: &str, action: &str) -> String {
158        let model = sanitize_model(model);
159        self.endpoint(&format!("{model}:{action}"))
160    }
161}
162
163pub fn sanitize_model(model: impl Into<String>) -> String {
164    let model = model.into();
165    if model.starts_with("models/") {
166        model
167    } else {
168        format!("models/{model}")
169    }
170}