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 Function,
13}
14
15#[derive(Debug, Clone, Serialize, Deserialize)]
17#[serde(untagged)]
18pub enum Part {
19 Text {
21 text: String,
23 },
24 FunctionCall {
26 #[serde(rename = "functionCall")]
28 function_call: super::tools::FunctionCall,
29 },
30 FunctionResponse {
32 #[serde(rename = "functionResponse")]
34 function_response: super::tools::FunctionResponse,
35 },
36}
37
38#[derive(Debug, Default, Clone, Serialize, Deserialize)]
40pub struct Content {
41 pub parts: Vec<Part>,
43 #[serde(skip_serializing_if = "Option::is_none")]
45 pub role: Option<Role>,
46}
47
48impl Content {
49 pub fn text(text: impl Into<String>) -> Self {
51 Self {
52 parts: vec![Part::Text { text: text.into() }],
53 role: None,
54 }
55 }
56
57 pub fn function_call(function_call: super::tools::FunctionCall) -> Self {
59 Self {
60 parts: vec![Part::FunctionCall { function_call }],
61 role: None,
62 }
63 }
64
65 pub fn function_response(function_response: super::tools::FunctionResponse) -> Self {
67 Self {
68 parts: vec![Part::FunctionResponse { function_response }],
69 role: None,
70 }
71 }
72
73 pub fn function_response_json(name: impl Into<String>, response: serde_json::Value) -> Self {
75 Self {
76 parts: vec![Part::FunctionResponse {
77 function_response: super::tools::FunctionResponse::new(name, response),
78 }],
79 role: None,
80 }
81 }
82
83 pub fn with_role(mut self, role: Role) -> Self {
85 self.role = Some(role);
86 self
87 }
88}
89
90#[derive(Debug, Clone, Serialize, Deserialize)]
92pub struct Message {
93 pub content: Content,
95 pub role: Role,
97}
98
99impl Message {
100 pub fn user(text: impl Into<String>) -> Self {
102 Self {
103 content: Content::text(text),
104 role: Role::User,
105 }
106 }
107
108 pub fn model(text: impl Into<String>) -> Self {
110 Self {
111 content: Content::text(text),
112 role: Role::Model,
113 }
114 }
115
116 pub fn function(name: impl Into<String>, response: serde_json::Value) -> Self {
118 Self {
119 content: Content::function_response_json(name, response),
120 role: Role::Function,
121 }
122 }
123
124 pub fn function_str(
126 name: impl Into<String>,
127 response: impl Into<String>,
128 ) -> Result<Self, serde_json::Error> {
129 let response_str = response.into();
130 let json = serde_json::from_str(&response_str)?;
131 Ok(Self {
132 content: Content::function_response_json(name, json),
133 role: Role::Function,
134 })
135 }
136}
137
138#[derive(Debug, Clone, Serialize, Deserialize)]
140pub struct SafetyRating {
141 pub category: String,
143 pub probability: String,
145}
146
147#[derive(Debug, Clone, Serialize, Deserialize)]
149pub struct CitationMetadata {
150 pub citation_sources: Vec<CitationSource>,
152}
153
154#[derive(Debug, Clone, Serialize, Deserialize)]
156pub struct CitationSource {
157 pub uri: Option<String>,
159 pub title: Option<String>,
161 pub start_index: Option<i32>,
163 pub end_index: Option<i32>,
165 pub license: Option<String>,
167 pub publication_date: Option<String>,
169}
170
171#[derive(Debug, Clone, Serialize, Deserialize)]
173pub struct Candidate {
174 pub content: Content,
176 #[serde(skip_serializing_if = "Option::is_none")]
178 pub safety_ratings: Option<Vec<SafetyRating>>,
179 #[serde(skip_serializing_if = "Option::is_none")]
181 pub citation_metadata: Option<CitationMetadata>,
182 #[serde(skip_serializing_if = "Option::is_none")]
184 pub finish_reason: Option<String>,
185 #[serde(skip_serializing_if = "Option::is_none")]
187 pub usage_metadata: Option<UsageMetadata>,
188}
189
190#[derive(Debug, Clone, Serialize, Deserialize)]
192pub struct UsageMetadata {
193 pub prompt_token_count: i32,
195 pub candidates_token_count: i32,
197 pub total_token_count: i32,
199}
200
201#[derive(Debug, Clone, Serialize, Deserialize)]
203pub struct GenerationResponse {
204 pub candidates: Vec<Candidate>,
206 #[serde(skip_serializing_if = "Option::is_none")]
208 pub prompt_feedback: Option<PromptFeedback>,
209 #[serde(skip_serializing_if = "Option::is_none")]
211 pub usage_metadata: Option<UsageMetadata>,
212}
213
214#[derive(Debug, Clone, Serialize, Deserialize)]
216pub struct PromptFeedback {
217 pub safety_ratings: Vec<SafetyRating>,
219 #[serde(skip_serializing_if = "Option::is_none")]
221 pub block_reason: Option<String>,
222}
223
224impl GenerationResponse {
225 pub fn text(&self) -> String {
227 self.candidates
228 .first()
229 .and_then(|c| {
230 c.content.parts.first().and_then(|p| match p {
231 Part::Text { text } => Some(text.clone()),
232 _ => None,
233 })
234 })
235 .unwrap_or_default()
236 }
237
238 pub fn function_calls(&self) -> Vec<&super::tools::FunctionCall> {
240 self.candidates
241 .iter()
242 .flat_map(|c| {
243 c.content.parts.iter().filter_map(|p| match p {
244 Part::FunctionCall { function_call } => Some(function_call),
245 _ => None,
246 })
247 })
248 .collect()
249 }
250}
251
252#[derive(Debug, Clone, Serialize, Deserialize)]
254pub struct GenerateContentRequest {
255 pub contents: Vec<Content>,
257 #[serde(skip_serializing_if = "Option::is_none")]
259 pub generation_config: Option<GenerationConfig>,
260 #[serde(skip_serializing_if = "Option::is_none")]
262 pub safety_settings: Option<Vec<SafetySetting>>,
263 #[serde(skip_serializing_if = "Option::is_none")]
265 pub tools: Option<Vec<super::tools::Tool>>,
266 #[serde(skip_serializing_if = "Option::is_none")]
268 pub tool_config: Option<ToolConfig>,
269 #[serde(skip_serializing_if = "Option::is_none")]
271 pub system_instruction: Option<Content>,
272}
273
274#[derive(Debug, Clone, Serialize, Deserialize)]
276pub struct GenerationConfig {
277 #[serde(skip_serializing_if = "Option::is_none")]
282 pub temperature: Option<f32>,
283
284 #[serde(skip_serializing_if = "Option::is_none")]
290 pub top_p: Option<f32>,
291
292 #[serde(skip_serializing_if = "Option::is_none")]
297 pub top_k: Option<i32>,
298
299 #[serde(skip_serializing_if = "Option::is_none")]
303 pub max_output_tokens: Option<i32>,
304
305 #[serde(skip_serializing_if = "Option::is_none")]
309 pub candidate_count: Option<i32>,
310
311 #[serde(skip_serializing_if = "Option::is_none")]
315 pub stop_sequences: Option<Vec<String>>,
316
317 #[serde(skip_serializing_if = "Option::is_none")]
321 pub response_mime_type: Option<String>,
322
323 #[serde(skip_serializing_if = "Option::is_none")]
327 pub response_schema: Option<serde_json::Value>,
328}
329
330impl Default for GenerationConfig {
331 fn default() -> Self {
332 Self {
333 temperature: Some(0.7),
334 top_p: Some(0.95),
335 top_k: Some(40),
336 max_output_tokens: Some(1024),
337 candidate_count: Some(1),
338 stop_sequences: None,
339 response_mime_type: None,
340 response_schema: None,
341 }
342 }
343}
344
345#[derive(Debug, Clone, Serialize, Deserialize)]
347pub struct ToolConfig {
348 #[serde(skip_serializing_if = "Option::is_none")]
350 pub function_calling_config: Option<FunctionCallingConfig>,
351}
352
353#[derive(Debug, Clone, Serialize, Deserialize)]
355pub struct FunctionCallingConfig {
356 pub mode: FunctionCallingMode,
358}
359
360#[derive(Debug, Clone, Serialize, Deserialize)]
362#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
363pub enum FunctionCallingMode {
364 Auto,
366 Any,
368 None,
370}
371
372#[derive(Debug, Clone, Serialize, Deserialize)]
374pub struct SafetySetting {
375 pub category: HarmCategory,
377 pub threshold: HarmBlockThreshold,
379}
380
381#[derive(Debug, Clone, Serialize, Deserialize)]
383#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
384pub enum HarmCategory {
385 Dangerous,
387 Harassment,
389 HateSpeech,
391 SexuallyExplicit,
393}
394
395#[derive(Debug, Clone, Serialize, Deserialize)]
397#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
398pub enum HarmBlockThreshold {
399 BlockLowAndAbove,
401 BlockMediumAndAbove,
403 BlockHighAndAbove,
405 BlockOnlyHigh,
407 BlockNone,
409}