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 System,
13 Function,
15}
16
17#[derive(Debug, Clone, Serialize, Deserialize)]
19#[serde(untagged)]
20pub enum Part {
21 Text {
23 text: String,
25 },
26 FunctionCall {
28 #[serde(rename = "functionCall")]
30 function_call: super::tools::FunctionCall,
31 },
32 FunctionResponse {
34 #[serde(rename = "functionResponse")]
36 function_response: super::tools::FunctionResponse,
37 },
38}
39
40#[derive(Debug, Default, Clone, Serialize, Deserialize)]
42pub struct Content {
43 pub parts: Vec<Part>,
45 #[serde(skip_serializing_if = "Option::is_none")]
47 pub role: Option<Role>,
48}
49
50impl Content {
51 pub fn text(text: impl Into<String>) -> Self {
53 Self {
54 parts: vec![Part::Text { text: text.into() }],
55 role: None,
56 }
57 }
58
59 pub fn function_call(function_call: super::tools::FunctionCall) -> Self {
61 Self {
62 parts: vec![Part::FunctionCall { function_call }],
63 role: None,
64 }
65 }
66
67 pub fn function_response(function_response: super::tools::FunctionResponse) -> Self {
69 Self {
70 parts: vec![Part::FunctionResponse { function_response }],
71 role: None,
72 }
73 }
74
75 pub fn function_response_json(
77 name: impl Into<String>,
78 response: serde_json::Value
79 ) -> Self {
80 Self {
81 parts: vec![Part::FunctionResponse {
82 function_response: super::tools::FunctionResponse::new(name, response)
83 }],
84 role: None,
85 }
86 }
87
88 pub fn with_role(mut self, role: Role) -> Self {
90 self.role = Some(role);
91 self
92 }
93}
94
95#[derive(Debug, Clone, Serialize, Deserialize)]
97pub struct Message {
98 pub content: Content,
100 pub role: Role,
102}
103
104impl Message {
105 pub fn user(text: impl Into<String>) -> Self {
107 Self {
108 content: Content::text(text),
109 role: Role::User,
110 }
111 }
112
113 pub fn model(text: impl Into<String>) -> Self {
115 Self {
116 content: Content::text(text),
117 role: Role::Model,
118 }
119 }
120
121 pub fn system(text: impl Into<String>) -> Self {
123 Self {
124 content: Content::text(text),
125 role: Role::System,
126 }
127 }
128
129 pub fn function(name: impl Into<String>, response: serde_json::Value) -> Self {
131 Self {
132 content: Content::function_response_json(name, response),
133 role: Role::Function,
134 }
135 }
136
137 pub fn function_str(name: impl Into<String>, response: impl Into<String>) -> Result<Self, serde_json::Error> {
139 let response_str = response.into();
140 let json = serde_json::from_str(&response_str)?;
141 Ok(Self {
142 content: Content::function_response_json(name, json),
143 role: Role::Function,
144 })
145 }
146}
147
148#[derive(Debug, Clone, Serialize, Deserialize)]
150pub struct SafetyRating {
151 pub category: String,
153 pub probability: String,
155}
156
157#[derive(Debug, Clone, Serialize, Deserialize)]
159pub struct CitationMetadata {
160 pub citation_sources: Vec<CitationSource>,
162}
163
164#[derive(Debug, Clone, Serialize, Deserialize)]
166pub struct CitationSource {
167 pub uri: Option<String>,
169 pub title: Option<String>,
171 pub start_index: Option<i32>,
173 pub end_index: Option<i32>,
175 pub license: Option<String>,
177 pub publication_date: Option<String>,
179}
180
181#[derive(Debug, Clone, Serialize, Deserialize)]
183pub struct Candidate {
184 pub content: Content,
186 #[serde(skip_serializing_if = "Option::is_none")]
188 pub safety_ratings: Option<Vec<SafetyRating>>,
189 #[serde(skip_serializing_if = "Option::is_none")]
191 pub citation_metadata: Option<CitationMetadata>,
192 #[serde(skip_serializing_if = "Option::is_none")]
194 pub finish_reason: Option<String>,
195 #[serde(skip_serializing_if = "Option::is_none")]
197 pub usage_metadata: Option<UsageMetadata>,
198}
199
200#[derive(Debug, Clone, Serialize, Deserialize)]
202pub struct UsageMetadata {
203 pub prompt_token_count: i32,
205 pub candidates_token_count: i32,
207 pub total_token_count: i32,
209}
210
211#[derive(Debug, Clone, Serialize, Deserialize)]
213pub struct GenerationResponse {
214 pub candidates: Vec<Candidate>,
216 #[serde(skip_serializing_if = "Option::is_none")]
218 pub prompt_feedback: Option<PromptFeedback>,
219 #[serde(skip_serializing_if = "Option::is_none")]
221 pub usage_metadata: Option<UsageMetadata>,
222}
223
224#[derive(Debug, Clone, Serialize, Deserialize)]
226pub struct PromptFeedback {
227 pub safety_ratings: Vec<SafetyRating>,
229 #[serde(skip_serializing_if = "Option::is_none")]
231 pub block_reason: Option<String>,
232}
233
234impl GenerationResponse {
235 pub fn text(&self) -> String {
237 self.candidates
238 .first()
239 .and_then(|c| {
240 c.content.parts.first().and_then(|p| match p {
241 Part::Text { text } => Some(text.clone()),
242 _ => None,
243 })
244 })
245 .unwrap_or_default()
246 }
247
248 pub fn function_calls(&self) -> Vec<&super::tools::FunctionCall> {
250 self.candidates
251 .iter()
252 .flat_map(|c| {
253 c.content.parts.iter().filter_map(|p| match p {
254 Part::FunctionCall { function_call } => Some(function_call),
255 _ => None,
256 })
257 })
258 .collect()
259 }
260}
261
262#[derive(Debug, Clone, Serialize, Deserialize)]
264pub struct GenerateContentRequest {
265 pub contents: Vec<Content>,
267 #[serde(skip_serializing_if = "Option::is_none")]
269 pub generation_config: Option<GenerationConfig>,
270 #[serde(skip_serializing_if = "Option::is_none")]
272 pub safety_settings: Option<Vec<SafetySetting>>,
273 #[serde(skip_serializing_if = "Option::is_none")]
275 pub tools: Option<Vec<super::tools::Tool>>,
276 #[serde(skip_serializing_if = "Option::is_none")]
278 pub tool_config: Option<ToolConfig>,
279}
280
281#[derive(Debug, Clone, Serialize, Deserialize)]
283pub struct GenerationConfig {
284 #[serde(skip_serializing_if = "Option::is_none")]
289 pub temperature: Option<f32>,
290
291 #[serde(skip_serializing_if = "Option::is_none")]
297 pub top_p: Option<f32>,
298
299 #[serde(skip_serializing_if = "Option::is_none")]
304 pub top_k: Option<i32>,
305
306 #[serde(skip_serializing_if = "Option::is_none")]
310 pub max_output_tokens: Option<i32>,
311
312 #[serde(skip_serializing_if = "Option::is_none")]
316 pub candidate_count: Option<i32>,
317
318 #[serde(skip_serializing_if = "Option::is_none")]
322 pub stop_sequences: Option<Vec<String>>,
323
324 #[serde(skip_serializing_if = "Option::is_none")]
328 pub response_mime_type: Option<String>,
329
330 #[serde(skip_serializing_if = "Option::is_none")]
334 pub response_schema: Option<serde_json::Value>,
335}
336
337impl Default for GenerationConfig {
338 fn default() -> Self {
339 Self {
340 temperature: Some(0.7),
341 top_p: Some(0.95),
342 top_k: Some(40),
343 max_output_tokens: Some(1024),
344 candidate_count: Some(1),
345 stop_sequences: None,
346 response_mime_type: None,
347 response_schema: None,
348 }
349 }
350}
351
352#[derive(Debug, Clone, Serialize, Deserialize)]
354pub struct ToolConfig {
355 #[serde(skip_serializing_if = "Option::is_none")]
357 pub function_calling_config: Option<FunctionCallingConfig>,
358}
359
360#[derive(Debug, Clone, Serialize, Deserialize)]
362pub struct FunctionCallingConfig {
363 pub mode: FunctionCallingMode,
365}
366
367#[derive(Debug, Clone, Serialize, Deserialize)]
369#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
370pub enum FunctionCallingMode {
371 Auto,
373 Any,
375 None,
377}
378
379#[derive(Debug, Clone, Serialize, Deserialize)]
381pub struct SafetySetting {
382 pub category: HarmCategory,
384 pub threshold: HarmBlockThreshold,
386}
387
388#[derive(Debug, Clone, Serialize, Deserialize)]
390#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
391pub enum HarmCategory {
392 Dangerous,
394 Harassment,
396 HateSpeech,
398 SexuallyExplicit,
400}
401
402#[derive(Debug, Clone, Serialize, Deserialize)]
404#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
405pub enum HarmBlockThreshold {
406 BlockLowAndAbove,
408 BlockMediumAndAbove,
410 BlockHighAndAbove,
412 BlockOnlyHigh,
414 BlockNone,
416}