gemini_client_api/gemini/
ask.rs1use super::error::GeminiResponseError;
2use super::types::caching::{CachedContent, CachedContentList, CachedContentUpdate};
3use super::types::request::*;
4use super::types::response::*;
5use super::types::sessions::Session;
6use reqwest::Client;
7use serde_json::{Value, json};
8use std::time::Duration;
9
10pub const BASE_URL: &str = "https://generativelanguage.googleapis.com/v1beta/models";
11
12#[derive(Clone, Default, Debug)]
14pub struct Gemini {
15 client: Client,
16 api_key: String,
17 model: String,
18 sys_prompt: Option<SystemInstruction>,
19 generation_config: Option<Value>,
20 safety_settings: Option<Vec<SafetySetting>>,
21 tools: Option<Vec<Tool>>,
22 tool_config: Option<ToolConfig>,
23 cached_content: Option<String>,
24}
25
26impl Gemini {
27 pub fn new(
34 api_key: impl Into<String>,
35 model: impl Into<String>,
36 sys_prompt: Option<SystemInstruction>,
37 ) -> Self {
38 Self {
39 client: Client::default(),
40 api_key: api_key.into(),
41 model: model.into(),
42 sys_prompt,
43 generation_config: None,
44 safety_settings: None,
45 tools: None,
46 tool_config: None,
47 cached_content: None,
48 }
49 }
50 #[deprecated]
58 pub fn new_with_timeout(
59 api_key: impl Into<String>,
60 model: impl Into<String>,
61 sys_prompt: Option<SystemInstruction>,
62 api_timeout: Duration,
63 ) -> Self {
64 Self {
65 client: Client::builder().timeout(api_timeout).build().unwrap(),
66 api_key: api_key.into(),
67 model: model.into(),
68 sys_prompt,
69 generation_config: None,
70 safety_settings: None,
71 tools: None,
72 tool_config: None,
73 cached_content: None,
74 }
75 }
76 pub fn new_with_client(
84 api_key: impl Into<String>,
85 model: impl Into<String>,
86 sys_prompt: Option<SystemInstruction>,
87 client: Client,
88 ) -> Self {
89 Self {
90 client,
91 api_key: api_key.into(),
92 model: model.into(),
93 sys_prompt,
94 generation_config: None,
95 safety_settings: None,
96 tools: None,
97 tool_config: None,
98 cached_content: None,
99 }
100 }
101 pub fn set_generation_config(&mut self) -> &mut Value {
106 if let None = self.generation_config {
107 self.generation_config = Some(json!({}));
108 }
109 self.generation_config.as_mut().unwrap()
110 }
111 pub fn set_tool_config(mut self, config: ToolConfig) -> Self {
112 self.tool_config = Some(config);
113 self
114 }
115 pub fn set_thinking_config(mut self, config: ThinkingConfig) -> Self {
116 if let Value::Object(map) = self.set_generation_config() {
117 if let Ok(thinking_value) = serde_json::to_value(config) {
118 map.insert("thinking_config".to_string(), thinking_value);
119 }
120 }
121 self
122 }
123 pub fn set_model(mut self, model: impl Into<String>) -> Self {
124 self.model = model.into();
125 self
126 }
127 pub fn set_sys_prompt(mut self, sys_prompt: Option<SystemInstruction>) -> Self {
130 self.sys_prompt = sys_prompt;
131 self
132 }
133 pub fn set_safety_settings(mut self, settings: Option<Vec<SafetySetting>>) -> Self {
134 self.safety_settings = settings;
135 self
136 }
137 pub fn set_api_key(mut self, api_key: impl Into<String>) -> Self {
138 self.api_key = api_key.into();
139 self
140 }
141 pub fn set_json_mode(mut self, schema: Value) -> Self {
149 let config = self.set_generation_config();
150 config["response_mime_type"] = "application/json".into();
151 config["response_schema"] = schema.into();
152 self
153 }
154 pub fn remove_json_mode(mut self) -> Self {
155 if let Some(ref mut generation_config) = self.generation_config {
156 generation_config["response_schema"] = None::<Value>.into();
157 generation_config["response_mime_type"] = None::<Value>.into();
158 }
159 self
160 }
161 pub fn set_tools(mut self, tools: Vec<Tool>) -> Self {
163 self.tools = Some(tools);
164 self
165 }
166 pub fn remove_tools(mut self) -> Self {
168 self.tools = None;
169 self
170 }
171 pub fn set_cached_content(mut self, name: impl Into<String>) -> Self {
172 self.cached_content = Some(name.into());
173 self
174 }
175 pub fn remove_cached_content(mut self) -> Self {
176 self.cached_content = None;
177 self
178 }
179
180 pub async fn create_cache(
183 &self,
184 cached_content: &CachedContent,
185 ) -> Result<CachedContent, GeminiResponseError> {
186 let req_url = format!(
187 "https://generativelanguage.googleapis.com/v1beta/cachedContents?key={}",
188 self.api_key
189 );
190
191 let response = self
192 .client
193 .post(req_url)
194 .json(cached_content)
195 .send()
196 .await
197 .map_err(|e| GeminiResponseError::ReqwestError(e))?;
198
199 if !response.status().is_success() {
200 let error = response
201 .json()
202 .await
203 .map_err(|e| GeminiResponseError::ReqwestError(e))?;
204 return Err(GeminiResponseError::StatusNotOk(error));
205 }
206
207 let cached_content: CachedContent = response
208 .json()
209 .await
210 .map_err(|e| GeminiResponseError::ReqwestError(e))?;
211 Ok(cached_content)
212 }
213
214 pub async fn list_caches(&self) -> Result<CachedContentList, GeminiResponseError> {
215 let req_url = format!(
216 "https://generativelanguage.googleapis.com/v1beta/cachedContents?key={}",
217 self.api_key
218 );
219
220 let response = self
221 .client
222 .get(req_url)
223 .send()
224 .await
225 .map_err(|e| GeminiResponseError::ReqwestError(e))?;
226
227 if !response.status().is_success() {
228 let error = response
229 .json()
230 .await
231 .map_err(|e| GeminiResponseError::ReqwestError(e))?;
232 return Err(GeminiResponseError::StatusNotOk(error));
233 }
234
235 let list: CachedContentList = response
236 .json()
237 .await
238 .map_err(|e| GeminiResponseError::ReqwestError(e))?;
239 Ok(list)
240 }
241
242 pub async fn get_cache(&self, name: &str) -> Result<CachedContent, GeminiResponseError> {
243 let req_url = format!(
244 "https://generativelanguage.googleapis.com/v1beta/{}?key={}",
245 name, self.api_key
246 );
247
248 let response = self
249 .client
250 .get(req_url)
251 .send()
252 .await
253 .map_err(|e| GeminiResponseError::ReqwestError(e))?;
254
255 if !response.status().is_success() {
256 let error = response
257 .json()
258 .await
259 .map_err(|e| GeminiResponseError::ReqwestError(e))?;
260 return Err(GeminiResponseError::StatusNotOk(error));
261 }
262
263 let cached_content: CachedContent = response
264 .json()
265 .await
266 .map_err(|e| GeminiResponseError::ReqwestError(e))?;
267 Ok(cached_content)
268 }
269
270 pub async fn update_cache(
271 &self,
272 name: &str,
273 update: &CachedContentUpdate,
274 ) -> Result<CachedContent, GeminiResponseError> {
275 let req_url = format!(
276 "https://generativelanguage.googleapis.com/v1beta/{}?key={}",
277 name, self.api_key
278 );
279
280 let response = self
281 .client
282 .patch(req_url)
283 .json(update)
284 .send()
285 .await
286 .map_err(|e| GeminiResponseError::ReqwestError(e))?;
287
288 if !response.status().is_success() {
289 let error = response
290 .json()
291 .await
292 .map_err(|e| GeminiResponseError::ReqwestError(e))?;
293 return Err(GeminiResponseError::StatusNotOk(error));
294 }
295
296 let cached_content: CachedContent = response
297 .json()
298 .await
299 .map_err(|e| GeminiResponseError::ReqwestError(e))?;
300 Ok(cached_content)
301 }
302
303 pub async fn delete_cache(&self, name: &str) -> Result<(), GeminiResponseError> {
304 let req_url = format!(
305 "https://generativelanguage.googleapis.com/v1beta/{}?key={}",
306 name, self.api_key
307 );
308
309 let response = self
310 .client
311 .delete(req_url)
312 .send()
313 .await
314 .map_err(|e| GeminiResponseError::ReqwestError(e))?;
315
316 if !response.status().is_success() {
317 let error = response
318 .json()
319 .await
320 .map_err(|e| GeminiResponseError::ReqwestError(e))?;
321 return Err(GeminiResponseError::StatusNotOk(error));
322 }
323
324 Ok(())
325 }
326
327 pub async fn ask(&self, session: &mut Session) -> Result<GeminiResponse, GeminiResponseError> {
334 if session
335 .get_last_chat()
336 .is_some_and(|chat| *chat.role() == Role::Model)
337 {
338 return Err(GeminiResponseError::NothingToRespond);
339 }
340 let req_url = format!(
341 "{BASE_URL}/{}:generateContent?key={}",
342 self.model, self.api_key
343 );
344
345 let response = self
346 .client
347 .post(req_url)
348 .json(&GeminiRequestBody::new(
349 self.sys_prompt.as_ref(),
350 self.tools.as_deref(),
351 &session.get_history().as_slice(),
352 self.generation_config.as_ref(),
353 self.safety_settings.as_deref(),
354 self.tool_config.as_ref(),
355 self.cached_content.clone(),
356 ))
357 .send()
358 .await
359 .map_err(|e| GeminiResponseError::ReqwestError(e))?;
360
361 if !response.status().is_success() {
362 let error = response
363 .json()
364 .await
365 .map_err(|e| GeminiResponseError::ReqwestError(e))?;
366 return Err(GeminiResponseError::StatusNotOk(error));
367 }
368
369 let reply = GeminiResponse::new(response)
370 .await
371 .map_err(|e| GeminiResponseError::ReqwestError(e))?;
372 session.update(&reply);
373 Ok(reply)
374 }
375 pub async fn ask_as_stream_with_extractor<F, StreamType>(
389 &self,
390 session: Session,
391 data_extractor: F,
392 ) -> Result<ResponseStream<F, StreamType>, (Session, GeminiResponseError)>
393 where
394 F: FnMut(&Session, GeminiResponse) -> StreamType,
395 {
396 if session
397 .get_last_chat()
398 .is_some_and(|chat| *chat.role() == Role::Model)
399 {
400 return Err((session, GeminiResponseError::NothingToRespond));
401 }
402 let req_url = format!(
403 "{BASE_URL}/{}:streamGenerateContent?alt=sse&key={}",
404 self.model, self.api_key
405 );
406
407 let request = self
408 .client
409 .post(req_url)
410 .json(&GeminiRequestBody::new(
411 self.sys_prompt.as_ref(),
412 self.tools.as_deref(),
413 session.get_history().as_slice(),
414 self.generation_config.as_ref(),
415 self.safety_settings.as_deref(),
416 self.tool_config.as_ref(),
417 self.cached_content.clone(),
418 ))
419 .send()
420 .await;
421 let response = match request {
422 Ok(response) => response,
423 Err(e) => return Err((session, GeminiResponseError::ReqwestError(e))),
424 };
425
426 if !response.status().is_success() {
427 let error = match response.json().await {
428 Ok(response) => response,
429 Err(e) => return Err((session, GeminiResponseError::ReqwestError(e))),
430 };
431 return Err((session, GeminiResponseError::StatusNotOk(error)));
432 }
433
434 Ok(ResponseStream::new(
435 Box::new(response.bytes_stream()),
436 session,
437 data_extractor,
438 ))
439 }
440 pub async fn ask_as_stream(
459 &self,
460 session: Session,
461 ) -> Result<GeminiResponseStream, (Session, GeminiResponseError)> {
462 self.ask_as_stream_with_extractor(
463 session,
464 (|_, gemini_response| gemini_response)
465 as fn(&Session, GeminiResponse) -> GeminiResponse,
466 )
467 .await
468 }
469}