gemini_rust/generation/
builder.rs

1use std::sync::Arc;
2use tracing::instrument;
3
4use crate::{
5    cache::CachedContentHandle,
6    client::{Error as ClientError, GeminiClient, GenerationStream},
7    files::Error as FilesError,
8    generation::{
9        GenerateContentRequest, ImageConfig, MediaResolutionLevel, SpeakerVoiceConfig,
10        SpeechConfig, ThinkingConfig, ThinkingLevel,
11    },
12    tools::{FunctionCallingConfig, ToolConfig},
13    Content, FileHandle, FunctionCallingMode, FunctionDeclaration, GenerationConfig,
14    GenerationResponse, Message, Role, SafetySetting, Tool,
15};
16
17/// Builder for content generation requests
18#[derive(Clone)]
19pub struct ContentBuilder {
20    client: Arc<GeminiClient>,
21    pub contents: Vec<Content>,
22    generation_config: Option<GenerationConfig>,
23    safety_settings: Option<Vec<SafetySetting>>,
24    tools: Option<Vec<Tool>>,
25    tool_config: Option<ToolConfig>,
26    system_instruction: Option<Content>,
27    cached_content: Option<String>,
28}
29
30impl ContentBuilder {
31    /// Creates a new `ContentBuilder`.
32    pub(crate) fn new(client: Arc<GeminiClient>) -> Self {
33        Self {
34            client,
35            contents: Vec::new(),
36            generation_config: None,
37            safety_settings: None,
38            tools: None,
39            tool_config: None,
40            system_instruction: None,
41            cached_content: None,
42        }
43    }
44
45    /// Sets the safety settings for the request.
46    pub fn with_safety_settings(mut self, safety_settings: Vec<SafetySetting>) -> Self {
47        self.safety_settings = Some(safety_settings);
48        self
49    }
50
51    /// Sets the system prompt for the request.
52    ///
53    /// This is an alias for [`with_system_instruction()`](Self::with_system_instruction).
54    pub fn with_system_prompt(self, text: impl Into<String>) -> Self {
55        self.with_system_instruction(text)
56    }
57
58    /// Sets the system instruction for the request.
59    ///
60    /// System instructions are used to provide high-level guidance to the model, such as
61    /// setting a persona, providing context, or defining the desired output format.
62    pub fn with_system_instruction(mut self, text: impl Into<String>) -> Self {
63        let content = Content::text(text);
64        self.system_instruction = Some(content);
65        self
66    }
67
68    /// Adds a user message to the conversation history.
69    pub fn with_user_message(mut self, text: impl Into<String>) -> Self {
70        let message = Message::user(text);
71        self.contents.push(message.content);
72        self
73    }
74
75    /// Adds a user message, together with coordinates for a previously uploaded file.
76    ///
77    /// Uploading a file and using it avoids encoding large files and sending them, in particular
78    /// when this would need to happen more than once with a file.
79    ///
80    /// # Errors
81    ///
82    /// Returns an error if the file metadata is incomplete (missing MIME type or URI).
83    pub fn with_user_message_and_file(
84        mut self,
85        text: impl Into<String>,
86        file_handle: &FileHandle,
87    ) -> Result<Self, FilesError> {
88        let content = Content::text_with_file(text, file_handle)?.with_role(Role::User);
89        self.contents.push(content);
90        Ok(self)
91    }
92
93    /// Adds a model message to the conversation history.
94    pub fn with_model_message(mut self, text: impl Into<String>) -> Self {
95        let message = Message::model(text);
96        self.contents.push(message.content);
97        self
98    }
99
100    /// Adds inline data (e.g., an image) to the request.
101    ///
102    /// The data should be base64-encoded.
103    pub fn with_inline_data(
104        mut self,
105        data: impl Into<String>,
106        mime_type: impl Into<String>,
107    ) -> Self {
108        let content = Content::inline_data(mime_type, data).with_role(Role::User);
109        self.contents.push(content);
110        self
111    }
112
113    /// Adds inline data with explicit media resolution control.
114    ///
115    /// This allows fine-grained control over the resolution used for processing
116    /// the inline data, which affects both quality and token consumption.
117    /// This method is useful for optimizing token usage.
118    /// The data should be base64-encoded.
119    pub fn with_inline_data_and_resolution(
120        mut self,
121        data: impl Into<String>,
122        mime_type: impl Into<String>,
123        resolution: MediaResolutionLevel,
124    ) -> Self {
125        let content =
126            Content::inline_data_with_resolution(mime_type, data, resolution).with_role(Role::User);
127        self.contents.push(content);
128        self
129    }
130
131    /// Adds a function response to the request using a `Serialize` response.
132    ///
133    /// This is used to provide the model with the result of a function call it has requested.
134    pub fn with_function_response<Response>(
135        mut self,
136        name: impl Into<String>,
137        response: Response,
138    ) -> std::result::Result<Self, serde_json::Error>
139    where
140        Response: serde::Serialize,
141    {
142        let content = Content::function_response_json(name, serde_json::to_value(response)?)
143            .with_role(Role::User);
144        self.contents.push(content);
145        Ok(self)
146    }
147
148    /// Adds a function response to the request using a JSON string.
149    ///
150    /// This is a convenience method that parses the string into a `serde_json::Value`.
151    pub fn with_function_response_str(
152        mut self,
153        name: impl Into<String>,
154        response: impl Into<String>,
155    ) -> std::result::Result<Self, serde_json::Error> {
156        let response_str = response.into();
157        let json = serde_json::from_str(&response_str)?;
158        let content = Content::function_response_json(name, json).with_role(Role::User);
159        self.contents.push(content);
160        Ok(self)
161    }
162
163    /// Adds a `Message` to the conversation history.
164    pub fn with_message(mut self, message: Message) -> Self {
165        let content = message.content.clone();
166        let role = content.role.clone().unwrap_or(message.role);
167        self.contents.push(content.with_role(role));
168        self
169    }
170
171    /// Uses cached content for this request.
172    ///
173    /// This allows reusing previously cached system instructions and conversation history,
174    /// which can reduce latency and cost.
175    pub fn with_cached_content(mut self, cached_content: &CachedContentHandle) -> Self {
176        self.cached_content = Some(cached_content.name().to_string());
177        self
178    }
179
180    /// Adds multiple messages to the conversation history.
181    pub fn with_messages(mut self, messages: impl IntoIterator<Item = Message>) -> Self {
182        for message in messages {
183            self = self.with_message(message);
184        }
185        self
186    }
187
188    /// Sets the generation configuration for the request.
189    pub fn with_generation_config(mut self, config: GenerationConfig) -> Self {
190        self.generation_config = Some(config);
191        self
192    }
193
194    /// Sets the temperature for the request.
195    ///
196    /// Temperature controls the randomness of the output. Higher values (e.g., 1.0) produce
197    /// more creative results, while lower values (e.g., 0.2) produce more deterministic results.
198    pub fn with_temperature(mut self, temperature: f32) -> Self {
199        self.generation_config
200            .get_or_insert_with(Default::default)
201            .temperature = Some(temperature);
202        self
203    }
204
205    /// Sets the top-p value for the request.
206    ///
207    /// Top-p is a sampling method that selects the next token from a cumulative probability
208    /// distribution. It can be used to control the diversity of the output.
209    pub fn with_top_p(mut self, top_p: f32) -> Self {
210        self.generation_config
211            .get_or_insert_with(Default::default)
212            .top_p = Some(top_p);
213        self
214    }
215
216    /// Sets the top-k value for the request.
217    ///
218    /// Top-k is a sampling method that selects the next token from the `k` most likely candidates.
219    pub fn with_top_k(mut self, top_k: i32) -> Self {
220        self.generation_config
221            .get_or_insert_with(Default::default)
222            .top_k = Some(top_k);
223        self
224    }
225
226    /// Sets the seed for the request.
227    pub fn with_seed(mut self, seed: i32) -> Self {
228        self.generation_config
229            .get_or_insert_with(Default::default)
230            .seed = Some(seed);
231        self
232    }
233
234    /// Sets the maximum number of output tokens for the request.
235    pub fn with_max_output_tokens(mut self, max_output_tokens: i32) -> Self {
236        self.generation_config
237            .get_or_insert_with(Default::default)
238            .max_output_tokens = Some(max_output_tokens);
239        self
240    }
241
242    /// Sets the number of candidate responses to generate.
243    pub fn with_candidate_count(mut self, candidate_count: i32) -> Self {
244        self.generation_config
245            .get_or_insert_with(Default::default)
246            .candidate_count = Some(candidate_count);
247        self
248    }
249
250    /// Sets the stop sequences for the request.
251    ///
252    /// The model will stop generating text when it encounters one of these sequences.
253    pub fn with_stop_sequences(mut self, stop_sequences: Vec<String>) -> Self {
254        self.generation_config
255            .get_or_insert_with(Default::default)
256            .stop_sequences = Some(stop_sequences);
257        self
258    }
259
260    /// Sets the response MIME type for the request.
261    ///
262    /// This can be used to request structured output, such as JSON.
263    pub fn with_response_mime_type(mut self, mime_type: impl Into<String>) -> Self {
264        self.generation_config
265            .get_or_insert_with(Default::default)
266            .response_mime_type = Some(mime_type.into());
267        self
268    }
269
270    /// Sets the response schema for structured output.
271    ///
272    /// When used with a JSON MIME type, this schema will be used to validate the model's
273    /// output.
274    pub fn with_response_schema(mut self, schema: serde_json::Value) -> Self {
275        self.generation_config
276            .get_or_insert_with(Default::default)
277            .response_schema = Some(schema);
278        self
279    }
280
281    /// Adds a tool to the request.
282    ///
283    /// Tools allow the model to interact with external systems, such as APIs or databases.
284    pub fn with_tool(mut self, tool: Tool) -> Self {
285        self.tools.get_or_insert_with(Vec::new).push(tool);
286        self
287    }
288
289    /// Adds a function declaration as a tool.
290    ///
291    /// This is a convenience method for creating a `Tool` from a `FunctionDeclaration`.
292    pub fn with_function(mut self, function: FunctionDeclaration) -> Self {
293        let tool = Tool::new(function);
294        self = self.with_tool(tool);
295        self
296    }
297
298    /// Sets the function calling mode for the request.
299    pub fn with_function_calling_mode(mut self, mode: FunctionCallingMode) -> Self {
300        self.tool_config
301            .get_or_insert_with(Default::default)
302            .function_calling_config = Some(FunctionCallingConfig { mode });
303        self
304    }
305
306    /// Sets the tool configuration for the request.
307    pub fn with_tool_config(mut self, tool_config: ToolConfig) -> Self {
308        self.tool_config = Some(tool_config);
309        self
310    }
311
312    /// Sets the thinking configuration for the request (Gemini 2.5 series only).
313    pub fn with_thinking_config(mut self, thinking_config: ThinkingConfig) -> Self {
314        self.generation_config
315            .get_or_insert_with(Default::default)
316            .thinking_config = Some(thinking_config);
317        self
318    }
319
320    /// Sets the thinking budget for the request (Gemini 2.5 series only).
321    ///
322    /// A budget of -1 enables dynamic thinking.
323    /// This is mutually exclusive with `thinking_level` (Gemini 3 models).
324    pub fn with_thinking_budget(mut self, budget: i32) -> Self {
325        let config = self
326            .generation_config
327            .get_or_insert_with(Default::default)
328            .thinking_config
329            .get_or_insert_with(Default::default);
330        config.thinking_budget = Some(budget);
331        config.thinking_level = None;
332        self
333    }
334
335    /// Enables dynamic thinking, which allows the model to decide its own thinking budget
336    /// (Gemini 2.5 series only).
337    ///
338    /// Note: This only enables the *capability*. To receive thoughts in the response,
339    /// you must also call `[.with_thoughts_included(true)](Self::with_thoughts_included)`.
340    pub fn with_dynamic_thinking(self) -> Self {
341        self.with_thinking_budget(-1)
342    }
343
344    /// Includes thought summaries in the response (Gemini 2.5 series only).
345    ///
346    /// This requires `with_dynamic_thinking()` or `with_thinking_budget()` to be enabled.
347    pub fn with_thoughts_included(mut self, include: bool) -> Self {
348        self.generation_config
349            .get_or_insert_with(Default::default)
350            .thinking_config
351            .get_or_insert_with(Default::default)
352            .include_thoughts = Some(include);
353        self
354    }
355
356    /// Sets the thinking level for Gemini 3 Pro.
357    ///
358    /// This controls the depth of reasoning the model applies. Use `Low` for simpler
359    /// queries requiring faster responses, or `High` for complex problems requiring
360    /// deeper analysis.
361    ///
362    /// Note: This is mutually exclusive with `thinking_budget` (used by Gemini 2.5 models).
363    /// Setting this will be ignored by Gemini 2.5 models.
364    pub fn with_thinking_level(mut self, level: ThinkingLevel) -> Self {
365        let config = self
366            .generation_config
367            .get_or_insert_with(Default::default)
368            .thinking_config
369            .get_or_insert_with(Default::default);
370        config.thinking_level = Some(level);
371        config.thinking_budget = None;
372        self
373    }
374
375    /// Sets the global media resolution level.
376    ///
377    /// This controls the token usage for all images and PDFs in the request.
378    /// Individual parts can override this setting using `with_inline_data_and_resolution()`.
379    /// Higher resolutions provide better quality but consume more tokens.
380    pub fn with_media_resolution(mut self, level: MediaResolutionLevel) -> Self {
381        self.generation_config
382            .get_or_insert_with(Default::default)
383            .media_resolution = Some(level);
384        self
385    }
386
387    /// Adds the code execution tool to the request.
388    ///
389    /// This allows the model to generate and execute Python code as part of the
390    /// generation process. Useful for mathematical calculations, data analysis,
391    /// and other computational tasks. Currently supports Python only.
392    pub fn with_code_execution(self) -> Self {
393        self.with_tool(Tool::code_execution())
394    }
395
396    /// Enables audio output (text-to-speech).
397    pub fn with_audio_output(mut self) -> Self {
398        self.generation_config
399            .get_or_insert_with(Default::default)
400            .response_modalities = Some(vec!["AUDIO".to_string()]);
401        self
402    }
403
404    /// Sets the image generation configuration.
405    pub fn with_image_config(mut self, image_config: ImageConfig) -> Self {
406        self.generation_config
407            .get_or_insert_with(Default::default)
408            .image_config = Some(image_config);
409        self
410    }
411
412    /// Sets the speech configuration for text-to-speech generation.
413    pub fn with_speech_config(mut self, speech_config: SpeechConfig) -> Self {
414        self.generation_config
415            .get_or_insert_with(Default::default)
416            .speech_config = Some(speech_config);
417        self
418    }
419
420    /// Sets a single voice for text-to-speech generation.
421    pub fn with_voice(self, voice_name: impl Into<String>) -> Self {
422        let speech_config = SpeechConfig::single_voice(voice_name);
423        self.with_speech_config(speech_config).with_audio_output()
424    }
425
426    /// Sets multi-speaker configuration for text-to-speech generation.
427    pub fn with_multi_speaker_config(self, speakers: Vec<SpeakerVoiceConfig>) -> Self {
428        let speech_config = SpeechConfig::multi_speaker(speakers);
429        self.with_speech_config(speech_config).with_audio_output()
430    }
431
432    /// Builds the `GenerateContentRequest`.
433    pub fn build(self) -> GenerateContentRequest {
434        GenerateContentRequest {
435            contents: self.contents,
436            generation_config: self.generation_config,
437            safety_settings: self.safety_settings,
438            tools: self.tools,
439            tool_config: self.tool_config,
440            system_instruction: self.system_instruction,
441            cached_content: self.cached_content,
442        }
443    }
444
445    /// Executes the content generation request.
446    #[instrument(skip_all, fields(
447        messages.parts.count = self.contents.len(),
448        tools.present = self.tools.is_some(),
449        system.instruction.present = self.system_instruction.is_some(),
450        cached.content.present = self.cached_content.is_some(),
451    ))]
452    pub async fn execute(self) -> Result<GenerationResponse, ClientError> {
453        let client = self.client.clone();
454        let request = self.build();
455        client.generate_content_raw(request).await
456    }
457
458    /// Executes the content generation request as a stream.
459    #[instrument(skip_all, fields(
460        messages.parts.count = self.contents.len(),
461        tools.present = self.tools.is_some(),
462        system.instruction.present = self.system_instruction.is_some(),
463        cached.content.present = self.cached_content.is_some(),
464    ))]
465    pub async fn execute_stream(self) -> Result<GenerationStream, ClientError> {
466        let client = self.client.clone();
467        let request = self.build();
468        client.generate_content_stream(request).await
469    }
470
471    /// Counts the number of tokens in the content generation request.
472    #[instrument(skip_all, fields(
473        messages.parts.count = self.contents.len(),
474    ))]
475    pub async fn count_tokens(self) -> Result<super::model::CountTokensResponse, ClientError> {
476        let client = self.client.clone();
477        let request = self.build();
478        client.count_tokens(request).await
479    }
480}