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}
12
13/// Content part that can be included in a message
14#[derive(Debug, Clone, Serialize, Deserialize)]
15#[serde(untagged)]
16pub enum Part {
17    /// Text content
18    Text {
19        /// The text content
20        text: String,
21        /// Whether this is a thought summary (Gemini 2.5 series only)
22        #[serde(skip_serializing_if = "Option::is_none")]
23        thought: Option<bool>,
24    },
25    InlineData {
26        /// The blob data
27        #[serde(rename = "inlineData")]
28        inline_data: Blob,
29    },
30    /// Function call from the model
31    FunctionCall {
32        /// The function call details
33        #[serde(rename = "functionCall")]
34        function_call: super::tools::FunctionCall,
35    },
36    /// Function response (results from executing a function call)
37    FunctionResponse {
38        /// The function response details
39        #[serde(rename = "functionResponse")]
40        function_response: super::tools::FunctionResponse,
41    },
42}
43
44/// Blob for a message part
45#[derive(Debug, Clone, Serialize, Deserialize)]
46#[serde(rename_all = "camelCase")]
47pub struct Blob {
48    pub mime_type: String,
49    pub data: String, // Base64 encoded data
50}
51
52impl Blob {
53    /// Create a new blob with mime type and data
54    pub fn new(mime_type: impl Into<String>, data: impl Into<String>) -> Self {
55        Self {
56            mime_type: mime_type.into(),
57            data: data.into(),
58        }
59    }
60}
61
62/// Content of a message
63#[derive(Debug, Default, Clone, Serialize, Deserialize)]
64pub struct Content {
65    /// Parts of the content
66    pub parts: Vec<Part>,
67    /// Role of the content
68    #[serde(skip_serializing_if = "Option::is_none")]
69    pub role: Option<Role>,
70}
71
72impl Content {
73    /// Create a new text content
74    pub fn text(text: impl Into<String>) -> Self {
75        Self {
76            parts: vec![Part::Text {
77                text: text.into(),
78                thought: None,
79            }],
80            role: None,
81        }
82    }
83
84    /// Create a new content with a function call
85    pub fn function_call(function_call: super::tools::FunctionCall) -> Self {
86        Self {
87            parts: vec![Part::FunctionCall { function_call }],
88            role: None,
89        }
90    }
91
92    /// Create a new content with a function response
93    pub fn function_response(function_response: super::tools::FunctionResponse) -> Self {
94        Self {
95            parts: vec![Part::FunctionResponse { function_response }],
96            role: None,
97        }
98    }
99
100    /// Create a new content with a function response from name and JSON value
101    pub fn function_response_json(name: impl Into<String>, response: serde_json::Value) -> Self {
102        Self {
103            parts: vec![Part::FunctionResponse {
104                function_response: super::tools::FunctionResponse::new(name, response),
105            }],
106            role: None,
107        }
108    }
109
110    /// Create a new content with inline data (blob data)
111    pub fn inline_data(mime_type: impl Into<String>, data: impl Into<String>) -> Self {
112        Self {
113            parts: vec![Part::InlineData {
114                inline_data: Blob::new(mime_type, data),
115            }],
116            role: None,
117        }
118    }
119
120    /// Add a role to this content
121    pub fn with_role(mut self, role: Role) -> Self {
122        self.role = Some(role);
123        self
124    }
125}
126
127/// Message in a conversation
128#[derive(Debug, Clone, Serialize, Deserialize)]
129pub struct Message {
130    /// Content of the message
131    pub content: Content,
132    /// Role of the message
133    pub role: Role,
134}
135
136impl Message {
137    /// Create a new user message with text content
138    pub fn user(text: impl Into<String>) -> Self {
139        Self {
140            content: Content::text(text).with_role(Role::User),
141            role: Role::User,
142        }
143    }
144
145    /// Create a new model message with text content
146    pub fn model(text: impl Into<String>) -> Self {
147        Self {
148            content: Content::text(text).with_role(Role::Model),
149            role: Role::Model,
150        }
151    }
152
153    pub fn embed(text: impl Into<String>) -> Self {
154        Self {
155            content: Content::text(text),
156            role: Role::Model,
157        }
158    }
159
160    /// Create a new function message with function response content from JSON
161    pub fn function(name: impl Into<String>, response: serde_json::Value) -> Self {
162        Self {
163            content: Content::function_response_json(name, response).with_role(Role::Model),
164            role: Role::Model,
165        }
166    }
167
168    /// Create a new function message with function response from a JSON string
169    pub fn function_str(
170        name: impl Into<String>,
171        response: impl Into<String>,
172    ) -> Result<Self, serde_json::Error> {
173        let response_str = response.into();
174        let json = serde_json::from_str(&response_str)?;
175        Ok(Self {
176            content: Content::function_response_json(name, json).with_role(Role::Model),
177            role: Role::Model,
178        })
179    }
180}
181
182/// Safety rating for content
183#[derive(Debug, Clone, Serialize, Deserialize)]
184pub struct SafetyRating {
185    /// The category of the safety rating
186    pub category: String,
187    /// The probability that the content is harmful
188    pub probability: String,
189}
190
191/// Citation metadata for content
192#[derive(Debug, Clone, Serialize, Deserialize)]
193#[serde(rename_all = "camelCase")]
194pub struct CitationMetadata {
195    /// The citation sources
196    pub citation_sources: Vec<CitationSource>,
197}
198
199/// Citation source
200#[derive(Debug, Clone, Serialize, Deserialize)]
201#[serde(rename_all = "camelCase")]
202pub struct CitationSource {
203    /// The URI of the citation source
204    pub uri: Option<String>,
205    /// The title of the citation source
206    pub title: Option<String>,
207    /// The start index of the citation in the response
208    pub start_index: Option<i32>,
209    /// The end index of the citation in the response
210    pub end_index: Option<i32>,
211    /// The license of the citation source
212    pub license: Option<String>,
213    /// The publication date of the citation source
214    pub publication_date: Option<String>,
215}
216
217/// A candidate response
218#[derive(Debug, Clone, Serialize, Deserialize)]
219#[serde(rename_all = "camelCase")]
220pub struct Candidate {
221    /// The content of the candidate
222    pub content: Content,
223    /// The safety ratings for the candidate
224    #[serde(skip_serializing_if = "Option::is_none")]
225    pub safety_ratings: Option<Vec<SafetyRating>>,
226    /// The citation metadata for the candidate
227    #[serde(skip_serializing_if = "Option::is_none")]
228    pub citation_metadata: Option<CitationMetadata>,
229    /// The finish reason for the candidate
230    #[serde(skip_serializing_if = "Option::is_none")]
231    pub finish_reason: Option<String>,
232    /// The index of the candidate
233    #[serde(skip_serializing_if = "Option::is_none")]
234    pub index: Option<i32>,
235}
236
237/// Metadata about token usage
238#[derive(Debug, Clone, Serialize, Deserialize)]
239#[serde(rename_all = "camelCase")]
240pub struct UsageMetadata {
241    /// The number of prompt tokens
242    pub prompt_token_count: i32,
243    /// The number of response tokens
244    #[serde(skip_serializing_if = "Option::is_none")]
245    pub candidates_token_count: Option<i32>,
246    /// The total number of tokens
247    pub total_token_count: i32,
248    /// The number of thinking tokens (Gemini 2.5 series only)
249    #[serde(skip_serializing_if = "Option::is_none")]
250    pub thoughts_token_count: Option<i32>,
251}
252
253/// Response from the Gemini API for content generation
254#[derive(Debug, Clone, Serialize, Deserialize)]
255#[serde(rename_all = "camelCase")]
256pub struct GenerationResponse {
257    /// The candidates generated
258    pub candidates: Vec<Candidate>,
259    /// The prompt feedback
260    #[serde(skip_serializing_if = "Option::is_none")]
261    pub prompt_feedback: Option<PromptFeedback>,
262    /// Usage metadata
263    #[serde(skip_serializing_if = "Option::is_none")]
264    pub usage_metadata: Option<UsageMetadata>,
265}
266
267/// Content of the embedding
268#[derive(Debug, Clone, Serialize, Deserialize)]
269pub struct ContentEmbedding {
270    /// The values generated
271    pub values: Vec<f32>, //Maybe Quantize this
272}
273
274/// Response from the Gemini API for content embedding
275#[derive(Debug, Clone, Serialize, Deserialize)]
276pub struct ContentEmbeddingResponse {
277    /// The embeddings generated
278    pub embedding: ContentEmbedding,
279}
280
281/// Response from the Gemini API for batch content embedding
282#[derive(Debug, Clone, Serialize, Deserialize)]
283pub struct BatchContentEmbeddingResponse {
284    /// The embeddings generated
285    pub embeddings: Vec<ContentEmbedding>,
286}
287
288/// Feedback about the prompt
289#[derive(Debug, Clone, Serialize, Deserialize)]
290#[serde(rename_all = "camelCase")]
291pub struct PromptFeedback {
292    /// The safety ratings for the prompt
293    pub safety_ratings: Vec<SafetyRating>,
294    /// The block reason if the prompt was blocked
295    #[serde(skip_serializing_if = "Option::is_none")]
296    pub block_reason: Option<String>,
297}
298
299impl GenerationResponse {
300    /// Get the text of the first candidate
301    pub fn text(&self) -> String {
302        self.candidates
303            .first()
304            .and_then(|c| {
305                c.content.parts.first().and_then(|p| match p {
306                    Part::Text { text, thought: _ } => Some(text.clone()),
307                    _ => None,
308                })
309            })
310            .unwrap_or_default()
311    }
312
313    /// Get function calls from the response
314    pub fn function_calls(&self) -> Vec<&super::tools::FunctionCall> {
315        self.candidates
316            .iter()
317            .flat_map(|c| {
318                c.content.parts.iter().filter_map(|p| match p {
319                    Part::FunctionCall { function_call } => Some(function_call),
320                    _ => None,
321                })
322            })
323            .collect()
324    }
325
326    /// Get thought summaries from the response
327    pub fn thoughts(&self) -> Vec<String> {
328        self.candidates
329            .iter()
330            .flat_map(|c| {
331                c.content.parts.iter().filter_map(|p| match p {
332                    Part::Text {
333                        text,
334                        thought: Some(true),
335                    } => Some(text.clone()),
336                    _ => None,
337                })
338            })
339            .collect()
340    }
341
342    /// Get all text parts (both regular text and thoughts)
343    pub fn all_text(&self) -> Vec<(String, bool)> {
344        self.candidates
345            .iter()
346            .flat_map(|c| {
347                c.content.parts.iter().filter_map(|p| match p {
348                    Part::Text { text, thought } => Some((text.clone(), thought.unwrap_or(false))),
349                    _ => None,
350                })
351            })
352            .collect()
353    }
354}
355
356/// Request to generate content
357#[derive(Debug, Clone, Serialize, Deserialize)]
358pub struct GenerateContentRequest {
359    /// The contents to generate content from
360    pub contents: Vec<Content>,
361    /// The generation config
362    #[serde(skip_serializing_if = "Option::is_none")]
363    pub generation_config: Option<GenerationConfig>,
364    /// The safety settings
365    #[serde(skip_serializing_if = "Option::is_none")]
366    pub safety_settings: Option<Vec<SafetySetting>>,
367    /// The tools that the model can use
368    #[serde(skip_serializing_if = "Option::is_none")]
369    pub tools: Option<Vec<super::tools::Tool>>,
370    /// The tool config
371    #[serde(skip_serializing_if = "Option::is_none")]
372    pub tool_config: Option<ToolConfig>,
373    /// The system instruction
374    #[serde(skip_serializing_if = "Option::is_none")]
375    pub system_instruction: Option<Content>,
376}
377
378/// Request to embed words
379#[derive(Debug, Clone, Serialize, Deserialize)]
380pub struct EmbedContentRequest {
381    /// The specified embedding model
382    pub model: String,
383    /// The chunks content to generate embeddings
384    pub content: Content,
385    /// The embedding task type (optional)
386    #[serde(skip_serializing_if = "Option::is_none")]
387    pub task_type: Option<TaskType>,
388    /// The title of the document (optional)
389    #[serde(skip_serializing_if = "Option::is_none")]
390    pub title: Option<String>,
391    /// The output_dimensionality (optional)
392    #[serde(skip_serializing_if = "Option::is_none")]
393    pub output_dimensionality: Option<i32>,
394}
395
396/// Request to batch embed requests
397#[derive(Debug, Clone, Serialize, Deserialize)]
398pub struct BatchEmbedContentsRequest {
399    /// The list of embed requests
400    pub requests: Vec<EmbedContentRequest>,
401}
402
403/// Configuration for thinking (Gemini 2.5 series only)
404#[derive(Debug, Clone, Serialize, Deserialize)]
405#[serde(rename_all = "camelCase")]
406pub struct ThinkingConfig {
407    /// The thinking budget (number of thinking tokens)
408    ///
409    /// - Set to 0 to disable thinking
410    /// - Set to -1 for dynamic thinking (model decides)
411    /// - Set to a positive number for a specific token budget
412    ///
413    /// Model-specific ranges:
414    /// - 2.5 Pro: 128 to 32768 (cannot disable thinking)
415    /// - 2.5 Flash: 0 to 24576
416    /// - 2.5 Flash Lite: 512 to 24576
417    #[serde(skip_serializing_if = "Option::is_none")]
418    pub thinking_budget: Option<i32>,
419
420    /// Whether to include thought summaries in the response
421    ///
422    /// When enabled, the response will include synthesized versions of the model's
423    /// raw thoughts, providing insights into the reasoning process.
424    #[serde(skip_serializing_if = "Option::is_none")]
425    pub include_thoughts: Option<bool>,
426}
427
428impl ThinkingConfig {
429    /// Create a new thinking config with default settings
430    pub fn new() -> Self {
431        Self {
432            thinking_budget: None,
433            include_thoughts: None,
434        }
435    }
436
437    /// Set the thinking budget
438    pub fn with_thinking_budget(mut self, budget: i32) -> Self {
439        self.thinking_budget = Some(budget);
440        self
441    }
442
443    /// Enable dynamic thinking (model decides the budget)
444    pub fn with_dynamic_thinking(mut self) -> Self {
445        self.thinking_budget = Some(-1);
446        self
447    }
448
449    /// Include thought summaries in the response
450    pub fn with_thoughts_included(mut self, include: bool) -> Self {
451        self.include_thoughts = Some(include);
452        self
453    }
454}
455
456impl Default for ThinkingConfig {
457    fn default() -> Self {
458        Self::new()
459    }
460}
461
462/// Configuration for generation
463#[derive(Debug, Clone, Serialize, Deserialize)]
464pub struct GenerationConfig {
465    /// The temperature for the model (0.0 to 1.0)
466    ///
467    /// Controls the randomness of the output. Higher values (e.g., 0.9) make output
468    /// more random, lower values (e.g., 0.1) make output more deterministic.
469    #[serde(skip_serializing_if = "Option::is_none")]
470    pub temperature: Option<f32>,
471
472    /// The top-p value for the model (0.0 to 1.0)
473    ///
474    /// For each token generation step, the model considers the top_p percentage of
475    /// probability mass for potential token choices. Lower values are more selective,
476    /// higher values allow more variety.
477    #[serde(skip_serializing_if = "Option::is_none")]
478    pub top_p: Option<f32>,
479
480    /// The top-k value for the model
481    ///
482    /// For each token generation step, the model considers the top_k most likely tokens.
483    /// Lower values are more selective, higher values allow more variety.
484    #[serde(skip_serializing_if = "Option::is_none")]
485    pub top_k: Option<i32>,
486
487    /// The maximum number of tokens to generate
488    ///
489    /// Limits the length of the generated content. One token is roughly 4 characters.
490    #[serde(skip_serializing_if = "Option::is_none")]
491    pub max_output_tokens: Option<i32>,
492
493    /// The candidate count
494    ///
495    /// Number of alternative responses to generate.
496    #[serde(skip_serializing_if = "Option::is_none")]
497    pub candidate_count: Option<i32>,
498
499    /// Whether to stop on specific sequences
500    ///
501    /// The model will stop generating content when it encounters any of these sequences.
502    #[serde(skip_serializing_if = "Option::is_none")]
503    pub stop_sequences: Option<Vec<String>>,
504
505    /// The response mime type
506    ///
507    /// Specifies the format of the model's response.
508    #[serde(skip_serializing_if = "Option::is_none")]
509    pub response_mime_type: Option<String>,
510
511    /// The response schema
512    ///
513    /// Specifies the JSON schema for structured responses.
514    #[serde(skip_serializing_if = "Option::is_none")]
515    pub response_schema: Option<serde_json::Value>,
516
517    /// The thinking configuration
518    ///
519    /// Configuration for the model's thinking process (Gemini 2.5 series only).
520    #[serde(skip_serializing_if = "Option::is_none")]
521    pub thinking_config: Option<ThinkingConfig>,
522}
523
524impl Default for GenerationConfig {
525    fn default() -> Self {
526        Self {
527            temperature: Some(0.7),
528            top_p: Some(0.95),
529            top_k: Some(40),
530            max_output_tokens: Some(1024),
531            candidate_count: Some(1),
532            stop_sequences: None,
533            response_mime_type: None,
534            response_schema: None,
535            thinking_config: None,
536        }
537    }
538}
539
540/// Configuration for tools
541#[derive(Debug, Clone, Serialize, Deserialize)]
542pub struct ToolConfig {
543    /// The function calling config
544    #[serde(skip_serializing_if = "Option::is_none")]
545    pub function_calling_config: Option<FunctionCallingConfig>,
546}
547
548/// Configuration for function calling
549#[derive(Debug, Clone, Serialize, Deserialize)]
550pub struct FunctionCallingConfig {
551    /// The mode for function calling
552    pub mode: FunctionCallingMode,
553}
554
555/// Mode for function calling
556#[derive(Debug, Clone, Serialize, Deserialize)]
557#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
558pub enum FunctionCallingMode {
559    /// The model may use function calling
560    Auto,
561    /// The model must use function calling
562    Any,
563    /// The model must not use function calling
564    None,
565}
566
567/// Setting for safety
568#[derive(Debug, Clone, Serialize, Deserialize)]
569pub struct SafetySetting {
570    /// The category of content to filter
571    pub category: HarmCategory,
572    /// The threshold for filtering
573    pub threshold: HarmBlockThreshold,
574}
575
576/// Category of harmful content
577#[derive(Debug, Clone, Serialize, Deserialize)]
578#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
579pub enum HarmCategory {
580    /// Dangerous content
581    Dangerous,
582    /// Harassment content
583    Harassment,
584    /// Hate speech
585    HateSpeech,
586    /// Sexually explicit content
587    SexuallyExplicit,
588}
589
590/// Threshold for blocking harmful content
591#[allow(clippy::enum_variant_names)]
592#[derive(Debug, Clone, Serialize, Deserialize)]
593#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
594pub enum HarmBlockThreshold {
595    /// Block content with low probability of harm
596    BlockLowAndAbove,
597    /// Block content with medium probability of harm
598    BlockMediumAndAbove,
599    /// Block content with high probability of harm
600    BlockHighAndAbove,
601    /// Block content with maximum probability of harm
602    BlockOnlyHigh,
603    /// Never block content
604    BlockNone,
605}
606
607/// Embedding Task types
608#[derive(Debug, Clone, Serialize, Deserialize)]
609#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
610pub enum TaskType {
611    ///Used to generate embeddings that are optimized to assess text similarity
612    SemanticSimilarity,
613    ///Used to generate embeddings that are optimized to classify texts according to preset labels
614    Classification,
615    ///Used to generate embeddings that are optimized to cluster texts based on their similarities
616    Clustering,
617
618    ///Used to generate embeddings that are optimized for document search or information retrieval.
619    RetrievalDocument,
620    RetrievalQuery,
621    QuestionAnswering,
622    FactVerification,
623
624    /// Used to retrieve a code block based on a natural language query, such as sort an array or reverse a linked list.
625    /// Embeddings of the code blocks are computed using RETRIEVAL_DOCUMENT.
626    CodeRetrievalQuery,
627}