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#[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 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 pub fn with_safety_settings(mut self, safety_settings: Vec<SafetySetting>) -> Self {
47 self.safety_settings = Some(safety_settings);
48 self
49 }
50
51 pub fn with_system_prompt(self, text: impl Into<String>) -> Self {
55 self.with_system_instruction(text)
56 }
57
58 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 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 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 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 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 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 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 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 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 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 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 pub fn with_generation_config(mut self, config: GenerationConfig) -> Self {
190 self.generation_config = Some(config);
191 self
192 }
193
194 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 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 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 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 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 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 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 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 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 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 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 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 pub fn with_tool_config(mut self, tool_config: ToolConfig) -> Self {
308 self.tool_config = Some(tool_config);
309 self
310 }
311
312 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 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 pub fn with_dynamic_thinking(self) -> Self {
341 self.with_thinking_budget(-1)
342 }
343
344 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 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 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 pub fn with_code_execution(self) -> Self {
393 self.with_tool(Tool::code_execution())
394 }
395
396 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 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 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 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 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 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 #[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 #[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 #[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}