Skip to main content

gemini_rs/
chat.rs

1use std::marker::PhantomData;
2
3use crate::{
4    Client, Result,
5    types::{self, Response},
6};
7
8/// Simplest way to use gemini-rs, and covers 80% of use cases
9pub struct Chat<T> {
10    model: Box<str>,
11    client: Client,
12    system_instruction: Option<Box<str>>,
13    safety_settings: Vec<types::SafetySettings>,
14    history: Vec<types::Content>,
15    config: Option<types::GenerationConfig>,
16    phantom: PhantomData<T>,
17}
18
19impl<T> Chat<T> {
20    pub fn new(client: &Client, model: &str) -> Self {
21        Self {
22            model: model.into(),
23            client: client.clone(),
24            system_instruction: None,
25            safety_settings: Vec::new(),
26            history: Vec::new(),
27            config: None,
28            phantom: PhantomData,
29        }
30    }
31
32    pub fn config(&mut self) -> &types::GenerationConfig {
33        self.config.get_or_insert_default()
34    }
35
36    pub fn to_json(mut self) -> Chat<Json> {
37        self.config_mut().response_mime_type = Some("application/json".into());
38        Chat {
39            model: self.model,
40            client: self.client,
41            system_instruction: self.system_instruction,
42            safety_settings: self.safety_settings,
43            history: self.history,
44            config: self.config,
45            phantom: PhantomData,
46        }
47    }
48
49    pub fn config_mut(&mut self) -> &mut types::GenerationConfig {
50        self.config.get_or_insert_default()
51    }
52
53    pub fn history(&self) -> &[types::Content] {
54        &self.history
55    }
56
57    pub fn history_mut(&mut self) -> &mut Vec<types::Content> {
58        &mut self.history
59    }
60
61    pub fn safety_settings(&mut self, safety_settings: Vec<types::SafetySettings>) {
62        self.safety_settings = safety_settings;
63    }
64
65    pub fn system_instruction(mut self, instruction: &str) -> Self {
66        self.system_instruction = Some(Box::from(instruction));
67        self
68    }
69
70    pub async fn generate_content(&mut self) -> Result<Response> {
71        let mut generate_content = self.client.generate_content(&self.model);
72
73        if let Some(system_instruction) = &self.system_instruction {
74            generate_content.system_instruction(system_instruction);
75        }
76
77        if let Some(config) = &self.config {
78            generate_content.config(config.clone());
79        }
80
81        generate_content.contents(self.history.clone());
82        generate_content.safety_settings(self.safety_settings.clone());
83        generate_content.await
84    }
85
86    pub async fn send_message(&mut self, message: &str) -> Result<Response> {
87        self.history.push(types::Content {
88            role: types::Role::User,
89            parts: vec![types::Part::text(message)],
90        });
91
92        self.generate_content().await
93    }
94    pub async fn send_parted_messages(&mut self, parts: Vec<types::Part>) -> Result<Response> {
95        self.history.push(types::Content {
96            role: types::Role::User,
97            parts,
98        });
99
100        self.generate_content().await
101    }
102}
103
104impl Chat<Json> {
105    pub fn response_schema(mut self, schema: types::Schema) -> Self {
106        self.config_mut().response_schema = Some(schema);
107        self
108    }
109
110    pub async fn json<T: serde::de::DeserializeOwned>(&mut self, message: &str) -> Result<T> {
111        let response = self.send_message(message).await?;
112        let json = format!("{response}");
113        serde_json::from_str(&json).map_err(Into::into)
114    }
115}
116
117pub struct Text {}
118
119pub struct Json {}