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_max_output_tokens(mut self, max_output_tokens: i32) -> Self {
228 self.generation_config
229 .get_or_insert_with(Default::default)
230 .max_output_tokens = Some(max_output_tokens);
231 self
232 }
233
234 pub fn with_candidate_count(mut self, candidate_count: i32) -> Self {
236 self.generation_config
237 .get_or_insert_with(Default::default)
238 .candidate_count = Some(candidate_count);
239 self
240 }
241
242 pub fn with_stop_sequences(mut self, stop_sequences: Vec<String>) -> Self {
246 self.generation_config
247 .get_or_insert_with(Default::default)
248 .stop_sequences = Some(stop_sequences);
249 self
250 }
251
252 pub fn with_response_mime_type(mut self, mime_type: impl Into<String>) -> Self {
256 self.generation_config
257 .get_or_insert_with(Default::default)
258 .response_mime_type = Some(mime_type.into());
259 self
260 }
261
262 pub fn with_response_schema(mut self, schema: serde_json::Value) -> Self {
267 self.generation_config
268 .get_or_insert_with(Default::default)
269 .response_schema = Some(schema);
270 self
271 }
272
273 pub fn with_tool(mut self, tool: Tool) -> Self {
277 self.tools.get_or_insert_with(Vec::new).push(tool);
278 self
279 }
280
281 pub fn with_function(mut self, function: FunctionDeclaration) -> Self {
285 let tool = Tool::new(function);
286 self = self.with_tool(tool);
287 self
288 }
289
290 pub fn with_function_calling_mode(mut self, mode: FunctionCallingMode) -> Self {
292 self.tool_config
293 .get_or_insert_with(Default::default)
294 .function_calling_config = Some(FunctionCallingConfig { mode });
295 self
296 }
297
298 pub fn with_tool_config(mut self, tool_config: ToolConfig) -> Self {
300 self.tool_config = Some(tool_config);
301 self
302 }
303
304 pub fn with_thinking_config(mut self, thinking_config: ThinkingConfig) -> Self {
306 self.generation_config
307 .get_or_insert_with(Default::default)
308 .thinking_config = Some(thinking_config);
309 self
310 }
311
312 pub fn with_thinking_budget(mut self, budget: i32) -> Self {
317 let config = self
318 .generation_config
319 .get_or_insert_with(Default::default)
320 .thinking_config
321 .get_or_insert_with(Default::default);
322 config.thinking_budget = Some(budget);
323 config.thinking_level = None;
324 self
325 }
326
327 pub fn with_dynamic_thinking(self) -> Self {
333 self.with_thinking_budget(-1)
334 }
335
336 pub fn with_thoughts_included(mut self, include: bool) -> Self {
340 self.generation_config
341 .get_or_insert_with(Default::default)
342 .thinking_config
343 .get_or_insert_with(Default::default)
344 .include_thoughts = Some(include);
345 self
346 }
347
348 pub fn with_thinking_level(mut self, level: ThinkingLevel) -> Self {
357 let config = self
358 .generation_config
359 .get_or_insert_with(Default::default)
360 .thinking_config
361 .get_or_insert_with(Default::default);
362 config.thinking_level = Some(level);
363 config.thinking_budget = None;
364 self
365 }
366
367 pub fn with_media_resolution(mut self, level: MediaResolutionLevel) -> Self {
373 self.generation_config
374 .get_or_insert_with(Default::default)
375 .media_resolution = Some(level);
376 self
377 }
378
379 pub fn with_code_execution(self) -> Self {
385 self.with_tool(Tool::code_execution())
386 }
387
388 pub fn with_audio_output(mut self) -> Self {
390 self.generation_config
391 .get_or_insert_with(Default::default)
392 .response_modalities = Some(vec!["AUDIO".to_string()]);
393 self
394 }
395
396 pub fn with_image_config(mut self, image_config: ImageConfig) -> Self {
398 self.generation_config
399 .get_or_insert_with(Default::default)
400 .image_config = Some(image_config);
401 self
402 }
403
404 pub fn with_speech_config(mut self, speech_config: SpeechConfig) -> Self {
406 self.generation_config
407 .get_or_insert_with(Default::default)
408 .speech_config = Some(speech_config);
409 self
410 }
411
412 pub fn with_voice(self, voice_name: impl Into<String>) -> Self {
414 let speech_config = SpeechConfig::single_voice(voice_name);
415 self.with_speech_config(speech_config).with_audio_output()
416 }
417
418 pub fn with_multi_speaker_config(self, speakers: Vec<SpeakerVoiceConfig>) -> Self {
420 let speech_config = SpeechConfig::multi_speaker(speakers);
421 self.with_speech_config(speech_config).with_audio_output()
422 }
423
424 pub fn build(self) -> GenerateContentRequest {
426 GenerateContentRequest {
427 contents: self.contents,
428 generation_config: self.generation_config,
429 safety_settings: self.safety_settings,
430 tools: self.tools,
431 tool_config: self.tool_config,
432 system_instruction: self.system_instruction,
433 cached_content: self.cached_content,
434 }
435 }
436
437 #[instrument(skip_all, fields(
439 messages.parts.count = self.contents.len(),
440 tools.present = self.tools.is_some(),
441 system.instruction.present = self.system_instruction.is_some(),
442 cached.content.present = self.cached_content.is_some(),
443 ))]
444 pub async fn execute(self) -> Result<GenerationResponse, ClientError> {
445 let client = self.client.clone();
446 let request = self.build();
447 client.generate_content_raw(request).await
448 }
449
450 #[instrument(skip_all, fields(
452 messages.parts.count = self.contents.len(),
453 tools.present = self.tools.is_some(),
454 system.instruction.present = self.system_instruction.is_some(),
455 cached.content.present = self.cached_content.is_some(),
456 ))]
457 pub async fn execute_stream(self) -> Result<GenerationStream, ClientError> {
458 let client = self.client.clone();
459 let request = self.build();
460 client.generate_content_stream(request).await
461 }
462
463 #[instrument(skip_all, fields(
465 messages.parts.count = self.contents.len(),
466 ))]
467 pub async fn count_tokens(self) -> Result<super::model::CountTokensResponse, ClientError> {
468 let client = self.client.clone();
469 let request = self.build();
470 client.count_tokens(request).await
471 }
472}