1use serde::{Deserialize, Serialize};
2
3#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
5#[serde(rename_all = "lowercase")]
6pub enum Role {
7 User,
9 Model,
11}
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
15#[serde(untagged)]
16pub enum Part {
17 Text {
19 text: String,
21 #[serde(skip_serializing_if = "Option::is_none")]
23 thought: Option<bool>,
24 },
25 InlineData {
26 #[serde(rename = "inlineData")]
28 inline_data: Blob,
29 },
30 FunctionCall {
32 #[serde(rename = "functionCall")]
34 function_call: super::tools::FunctionCall,
35 },
36 FunctionResponse {
38 #[serde(rename = "functionResponse")]
40 function_response: super::tools::FunctionResponse,
41 },
42}
43
44#[derive(Debug, Clone, Serialize, Deserialize)]
46#[serde(rename_all = "camelCase")]
47pub struct Blob {
48 pub mime_type: String,
49 pub data: String, }
51
52impl Blob {
53 pub fn new(mime_type: impl Into<String>, data: impl Into<String>) -> Self {
55 Self {
56 mime_type: mime_type.into(),
57 data: data.into(),
58 }
59 }
60}
61
62#[derive(Debug, Default, Clone, Serialize, Deserialize)]
64pub struct Content {
65 pub parts: Vec<Part>,
67 #[serde(skip_serializing_if = "Option::is_none")]
69 pub role: Option<Role>,
70}
71
72impl Content {
73 pub fn text(text: impl Into<String>) -> Self {
75 Self {
76 parts: vec![Part::Text {
77 text: text.into(),
78 thought: None,
79 }],
80 role: None,
81 }
82 }
83
84 pub fn function_call(function_call: super::tools::FunctionCall) -> Self {
86 Self {
87 parts: vec![Part::FunctionCall { function_call }],
88 role: None,
89 }
90 }
91
92 pub fn function_response(function_response: super::tools::FunctionResponse) -> Self {
94 Self {
95 parts: vec![Part::FunctionResponse { function_response }],
96 role: None,
97 }
98 }
99
100 pub fn function_response_json(name: impl Into<String>, response: serde_json::Value) -> Self {
102 Self {
103 parts: vec![Part::FunctionResponse {
104 function_response: super::tools::FunctionResponse::new(name, response),
105 }],
106 role: None,
107 }
108 }
109
110 pub fn inline_data(mime_type: impl Into<String>, data: impl Into<String>) -> Self {
112 Self {
113 parts: vec![Part::InlineData {
114 inline_data: Blob::new(mime_type, data),
115 }],
116 role: None,
117 }
118 }
119
120 pub fn with_role(mut self, role: Role) -> Self {
122 self.role = Some(role);
123 self
124 }
125}
126
127#[derive(Debug, Clone, Serialize, Deserialize)]
129pub struct Message {
130 pub content: Content,
132 pub role: Role,
134}
135
136impl Message {
137 pub fn user(text: impl Into<String>) -> Self {
139 Self {
140 content: Content::text(text).with_role(Role::User),
141 role: Role::User,
142 }
143 }
144
145 pub fn model(text: impl Into<String>) -> Self {
147 Self {
148 content: Content::text(text).with_role(Role::Model),
149 role: Role::Model,
150 }
151 }
152
153 pub fn embed(text: impl Into<String>) -> Self {
154 Self {
155 content: Content::text(text),
156 role: Role::Model,
157 }
158 }
159
160 pub fn function(name: impl Into<String>, response: serde_json::Value) -> Self {
162 Self {
163 content: Content::function_response_json(name, response).with_role(Role::Model),
164 role: Role::Model,
165 }
166 }
167
168 pub fn function_str(
170 name: impl Into<String>,
171 response: impl Into<String>,
172 ) -> Result<Self, serde_json::Error> {
173 let response_str = response.into();
174 let json = serde_json::from_str(&response_str)?;
175 Ok(Self {
176 content: Content::function_response_json(name, json).with_role(Role::Model),
177 role: Role::Model,
178 })
179 }
180}
181
182#[derive(Debug, Clone, Serialize, Deserialize)]
184pub struct SafetyRating {
185 pub category: String,
187 pub probability: String,
189}
190
191#[derive(Debug, Clone, Serialize, Deserialize)]
193#[serde(rename_all = "camelCase")]
194pub struct CitationMetadata {
195 pub citation_sources: Vec<CitationSource>,
197}
198
199#[derive(Debug, Clone, Serialize, Deserialize)]
201#[serde(rename_all = "camelCase")]
202pub struct CitationSource {
203 pub uri: Option<String>,
205 pub title: Option<String>,
207 pub start_index: Option<i32>,
209 pub end_index: Option<i32>,
211 pub license: Option<String>,
213 pub publication_date: Option<String>,
215}
216
217#[derive(Debug, Clone, Serialize, Deserialize)]
219#[serde(rename_all = "camelCase")]
220pub struct Candidate {
221 pub content: Content,
223 #[serde(skip_serializing_if = "Option::is_none")]
225 pub safety_ratings: Option<Vec<SafetyRating>>,
226 #[serde(skip_serializing_if = "Option::is_none")]
228 pub citation_metadata: Option<CitationMetadata>,
229 #[serde(skip_serializing_if = "Option::is_none")]
231 pub finish_reason: Option<String>,
232 #[serde(skip_serializing_if = "Option::is_none")]
234 pub index: Option<i32>,
235}
236
237#[derive(Debug, Clone, Serialize, Deserialize)]
239#[serde(rename_all = "camelCase")]
240pub struct UsageMetadata {
241 pub prompt_token_count: i32,
243 #[serde(skip_serializing_if = "Option::is_none")]
245 pub candidates_token_count: Option<i32>,
246 pub total_token_count: i32,
248 #[serde(skip_serializing_if = "Option::is_none")]
250 pub thoughts_token_count: Option<i32>,
251}
252
253#[derive(Debug, Clone, Serialize, Deserialize)]
255#[serde(rename_all = "camelCase")]
256pub struct GenerationResponse {
257 pub candidates: Vec<Candidate>,
259 #[serde(skip_serializing_if = "Option::is_none")]
261 pub prompt_feedback: Option<PromptFeedback>,
262 #[serde(skip_serializing_if = "Option::is_none")]
264 pub usage_metadata: Option<UsageMetadata>,
265}
266
267#[derive(Debug, Clone, Serialize, Deserialize)]
269pub struct ContentEmbedding {
270 pub values: Vec<f32>, }
273
274#[derive(Debug, Clone, Serialize, Deserialize)]
276pub struct ContentEmbeddingResponse {
277 pub embedding: ContentEmbedding,
279}
280
281#[derive(Debug, Clone, Serialize, Deserialize)]
283pub struct BatchContentEmbeddingResponse {
284 pub embeddings: Vec<ContentEmbedding>,
286}
287
288#[derive(Debug, Clone, Serialize, Deserialize)]
290#[serde(rename_all = "camelCase")]
291pub struct PromptFeedback {
292 pub safety_ratings: Vec<SafetyRating>,
294 #[serde(skip_serializing_if = "Option::is_none")]
296 pub block_reason: Option<String>,
297}
298
299impl GenerationResponse {
300 pub fn text(&self) -> String {
302 self.candidates
303 .first()
304 .and_then(|c| {
305 c.content.parts.first().and_then(|p| match p {
306 Part::Text { text, thought: _ } => Some(text.clone()),
307 _ => None,
308 })
309 })
310 .unwrap_or_default()
311 }
312
313 pub fn function_calls(&self) -> Vec<&super::tools::FunctionCall> {
315 self.candidates
316 .iter()
317 .flat_map(|c| {
318 c.content.parts.iter().filter_map(|p| match p {
319 Part::FunctionCall { function_call } => Some(function_call),
320 _ => None,
321 })
322 })
323 .collect()
324 }
325
326 pub fn thoughts(&self) -> Vec<String> {
328 self.candidates
329 .iter()
330 .flat_map(|c| {
331 c.content.parts.iter().filter_map(|p| match p {
332 Part::Text {
333 text,
334 thought: Some(true),
335 } => Some(text.clone()),
336 _ => None,
337 })
338 })
339 .collect()
340 }
341
342 pub fn all_text(&self) -> Vec<(String, bool)> {
344 self.candidates
345 .iter()
346 .flat_map(|c| {
347 c.content.parts.iter().filter_map(|p| match p {
348 Part::Text { text, thought } => Some((text.clone(), thought.unwrap_or(false))),
349 _ => None,
350 })
351 })
352 .collect()
353 }
354}
355
356#[derive(Debug, Clone, Serialize, Deserialize)]
358pub struct GenerateContentRequest {
359 pub contents: Vec<Content>,
361 #[serde(skip_serializing_if = "Option::is_none")]
363 pub generation_config: Option<GenerationConfig>,
364 #[serde(skip_serializing_if = "Option::is_none")]
366 pub safety_settings: Option<Vec<SafetySetting>>,
367 #[serde(skip_serializing_if = "Option::is_none")]
369 pub tools: Option<Vec<super::tools::Tool>>,
370 #[serde(skip_serializing_if = "Option::is_none")]
372 pub tool_config: Option<ToolConfig>,
373 #[serde(skip_serializing_if = "Option::is_none")]
375 pub system_instruction: Option<Content>,
376}
377
378#[derive(Debug, Clone, Serialize, Deserialize)]
380pub struct EmbedContentRequest {
381 pub model: String,
383 pub content: Content,
385 #[serde(skip_serializing_if = "Option::is_none")]
387 pub task_type: Option<TaskType>,
388 #[serde(skip_serializing_if = "Option::is_none")]
390 pub title: Option<String>,
391 #[serde(skip_serializing_if = "Option::is_none")]
393 pub output_dimensionality: Option<i32>,
394}
395
396#[derive(Debug, Clone, Serialize, Deserialize)]
398pub struct BatchEmbedContentsRequest {
399 pub requests: Vec<EmbedContentRequest>,
401}
402
403#[derive(Debug, Clone, Serialize, Deserialize)]
405#[serde(rename_all = "camelCase")]
406pub struct ThinkingConfig {
407 #[serde(skip_serializing_if = "Option::is_none")]
418 pub thinking_budget: Option<i32>,
419
420 #[serde(skip_serializing_if = "Option::is_none")]
425 pub include_thoughts: Option<bool>,
426}
427
428impl ThinkingConfig {
429 pub fn new() -> Self {
431 Self {
432 thinking_budget: None,
433 include_thoughts: None,
434 }
435 }
436
437 pub fn with_thinking_budget(mut self, budget: i32) -> Self {
439 self.thinking_budget = Some(budget);
440 self
441 }
442
443 pub fn with_dynamic_thinking(mut self) -> Self {
445 self.thinking_budget = Some(-1);
446 self
447 }
448
449 pub fn with_thoughts_included(mut self, include: bool) -> Self {
451 self.include_thoughts = Some(include);
452 self
453 }
454}
455
456impl Default for ThinkingConfig {
457 fn default() -> Self {
458 Self::new()
459 }
460}
461
462#[derive(Debug, Clone, Serialize, Deserialize)]
464pub struct GenerationConfig {
465 #[serde(skip_serializing_if = "Option::is_none")]
470 pub temperature: Option<f32>,
471
472 #[serde(skip_serializing_if = "Option::is_none")]
478 pub top_p: Option<f32>,
479
480 #[serde(skip_serializing_if = "Option::is_none")]
485 pub top_k: Option<i32>,
486
487 #[serde(skip_serializing_if = "Option::is_none")]
491 pub max_output_tokens: Option<i32>,
492
493 #[serde(skip_serializing_if = "Option::is_none")]
497 pub candidate_count: Option<i32>,
498
499 #[serde(skip_serializing_if = "Option::is_none")]
503 pub stop_sequences: Option<Vec<String>>,
504
505 #[serde(skip_serializing_if = "Option::is_none")]
509 pub response_mime_type: Option<String>,
510
511 #[serde(skip_serializing_if = "Option::is_none")]
515 pub response_schema: Option<serde_json::Value>,
516
517 #[serde(skip_serializing_if = "Option::is_none")]
521 pub thinking_config: Option<ThinkingConfig>,
522}
523
524impl Default for GenerationConfig {
525 fn default() -> Self {
526 Self {
527 temperature: Some(0.7),
528 top_p: Some(0.95),
529 top_k: Some(40),
530 max_output_tokens: Some(1024),
531 candidate_count: Some(1),
532 stop_sequences: None,
533 response_mime_type: None,
534 response_schema: None,
535 thinking_config: None,
536 }
537 }
538}
539
540#[derive(Debug, Clone, Serialize, Deserialize)]
542pub struct ToolConfig {
543 #[serde(skip_serializing_if = "Option::is_none")]
545 pub function_calling_config: Option<FunctionCallingConfig>,
546}
547
548#[derive(Debug, Clone, Serialize, Deserialize)]
550pub struct FunctionCallingConfig {
551 pub mode: FunctionCallingMode,
553}
554
555#[derive(Debug, Clone, Serialize, Deserialize)]
557#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
558pub enum FunctionCallingMode {
559 Auto,
561 Any,
563 None,
565}
566
567#[derive(Debug, Clone, Serialize, Deserialize)]
569pub struct SafetySetting {
570 pub category: HarmCategory,
572 pub threshold: HarmBlockThreshold,
574}
575
576#[derive(Debug, Clone, Serialize, Deserialize)]
578#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
579pub enum HarmCategory {
580 Dangerous,
582 Harassment,
584 HateSpeech,
586 SexuallyExplicit,
588}
589
590#[allow(clippy::enum_variant_names)]
592#[derive(Debug, Clone, Serialize, Deserialize)]
593#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
594pub enum HarmBlockThreshold {
595 BlockLowAndAbove,
597 BlockMediumAndAbove,
599 BlockHighAndAbove,
601 BlockOnlyHigh,
603 BlockNone,
605}
606
607#[derive(Debug, Clone, Serialize, Deserialize)]
609#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
610pub enum TaskType {
611 SemanticSimilarity,
613 Classification,
615 Clustering,
617
618 RetrievalDocument,
620 RetrievalQuery,
621 QuestionAnswering,
622 FactVerification,
623
624 CodeRetrievalQuery,
627}