1use aither_core::llm::model::Ability;
2
3pub const GEMINI_API_BASE_URL: &str = "https://generativelanguage.googleapis.com/v1beta";
5
6pub const USER_AGENT: &str = "aither-gemini/0.1";
7pub const DEFAULT_MODEL: &str = "gemini-3.7-flash";
10pub const DEFAULT_EMBEDDING_MODEL: &str = "gemini-embedding-001";
11pub 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub enum AuthMode {
19 Query,
21 Header,
23}
24
25#[derive(Clone, Debug)]
27pub struct Gemini {
28 inner: GeminiConfig,
29}
30
31impl Gemini {
32 #[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 #[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 #[must_use]
60 pub const fn with_auth_mode(mut self, mode: AuthMode) -> Self {
61 self.inner.auth = mode;
62 self
63 }
64
65 #[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 #[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 #[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 #[must_use]
89 pub fn without_image_model(mut self) -> Self {
90 self.inner.image_model = None;
91 self
92 }
93
94 #[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 #[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 #[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}