gemini_rust/
models.rs

1use serde::{Deserialize, Serialize};
2
3/// Role of a message in a conversation
4#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
5#[serde(rename_all = "lowercase")]
6pub enum Role {
7    /// Message from the user
8    User,
9    /// Message from the model
10    Model,
11    /// System prompt
12    System,
13    /// Function response
14    Function,
15}
16
17/// Content part that can be included in a message
18#[derive(Debug, Clone, Serialize, Deserialize)]
19#[serde(untagged)]
20pub enum Part {
21    /// Text content
22    Text {
23        /// The text content
24        text: String,
25    },
26    /// Function call from the model
27    FunctionCall {
28        /// The function call details
29        #[serde(rename = "functionCall")]
30        function_call: super::tools::FunctionCall,
31    },
32    /// Function response (results from executing a function call)
33    FunctionResponse {
34        /// The function response details
35        #[serde(rename = "functionResponse")]
36        function_response: super::tools::FunctionResponse,
37    },
38}
39
40/// Content of a message
41#[derive(Debug, Default, Clone, Serialize, Deserialize)]
42pub struct Content {
43    /// Parts of the content
44    pub parts: Vec<Part>,
45    /// Role of the content
46    #[serde(skip_serializing_if = "Option::is_none")]
47    pub role: Option<Role>,
48}
49
50impl Content {
51    /// Create a new text content
52    pub fn text(text: impl Into<String>) -> Self {
53        Self {
54            parts: vec![Part::Text { text: text.into() }],
55            role: None,
56        }
57    }
58
59    /// Create a new content with a function call
60    pub fn function_call(function_call: super::tools::FunctionCall) -> Self {
61        Self {
62            parts: vec![Part::FunctionCall { function_call }],
63            role: None,
64        }
65    }
66
67    /// Create a new content with a function response
68    pub fn function_response(function_response: super::tools::FunctionResponse) -> Self {
69        Self {
70            parts: vec![Part::FunctionResponse { function_response }],
71            role: None,
72        }
73    }
74    
75    /// Create a new content with a function response from name and JSON value
76    pub fn function_response_json(
77        name: impl Into<String>, 
78        response: serde_json::Value
79    ) -> Self {
80        Self {
81            parts: vec![Part::FunctionResponse { 
82                function_response: super::tools::FunctionResponse::new(name, response) 
83            }],
84            role: None,
85        }
86    }
87
88    /// Add a role to this content
89    pub fn with_role(mut self, role: Role) -> Self {
90        self.role = Some(role);
91        self
92    }
93}
94
95/// Message in a conversation
96#[derive(Debug, Clone, Serialize, Deserialize)]
97pub struct Message {
98    /// Content of the message
99    pub content: Content,
100    /// Role of the message
101    pub role: Role,
102}
103
104impl Message {
105    /// Create a new user message with text content
106    pub fn user(text: impl Into<String>) -> Self {
107        Self {
108            content: Content::text(text),
109            role: Role::User,
110        }
111    }
112
113    /// Create a new model message with text content
114    pub fn model(text: impl Into<String>) -> Self {
115        Self {
116            content: Content::text(text),
117            role: Role::Model,
118        }
119    }
120
121    /// Create a new system message with text content
122    pub fn system(text: impl Into<String>) -> Self {
123        Self {
124            content: Content::text(text),
125            role: Role::System,
126        }
127    }
128
129    /// Create a new function message with function response content from JSON
130    pub fn function(name: impl Into<String>, response: serde_json::Value) -> Self {
131        Self {
132            content: Content::function_response_json(name, response),
133            role: Role::Function,
134        }
135    }
136    
137    /// Create a new function message with function response from a JSON string
138    pub fn function_str(name: impl Into<String>, response: impl Into<String>) -> Result<Self, serde_json::Error> {
139        let response_str = response.into();
140        let json = serde_json::from_str(&response_str)?;
141        Ok(Self {
142            content: Content::function_response_json(name, json),
143            role: Role::Function,
144        })
145    }
146}
147
148/// Safety rating for content
149#[derive(Debug, Clone, Serialize, Deserialize)]
150pub struct SafetyRating {
151    /// The category of the safety rating
152    pub category: String,
153    /// The probability that the content is harmful
154    pub probability: String,
155}
156
157/// Citation metadata for content
158#[derive(Debug, Clone, Serialize, Deserialize)]
159pub struct CitationMetadata {
160    /// The citation sources
161    pub citation_sources: Vec<CitationSource>,
162}
163
164/// Citation source
165#[derive(Debug, Clone, Serialize, Deserialize)]
166pub struct CitationSource {
167    /// The URI of the citation source
168    pub uri: Option<String>,
169    /// The title of the citation source
170    pub title: Option<String>,
171    /// The start index of the citation in the response
172    pub start_index: Option<i32>,
173    /// The end index of the citation in the response
174    pub end_index: Option<i32>,
175    /// The license of the citation source
176    pub license: Option<String>,
177    /// The publication date of the citation source
178    pub publication_date: Option<String>,
179}
180
181/// A candidate response
182#[derive(Debug, Clone, Serialize, Deserialize)]
183pub struct Candidate {
184    /// The content of the candidate
185    pub content: Content,
186    /// The safety ratings for the candidate
187    #[serde(skip_serializing_if = "Option::is_none")]
188    pub safety_ratings: Option<Vec<SafetyRating>>,
189    /// The citation metadata for the candidate
190    #[serde(skip_serializing_if = "Option::is_none")]
191    pub citation_metadata: Option<CitationMetadata>,
192    /// The finish reason for the candidate
193    #[serde(skip_serializing_if = "Option::is_none")]
194    pub finish_reason: Option<String>,
195    /// The tokens used in the response
196    #[serde(skip_serializing_if = "Option::is_none")]
197    pub usage_metadata: Option<UsageMetadata>,
198}
199
200/// Metadata about token usage
201#[derive(Debug, Clone, Serialize, Deserialize)]
202pub struct UsageMetadata {
203    /// The number of prompt tokens
204    pub prompt_token_count: i32,
205    /// The number of response tokens
206    pub candidates_token_count: i32,
207    /// The total number of tokens
208    pub total_token_count: i32,
209}
210
211/// Response from the Gemini API for content generation
212#[derive(Debug, Clone, Serialize, Deserialize)]
213pub struct GenerationResponse {
214    /// The candidates generated
215    pub candidates: Vec<Candidate>,
216    /// The prompt feedback
217    #[serde(skip_serializing_if = "Option::is_none")]
218    pub prompt_feedback: Option<PromptFeedback>,
219    /// Usage metadata
220    #[serde(skip_serializing_if = "Option::is_none")]
221    pub usage_metadata: Option<UsageMetadata>,
222}
223
224/// Feedback about the prompt
225#[derive(Debug, Clone, Serialize, Deserialize)]
226pub struct PromptFeedback {
227    /// The safety ratings for the prompt
228    pub safety_ratings: Vec<SafetyRating>,
229    /// The block reason if the prompt was blocked
230    #[serde(skip_serializing_if = "Option::is_none")]
231    pub block_reason: Option<String>,
232}
233
234impl GenerationResponse {
235    /// Get the text of the first candidate
236    pub fn text(&self) -> String {
237        self.candidates
238            .first()
239            .and_then(|c| {
240                c.content.parts.first().and_then(|p| match p {
241                    Part::Text { text } => Some(text.clone()),
242                    _ => None,
243                })
244            })
245            .unwrap_or_default()
246    }
247
248    /// Get function calls from the response
249    pub fn function_calls(&self) -> Vec<&super::tools::FunctionCall> {
250        self.candidates
251            .iter()
252            .flat_map(|c| {
253                c.content.parts.iter().filter_map(|p| match p {
254                    Part::FunctionCall { function_call } => Some(function_call),
255                    _ => None,
256                })
257            })
258            .collect()
259    }
260}
261
262/// Request to generate content
263#[derive(Debug, Clone, Serialize, Deserialize)]
264pub struct GenerateContentRequest {
265    /// The contents to generate content from
266    pub contents: Vec<Content>,
267    /// The generation config
268    #[serde(skip_serializing_if = "Option::is_none")]
269    pub generation_config: Option<GenerationConfig>,
270    /// The safety settings
271    #[serde(skip_serializing_if = "Option::is_none")]
272    pub safety_settings: Option<Vec<SafetySetting>>,
273    /// The tools that the model can use
274    #[serde(skip_serializing_if = "Option::is_none")]
275    pub tools: Option<Vec<super::tools::Tool>>,
276    /// The tool config
277    #[serde(skip_serializing_if = "Option::is_none")]
278    pub tool_config: Option<ToolConfig>,
279}
280
281/// Configuration for generation
282#[derive(Debug, Clone, Serialize, Deserialize)]
283pub struct GenerationConfig {
284    /// The temperature for the model (0.0 to 1.0)
285    /// 
286    /// Controls the randomness of the output. Higher values (e.g., 0.9) make output 
287    /// more random, lower values (e.g., 0.1) make output more deterministic.
288    #[serde(skip_serializing_if = "Option::is_none")]
289    pub temperature: Option<f32>,
290    
291    /// The top-p value for the model (0.0 to 1.0)
292    /// 
293    /// For each token generation step, the model considers the top_p percentage of 
294    /// probability mass for potential token choices. Lower values are more selective, 
295    /// higher values allow more variety.
296    #[serde(skip_serializing_if = "Option::is_none")]
297    pub top_p: Option<f32>,
298    
299    /// The top-k value for the model
300    /// 
301    /// For each token generation step, the model considers the top_k most likely tokens.
302    /// Lower values are more selective, higher values allow more variety.
303    #[serde(skip_serializing_if = "Option::is_none")]
304    pub top_k: Option<i32>,
305    
306    /// The maximum number of tokens to generate
307    /// 
308    /// Limits the length of the generated content. One token is roughly 4 characters.
309    #[serde(skip_serializing_if = "Option::is_none")]
310    pub max_output_tokens: Option<i32>,
311    
312    /// The candidate count
313    /// 
314    /// Number of alternative responses to generate.
315    #[serde(skip_serializing_if = "Option::is_none")]
316    pub candidate_count: Option<i32>,
317    
318    /// Whether to stop on specific sequences
319    /// 
320    /// The model will stop generating content when it encounters any of these sequences.
321    #[serde(skip_serializing_if = "Option::is_none")]
322    pub stop_sequences: Option<Vec<String>>,
323    
324    /// The response mime type
325    /// 
326    /// Specifies the format of the model's response.
327    #[serde(skip_serializing_if = "Option::is_none")]
328    pub response_mime_type: Option<String>,
329    
330    /// The response schema
331    /// 
332    /// Specifies the JSON schema for structured responses.
333    #[serde(skip_serializing_if = "Option::is_none")]
334    pub response_schema: Option<serde_json::Value>,
335}
336
337impl Default for GenerationConfig {
338    fn default() -> Self {
339        Self {
340            temperature: Some(0.7),
341            top_p: Some(0.95),
342            top_k: Some(40),
343            max_output_tokens: Some(1024),
344            candidate_count: Some(1),
345            stop_sequences: None,
346            response_mime_type: None,
347            response_schema: None,
348        }
349    }
350}
351
352/// Configuration for tools
353#[derive(Debug, Clone, Serialize, Deserialize)]
354pub struct ToolConfig {
355    /// The function calling config
356    #[serde(skip_serializing_if = "Option::is_none")]
357    pub function_calling_config: Option<FunctionCallingConfig>,
358}
359
360/// Configuration for function calling
361#[derive(Debug, Clone, Serialize, Deserialize)]
362pub struct FunctionCallingConfig {
363    /// The mode for function calling
364    pub mode: FunctionCallingMode,
365}
366
367/// Mode for function calling
368#[derive(Debug, Clone, Serialize, Deserialize)]
369#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
370pub enum FunctionCallingMode {
371    /// The model may use function calling
372    Auto,
373    /// The model must use function calling
374    Any,
375    /// The model must not use function calling
376    None,
377}
378
379/// Setting for safety
380#[derive(Debug, Clone, Serialize, Deserialize)]
381pub struct SafetySetting {
382    /// The category of content to filter
383    pub category: HarmCategory,
384    /// The threshold for filtering
385    pub threshold: HarmBlockThreshold,
386}
387
388/// Category of harmful content
389#[derive(Debug, Clone, Serialize, Deserialize)]
390#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
391pub enum HarmCategory {
392    /// Dangerous content
393    Dangerous,
394    /// Harassment content
395    Harassment,
396    /// Hate speech
397    HateSpeech,
398    /// Sexually explicit content
399    SexuallyExplicit,
400}
401
402/// Threshold for blocking harmful content
403#[derive(Debug, Clone, Serialize, Deserialize)]
404#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
405pub enum HarmBlockThreshold {
406    /// Block content with low probability of harm
407    BlockLowAndAbove,
408    /// Block content with medium probability of harm
409    BlockMediumAndAbove,
410    /// Block content with high probability of harm
411    BlockHighAndAbove,
412    /// Block content with maximum probability of harm
413    BlockOnlyHigh,
414    /// Never block content
415    BlockNone,
416}