gemini_rust/
content_builder.rs

1use std::{pin::Pin, sync::Arc};
2
3use futures::Stream;
4
5use crate::{
6    client::GeminiClient,
7    models::{
8        FunctionCallingConfig, GenerateContentRequest, SpeakerVoiceConfig, SpeechConfig,
9        ThinkingConfig, ToolConfig,
10    },
11    Content, FunctionCallingMode, FunctionDeclaration, GenerationConfig, GenerationResponse,
12    Message, Result, Role, Tool,
13};
14
15/// Builder for content generation requests
16pub struct ContentBuilder {
17    client: Arc<GeminiClient>,
18    pub contents: Vec<Content>,
19    generation_config: Option<GenerationConfig>,
20    tools: Option<Vec<Tool>>,
21    tool_config: Option<ToolConfig>,
22    system_instruction: Option<Content>,
23}
24
25impl ContentBuilder {
26    /// Create a new content builder
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        }
36    }
37
38    /// Add a system prompt to the request
39    pub fn with_system_prompt(self, text: impl Into<String>) -> Self {
40        // Create a Content with text parts specifically for system_instruction field
41        self.with_system_instruction(text)
42    }
43
44    /// Set the system instruction directly (matching the API format in the curl example)
45    pub fn with_system_instruction(mut self, text: impl Into<String>) -> Self {
46        // Create a Content with text parts specifically for system_instruction field
47        let content = Content::text(text);
48        self.system_instruction = Some(content);
49        self
50    }
51
52    /// Add a user message to the request
53    pub fn with_user_message(mut self, text: impl Into<String>) -> Self {
54        let message = Message::user(text);
55        let content = message.content;
56        self.contents.push(content);
57        self
58    }
59
60    /// Add a model message to the request
61    pub fn with_model_message(mut self, text: impl Into<String>) -> Self {
62        let message = Message::model(text);
63        let content = message.content;
64        self.contents.push(content);
65        self
66    }
67
68    /// Add a inline data (blob data) to the request
69    pub fn with_inline_data(
70        mut self,
71        data: impl Into<String>,
72        mime_type: impl Into<String>,
73    ) -> Self {
74        let content = Content::inline_data(mime_type, data).with_role(Role::User);
75        self.contents.push(content);
76        self
77    }
78
79    /// Add a function response to the request using a JSON value
80    pub fn with_function_response(
81        mut self,
82        name: impl Into<String>,
83        response: serde_json::Value,
84    ) -> Self {
85        let content = Content::function_response_json(name, response).with_role(Role::User);
86        self.contents.push(content);
87        self
88    }
89
90    /// Add a function response to the request using a JSON string
91    pub fn with_function_response_str(
92        mut self,
93        name: impl Into<String>,
94        response: impl Into<String>,
95    ) -> std::result::Result<Self, serde_json::Error> {
96        let response_str = response.into();
97        let json = serde_json::from_str(&response_str)?;
98        let content = Content::function_response_json(name, json).with_role(Role::User);
99        self.contents.push(content);
100        Ok(self)
101    }
102
103    /// Add a message to the request
104    pub fn with_message(mut self, message: Message) -> Self {
105        let content = message.content.clone();
106        match &content.role {
107            Some(role) => {
108                let role_clone = role.clone();
109                self.contents.push(content.with_role(role_clone));
110            }
111            None => {
112                self.contents.push(content.with_role(message.role));
113            }
114        }
115        self
116    }
117
118    /// Add multiple messages to the request
119    pub fn with_messages(mut self, messages: impl IntoIterator<Item = Message>) -> Self {
120        for message in messages {
121            self = self.with_message(message);
122        }
123        self
124    }
125
126    /// Set the generation config for the request
127    pub fn with_generation_config(mut self, config: GenerationConfig) -> Self {
128        self.generation_config = Some(config);
129        self
130    }
131
132    /// Set the temperature for the request
133    pub fn with_temperature(mut self, temperature: f32) -> Self {
134        if self.generation_config.is_none() {
135            self.generation_config = Some(GenerationConfig::default());
136        }
137        if let Some(config) = &mut self.generation_config {
138            config.temperature = Some(temperature);
139        }
140        self
141    }
142
143    /// Set the top-p value for the request
144    pub fn with_top_p(mut self, top_p: f32) -> Self {
145        if self.generation_config.is_none() {
146            self.generation_config = Some(GenerationConfig::default());
147        }
148        if let Some(config) = &mut self.generation_config {
149            config.top_p = Some(top_p);
150        }
151        self
152    }
153
154    /// Set the top-k value for the request
155    pub fn with_top_k(mut self, top_k: i32) -> Self {
156        if self.generation_config.is_none() {
157            self.generation_config = Some(GenerationConfig::default());
158        }
159        if let Some(config) = &mut self.generation_config {
160            config.top_k = Some(top_k);
161        }
162        self
163    }
164
165    /// Set the maximum output tokens for the request
166    pub fn with_max_output_tokens(mut self, max_output_tokens: i32) -> Self {
167        if self.generation_config.is_none() {
168            self.generation_config = Some(GenerationConfig::default());
169        }
170        if let Some(config) = &mut self.generation_config {
171            config.max_output_tokens = Some(max_output_tokens);
172        }
173        self
174    }
175
176    /// Set the candidate count for the request
177    pub fn with_candidate_count(mut self, candidate_count: i32) -> Self {
178        if self.generation_config.is_none() {
179            self.generation_config = Some(GenerationConfig::default());
180        }
181        if let Some(config) = &mut self.generation_config {
182            config.candidate_count = Some(candidate_count);
183        }
184        self
185    }
186
187    /// Set the stop sequences for the request
188    pub fn with_stop_sequences(mut self, stop_sequences: Vec<String>) -> Self {
189        if self.generation_config.is_none() {
190            self.generation_config = Some(GenerationConfig::default());
191        }
192        if let Some(config) = &mut self.generation_config {
193            config.stop_sequences = Some(stop_sequences);
194        }
195        self
196    }
197
198    /// Set the response mime type for the request
199    pub fn with_response_mime_type(mut self, mime_type: impl Into<String>) -> Self {
200        if self.generation_config.is_none() {
201            self.generation_config = Some(GenerationConfig::default());
202        }
203        if let Some(config) = &mut self.generation_config {
204            config.response_mime_type = Some(mime_type.into());
205        }
206        self
207    }
208
209    /// Set the response schema for structured output
210    pub fn with_response_schema(mut self, schema: serde_json::Value) -> Self {
211        if self.generation_config.is_none() {
212            self.generation_config = Some(GenerationConfig::default());
213        }
214        if let Some(config) = &mut self.generation_config {
215            config.response_schema = Some(schema);
216        }
217        self
218    }
219
220    /// Add a tool to the request
221    pub fn with_tool(mut self, tool: Tool) -> Self {
222        if self.tools.is_none() {
223            self.tools = Some(Vec::new());
224        }
225        if let Some(tools) = &mut self.tools {
226            tools.push(tool);
227        }
228        self
229    }
230
231    /// Add a function declaration as a tool
232    pub fn with_function(mut self, function: FunctionDeclaration) -> Self {
233        let tool = Tool::new(function);
234        self = self.with_tool(tool);
235        self
236    }
237
238    /// Set the function calling mode for the request
239    pub fn with_function_calling_mode(mut self, mode: FunctionCallingMode) -> Self {
240        if self.tool_config.is_none() {
241            self.tool_config = Some(ToolConfig {
242                function_calling_config: Some(FunctionCallingConfig { mode }),
243            });
244        } else if let Some(tool_config) = &mut self.tool_config {
245            tool_config.function_calling_config = Some(FunctionCallingConfig { mode });
246        }
247        self
248    }
249
250    /// Set the thinking configuration for the request (Gemini 2.5 series only)
251    pub fn with_thinking_config(mut self, thinking_config: ThinkingConfig) -> Self {
252        if self.generation_config.is_none() {
253            self.generation_config = Some(GenerationConfig::default());
254        }
255        if let Some(config) = &mut self.generation_config {
256            config.thinking_config = Some(thinking_config);
257        }
258        self
259    }
260
261    /// Set the thinking budget for the request (Gemini 2.5 series only)
262    pub fn with_thinking_budget(mut self, budget: i32) -> Self {
263        if self.generation_config.is_none() {
264            self.generation_config = Some(GenerationConfig::default());
265        }
266        if let Some(config) = &mut self.generation_config {
267            if config.thinking_config.is_none() {
268                config.thinking_config = Some(ThinkingConfig::default());
269            }
270            if let Some(thinking_config) = &mut config.thinking_config {
271                thinking_config.thinking_budget = Some(budget);
272            }
273        }
274        self
275    }
276
277    /// Enable dynamic thinking (model decides the budget) (Gemini 2.5 series only)
278    pub fn with_dynamic_thinking(self) -> Self {
279        self.with_thinking_budget(-1)
280    }
281
282    /// Include thought summaries in the response (Gemini 2.5 series only)
283    pub fn with_thoughts_included(mut self, include: bool) -> Self {
284        if self.generation_config.is_none() {
285            self.generation_config = Some(GenerationConfig::default());
286        }
287        if let Some(config) = &mut self.generation_config {
288            if config.thinking_config.is_none() {
289                config.thinking_config = Some(ThinkingConfig::default());
290            }
291            if let Some(thinking_config) = &mut config.thinking_config {
292                thinking_config.include_thoughts = Some(include);
293            }
294        }
295        self
296    }
297
298    /// Enable audio output (text-to-speech)
299    pub fn with_audio_output(mut self) -> Self {
300        if self.generation_config.is_none() {
301            self.generation_config = Some(GenerationConfig::default());
302        }
303        if let Some(config) = &mut self.generation_config {
304            config.response_modalities = Some(vec!["AUDIO".to_string()]);
305        }
306        self
307    }
308
309    /// Set speech configuration for text-to-speech generation
310    pub fn with_speech_config(mut self, speech_config: SpeechConfig) -> Self {
311        if self.generation_config.is_none() {
312            self.generation_config = Some(GenerationConfig::default());
313        }
314        if let Some(config) = &mut self.generation_config {
315            config.speech_config = Some(speech_config);
316        }
317        self
318    }
319
320    /// Set a single voice for text-to-speech generation
321    pub fn with_voice(self, voice_name: impl Into<String>) -> Self {
322        let speech_config = SpeechConfig::single_voice(voice_name);
323        self.with_speech_config(speech_config).with_audio_output()
324    }
325
326    /// Set multi-speaker configuration for text-to-speech generation
327    pub fn with_multi_speaker_config(self, speakers: Vec<SpeakerVoiceConfig>) -> Self {
328        let speech_config = SpeechConfig::multi_speaker(speakers);
329        self.with_speech_config(speech_config).with_audio_output()
330    }
331
332    pub fn build(self) -> GenerateContentRequest {
333        GenerateContentRequest {
334            contents: self.contents,
335            generation_config: self.generation_config,
336            safety_settings: None,
337            tools: self.tools,
338            tool_config: self.tool_config,
339            system_instruction: self.system_instruction,
340        }
341    }
342
343    /// Execute the request
344    pub async fn execute(self) -> Result<GenerationResponse> {
345        let client = self.client.clone();
346        let request = self.build();
347        client.generate_content_raw(request).await
348    }
349
350    /// Execute the request with streaming
351    pub async fn execute_stream(
352        self,
353    ) -> Result<Pin<Box<dyn Stream<Item = Result<GenerationResponse>> + Send>>> {
354        let request = GenerateContentRequest {
355            contents: self.contents,
356            generation_config: self.generation_config,
357            safety_settings: None,
358            tools: self.tools,
359            tool_config: self.tool_config,
360            system_instruction: self.system_instruction,
361        };
362
363        self.client.generate_content_stream(request).await
364    }
365}
gemini_rust/content_builder.rs

gemini_rust/
content_builder.rs
