gemini_rust/generation/
builder.rs

1use futures::TryStream;
2use std::sync::Arc;
3use tracing::instrument;
4
5use crate::{
6    cache::CachedContentHandle,
7    client::{Error as ClientError, GeminiClient},
8    generation::{GenerateContentRequest, SpeakerVoiceConfig, SpeechConfig, ThinkingConfig},
9    tools::{FunctionCallingConfig, ToolConfig},
10    Content, FunctionCallingMode, FunctionDeclaration, GenerationConfig, GenerationResponse,
11    Message, Role, Tool,
12};
13
14/// Builder for content generation requests
15pub struct ContentBuilder {
16    client: Arc<GeminiClient>,
17    pub contents: Vec<Content>,
18    generation_config: Option<GenerationConfig>,
19    tools: Option<Vec<Tool>>,
20    tool_config: Option<ToolConfig>,
21    system_instruction: Option<Content>,
22    cached_content: Option<String>,
23}
24
25impl ContentBuilder {
26    /// Creates a new `ContentBuilder`.
27    pub(crate) fn new(client: Arc<GeminiClient>) -> Self {
28        Self {
29            client,
30            contents: Vec::new(),
31            generation_config: None,
32            tools: None,
33            tool_config: None,
34            system_instruction: None,
35            cached_content: None,
36        }
37    }
38
39    /// Sets the system prompt for the request.
40    ///
41    /// This is an alias for [`with_system_instruction()`](Self::with_system_instruction).
42    pub fn with_system_prompt(self, text: impl Into<String>) -> Self {
43        self.with_system_instruction(text)
44    }
45
46    /// Sets the system instruction for the request.
47    ///
48    /// System instructions are used to provide high-level guidance to the model, such as
49    /// setting a persona, providing context, or defining the desired output format.
50    pub fn with_system_instruction(mut self, text: impl Into<String>) -> Self {
51        let content = Content::text(text);
52        self.system_instruction = Some(content);
53        self
54    }
55
56    /// Adds a user message to the conversation history.
57    pub fn with_user_message(mut self, text: impl Into<String>) -> Self {
58        let message = Message::user(text);
59        self.contents.push(message.content);
60        self
61    }
62
63    /// Adds a model message to the conversation history.
64    pub fn with_model_message(mut self, text: impl Into<String>) -> Self {
65        let message = Message::model(text);
66        self.contents.push(message.content);
67        self
68    }
69
70    /// Adds inline data (e.g., an image) to the request.
71    ///
72    /// The data should be base64-encoded.
73    pub fn with_inline_data(
74        mut self,
75        data: impl Into<String>,
76        mime_type: impl Into<String>,
77    ) -> Self {
78        let content = Content::inline_data(mime_type, data).with_role(Role::User);
79        self.contents.push(content);
80        self
81    }
82
83    /// Adds a function response to the request using a `Serialize` response.
84    ///
85    /// This is used to provide the model with the result of a function call it has requested.
86    pub fn with_function_response<Response>(
87        mut self,
88        name: impl Into<String>,
89        response: Response,
90    ) -> std::result::Result<Self, serde_json::Error>
91    where
92        Response: serde::Serialize,
93    {
94        let content = Content::function_response_json(name, serde_json::to_value(response)?)
95            .with_role(Role::User);
96        self.contents.push(content);
97        Ok(self)
98    }
99
100    /// Adds a function response to the request using a JSON string.
101    ///
102    /// This is a convenience method that parses the string into a `serde_json::Value`.
103    pub fn with_function_response_str(
104        mut self,
105        name: impl Into<String>,
106        response: impl Into<String>,
107    ) -> std::result::Result<Self, serde_json::Error> {
108        let response_str = response.into();
109        let json = serde_json::from_str(&response_str)?;
110        let content = Content::function_response_json(name, json).with_role(Role::User);
111        self.contents.push(content);
112        Ok(self)
113    }
114
115    /// Adds a `Message` to the conversation history.
116    pub fn with_message(mut self, message: Message) -> Self {
117        let content = message.content.clone();
118        let role = content.role.clone().unwrap_or(message.role);
119        self.contents.push(content.with_role(role));
120        self
121    }
122
123    /// Uses cached content for this request.
124    ///
125    /// This allows reusing previously cached system instructions and conversation history,
126    /// which can reduce latency and cost.
127    pub fn with_cached_content(mut self, cached_content: &CachedContentHandle) -> Self {
128        self.cached_content = Some(cached_content.name().to_string());
129        self
130    }
131
132    /// Adds multiple messages to the conversation history.
133    pub fn with_messages(mut self, messages: impl IntoIterator<Item = Message>) -> Self {
134        for message in messages {
135            self = self.with_message(message);
136        }
137        self
138    }
139
140    /// Sets the generation configuration for the request.
141    pub fn with_generation_config(mut self, config: GenerationConfig) -> Self {
142        self.generation_config = Some(config);
143        self
144    }
145
146    /// Sets the temperature for the request.
147    ///
148    /// Temperature controls the randomness of the output. Higher values (e.g., 1.0) produce
149    /// more creative results, while lower values (e.g., 0.2) produce more deterministic results.
150    pub fn with_temperature(mut self, temperature: f32) -> Self {
151        self.generation_config
152            .get_or_insert_with(Default::default)
153            .temperature = Some(temperature);
154        self
155    }
156
157    /// Sets the top-p value for the request.
158    ///
159    /// Top-p is a sampling method that selects the next token from a cumulative probability
160    /// distribution. It can be used to control the diversity of the output.
161    pub fn with_top_p(mut self, top_p: f32) -> Self {
162        self.generation_config
163            .get_or_insert_with(Default::default)
164            .top_p = Some(top_p);
165        self
166    }
167
168    /// Sets the top-k value for the request.
169    ///
170    /// Top-k is a sampling method that selects the next token from the `k` most likely candidates.
171    pub fn with_top_k(mut self, top_k: i32) -> Self {
172        self.generation_config
173            .get_or_insert_with(Default::default)
174            .top_k = Some(top_k);
175        self
176    }
177
178    /// Sets the maximum number of output tokens for the request.
179    pub fn with_max_output_tokens(mut self, max_output_tokens: i32) -> Self {
180        self.generation_config
181            .get_or_insert_with(Default::default)
182            .max_output_tokens = Some(max_output_tokens);
183        self
184    }
185
186    /// Sets the number of candidate responses to generate.
187    pub fn with_candidate_count(mut self, candidate_count: i32) -> Self {
188        self.generation_config
189            .get_or_insert_with(Default::default)
190            .candidate_count = Some(candidate_count);
191        self
192    }
193
194    /// Sets the stop sequences for the request.
195    ///
196    /// The model will stop generating text when it encounters one of these sequences.
197    pub fn with_stop_sequences(mut self, stop_sequences: Vec<String>) -> Self {
198        self.generation_config
199            .get_or_insert_with(Default::default)
200            .stop_sequences = Some(stop_sequences);
201        self
202    }
203
204    /// Sets the response MIME type for the request.
205    ///
206    /// This can be used to request structured output, such as JSON.
207    pub fn with_response_mime_type(mut self, mime_type: impl Into<String>) -> Self {
208        self.generation_config
209            .get_or_insert_with(Default::default)
210            .response_mime_type = Some(mime_type.into());
211        self
212    }
213
214    /// Sets the response schema for structured output.
215    ///
216    /// When used with a JSON MIME type, this schema will be used to validate the model's
217    /// output.
218    pub fn with_response_schema(mut self, schema: serde_json::Value) -> Self {
219        self.generation_config
220            .get_or_insert_with(Default::default)
221            .response_schema = Some(schema);
222        self
223    }
224
225    /// Adds a tool to the request.
226    ///
227    /// Tools allow the model to interact with external systems, such as APIs or databases.
228    pub fn with_tool(mut self, tool: Tool) -> Self {
229        self.tools.get_or_insert_with(Vec::new).push(tool);
230        self
231    }
232
233    /// Adds a function declaration as a tool.
234    ///
235    /// This is a convenience method for creating a `Tool` from a `FunctionDeclaration`.
236    pub fn with_function(mut self, function: FunctionDeclaration) -> Self {
237        let tool = Tool::new(function);
238        self = self.with_tool(tool);
239        self
240    }
241
242    /// Sets the function calling mode for the request.
243    pub fn with_function_calling_mode(mut self, mode: FunctionCallingMode) -> Self {
244        self.tool_config
245            .get_or_insert_with(Default::default)
246            .function_calling_config = Some(FunctionCallingConfig { mode });
247        self
248    }
249
250    /// Sets the thinking configuration for the request (Gemini 2.5 series only).
251    pub fn with_thinking_config(mut self, thinking_config: ThinkingConfig) -> Self {
252        self.generation_config
253            .get_or_insert_with(Default::default)
254            .thinking_config = Some(thinking_config);
255        self
256    }
257
258    /// Sets the thinking budget for the request (Gemini 2.5 series only).
259    ///
260    /// A budget of -1 enables dynamic thinking.
261    pub fn with_thinking_budget(mut self, budget: i32) -> Self {
262        self.generation_config
263            .get_or_insert_with(Default::default)
264            .thinking_config
265            .get_or_insert_with(Default::default)
266            .thinking_budget = Some(budget);
267        self
268    }
269
270    /// Enables dynamic thinking, which allows the model to decide its own thinking budget
271    /// (Gemini 2.5 series only).
272    ///
273    /// Note: This only enables the *capability*. To receive thoughts in the response,
274    /// you must also call `[.with_thoughts_included(true)](Self::with_thoughts_included)`.
275    pub fn with_dynamic_thinking(self) -> Self {
276        self.with_thinking_budget(-1)
277    }
278
279    /// Includes thought summaries in the response (Gemini 2.5 series only).
280    ///
281    /// This requires `with_dynamic_thinking()` or `with_thinking_budget()` to be enabled.
282    pub fn with_thoughts_included(mut self, include: bool) -> Self {
283        self.generation_config
284            .get_or_insert_with(Default::default)
285            .thinking_config
286            .get_or_insert_with(Default::default)
287            .include_thoughts = Some(include);
288        self
289    }
290
291    /// Enables audio output (text-to-speech).
292    pub fn with_audio_output(mut self) -> Self {
293        self.generation_config
294            .get_or_insert_with(Default::default)
295            .response_modalities = Some(vec!["AUDIO".to_string()]);
296        self
297    }
298
299    /// Sets the speech configuration for text-to-speech generation.
300    pub fn with_speech_config(mut self, speech_config: SpeechConfig) -> Self {
301        self.generation_config
302            .get_or_insert_with(Default::default)
303            .speech_config = Some(speech_config);
304        self
305    }
306
307    /// Sets a single voice for text-to-speech generation.
308    pub fn with_voice(self, voice_name: impl Into<String>) -> Self {
309        let speech_config = SpeechConfig::single_voice(voice_name);
310        self.with_speech_config(speech_config).with_audio_output()
311    }
312
313    /// Sets multi-speaker configuration for text-to-speech generation.
314    pub fn with_multi_speaker_config(self, speakers: Vec<SpeakerVoiceConfig>) -> Self {
315        let speech_config = SpeechConfig::multi_speaker(speakers);
316        self.with_speech_config(speech_config).with_audio_output()
317    }
318
319    /// Builds the `GenerateContentRequest`.
320    pub fn build(self) -> GenerateContentRequest {
321        GenerateContentRequest {
322            contents: self.contents,
323            generation_config: self.generation_config,
324            safety_settings: None,
325            tools: self.tools,
326            tool_config: self.tool_config,
327            system_instruction: self.system_instruction,
328            cached_content: self.cached_content,
329        }
330    }
331
332    /// Executes the content generation request.
333    #[instrument(skip_all, fields(
334        messages.parts.count = self.contents.len(),
335        tools.present = self.tools.is_some(),
336        system.instruction.present = self.system_instruction.is_some(),
337        cached.content.present = self.cached_content.is_some(),
338    ))]
339    pub async fn execute(self) -> Result<GenerationResponse, ClientError> {
340        let client = self.client.clone();
341        let request = self.build();
342        client.generate_content_raw(request).await
343    }
344
345    /// Executes the content generation request as a stream.
346    #[instrument(skip_all, fields(
347        messages.parts.count = self.contents.len(),
348        tools.present = self.tools.is_some(),
349        system.instruction.present = self.system_instruction.is_some(),
350        cached.content.present = self.cached_content.is_some(),
351    ))]
352    pub async fn execute_stream(
353        self,
354    ) -> Result<impl TryStream<Ok = GenerationResponse, Error = ClientError> + Send, ClientError>
355    {
356        let client = self.client.clone();
357        let request = self.build();
358        client.generate_content_stream(request).await
359    }
360}
gemini_rust/generation/builder.rs

gemini_rust/generation/
builder.rs
