gemini_rust/
content_builder.rs

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

gemini_rust/
content_builder.rs
