1pub const GEMINI_3_1_FLASH_LITE_PREVIEW: &str = "gemini-3.1-flash-lite-preview";
7pub const GEMINI_3_FLASH_PREVIEW: &str = "gemini-3-flash-preview";
9pub const GEMINI_2_5_PRO_PREVIEW_06_05: &str = "gemini-2.5-pro-preview-06-05";
11pub const GEMINI_2_5_PRO_PREVIEW_05_06: &str = "gemini-2.5-pro-preview-05-06";
13pub const GEMINI_2_5_PRO_PREVIEW_03_25: &str = "gemini-2.5-pro-preview-03-25";
15pub const GEMINI_2_5_FLASH_PREVIEW_04_17: &str = "gemini-2.5-flash-preview-04-17";
17pub const GEMINI_2_5_PRO_EXP_03_25: &str = "gemini-2.5-pro-exp-03-25";
19pub const GEMINI_2_5_FLASH: &str = "gemini-2.5-flash";
21#[cfg(feature = "image")]
23#[cfg_attr(docsrs, doc(cfg(feature = "image")))]
24pub const GEMINI_2_5_FLASH_IMAGE: &str = "gemini-2.5-flash-image";
25pub const GEMINI_2_0_FLASH_LITE: &str = "gemini-2.0-flash-lite";
27pub const GEMINI_2_0_FLASH: &str = "gemini-2.0-flash";
29
30use self::gemini_api_types::tool_parameters_to_schema;
31use crate::completion::{self, CompletionError, CompletionRequest};
32use crate::http_client::HttpClientExt;
33use crate::message::{self, MimeType, Reasoning};
34use crate::providers::gemini::completion::gemini_api_types::{
35 AdditionalParameters, FunctionCallingMode, ToolConfig,
36};
37use crate::providers::internal::completion_send::send_completion;
38use crate::providers::internal::envelope::DirectPayload;
39use crate::telemetry::{CompletionOperation, CompletionSpanBuilder, SpanCombinator};
40use gemini_api_types::{
41 Content, FinishReason, FunctionDeclaration, GenerateContentRequest, GenerateContentResponse,
42 GenerationConfig, Part, PartKind, Role, Tool, map_finish_reason,
43};
44use serde_json::{Map, Value};
45use std::convert::TryFrom;
46use tracing_futures::Instrument;
47
48use super::Client;
49
50pub(crate) const PROVIDER_NAME: &str = "gcp.gemini";
59
60#[derive(Clone, Debug)]
61pub struct CompletionModel<T = reqwest::Client> {
62 pub(crate) client: Client<T>,
63 pub model: String,
64}
65
66impl<T> CompletionModel<T> {
67 pub fn new(client: Client<T>, model: impl Into<String>) -> Self {
68 Self {
69 client,
70 model: model.into(),
71 }
72 }
73
74 pub fn with_model(client: Client<T>, model: &str) -> Self {
75 Self {
76 client,
77 model: model.into(),
78 }
79 }
80}
81
82impl<T> CompletionModel<T>
83where
84 T: HttpClientExt + Clone + 'static,
85{
86 pub async fn raw_completion(
95 &self,
96 completion_request: CompletionRequest,
97 ) -> Result<GenerateContentResponse, CompletionError> {
98 let request_model = resolve_request_model(&self.model, &completion_request);
99 let span = CompletionSpanBuilder::new(
100 PROVIDER_NAME,
101 &request_model,
102 CompletionOperation::GenerateContent,
103 )
104 .system_instructions(
105 completion_request.preamble.as_deref(),
106 completion_request.record_telemetry_content,
107 )
108 .build();
109
110 let request = create_request_body(completion_request)?;
111
112 crate::providers::internal::trace_json(
113 crate::providers::internal::LogTarget::Completions,
114 "Gemini completion request",
115 &request,
116 );
117
118 let body = serde_json::to_vec(&request)?;
119
120 let path = completion_endpoint(&request_model);
121
122 let request = self
123 .client
124 .post(path.as_str())?
125 .body(body)
126 .map_err(|e| CompletionError::HttpError(e.into()))?;
127
128 send_completion::<_, DirectPayload<GenerateContentResponse>, _>(
129 &self.client,
130 request,
131 "Gemini completion",
132 None,
135 |response| {
136 let span = tracing::Span::current();
137 span.record_response_metadata(response);
138 let usage = response
139 .usage_metadata
140 .as_ref()
141 .map(crate::completion::Usage::from)
142 .unwrap_or_default();
143 span.record_token_usage(&usage);
144 },
145 )
146 .instrument(span)
147 .await
148 .map(|(payload, _)| payload)
149 }
150}
151
152impl<T> completion::CompletionModel for CompletionModel<T>
153where
154 T: HttpClientExt + Clone + 'static,
155{
156 async fn completion(
157 &self,
158 completion_request: CompletionRequest,
159 ) -> Result<completion::CompletionResponse, CompletionError> {
160 let raw = self.raw_completion(completion_request).await?;
162 let captured = serde_json::to_value(&raw)?;
163 let response: completion::CompletionResponse = raw.try_into()?;
164 Ok(response.with_raw(captured))
165 }
166
167 async fn stream(
168 &self,
169 request: CompletionRequest,
170 ) -> Result<crate::streaming::StreamingCompletionResponse, CompletionError> {
171 CompletionModel::stream(self, request).await
172 }
173}
174
175impl<T> crate::client::ConstructCompletionModel<Client<T>> for CompletionModel<T>
176where
177 Client<T>: Clone,
178{
179 fn construct(client: &Client<T>, model: String) -> Self {
180 Self::new(client.clone(), model)
181 }
182}
183
184pub(crate) fn create_request_body(
185 completion_request: CompletionRequest,
186) -> Result<GenerateContentRequest, CompletionError> {
187 let chat_history = completion_request.chat_history_with_documents();
188
189 let CompletionRequest {
190 model: _,
191 preamble,
192 chat_history: _,
193 documents: _,
194 tools: function_tools,
195 temperature,
196 max_tokens,
197 tool_choice,
198 mut additional_params,
199 output_schema,
200 record_telemetry_content: _,
201 } = completion_request;
202
203 let mut full_history = Vec::new();
204 full_history.extend(chat_history);
205 crate::providers::internal::resolve_empty_tool_result_names(&mut full_history);
208 let (history_system, full_history) = split_system_messages_from_history(full_history);
209
210 let mut additional_params_payload = additional_params
211 .take()
212 .unwrap_or_else(|| Value::Object(Map::new()));
213 let mut additional_tools =
214 extract_tools_from_additional_params(&mut additional_params_payload)?;
215
216 let AdditionalParameters {
217 mut generation_config,
218 additional_params,
219 } = serde_json::from_value::<AdditionalParameters>(additional_params_payload)?;
220
221 if let Some(schema) = output_schema {
223 let cfg = generation_config.get_or_insert_with(GenerationConfig::default);
224 cfg.response_mime_type = Some("application/json".to_string());
225 cfg.response_json_schema = Some(schema.to_value());
226 }
227
228 if temperature.is_some() || max_tokens.is_some() {
239 let cfg = generation_config.get_or_insert_with(GenerationConfig::default);
240
241 if let Some(temp) = temperature {
242 cfg.temperature = Some(temp);
243 }
244
245 if let Some(max_tokens) = max_tokens {
246 cfg.max_output_tokens = Some(max_tokens);
247 }
248 }
249
250 let mut system_parts: Vec<Part> = Vec::new();
251 if let Some(preamble) = preamble.filter(|preamble| !preamble.is_empty()) {
252 system_parts.push(preamble.into());
253 }
254 for content in history_system {
255 if !content.is_empty() {
256 system_parts.push(content.into());
257 }
258 }
259 let system_instruction = if system_parts.is_empty() {
260 None
261 } else {
262 Some(Content {
263 parts: system_parts,
264 role: Some(Role::Model),
265 })
266 };
267
268 let mut tools = if function_tools.is_empty() {
269 Vec::new()
270 } else {
271 vec![serde_json::to_value(Tool::try_from(function_tools)?)?]
272 };
273 tools.append(&mut additional_tools);
274 let tools = if tools.is_empty() { None } else { Some(tools) };
275
276 let tool_config = if let Some(cfg) = tool_choice {
277 Some(ToolConfig {
278 function_calling_config: Some(FunctionCallingMode::try_from(cfg)?),
279 })
280 } else {
281 None
282 };
283
284 let request = GenerateContentRequest {
285 contents: full_history
286 .into_iter()
287 .map(|msg| {
288 msg.try_into()
289 .map_err(|e| CompletionError::RequestError(Box::new(e)))
290 })
291 .collect::<Result<Vec<_>, _>>()?,
292 generation_config,
293 safety_settings: None,
294 tools,
295 tool_config,
296 system_instruction,
297 additional_params,
298 };
299
300 Ok(request)
301}
302
303pub fn split_system_messages_from_history(
306 history: Vec<completion::Message>,
307) -> (Vec<String>, Vec<completion::Message>) {
308 let mut system = Vec::new();
309 let mut remaining = Vec::new();
310
311 for message in history {
312 match message {
313 completion::Message::System { content } => system.push(content),
314 other => remaining.push(other),
315 }
316 }
317
318 (system, remaining)
319}
320
321fn extract_tools_from_additional_params(
322 additional_params: &mut Value,
323) -> Result<Vec<Value>, CompletionError> {
324 if let Some(map) = additional_params.as_object_mut()
325 && let Some(raw_tools) = map.remove("tools")
326 {
327 return serde_json::from_value::<Vec<Value>>(raw_tools).map_err(|err| {
328 CompletionError::RequestError(
329 format!("Invalid Gemini `additional_params.tools` payload: {err}").into(),
330 )
331 });
332 }
333
334 Ok(Vec::new())
335}
336
337pub(crate) fn resolve_request_model(
338 default_model: &str,
339 completion_request: &CompletionRequest,
340) -> String {
341 completion_request
342 .model
343 .clone()
344 .unwrap_or_else(|| default_model.to_string())
345}
346
347pub(crate) fn completion_endpoint(model: &str) -> String {
348 format!("/v1beta/models/{model}:generateContent")
349}
350
351pub(crate) fn streaming_endpoint(model: &str) -> String {
352 format!("/v1beta/models/{model}:streamGenerateContent")
353}
354
355impl TryFrom<completion::ToolDefinition> for Tool {
356 type Error = CompletionError;
357
358 fn try_from(tool: completion::ToolDefinition) -> Result<Self, Self::Error> {
359 let parameters = tool_parameters_to_schema(tool.parameters)?;
360
361 Ok(Self {
362 function_declarations: vec![FunctionDeclaration {
363 name: tool.name,
364 description: tool.description,
365 parameters,
366 }],
367 code_execution: None,
368 })
369 }
370}
371
372impl TryFrom<Vec<completion::ToolDefinition>> for Tool {
373 type Error = CompletionError;
374
375 fn try_from(tools: Vec<completion::ToolDefinition>) -> Result<Self, Self::Error> {
376 let mut function_declarations = Vec::new();
377
378 for tool in tools {
379 let parameters = tool_parameters_to_schema(tool.parameters).map_err(|e| {
380 CompletionError::ProviderError(format!(
381 "Tool '{}' could not be converted to a schema: {:?}",
382 tool.name, e,
383 ))
384 })?;
385
386 function_declarations.push(FunctionDeclaration {
387 name: tool.name,
388 description: tool.description,
389 parameters,
390 });
391 }
392
393 Ok(Self {
394 function_declarations,
395 code_execution: None,
396 })
397 }
398}
399
400pub(crate) fn function_call_finish_reason_error(
401 reason: &FinishReason,
402 finish_message: Option<&str>,
403) -> Option<CompletionError> {
404 match reason {
405 FinishReason::MalformedFunctionCall
406 | FinishReason::UnexpectedToolCall
407 | FinishReason::MissingThoughtSignature
408 | FinishReason::TooManyToolCalls
409 | FinishReason::MalformedResponse => {
410 let message = finish_message.unwrap_or("no finish message provided");
411 Some(CompletionError::ResponseError(format!(
412 "Gemini stopped with finish_reason={reason:?}: {message}"
413 )))
414 }
415 _ => None,
416 }
417}
418
419fn map_response_part(part: &Part) -> Result<Vec<completion::AssistantContent>, CompletionError> {
428 let Part {
429 thought,
430 thought_signature,
431 part,
432 ..
433 } = part;
434
435 Ok(vec![match part {
436 PartKind::Text(text) => {
437 if let Some(thought) = thought
438 && *thought
439 {
440 completion::AssistantContent::Reasoning(Reasoning::new_with_signature(
441 text,
442 thought_signature.clone(),
443 ))
444 } else if thought_signature.is_some() {
445 return Ok(vec![completion::AssistantContent::text(text)]);
449 } else {
450 completion::AssistantContent::text(text)
451 }
452 }
453 PartKind::InlineData(inline_data) => {
454 let mime_type = message::MediaType::from_mime_type(&inline_data.mime_type);
455
456 match mime_type {
457 Some(message::MediaType::Image(media_type)) => {
458 message::AssistantContent::image_base64(
459 &inline_data.data,
460 Some(media_type),
461 Some(message::ImageDetail::default()),
462 )
463 }
464 _ => {
465 return Err(CompletionError::ResponseError(format!(
466 "Unsupported media type {mime_type:?}"
467 )));
468 }
469 }
470 }
471 PartKind::FunctionCall(function_call) => {
472 let tool_call = message::ToolCall::from_wire(
473 function_call.id.clone().unwrap_or_default(),
474 message::ToolFunction::new(function_call.name.clone(), function_call.args.clone()),
475 )
476 .with_signature(thought_signature.clone());
477 completion::AssistantContent::ToolCall(tool_call)
478 }
479 PartKind::ExecutableCode(_) | PartKind::CodeExecutionResult(_) => return Ok(Vec::new()),
493 other => {
494 return Err(CompletionError::ResponseError(format!(
495 "Gemini response part kind {} carries no assistant content rig can account for",
496 part_kind_name(other)
497 )));
498 }
499 }])
500}
501
502pub fn attach_trailing_signature(
524 content: &mut Vec<completion::AssistantContent>,
525 signature: String,
526) {
527 let unsigned_reasoning = content.iter_mut().rev().find_map(|item| match item {
528 completion::AssistantContent::Reasoning(reasoning) => match reasoning.content.first_mut() {
529 Some(message::ReasoningContent::Text {
530 signature: slot @ None,
531 ..
532 }) => Some(slot),
533 _ => None,
534 },
535 _ => None,
536 });
537
538 match unsigned_reasoning {
539 Some(slot) => *slot = Some(signature),
540 None => content.push(completion::AssistantContent::Reasoning(
541 Reasoning::new_with_signature("", Some(signature)),
542 )),
543 }
544}
545
546fn part_kind_name(part: &PartKind) -> &'static str {
548 match part {
549 PartKind::Text(_) => "text",
550 PartKind::InlineData(_) => "inlineData",
551 PartKind::FunctionCall(_) => "functionCall",
552 PartKind::FunctionResponse(_) => "functionResponse",
553 PartKind::FileData(_) => "fileData",
554 PartKind::ExecutableCode(_) => "executableCode",
555 PartKind::CodeExecutionResult(_) => "codeExecutionResult",
556 }
557}
558
559impl TryFrom<GenerateContentResponse> for completion::CompletionResponse {
561 type Error = CompletionError;
562
563 fn try_from(response: GenerateContentResponse) -> Result<Self, Self::Error> {
564 let candidate = response.candidates.first().ok_or_else(|| {
565 CompletionError::ResponseError("No response candidates in response".into())
566 })?;
567
568 if let Some(reason) = candidate.finish_reason.as_ref()
569 && let Some(err) =
570 function_call_finish_reason_error(reason, candidate.finish_message.as_deref())
571 {
572 return Err(err);
573 }
574
575 let finish_reason = candidate.finish_reason.as_ref().and_then(map_finish_reason);
576
577 let parts = &candidate
578 .content
579 .as_ref()
580 .ok_or_else(|| {
581 let reason = candidate
582 .finish_reason
583 .as_ref()
584 .map(|r| format!("finish_reason={r:?}"))
585 .unwrap_or_else(|| "finish_reason=<unknown>".to_string());
586 let message = candidate
587 .finish_message
588 .as_deref()
589 .unwrap_or("no finish message provided");
590 CompletionError::ResponseError(format!(
591 "Gemini candidate missing content ({reason}, finish_message={message})"
592 ))
593 })?
594 .parts;
595
596 let mut content: Vec<completion::AssistantContent> = Vec::with_capacity(parts.len());
602 for part in parts {
603 content.extend(map_response_part(part)?);
604 if !part.thought.unwrap_or(false)
605 && matches!(part.part, PartKind::Text(_))
606 && let Some(signature) = part.thought_signature.clone()
607 {
608 attach_trailing_signature(&mut content, signature);
609 }
610 }
611
612 let choice = crate::message::require_non_empty_response(content)?;
613
614 let usage = response
615 .usage_metadata
616 .as_ref()
617 .map(crate::completion::Usage::from)
618 .unwrap_or_default();
619
620 Ok(
621 completion::CompletionResponse::new(choice, usage, PROVIDER_NAME)
622 .with_optional_response_id(
623 Some(response.response_id.as_str()).filter(|id| !id.is_empty()),
624 )
625 .with_optional_model(response.model_version.as_deref())
626 .with_optional_finish_reason(finish_reason),
627 )
628 }
629}
630
631pub mod gemini_api_types {
632 use crate::telemetry::ProviderResponseExt;
633 use std::{collections::HashMap, convert::Infallible, str::FromStr};
634
635 use serde::{Deserialize, Serialize};
639 use serde_json::{Value, json};
640
641 use crate::message::{DocumentSourceKind, ImageMediaType, MessageError, MimeType};
642 use crate::{
643 completion::CompletionError,
644 message::{self},
645 providers::gemini::gemini_api_types::{CodeExecutionResult, ExecutableCode},
646 };
647
648 #[derive(Debug, Deserialize, Serialize, Default)]
649 #[serde(rename_all = "camelCase")]
650 pub struct AdditionalParameters {
651 pub generation_config: Option<GenerationConfig>,
653 #[serde(flatten, skip_serializing_if = "Option::is_none")]
655 pub additional_params: Option<serde_json::Value>,
656 }
657
658 impl AdditionalParameters {
659 pub fn with_config(mut self, cfg: GenerationConfig) -> Self {
660 self.generation_config = Some(cfg);
661 self
662 }
663
664 pub fn with_params(mut self, params: serde_json::Value) -> Self {
665 self.additional_params = Some(params);
666 self
667 }
668 }
669
670 #[derive(Debug, Deserialize, Serialize)]
678 #[serde(rename_all = "camelCase")]
679 pub struct GenerateContentResponse {
680 #[serde(default)]
681 pub response_id: String,
682 #[serde(default)]
684 pub candidates: Vec<ContentCandidate>,
685 pub prompt_feedback: Option<PromptFeedback>,
687 pub usage_metadata: Option<UsageMetadata>,
689 pub model_version: Option<String>,
690 }
691
692 impl ProviderResponseExt for GenerateContentResponse {
693 type Usage = UsageMetadata;
694
695 fn get_response_id(&self) -> Option<String> {
696 Some(self.response_id.clone())
697 }
698
699 fn get_response_model_name(&self) -> Option<String> {
700 self.model_version.clone()
701 }
702
703 fn get_text_response(&self) -> Option<String> {
704 let str = self
705 .candidates
706 .iter()
707 .filter_map(|x| {
708 let content = x.content.as_ref()?;
709 if content.role.as_ref().is_none_or(|y| y != &Role::Model) {
710 return None;
711 }
712
713 Some(visible_text_parts(content).collect::<Vec<_>>().join("\n"))
714 })
715 .collect::<Vec<String>>()
716 .join("\n");
717
718 if str.is_empty() { None } else { Some(str) }
719 }
720
721 fn get_usage(&self) -> Option<Self::Usage> {
722 self.usage_metadata.clone()
723 }
724 }
725
726 pub(crate) fn visible_text_parts(content: &Content) -> impl Iterator<Item = &str> {
741 content.parts.iter().filter_map(|part| match &part.part {
742 PartKind::Text(text) if !part.thought.unwrap_or(false) => Some(text.as_str()),
743 _ => None,
744 })
745 }
746
747 #[derive(Clone, Debug, Deserialize, Serialize)]
749 #[serde(rename_all = "camelCase")]
750 pub struct ContentCandidate {
751 #[serde(skip_serializing_if = "Option::is_none")]
753 pub content: Option<Content>,
754 pub finish_reason: Option<FinishReason>,
757 pub safety_ratings: Option<Vec<SafetyRating>>,
760 pub citation_metadata: Option<CitationMetadata>,
764 pub token_count: Option<i32>,
766 pub avg_logprobs: Option<f64>,
768 pub logprobs_result: Option<LogprobsResult>,
770 pub index: Option<i32>,
772 pub finish_message: Option<String>,
774 }
775
776 #[derive(Clone, Debug, Deserialize, Serialize)]
777 pub struct Content {
778 #[serde(default)]
780 pub parts: Vec<Part>,
781 pub role: Option<Role>,
784 }
785
786 impl TryFrom<message::Message> for Content {
787 type Error = message::MessageError;
788
789 fn try_from(msg: message::Message) -> Result<Self, Self::Error> {
790 Ok(match msg {
791 message::Message::System { content } => Content {
792 parts: vec![content.into()],
793 role: Some(Role::User),
794 },
795 message::Message::User { content } => Content {
796 parts: content
797 .into_iter()
798 .map(|c| c.try_into())
799 .collect::<Result<Vec<_>, _>>()?,
800 role: Some(Role::User),
801 },
802 message::Message::Assistant { content, .. } => Content {
803 role: Some(Role::Model),
804 parts: content
805 .into_iter()
806 .map(|content| content.try_into())
807 .collect::<Result<Vec<_>, _>>()?,
808 },
809 })
810 }
811 }
812
813 #[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
814 #[serde(rename_all = "lowercase")]
815 pub enum Role {
816 User,
817 Model,
818 }
819
820 #[derive(Debug, Default, Deserialize, Serialize, Clone, PartialEq)]
821 #[serde(rename_all = "camelCase")]
822 pub struct Part {
823 #[serde(skip_serializing_if = "Option::is_none")]
825 pub thought: Option<bool>,
826 #[serde(skip_serializing_if = "Option::is_none")]
828 pub thought_signature: Option<String>,
829 #[serde(flatten)]
830 pub part: PartKind,
831 #[serde(flatten, skip_serializing_if = "Option::is_none")]
832 pub additional_params: Option<Value>,
833 }
834
835 #[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
839 #[serde(rename_all = "camelCase")]
840 pub enum PartKind {
841 Text(String),
842 InlineData(Blob),
843 FunctionCall(FunctionCall),
844 FunctionResponse(FunctionResponse),
845 FileData(FileData),
846 ExecutableCode(ExecutableCode),
847 CodeExecutionResult(CodeExecutionResult),
848 }
849
850 impl Default for PartKind {
853 fn default() -> Self {
854 Self::Text(String::new())
855 }
856 }
857
858 impl From<String> for Part {
859 fn from(text: String) -> Self {
860 Self {
861 thought: Some(false),
862 thought_signature: None,
863 part: PartKind::Text(text),
864 additional_params: None,
865 }
866 }
867 }
868
869 impl From<&str> for Part {
870 fn from(text: &str) -> Self {
871 Self::from(text.to_string())
872 }
873 }
874
875 impl FromStr for Part {
876 type Err = Infallible;
877
878 fn from_str(s: &str) -> Result<Self, Self::Err> {
879 Ok(s.into())
880 }
881 }
882
883 fn media_source_to_part_kind(
893 kind: &str,
894 mime_type: String,
895 source: DocumentSourceKind,
896 string_is_data: bool,
897 ) -> Result<PartKind, message::MessageError> {
898 match source {
899 DocumentSourceKind::Url(file_uri) => Ok(PartKind::FileData(FileData {
900 mime_type: Some(mime_type),
901 file_uri,
902 })),
903 DocumentSourceKind::Base64(data) => Ok(PartKind::InlineData(Blob { mime_type, data })),
904 DocumentSourceKind::String(data) if string_is_data => {
905 Ok(PartKind::InlineData(Blob { mime_type, data }))
906 }
907 DocumentSourceKind::String(_) => Err(message::MessageError::ConversionError(format!(
908 "Strings cannot be used as Gemini {kind} inputs"
909 ))),
910 DocumentSourceKind::Raw(_) => Err(message::MessageError::ConversionError(
911 "Raw files not supported, encode as base64 first".to_string(),
912 )),
913 DocumentSourceKind::FileId(_) => Err(message::MessageError::ConversionError(format!(
914 "Provider file IDs are not supported for Gemini {kind} inputs"
915 ))),
916 DocumentSourceKind::Unknown => Err(message::MessageError::ConversionError(format!(
917 "Gemini {kind} input has no body"
918 ))),
919 }
920 }
921
922 impl TryFrom<(ImageMediaType, DocumentSourceKind)> for PartKind {
923 type Error = message::MessageError;
924 fn try_from(
925 (mime_type, doc_src): (ImageMediaType, DocumentSourceKind),
926 ) -> Result<Self, Self::Error> {
927 media_source_to_part_kind("image", mime_type.to_mime_type().to_string(), doc_src, true)
928 }
929 }
930
931 fn image_to_part(image: message::Image) -> Result<Part, message::MessageError> {
936 let message::Image {
937 data, media_type, ..
938 } = image;
939
940 let Some(media_type) = media_type else {
941 return Err(message::MessageError::ConversionError(
942 "Media type for image is required for Gemini".to_string(),
943 ));
944 };
945
946 match media_type {
947 message::ImageMediaType::JPEG
948 | message::ImageMediaType::PNG
949 | message::ImageMediaType::WEBP
950 | message::ImageMediaType::HEIC
951 | message::ImageMediaType::HEIF => Ok(Part {
952 thought: Some(false),
953 thought_signature: None,
954 part: PartKind::try_from((media_type, data))?,
955 additional_params: None,
956 }),
957 _ => Err(message::MessageError::ConversionError(format!(
958 "Unsupported image media type {media_type:?}"
959 ))),
960 }
961 }
962
963 fn gemini_tool_result_image_mime_type(
964 media_type: Option<&ImageMediaType>,
965 ) -> Result<&'static str, MessageError> {
966 let media_type = media_type.ok_or_else(|| {
967 MessageError::ConversionError(
968 "Image media type is required for Gemini tool results".to_string(),
969 )
970 })?;
971
972 match media_type {
973 ImageMediaType::JPEG | ImageMediaType::PNG | ImageMediaType::WEBP => {
974 Ok(media_type.to_mime_type())
975 }
976 _ => Err(MessageError::ConversionError(format!(
977 "Unsupported image media type {media_type:?} for Gemini tool results; supported types are JPEG, PNG, and WEBP"
978 ))),
979 }
980 }
981
982 impl TryFrom<message::UserContent> for Part {
983 type Error = message::MessageError;
984
985 fn try_from(content: message::UserContent) -> Result<Self, Self::Error> {
986 match content {
987 message::UserContent::Text(message::Text { text, .. }) => Ok(Part {
988 thought: Some(false),
989 thought_signature: None,
990 part: PartKind::Text(text),
991 additional_params: None,
992 }),
993 message::UserContent::ToolResult(message::ToolResult {
994 call: _,
995 provider,
996 name,
997 content,
998 }) => {
999 let function_name = name;
1001 let mut response_values = Vec::new();
1002 let mut parts: Vec<FunctionResponsePart> = Vec::new();
1003
1004 for item in content.iter() {
1005 match item {
1006 message::ToolResultContent::Text(text) => {
1007 response_values.push(json!(&text.text));
1008 }
1009 message::ToolResultContent::Json { value } => {
1010 response_values.push(value.clone());
1011 }
1012 message::ToolResultContent::Image(image) => {
1013 let part = match &image.data {
1014 DocumentSourceKind::Base64(b64) => {
1015 let mime_type = gemini_tool_result_image_mime_type(
1016 image.media_type.as_ref(),
1017 )?;
1018
1019 FunctionResponsePart {
1026 inline_data: Some(FunctionResponseInlineData {
1027 mime_type: mime_type.to_string(),
1028 data: b64.clone(),
1029 display_name: None,
1030 }),
1031 file_data: None,
1032 }
1033 }
1034 DocumentSourceKind::Url(_) => {
1035 return Err(message::MessageError::ConversionError(
1036 "Gemini tool result images must use base64 inline data; URL-backed images are not supported"
1037 .to_string(),
1038 ));
1039 }
1040 _ => {
1041 return Err(message::MessageError::ConversionError(
1042 "Unsupported image source kind for tool results"
1043 .to_string(),
1044 ));
1045 }
1046 };
1047 parts.push(part);
1048 }
1049 }
1050 }
1051
1052 let response_json = if response_values.is_empty() {
1053 None
1054 } else {
1055 let result = if response_values.len() == 1 {
1056 response_values.remove(0)
1057 } else {
1058 serde_json::Value::Array(response_values)
1059 };
1060 Some(json!({ "result": result }))
1061 };
1062
1063 Ok(Part {
1064 thought: Some(false),
1065 thought_signature: None,
1066 part: PartKind::FunctionResponse(FunctionResponse {
1067 name: function_name,
1068 id: provider.map(|provider| provider.call_id),
1069 response: response_json,
1070 parts: if parts.is_empty() { None } else { Some(parts) },
1071 }),
1072 additional_params: None,
1073 })
1074 }
1075 message::UserContent::Image(image) => image_to_part(image),
1076 message::UserContent::Document(message::Document {
1077 data, media_type, ..
1078 }) => {
1079 let Some(media_type) = media_type else {
1080 return Err(MessageError::ConversionError(
1081 "A mime type is required for document inputs to Gemini".to_string(),
1082 ));
1083 };
1084
1085 if matches!(
1088 media_type,
1089 message::DocumentMediaType::TXT
1090 | message::DocumentMediaType::RTF
1091 | message::DocumentMediaType::HTML
1092 | message::DocumentMediaType::CSS
1093 | message::DocumentMediaType::MARKDOWN
1094 | message::DocumentMediaType::CSV
1095 | message::DocumentMediaType::XML
1096 | message::DocumentMediaType::Javascript
1097 | message::DocumentMediaType::Python
1098 ) {
1099 use base64::Engine;
1100 let part = match data {
1101 DocumentSourceKind::String(text) => PartKind::Text(text),
1102 DocumentSourceKind::Base64(data) => {
1103 let text = String::from_utf8(
1105 base64::engine::general_purpose::STANDARD
1106 .decode(&data)
1107 .map_err(|e| {
1108 MessageError::ConversionError(format!(
1109 "Failed to decode base64: {e}"
1110 ))
1111 })?,
1112 )
1113 .map_err(|e| {
1114 MessageError::ConversionError(format!(
1115 "Invalid UTF-8 in document: {e}"
1116 ))
1117 })?;
1118 PartKind::Text(text)
1119 }
1120 DocumentSourceKind::Url(file_uri) => PartKind::FileData(FileData {
1121 mime_type: Some(media_type.to_mime_type().to_string()),
1122 file_uri,
1123 }),
1124 DocumentSourceKind::Raw(_) => {
1125 return Err(MessageError::ConversionError(
1126 "Raw files not supported, encode as base64 first".to_string(),
1127 ));
1128 }
1129 DocumentSourceKind::FileId(_) => {
1130 return Err(MessageError::ConversionError(
1131 "Provider file IDs are not supported for Gemini documents"
1132 .to_string(),
1133 ));
1134 }
1135 DocumentSourceKind::Unknown => {
1136 return Err(MessageError::ConversionError(
1137 "Document has no body".to_string(),
1138 ));
1139 }
1140 };
1141
1142 Ok(Part {
1143 thought: Some(false),
1144 part,
1145 ..Default::default()
1146 })
1147 } else if !media_type.is_code() {
1148 let part = media_source_to_part_kind(
1149 "document",
1150 media_type.to_mime_type().to_string(),
1151 data,
1152 true,
1153 )?;
1154
1155 Ok(Part {
1156 thought: Some(false),
1157 part,
1158 ..Default::default()
1159 })
1160 } else {
1161 Err(message::MessageError::ConversionError(format!(
1162 "Unsupported document media type {media_type:?}"
1163 )))
1164 }
1165 }
1166
1167 message::UserContent::Audio(message::Audio {
1168 data, media_type, ..
1169 }) => {
1170 let Some(media_type) = media_type else {
1171 return Err(MessageError::ConversionError(
1172 "A mime type is required for audio inputs to Gemini".to_string(),
1173 ));
1174 };
1175
1176 let part = media_source_to_part_kind(
1177 "audio",
1178 media_type.to_mime_type().to_string(),
1179 data,
1180 false,
1181 )?;
1182
1183 Ok(Part {
1184 thought: Some(false),
1185 part,
1186 ..Default::default()
1187 })
1188 }
1189 message::UserContent::Video(message::Video {
1190 data,
1191 media_type,
1192 additional_params,
1193 ..
1194 }) => {
1195 let mime_type = media_type.map(|media_ty| media_ty.to_mime_type().to_string());
1196
1197 let part = match data {
1198 DocumentSourceKind::Url(file_uri)
1202 if file_uri.starts_with("https://www.youtube.com") =>
1203 {
1204 PartKind::FileData(FileData {
1205 mime_type,
1206 file_uri,
1207 })
1208 }
1209 data => {
1210 let mime_type = mime_type.ok_or_else(|| {
1211 MessageError::ConversionError(
1212 "A mime type is required for non-Youtube video inputs to Gemini"
1213 .to_string(),
1214 )
1215 })?;
1216
1217 media_source_to_part_kind("video", mime_type, data, false)?
1218 }
1219 };
1220
1221 Ok(Part {
1222 thought: Some(false),
1223 thought_signature: None,
1224 part,
1225 additional_params: additional_params.map(Into::into),
1226 })
1227 }
1228 }
1229 }
1230 }
1231
1232 impl TryFrom<message::AssistantContent> for Part {
1233 type Error = message::MessageError;
1234
1235 fn try_from(content: message::AssistantContent) -> Result<Self, Self::Error> {
1236 match content {
1237 message::AssistantContent::Text(message::Text { text, .. }) => Ok(text.into()),
1238 message::AssistantContent::Image(image) => image_to_part(image),
1239 message::AssistantContent::ToolCall(tool_call) => Ok(tool_call.into()),
1240 message::AssistantContent::Reasoning(reasoning) => Ok(Part {
1241 thought: Some(true),
1242 thought_signature: reasoning.first_signature().map(str::to_owned),
1243 part: PartKind::Text(reasoning.display_text()),
1244 additional_params: None,
1245 }),
1246 }
1247 }
1248 }
1249
1250 impl From<message::ToolCall> for Part {
1251 fn from(tool_call: message::ToolCall) -> Self {
1252 Self {
1253 thought: Some(false),
1254 thought_signature: tool_call.signature,
1255 part: PartKind::FunctionCall(FunctionCall {
1256 name: tool_call.function.name,
1257 args: tool_call.function.arguments,
1258 id: tool_call.provider.map(|provider| provider.call_id),
1261 }),
1262 additional_params: None,
1263 }
1264 }
1265 }
1266
1267 #[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
1270 #[serde(rename_all = "camelCase")]
1271 pub struct Blob {
1272 pub mime_type: String,
1275 pub data: String,
1277 }
1278
1279 #[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
1282 pub struct FunctionCall {
1283 pub name: String,
1286 pub args: serde_json::Value,
1288 #[serde(skip_serializing_if = "Option::is_none")]
1290 pub id: Option<String>,
1291 }
1292
1293 impl From<message::ToolCall> for FunctionCall {
1294 fn from(tool_call: message::ToolCall) -> Self {
1295 Self {
1296 name: tool_call.function.name,
1297 args: tool_call.function.arguments,
1298 id: tool_call.provider.map(|provider| provider.call_id),
1299 }
1300 }
1301 }
1302
1303 #[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
1307 pub struct FunctionResponse {
1308 pub name: String,
1311 #[serde(skip_serializing_if = "Option::is_none")]
1313 pub id: Option<String>,
1314 #[serde(skip_serializing_if = "Option::is_none")]
1316 pub response: Option<serde_json::Value>,
1317 #[serde(skip_serializing_if = "Option::is_none")]
1319 pub parts: Option<Vec<FunctionResponsePart>>,
1320 }
1321
1322 #[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
1324 #[serde(rename_all = "camelCase")]
1325 pub struct FunctionResponsePart {
1326 #[serde(skip_serializing_if = "Option::is_none")]
1328 pub inline_data: Option<FunctionResponseInlineData>,
1329 #[serde(skip_serializing_if = "Option::is_none")]
1331 pub file_data: Option<FileData>,
1332 }
1333
1334 #[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
1336 #[serde(rename_all = "camelCase")]
1337 pub struct FunctionResponseInlineData {
1338 pub mime_type: String,
1340 pub data: String,
1342 #[serde(skip_serializing_if = "Option::is_none")]
1344 pub display_name: Option<String>,
1345 }
1346
1347 #[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
1349 #[serde(rename_all = "camelCase")]
1350 pub struct FileData {
1351 pub mime_type: Option<String>,
1353 pub file_uri: String,
1355 }
1356
1357 #[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
1358 pub struct SafetyRating {
1359 pub category: HarmCategory,
1360 pub probability: HarmProbability,
1361 }
1362
1363 #[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
1364 #[serde(rename_all = "SCREAMING_SNAKE_CASE")]
1365 pub enum HarmProbability {
1366 HarmProbabilityUnspecified,
1367 Negligible,
1368 Low,
1369 Medium,
1370 High,
1371 }
1372
1373 #[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
1374 #[serde(rename_all = "SCREAMING_SNAKE_CASE")]
1375 pub enum HarmCategory {
1376 HarmCategoryUnspecified,
1377 HarmCategoryDerogatory,
1378 HarmCategoryToxicity,
1379 HarmCategoryViolence,
1380 HarmCategorySexually,
1381 HarmCategoryMedical,
1382 HarmCategoryDangerous,
1383 HarmCategoryHarassment,
1384 HarmCategoryHateSpeech,
1385 HarmCategorySexuallyExplicit,
1386 HarmCategoryDangerousContent,
1387 HarmCategoryCivicIntegrity,
1388 }
1389
1390 #[derive(Debug, Deserialize, Clone, Default, Serialize)]
1391 #[serde(rename_all = "camelCase")]
1392 pub struct UsageMetadata {
1393 #[serde(default)]
1394 pub prompt_token_count: i32,
1395 #[serde(skip_serializing_if = "Option::is_none")]
1396 pub cached_content_token_count: Option<i32>,
1397 #[serde(skip_serializing_if = "Option::is_none")]
1398 pub candidates_token_count: Option<i32>,
1399 #[serde(default)]
1400 pub total_token_count: i32,
1401 #[serde(skip_serializing_if = "Option::is_none")]
1402 pub thoughts_token_count: Option<i32>,
1403 #[serde(default, skip_serializing_if = "Option::is_none")]
1404 pub prompt_tokens_details: Option<Vec<ModalityTokenCount>>,
1405 #[serde(default, skip_serializing_if = "Option::is_none")]
1406 pub cache_tokens_details: Option<Vec<ModalityTokenCount>>,
1407 #[serde(default, skip_serializing_if = "Option::is_none")]
1408 pub candidates_tokens_details: Option<Vec<ModalityTokenCount>>,
1409 #[serde(default, skip_serializing_if = "Option::is_none")]
1410 pub tool_use_prompt_token_count: Option<i32>,
1411 #[serde(default, skip_serializing_if = "Option::is_none")]
1412 pub tool_use_prompt_tokens_details: Option<Vec<ModalityTokenCount>>,
1413 #[serde(default, skip_serializing_if = "Option::is_none")]
1414 pub traffic_type: Option<TrafficType>,
1415 }
1416
1417 #[derive(Clone, Debug, Deserialize, Serialize)]
1418 #[serde(rename_all = "camelCase")]
1419 pub struct ModalityTokenCount {
1420 pub modality: Modality,
1421 #[serde(default)]
1422 pub token_count: i32,
1423 }
1424
1425 #[derive(Clone, Debug, Deserialize, Serialize)]
1426 #[serde(rename_all = "SCREAMING_SNAKE_CASE")]
1427 pub enum Modality {
1428 ModalityUnspecified,
1429 Text,
1430 Image,
1431 Video,
1432 Audio,
1433 Document,
1434 }
1435
1436 #[derive(Clone, Debug, Deserialize, Serialize)]
1437 #[serde(rename_all = "SCREAMING_SNAKE_CASE")]
1438 pub enum TrafficType {
1439 TrafficTypeUnspecified,
1440 OnDemand,
1441 ProvisionedThroughput,
1442 }
1443
1444 impl std::fmt::Display for UsageMetadata {
1445 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1446 write!(
1447 f,
1448 "Prompt token count: {}\nCached content token count: {}\nCandidates token count: {}\nTotal token count: {}",
1449 self.prompt_token_count,
1450 match self.cached_content_token_count {
1451 Some(count) => count.to_string(),
1452 None => "n/a".to_string(),
1453 },
1454 match self.candidates_token_count {
1455 Some(count) => count.to_string(),
1456 None => "n/a".to_string(),
1457 },
1458 self.total_token_count
1459 )
1460 }
1461 }
1462
1463 impl From<&UsageMetadata> for crate::completion::Usage {
1464 fn from(value: &UsageMetadata) -> crate::completion::Usage {
1465 let mut usage = crate::completion::Usage::new();
1466
1467 usage.input_tokens = value.prompt_token_count as u64;
1468 usage.output_tokens = value.candidates_token_count.unwrap_or_default() as u64;
1469 usage.cached_input_tokens = value.cached_content_token_count.unwrap_or_default() as u64;
1470 usage.reasoning_tokens = value.thoughts_token_count.unwrap_or_default() as u64;
1471 usage.tool_use_prompt_tokens =
1472 value.tool_use_prompt_token_count.unwrap_or_default() as u64;
1473 usage.total_tokens = value.total_token_count as u64;
1474
1475 usage
1476 }
1477 }
1478
1479 impl From<UsageMetadata> for crate::completion::Usage {
1480 fn from(value: UsageMetadata) -> crate::completion::Usage {
1481 (&value).into()
1482 }
1483 }
1484
1485 #[derive(Debug, Deserialize, Serialize)]
1487 #[serde(rename_all = "camelCase")]
1488 pub struct PromptFeedback {
1489 pub block_reason: Option<BlockReason>,
1491 pub safety_ratings: Option<Vec<SafetyRating>>,
1493 }
1494
1495 #[derive(Debug, Deserialize, Serialize)]
1497 #[serde(rename_all = "SCREAMING_SNAKE_CASE")]
1498 pub enum BlockReason {
1499 BlockReasonUnspecified,
1501 Safety,
1503 Other,
1505 Blocklist,
1507 ProhibitedContent,
1509 #[serde(untagged)]
1513 Unknown(String),
1514 }
1515
1516 #[derive(Clone, Debug, Deserialize, Serialize)]
1517 #[serde(rename_all = "SCREAMING_SNAKE_CASE")]
1518 pub enum FinishReason {
1519 FinishReasonUnspecified,
1521 Stop,
1523 MaxTokens,
1525 Safety,
1527 Recitation,
1529 Language,
1531 Other,
1533 Blocklist,
1535 ProhibitedContent,
1537 Spii,
1539 MalformedFunctionCall,
1541 UnexpectedToolCall,
1543 MissingThoughtSignature,
1545 TooManyToolCalls,
1547 MalformedResponse,
1549 #[serde(untagged)]
1554 Unknown(String),
1555 }
1556
1557 impl FinishReason {
1558 pub fn as_wire_str(&self) -> &str {
1564 match self {
1565 Self::FinishReasonUnspecified => "FINISH_REASON_UNSPECIFIED",
1566 Self::Stop => "STOP",
1567 Self::MaxTokens => "MAX_TOKENS",
1568 Self::Safety => "SAFETY",
1569 Self::Recitation => "RECITATION",
1570 Self::Language => "LANGUAGE",
1571 Self::Other => "OTHER",
1572 Self::Blocklist => "BLOCKLIST",
1573 Self::ProhibitedContent => "PROHIBITED_CONTENT",
1574 Self::Spii => "SPII",
1575 Self::MalformedFunctionCall => "MALFORMED_FUNCTION_CALL",
1576 Self::UnexpectedToolCall => "UNEXPECTED_TOOL_CALL",
1577 Self::MissingThoughtSignature => "MISSING_THOUGHT_SIGNATURE",
1578 Self::TooManyToolCalls => "TOO_MANY_TOOL_CALLS",
1579 Self::MalformedResponse => "MALFORMED_RESPONSE",
1580 Self::Unknown(reason) => reason,
1581 }
1582 }
1583 }
1584
1585 pub fn map_google_finish_reason(wire_name: &str) -> Option<crate::completion::FinishReason> {
1599 Some(match wire_name {
1600 "FINISH_REASON_UNSPECIFIED" => return None,
1601 "STOP" => crate::completion::FinishReason::Stop,
1602 "MAX_TOKENS" => crate::completion::FinishReason::Length,
1603 "SAFETY" | "BLOCKLIST" | "PROHIBITED_CONTENT" | "SPII" => {
1604 crate::completion::FinishReason::ContentFilter
1605 }
1606 other => crate::completion::FinishReason::Other(other.to_owned()),
1607 })
1608 }
1609
1610 pub(crate) fn map_finish_reason(
1614 reason: &FinishReason,
1615 ) -> Option<crate::completion::FinishReason> {
1616 map_google_finish_reason(reason.as_wire_str())
1617 }
1618
1619 #[derive(Clone, Debug, Deserialize, Serialize)]
1620 #[serde(rename_all = "camelCase")]
1621 pub struct CitationMetadata {
1622 #[serde(default)]
1623 pub citation_sources: Vec<CitationSource>,
1624 }
1625
1626 #[derive(Clone, Debug, Deserialize, Serialize)]
1627 #[serde(rename_all = "camelCase")]
1628 pub struct CitationSource {
1629 #[serde(skip_serializing_if = "Option::is_none")]
1630 pub uri: Option<String>,
1631 #[serde(skip_serializing_if = "Option::is_none")]
1632 pub start_index: Option<i32>,
1633 #[serde(skip_serializing_if = "Option::is_none")]
1634 pub end_index: Option<i32>,
1635 #[serde(skip_serializing_if = "Option::is_none")]
1636 pub license: Option<String>,
1637 }
1638
1639 #[derive(Clone, Debug, Deserialize, Serialize)]
1640 #[serde(rename_all = "camelCase")]
1641 pub struct LogprobsResult {
1642 #[serde(default)]
1643 pub top_candidates: Vec<TopCandidate>,
1644 #[serde(skip_serializing_if = "Option::is_none")]
1645 pub log_probability_sum: Option<f64>,
1646 #[serde(default)]
1647 pub chosen_candidates: Vec<LogProbCandidate>,
1648 }
1649
1650 #[derive(Clone, Debug, Deserialize, Serialize)]
1651 pub struct TopCandidate {
1652 #[serde(default)]
1653 pub candidates: Vec<LogProbCandidate>,
1654 }
1655
1656 #[derive(Clone, Debug, Deserialize, Serialize)]
1657 #[serde(rename_all = "camelCase")]
1658 pub struct LogProbCandidate {
1659 #[serde(skip_serializing_if = "Option::is_none")]
1660 pub token: Option<String>,
1661 #[serde(skip_serializing_if = "Option::is_none")]
1662 pub token_id: Option<i32>,
1663 #[serde(skip_serializing_if = "Option::is_none")]
1664 pub log_probability: Option<f64>,
1665 }
1666
1667 #[derive(Debug, Default, Deserialize, Serialize)]
1683 #[serde(rename_all = "camelCase")]
1684 pub struct GenerationConfig {
1685 #[serde(skip_serializing_if = "Option::is_none")]
1688 pub stop_sequences: Option<Vec<String>>,
1689 #[serde(skip_serializing_if = "Option::is_none")]
1695 pub response_mime_type: Option<String>,
1696 #[serde(skip_serializing_if = "Option::is_none")]
1700 pub response_schema: Option<Schema>,
1701 #[serde(
1707 skip_serializing_if = "Option::is_none",
1708 rename = "_responseJsonSchema"
1709 )]
1710 pub _response_json_schema: Option<Value>,
1711 #[serde(skip_serializing_if = "Option::is_none")]
1713 pub response_json_schema: Option<Value>,
1714 #[serde(skip_serializing_if = "Option::is_none")]
1717 pub candidate_count: Option<i32>,
1718 #[serde(skip_serializing_if = "Option::is_none")]
1721 pub max_output_tokens: Option<u64>,
1722 #[serde(skip_serializing_if = "Option::is_none")]
1725 pub temperature: Option<f64>,
1726 #[serde(skip_serializing_if = "Option::is_none")]
1733 pub top_p: Option<f64>,
1734 #[serde(skip_serializing_if = "Option::is_none")]
1740 pub top_k: Option<i32>,
1741 #[serde(skip_serializing_if = "Option::is_none")]
1747 pub presence_penalty: Option<f64>,
1748 #[serde(skip_serializing_if = "Option::is_none")]
1756 pub frequency_penalty: Option<f64>,
1757 #[serde(skip_serializing_if = "Option::is_none")]
1759 pub response_logprobs: Option<bool>,
1760 #[serde(skip_serializing_if = "Option::is_none")]
1763 pub logprobs: Option<i32>,
1764 #[serde(skip_serializing_if = "Option::is_none")]
1766 pub thinking_config: Option<ThinkingConfig>,
1767 #[serde(skip_serializing_if = "Option::is_none")]
1769 pub response_modalities: Option<Vec<ResponseModality>>,
1770 #[serde(skip_serializing_if = "Option::is_none")]
1771 pub image_config: Option<ImageConfig>,
1772 }
1773
1774 #[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
1776 #[serde(rename_all = "SCREAMING_SNAKE_CASE")]
1777 pub enum ResponseModality {
1778 Text,
1779 Image,
1780 Audio,
1781 }
1782
1783 #[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
1785 #[serde(rename_all = "snake_case")]
1786 pub enum ThinkingLevel {
1787 Minimal,
1788 Low,
1789 Medium,
1790 High,
1791 }
1792
1793 #[derive(Debug, Deserialize, Serialize)]
1797 #[serde(rename_all = "camelCase")]
1798 pub struct ThinkingConfig {
1799 #[serde(skip_serializing_if = "Option::is_none")]
1801 pub thinking_budget: Option<u32>,
1802 #[serde(skip_serializing_if = "Option::is_none")]
1804 pub thinking_level: Option<ThinkingLevel>,
1805 #[serde(skip_serializing_if = "Option::is_none")]
1807 pub include_thoughts: Option<bool>,
1808 }
1809
1810 #[derive(Debug, Deserialize, Serialize)]
1811 #[serde(rename_all = "camelCase")]
1812 pub struct ImageConfig {
1813 #[serde(skip_serializing_if = "Option::is_none")]
1814 pub aspect_ratio: Option<String>,
1815 #[serde(skip_serializing_if = "Option::is_none")]
1816 pub image_size: Option<String>,
1817 }
1818
1819 #[derive(Debug, Deserialize, Serialize, Clone)]
1823 pub struct Schema {
1824 pub r#type: String,
1825 #[serde(skip_serializing_if = "Option::is_none")]
1826 pub format: Option<String>,
1827 #[serde(skip_serializing_if = "Option::is_none")]
1828 pub description: Option<String>,
1829 #[serde(skip_serializing_if = "Option::is_none")]
1830 pub nullable: Option<bool>,
1831 #[serde(skip_serializing_if = "Option::is_none")]
1832 pub r#enum: Option<Vec<String>>,
1833 #[serde(skip_serializing_if = "Option::is_none")]
1834 pub max_items: Option<i32>,
1835 #[serde(skip_serializing_if = "Option::is_none")]
1836 pub min_items: Option<i32>,
1837 #[serde(skip_serializing_if = "Option::is_none")]
1838 pub properties: Option<HashMap<String, Schema>>,
1839 #[serde(skip_serializing_if = "Option::is_none")]
1840 pub required: Option<Vec<String>>,
1841 #[serde(skip_serializing_if = "Option::is_none")]
1842 pub items: Option<Box<Schema>>,
1843 }
1844
1845 pub fn tool_parameters_to_schema(parameters: Value) -> Result<Option<Schema>, CompletionError> {
1851 if parameters.is_null() || parameters == json!({"type": "object", "properties": {}}) {
1852 Ok(None)
1853 } else {
1854 parameters.try_into().map(Some)
1855 }
1856 }
1857
1858 pub fn flatten_schema(mut schema: Value) -> Result<Value, CompletionError> {
1864 let defs = if let Some(obj) = schema.as_object() {
1866 obj.get("$defs").or_else(|| obj.get("definitions")).cloned()
1867 } else {
1868 None
1869 };
1870
1871 let Some(defs_value) = defs else {
1872 return Ok(schema);
1873 };
1874
1875 let Some(defs_obj) = defs_value.as_object() else {
1876 return Err(CompletionError::ResponseError(
1877 "$defs must be an object".into(),
1878 ));
1879 };
1880
1881 resolve_refs(&mut schema, defs_obj)?;
1882
1883 if let Some(obj) = schema.as_object_mut() {
1885 obj.remove("$defs");
1886 obj.remove("definitions");
1887 }
1888
1889 Ok(schema)
1890 }
1891
1892 fn resolve_refs(
1895 value: &mut Value,
1896 defs: &serde_json::Map<String, Value>,
1897 ) -> Result<(), CompletionError> {
1898 match value {
1899 Value::Object(obj) => {
1900 if let Some(ref_value) = obj.get("$ref")
1901 && let Some(ref_str) = ref_value.as_str()
1902 {
1903 let def_name = parse_ref_path(ref_str)?;
1905
1906 let def = defs.get(&def_name).ok_or_else(|| {
1907 CompletionError::ResponseError(format!("Reference not found: {}", ref_str))
1908 })?;
1909
1910 let mut resolved = def.clone();
1911 resolve_refs(&mut resolved, defs)?;
1912 *value = resolved;
1913 return Ok(());
1914 }
1915
1916 for (_, v) in obj.iter_mut() {
1917 resolve_refs(v, defs)?;
1918 }
1919 }
1920 Value::Array(arr) => {
1921 for item in arr.iter_mut() {
1922 resolve_refs(item, defs)?;
1923 }
1924 }
1925 _ => {}
1926 }
1927
1928 Ok(())
1929 }
1930
1931 fn parse_ref_path(ref_str: &str) -> Result<String, CompletionError> {
1937 if let Some(fragment) = ref_str.strip_prefix('#') {
1938 if let Some(name) = fragment.strip_prefix("/$defs/") {
1939 Ok(name.to_string())
1940 } else if let Some(name) = fragment.strip_prefix("/definitions/") {
1941 Ok(name.to_string())
1942 } else {
1943 Err(CompletionError::ResponseError(format!(
1944 "Unsupported reference format: {}",
1945 ref_str
1946 )))
1947 }
1948 } else {
1949 Err(CompletionError::ResponseError(format!(
1950 "Only fragment references (#/...) are supported: {}",
1951 ref_str
1952 )))
1953 }
1954 }
1955
1956 fn extract_type(type_value: &Value) -> Option<String> {
1959 if let Some(t) = type_value.as_str() {
1960 return Some(t.to_string());
1961 }
1962
1963 type_value.as_array().and_then(|arr| {
1964 arr.iter()
1965 .filter_map(|v| v.as_str())
1966 .find(|t| *t != "null")
1967 .or_else(|| arr.iter().find_map(|v| v.as_str()))
1968 .map(str::to_owned)
1969 })
1970 }
1971
1972 fn schema_is_null(obj: &serde_json::Map<String, Value>) -> bool {
1973 obj.get("type")
1974 .and_then(extract_type)
1975 .as_deref()
1976 .is_some_and(|t| t == "null")
1977 }
1978
1979 fn schema_is_nullable(obj: &serde_json::Map<String, Value>) -> bool {
1980 obj.get("nullable")
1981 .and_then(|v| v.as_bool())
1982 .unwrap_or(false)
1983 || obj
1984 .get("type")
1985 .and_then(|v| v.as_array())
1986 .is_some_and(|arr| arr.iter().any(|v| v.as_str() == Some("null")))
1987 || ["anyOf", "oneOf", "allOf"].iter().any(|key| {
1988 obj.get(*key).and_then(|v| v.as_array()).is_some_and(|arr| {
1989 arr.iter()
1990 .filter_map(|schema| schema.as_object())
1991 .any(schema_is_null)
1992 })
1993 })
1994 }
1995
1996 fn extract_type_from_composition(composition: &Value) -> Option<String> {
1999 composition.as_array().and_then(|arr| {
2000 arr.iter().find_map(|schema| {
2001 let obj = schema.as_object()?;
2002 if schema_is_null(obj) {
2003 return None;
2004 }
2005
2006 obj.get("type").and_then(extract_type).or_else(|| {
2007 if obj.contains_key("properties") {
2008 Some("object".to_string())
2009 } else if obj.contains_key("enum") {
2010 Some("string".to_string())
2012 } else {
2013 None
2014 }
2015 })
2016 })
2017 })
2018 }
2019
2020 fn extract_schema_from_composition(
2023 composition: &Value,
2024 ) -> Option<serde_json::Map<String, Value>> {
2025 composition.as_array().and_then(|arr| {
2026 arr.iter().find_map(|schema| {
2027 let obj = schema.as_object()?;
2028 if schema_is_null(obj) {
2029 None
2030 } else {
2031 Some(obj.clone())
2032 }
2033 })
2034 })
2035 }
2036
2037 fn extract_schema_from_composition_obj(
2038 obj: &serde_json::Map<String, Value>,
2039 ) -> Option<serde_json::Map<String, Value>> {
2040 obj.get("anyOf")
2041 .and_then(extract_schema_from_composition)
2042 .or_else(|| obj.get("oneOf").and_then(extract_schema_from_composition))
2043 .or_else(|| obj.get("allOf").and_then(extract_schema_from_composition))
2044 }
2045
2046 fn infer_type(obj: &serde_json::Map<String, Value>) -> String {
2049 if let Some(type_val) = obj.get("type")
2051 && let Some(type_str) = extract_type(type_val)
2052 {
2053 return type_str;
2054 }
2055
2056 if let Some(any_of) = obj.get("anyOf")
2058 && let Some(type_str) = extract_type_from_composition(any_of)
2059 {
2060 return type_str;
2061 }
2062
2063 if let Some(one_of) = obj.get("oneOf")
2064 && let Some(type_str) = extract_type_from_composition(one_of)
2065 {
2066 return type_str;
2067 }
2068
2069 if let Some(all_of) = obj.get("allOf")
2070 && let Some(type_str) = extract_type_from_composition(all_of)
2071 {
2072 return type_str;
2073 }
2074
2075 if obj.contains_key("properties") {
2077 "object".to_string()
2078 } else if obj.contains_key("enum") {
2079 "string".to_string()
2080 } else {
2081 String::new()
2082 }
2083 }
2084
2085 impl TryFrom<Value> for Schema {
2086 type Error = CompletionError;
2087
2088 fn try_from(value: Value) -> Result<Self, Self::Error> {
2089 let flattened_val = flatten_schema(value)?;
2090 if let Some(obj) = flattened_val.as_object() {
2091 let composition_source = extract_schema_from_composition_obj(obj);
2094 let props_source = if obj.get("properties").is_none() {
2095 composition_source.clone().unwrap_or(obj.clone())
2096 } else {
2097 obj.clone()
2098 };
2099
2100 let schema_type = infer_type(obj);
2101 let items = obj
2102 .get("items")
2103 .or_else(|| props_source.get("items"))
2104 .and_then(|v| v.clone().try_into().ok())
2105 .map(Box::new);
2106
2107 let items = if schema_type == "array" && items.is_none() {
2110 Some(Box::new(Schema {
2111 r#type: "string".to_string(),
2112 format: None,
2113 description: None,
2114 nullable: None,
2115 r#enum: None,
2116 max_items: None,
2117 min_items: None,
2118 properties: None,
2119 required: None,
2120 items: None,
2121 }))
2122 } else {
2123 items
2124 };
2125
2126 Ok(Schema {
2127 r#type: schema_type,
2128 format: obj
2129 .get("format")
2130 .or_else(|| props_source.get("format"))
2131 .and_then(|v| v.as_str())
2132 .map(String::from),
2133 description: obj
2134 .get("description")
2135 .or_else(|| props_source.get("description"))
2136 .and_then(|v| v.as_str())
2137 .map(String::from),
2138 nullable: if schema_is_nullable(obj)
2139 || composition_source.as_ref().is_some_and(schema_is_nullable)
2140 {
2141 Some(true)
2142 } else {
2143 None
2144 },
2145 r#enum: obj
2146 .get("enum")
2147 .or_else(|| props_source.get("enum"))
2148 .and_then(|v| v.as_array())
2149 .map(|arr| {
2150 arr.iter()
2151 .filter_map(|v| v.as_str().map(String::from))
2152 .collect()
2153 }),
2154 max_items: obj
2155 .get("maxItems")
2156 .and_then(|v| v.as_i64())
2157 .map(|v| v as i32),
2158 min_items: obj
2159 .get("minItems")
2160 .and_then(|v| v.as_i64())
2161 .map(|v| v as i32),
2162 properties: props_source
2163 .get("properties")
2164 .and_then(|v| v.as_object())
2165 .map(|map| {
2166 map.iter()
2167 .filter_map(|(k, v)| {
2168 v.clone().try_into().ok().map(|schema| (k.clone(), schema))
2169 })
2170 .collect()
2171 }),
2172 required: props_source
2173 .get("required")
2174 .and_then(|v| v.as_array())
2175 .map(|arr| {
2176 arr.iter()
2177 .filter_map(|v| v.as_str().map(String::from))
2178 .collect()
2179 }),
2180 items,
2181 })
2182 } else {
2183 Err(CompletionError::ResponseError(
2184 "Expected a JSON object for Schema".into(),
2185 ))
2186 }
2187 }
2188 }
2189
2190 #[derive(Debug, Serialize)]
2191 #[serde(rename_all = "camelCase")]
2192 pub struct GenerateContentRequest {
2193 pub contents: Vec<Content>,
2194 #[serde(skip_serializing_if = "Option::is_none")]
2195 pub tools: Option<Vec<Value>>,
2196 pub tool_config: Option<ToolConfig>,
2197 pub generation_config: Option<GenerationConfig>,
2199 pub safety_settings: Option<Vec<SafetySetting>>,
2213 pub system_instruction: Option<Content>,
2216 #[serde(flatten, skip_serializing_if = "Option::is_none")]
2219 pub additional_params: Option<serde_json::Value>,
2220 }
2221
2222 #[derive(Debug, Serialize)]
2223 #[serde(rename_all = "camelCase")]
2224 pub struct Tool {
2225 pub function_declarations: Vec<FunctionDeclaration>,
2226 pub code_execution: Option<CodeExecution>,
2227 }
2228
2229 #[derive(Debug, Serialize, Clone)]
2230 #[serde(rename_all = "camelCase")]
2231 pub struct FunctionDeclaration {
2232 pub name: String,
2233 pub description: String,
2234 #[serde(skip_serializing_if = "Option::is_none")]
2235 pub parameters: Option<Schema>,
2236 }
2237
2238 #[derive(Debug, Serialize, Deserialize)]
2239 #[serde(rename_all = "camelCase")]
2240 pub struct ToolConfig {
2241 pub function_calling_config: Option<FunctionCallingMode>,
2242 }
2243
2244 #[derive(Debug, Serialize, Deserialize, Default)]
2245 #[serde(tag = "mode", rename_all = "UPPERCASE")]
2246 pub enum FunctionCallingMode {
2247 #[default]
2248 Auto,
2249 None,
2250 Any {
2251 #[serde(skip_serializing_if = "Option::is_none")]
2252 allowed_function_names: Option<Vec<String>>,
2253 },
2254 }
2255
2256 impl TryFrom<message::ToolChoice> for FunctionCallingMode {
2257 type Error = CompletionError;
2258 fn try_from(value: message::ToolChoice) -> Result<Self, Self::Error> {
2259 let res = match value {
2260 message::ToolChoice::Auto => Self::Auto,
2261 message::ToolChoice::None => Self::None,
2262 message::ToolChoice::Required => Self::Any {
2263 allowed_function_names: None,
2264 },
2265 message::ToolChoice::Specific { function_names } => Self::Any {
2266 allowed_function_names: Some(function_names),
2267 },
2268 };
2269
2270 Ok(res)
2271 }
2272 }
2273
2274 #[derive(Debug, Serialize)]
2275 pub struct CodeExecution {}
2276
2277 #[derive(Debug, Serialize)]
2278 #[serde(rename_all = "camelCase")]
2279 pub struct SafetySetting {
2280 pub category: HarmCategory,
2281 pub threshold: HarmBlockThreshold,
2282 }
2283
2284 #[derive(Debug, Serialize)]
2285 #[serde(rename_all = "SCREAMING_SNAKE_CASE")]
2286 pub enum HarmBlockThreshold {
2287 HarmBlockThresholdUnspecified,
2288 BlockLowAndAbove,
2289 BlockMediumAndAbove,
2290 BlockOnlyHigh,
2291 BlockNone,
2292 Off,
2293 }
2294}
2295
2296#[cfg(test)]
2297mod tests {
2298 use crate::{
2299 message,
2300 providers::gemini::completion::gemini_api_types::{
2301 BlockReason, CitationMetadata, ContentCandidate, FinishReason, FunctionCall,
2302 GenerateContentResponse, LogprobsResult, ModalityTokenCount, PromptFeedback, Schema,
2303 TopCandidate, UsageMetadata, flatten_schema, tool_parameters_to_schema,
2304 },
2305 };
2306
2307 use super::*;
2308 use serde_json::json;
2309
2310 #[test]
2311 fn test_usage_metadata_deserializes_without_total_token_count() {
2312 let usage: UsageMetadata =
2315 serde_json::from_str(r#"{"promptTokenCount": 12}"#).expect("should deserialize");
2316 assert_eq!(usage.total_token_count, 0);
2317 assert_eq!(usage.prompt_token_count, 12);
2318 }
2319
2320 #[test]
2321 fn test_generate_content_response_deserializes_without_candidates_or_response_id() {
2322 let response: GenerateContentResponse = serde_json::from_value(json!({
2325 "promptFeedback": {
2326 "blockReason": "SAFETY"
2327 }
2328 }))
2329 .expect("blocked prompt response should deserialize");
2330
2331 assert!(response.response_id.is_empty());
2332 assert!(response.candidates.is_empty());
2333
2334 let error = completion::CompletionResponse::try_from(response)
2335 .expect_err("empty candidates should become a response error");
2336 assert!(error.to_string().contains("No response candidates"));
2337 }
2338
2339 #[test]
2340 fn test_modality_token_count_deserializes_without_zero_token_count() {
2341 let count: ModalityTokenCount = serde_json::from_value(json!({
2342 "modality": "TEXT"
2343 }))
2344 .expect("zero tokenCount may be omitted");
2345
2346 assert_eq!(count.token_count, 0);
2347 }
2348
2349 #[test]
2350 fn test_response_metadata_repeated_fields_deserialize_when_omitted() {
2351 let citation_metadata: CitationMetadata =
2352 serde_json::from_value(json!({})).expect("empty citation metadata should deserialize");
2353 assert!(citation_metadata.citation_sources.is_empty());
2354
2355 let logprobs: LogprobsResult =
2356 serde_json::from_value(json!({})).expect("empty logprobs result should deserialize");
2357 assert!(logprobs.top_candidates.is_empty());
2358 assert_eq!(logprobs.log_probability_sum, None);
2359 assert!(logprobs.chosen_candidates.is_empty());
2360
2361 let top_candidate: TopCandidate =
2362 serde_json::from_value(json!({})).expect("empty top candidate should deserialize");
2363 assert!(top_candidate.candidates.is_empty());
2364 }
2365
2366 #[test]
2367 fn test_logprobs_result_deserializes_official_json_field_names() {
2368 let logprobs: LogprobsResult = serde_json::from_value(json!({
2369 "topCandidates": [
2370 {
2371 "candidates": [
2372 {
2373 "token": "Hello",
2374 "tokenId": 123,
2375 "logProbability": -0.1
2376 },
2377 {
2378 "token": "Hi",
2379 "tokenId": 124,
2380 "logProbability": -1.25
2381 }
2382 ]
2383 }
2384 ],
2385 "logProbabilitySum": -0.1,
2386 "chosenCandidates": [
2387 {
2388 "token": "Hello",
2389 "tokenId": 123,
2390 "logProbability": -0.1
2391 }
2392 ]
2393 }))
2394 .expect("official Gemini logprobs result should deserialize");
2395
2396 assert_eq!(logprobs.top_candidates.len(), 1);
2397 assert_eq!(logprobs.top_candidates[0].candidates.len(), 2);
2398 assert_eq!(
2399 logprobs.top_candidates[0].candidates[0].token.as_deref(),
2400 Some("Hello")
2401 );
2402 assert_eq!(logprobs.top_candidates[0].candidates[0].token_id, Some(123));
2403 assert_eq!(
2404 logprobs.top_candidates[0].candidates[0].log_probability,
2405 Some(-0.1)
2406 );
2407 assert_eq!(logprobs.log_probability_sum, Some(-0.1));
2408 assert_eq!(logprobs.chosen_candidates.len(), 1);
2409 assert_eq!(
2410 logprobs.chosen_candidates[0].token.as_deref(),
2411 Some("Hello")
2412 );
2413 assert_eq!(logprobs.chosen_candidates[0].token_id, Some(123));
2414 assert_eq!(logprobs.chosen_candidates[0].log_probability, Some(-0.1));
2415 }
2416
2417 #[test]
2418 fn test_resolve_request_model_uses_override() {
2419 let request = CompletionRequest {
2420 model: Some("gemini-2.5-flash".to_string()),
2421 preamble: None,
2422 chat_history: vec!["Hello".into()],
2423 documents: vec![],
2424 tools: vec![],
2425 temperature: None,
2426 max_tokens: None,
2427 tool_choice: None,
2428 additional_params: None,
2429 output_schema: None,
2430 record_telemetry_content: false,
2431 };
2432
2433 let request_model = resolve_request_model("gemini-2.0-flash", &request);
2434 assert_eq!(request_model, "gemini-2.5-flash");
2435 assert_eq!(
2436 completion_endpoint(&request_model),
2437 "/v1beta/models/gemini-2.5-flash:generateContent"
2438 );
2439 assert_eq!(
2440 streaming_endpoint(&request_model),
2441 "/v1beta/models/gemini-2.5-flash:streamGenerateContent"
2442 );
2443 }
2444
2445 #[test]
2446 fn test_resolve_request_model_uses_default_when_unset() {
2447 let request = CompletionRequest {
2448 model: None,
2449 preamble: None,
2450 chat_history: vec!["Hello".into()],
2451 documents: vec![],
2452 tools: vec![],
2453 temperature: None,
2454 max_tokens: None,
2455 tool_choice: None,
2456 additional_params: None,
2457 output_schema: None,
2458 record_telemetry_content: false,
2459 };
2460
2461 assert_eq!(
2462 resolve_request_model("gemini-2.0-flash", &request),
2463 "gemini-2.0-flash"
2464 );
2465 }
2466
2467 #[test]
2468 fn test_deserialize_message_user() {
2469 let raw_message = r#"{
2470 "parts": [
2471 {"text": "Hello, world!"},
2472 {"inlineData": {"mimeType": "image/png", "data": "base64encodeddata"}},
2473 {"functionCall": {"name": "test_function", "args": {"arg1": "value1"}}},
2474 {"functionResponse": {"name": "test_function", "response": {"result": "success"}}},
2475 {"fileData": {"mimeType": "application/pdf", "fileUri": "http://example.com/file.pdf"}},
2476 {"executableCode": {"code": "print('Hello, world!')", "language": "PYTHON"}},
2477 {"codeExecutionResult": {"output": "Hello, world!", "outcome": "OUTCOME_OK"}}
2478 ],
2479 "role": "user"
2480 }"#;
2481
2482 let content: Content = {
2483 let jd = &mut serde_json::Deserializer::from_str(raw_message);
2484 serde_path_to_error::deserialize(jd).unwrap_or_else(|err| {
2485 panic!("Deserialization error at {}: {}", err.path(), err);
2486 })
2487 };
2488 assert_eq!(content.role, Some(Role::User));
2489 assert_eq!(content.parts.len(), 7);
2490
2491 let parts: Vec<Part> = content.parts.into_iter().collect();
2492
2493 if let Part {
2494 part: PartKind::Text(text),
2495 ..
2496 } = &parts[0]
2497 {
2498 assert_eq!(text, "Hello, world!");
2499 } else {
2500 panic!("Expected text part");
2501 }
2502
2503 if let Part {
2504 part: PartKind::InlineData(inline_data),
2505 ..
2506 } = &parts[1]
2507 {
2508 assert_eq!(inline_data.mime_type, "image/png");
2509 assert_eq!(inline_data.data, "base64encodeddata");
2510 } else {
2511 panic!("Expected inline data part");
2512 }
2513
2514 if let Part {
2515 part: PartKind::FunctionCall(function_call),
2516 ..
2517 } = &parts[2]
2518 {
2519 assert_eq!(function_call.name, "test_function");
2520 assert_eq!(
2521 function_call.args.as_object().unwrap().get("arg1").unwrap(),
2522 "value1"
2523 );
2524 } else {
2525 panic!("Expected function call part");
2526 }
2527
2528 if let Part {
2529 part: PartKind::FunctionResponse(function_response),
2530 ..
2531 } = &parts[3]
2532 {
2533 assert_eq!(function_response.name, "test_function");
2534 assert_eq!(
2535 function_response
2536 .response
2537 .as_ref()
2538 .unwrap()
2539 .get("result")
2540 .unwrap(),
2541 "success"
2542 );
2543 } else {
2544 panic!("Expected function response part");
2545 }
2546
2547 if let Part {
2548 part: PartKind::FileData(file_data),
2549 ..
2550 } = &parts[4]
2551 {
2552 assert_eq!(file_data.mime_type.as_ref().unwrap(), "application/pdf");
2553 assert_eq!(file_data.file_uri, "http://example.com/file.pdf");
2554 } else {
2555 panic!("Expected file data part");
2556 }
2557
2558 if let Part {
2559 part: PartKind::ExecutableCode(executable_code),
2560 ..
2561 } = &parts[5]
2562 {
2563 assert_eq!(executable_code.code, "print('Hello, world!')");
2564 } else {
2565 panic!("Expected executable code part");
2566 }
2567
2568 if let Part {
2569 part: PartKind::CodeExecutionResult(code_execution_result),
2570 ..
2571 } = &parts[6]
2572 {
2573 assert_eq!(
2574 code_execution_result.clone().output.unwrap(),
2575 "Hello, world!"
2576 );
2577 } else {
2578 panic!("Expected code execution result part");
2579 }
2580 }
2581
2582 #[test]
2583 fn test_deserialize_message_model() {
2584 let json_data = json!({
2585 "parts": [{"text": "Hello, user!"}],
2586 "role": "model"
2587 });
2588
2589 let content: Content = serde_json::from_value(json_data).unwrap();
2590 assert_eq!(content.role, Some(Role::Model));
2591 assert_eq!(content.parts.len(), 1);
2592 if let Some(Part {
2593 part: PartKind::Text(text),
2594 ..
2595 }) = content.parts.first()
2596 {
2597 assert_eq!(text, "Hello, user!");
2598 } else {
2599 panic!("Expected text part");
2600 }
2601 }
2602
2603 #[test]
2604 fn test_message_conversion_user() {
2605 let msg = message::Message::user("Hello, world!");
2606 let content: Content = msg.try_into().unwrap();
2607 assert_eq!(content.role, Some(Role::User));
2608 assert_eq!(content.parts.len(), 1);
2609 if let Some(Part {
2610 part: PartKind::Text(text),
2611 ..
2612 }) = &content.parts.first()
2613 {
2614 assert_eq!(text, "Hello, world!");
2615 } else {
2616 panic!("Expected text part");
2617 }
2618 }
2619
2620 #[test]
2621 fn test_message_conversion_model() {
2622 let msg = message::Message::assistant("Hello, user!");
2623
2624 let content: Content = msg.try_into().unwrap();
2625 assert_eq!(content.role, Some(Role::Model));
2626 assert_eq!(content.parts.len(), 1);
2627 if let Some(Part {
2628 part: PartKind::Text(text),
2629 ..
2630 }) = &content.parts.first()
2631 {
2632 assert_eq!(text, "Hello, user!");
2633 } else {
2634 panic!("Expected text part");
2635 }
2636 }
2637
2638 #[test]
2639 fn test_thought_signature_is_preserved_from_response_reasoning_part() {
2640 let response = GenerateContentResponse {
2641 response_id: "resp_1".to_string(),
2642 candidates: vec![ContentCandidate {
2643 content: Some(Content {
2644 parts: vec![Part {
2645 thought: Some(true),
2646 thought_signature: Some("thought_sig_123".to_string()),
2647 part: PartKind::Text("thinking text".to_string()),
2648 additional_params: None,
2649 }],
2650 role: Some(Role::Model),
2651 }),
2652 finish_reason: Some(FinishReason::Stop),
2653 safety_ratings: None,
2654 citation_metadata: None,
2655 token_count: None,
2656 avg_logprobs: None,
2657 logprobs_result: None,
2658 index: Some(0),
2659 finish_message: None,
2660 }],
2661 prompt_feedback: None,
2662 usage_metadata: None,
2663 model_version: None,
2664 };
2665
2666 let converted: crate::completion::CompletionResponse =
2667 response.try_into().expect("convert response");
2668 let first = converted.choice.first();
2669 assert!(matches!(
2670 first,
2671 Some(message::AssistantContent::Reasoning(message::Reasoning { content, .. }))
2672 if matches!(
2673 content.first(),
2674 Some(message::ReasoningContent::Text {
2675 text,
2676 signature: Some(signature)
2677 }) if text == "thinking text" && signature == "thought_sig_123"
2678 )
2679 ));
2680 }
2681
2682 #[test]
2683 fn test_tool_protocol_finish_reason_returns_response_error() {
2684 for (reason, finish_message) in [
2685 (
2686 FinishReason::MalformedFunctionCall,
2687 "malformed function call: default_api",
2688 ),
2689 (
2690 FinishReason::UnexpectedToolCall,
2691 "unexpected tool call: default_api",
2692 ),
2693 (
2694 FinishReason::MissingThoughtSignature,
2695 "missing thought signature for tool call",
2696 ),
2697 (
2698 FinishReason::TooManyToolCalls,
2699 "too many tool calls in response",
2700 ),
2701 (
2702 FinishReason::MalformedResponse,
2703 "malformed response from provider",
2704 ),
2705 ] {
2706 let reason_name = format!("{reason:?}");
2707 let response = GenerateContentResponse {
2708 response_id: "resp_tool_protocol_error".to_string(),
2709 candidates: vec![ContentCandidate {
2710 content: Some(Content {
2711 parts: vec![Part {
2712 thought: None,
2713 thought_signature: None,
2714 part: PartKind::FunctionCall(FunctionCall {
2715 name: "default_api".to_string(),
2716 args: json!({"x": 1}),
2717 id: None,
2718 }),
2719 additional_params: None,
2720 }],
2721 role: Some(Role::Model),
2722 }),
2723 finish_reason: Some(reason),
2724 safety_ratings: None,
2725 citation_metadata: None,
2726 token_count: None,
2727 avg_logprobs: None,
2728 logprobs_result: None,
2729 index: Some(0),
2730 finish_message: Some(finish_message.to_string()),
2731 }],
2732 prompt_feedback: None,
2733 usage_metadata: None,
2734 model_version: None,
2735 };
2736
2737 let err = crate::completion::CompletionResponse::try_from(response)
2738 .expect_err("tool protocol finish reason should fail");
2739
2740 assert!(matches!(
2741 err,
2742 CompletionError::ResponseError(message)
2743 if message.contains(&reason_name)
2744 && message.contains(finish_message)
2745 ));
2746 }
2747 }
2748
2749 #[test]
2750 fn test_completion_response_usage_preserves_cached_and_reasoning_tokens() {
2751 let response = GenerateContentResponse {
2752 response_id: "resp_1".to_string(),
2753 candidates: vec![ContentCandidate {
2754 content: Some(Content {
2755 parts: vec![Part {
2756 thought: None,
2757 thought_signature: None,
2758 part: PartKind::Text("answer".to_string()),
2759 additional_params: None,
2760 }],
2761 role: Some(Role::Model),
2762 }),
2763 finish_reason: Some(FinishReason::Stop),
2764 safety_ratings: None,
2765 citation_metadata: None,
2766 token_count: None,
2767 avg_logprobs: None,
2768 logprobs_result: None,
2769 index: Some(0),
2770 finish_message: None,
2771 }],
2772 prompt_feedback: None,
2773 usage_metadata: Some(UsageMetadata {
2774 prompt_token_count: 40,
2775 cached_content_token_count: Some(20),
2776 candidates_token_count: Some(30),
2777 total_token_count: 100,
2778 thoughts_token_count: Some(10),
2779 prompt_tokens_details: None,
2780 cache_tokens_details: None,
2781 candidates_tokens_details: None,
2782 tool_use_prompt_token_count: Some(12),
2783 tool_use_prompt_tokens_details: None,
2784 traffic_type: None,
2785 }),
2786 model_version: Some("gemini-2.0-flash-001".to_string()),
2787 };
2788
2789 let converted: crate::completion::CompletionResponse =
2790 response.try_into().expect("convert response");
2791
2792 assert_eq!(converted.usage.input_tokens, 40);
2793 assert_eq!(converted.usage.cached_input_tokens, 20);
2794 assert_eq!(converted.usage.output_tokens, 30);
2795 assert_eq!(converted.usage.reasoning_tokens, 10);
2796 assert_eq!(converted.usage.tool_use_prompt_tokens, 12);
2797 assert_eq!(converted.usage.total_tokens, 100);
2798 }
2799
2800 #[test]
2801 fn test_finish_reason_maps_every_wire_variant() {
2802 use crate::completion::FinishReason as Normalized;
2803
2804 for (wire, expected) in [
2805 (FinishReason::Stop, Normalized::Stop),
2806 (FinishReason::MaxTokens, Normalized::Length),
2807 (FinishReason::Safety, Normalized::ContentFilter),
2808 (FinishReason::Blocklist, Normalized::ContentFilter),
2809 (FinishReason::ProhibitedContent, Normalized::ContentFilter),
2810 (FinishReason::Spii, Normalized::ContentFilter),
2811 (
2814 FinishReason::Recitation,
2815 Normalized::Other("RECITATION".to_string()),
2816 ),
2817 (
2818 FinishReason::Language,
2819 Normalized::Other("LANGUAGE".to_string()),
2820 ),
2821 (FinishReason::Other, Normalized::Other("OTHER".to_string())),
2822 (
2823 FinishReason::MalformedFunctionCall,
2824 Normalized::Other("MALFORMED_FUNCTION_CALL".to_string()),
2825 ),
2826 (
2827 FinishReason::UnexpectedToolCall,
2828 Normalized::Other("UNEXPECTED_TOOL_CALL".to_string()),
2829 ),
2830 (
2831 FinishReason::MissingThoughtSignature,
2832 Normalized::Other("MISSING_THOUGHT_SIGNATURE".to_string()),
2833 ),
2834 (
2835 FinishReason::TooManyToolCalls,
2836 Normalized::Other("TOO_MANY_TOOL_CALLS".to_string()),
2837 ),
2838 (
2839 FinishReason::MalformedResponse,
2840 Normalized::Other("MALFORMED_RESPONSE".to_string()),
2841 ),
2842 ] {
2843 assert_eq!(
2844 map_finish_reason(&wire),
2845 Some(expected),
2846 "wire reason {wire:?}"
2847 );
2848 }
2849
2850 assert_eq!(
2853 map_finish_reason(&FinishReason::FinishReasonUnspecified),
2854 None
2855 );
2856 }
2857
2858 #[test]
2859 fn test_finish_reason_wire_spelling_matches_serde() {
2860 for reason in [
2863 FinishReason::FinishReasonUnspecified,
2864 FinishReason::Stop,
2865 FinishReason::MaxTokens,
2866 FinishReason::Safety,
2867 FinishReason::Recitation,
2868 FinishReason::Language,
2869 FinishReason::Other,
2870 FinishReason::Blocklist,
2871 FinishReason::ProhibitedContent,
2872 FinishReason::Spii,
2873 FinishReason::MalformedFunctionCall,
2874 FinishReason::UnexpectedToolCall,
2875 FinishReason::MissingThoughtSignature,
2876 FinishReason::TooManyToolCalls,
2877 FinishReason::MalformedResponse,
2878 ] {
2879 let serialized = serde_json::to_value(&reason).expect("reason should serialize");
2880 assert_eq!(serialized, json!(reason.as_wire_str()));
2881 }
2882 }
2883
2884 #[test]
2885 fn test_unknown_finish_reason_round_trips_verbatim() {
2886 let reason: FinishReason = serde_json::from_value(json!("FINISH_REASON_FUTURE"))
2890 .expect("unknown finish reason should deserialize");
2891 assert!(matches!(&reason, FinishReason::Unknown(s) if s == "FINISH_REASON_FUTURE"));
2892 assert_eq!(reason.as_wire_str(), "FINISH_REASON_FUTURE");
2893 assert_eq!(
2894 serde_json::to_value(&reason).expect("reason should serialize"),
2895 json!("FINISH_REASON_FUTURE")
2896 );
2897 assert_eq!(
2898 map_finish_reason(&reason),
2899 Some(crate::completion::FinishReason::Other(
2900 "FINISH_REASON_FUTURE".to_string()
2901 ))
2902 );
2903 }
2904
2905 #[test]
2906 fn test_unknown_block_reason_deserializes_verbatim() {
2907 let feedback: PromptFeedback = serde_json::from_value(json!({
2910 "blockReason": "BLOCK_REASON_FUTURE"
2911 }))
2912 .expect("unknown block reason should deserialize");
2913 assert!(matches!(
2914 feedback.block_reason,
2915 Some(BlockReason::Unknown(ref s)) if s == "BLOCK_REASON_FUTURE"
2916 ));
2917 }
2918
2919 #[test]
2920 fn test_unary_response_with_unknown_finish_reason_stays_parseable() {
2921 let response: GenerateContentResponse = serde_json::from_value(json!({
2925 "responseId": "resp-future",
2926 "candidates": [{
2927 "content": {
2928 "parts": [{"text": "hi"}],
2929 "role": "model"
2930 },
2931 "finishReason": "FINISH_REASON_FUTURE"
2932 }],
2933 "usageMetadata": {
2934 "promptTokenCount": 3,
2935 "candidatesTokenCount": 2,
2936 "totalTokenCount": 5
2937 }
2938 }))
2939 .expect("unknown finish reason should not fail the payload");
2940
2941 let converted: crate::completion::CompletionResponse =
2942 response.try_into().expect("convert response");
2943
2944 assert!(matches!(
2945 converted.choice.first(),
2946 Some(message::AssistantContent::Text(text)) if text.text == "hi"
2947 ));
2948 assert_eq!(converted.usage.total_tokens, 5);
2949 assert_eq!(
2950 converted.finish_reason(),
2951 Some(crate::completion::FinishReason::Other(
2952 "FINISH_REASON_FUTURE".to_string()
2953 ))
2954 );
2955 }
2956
2957 #[test]
2958 fn test_streaming_candidate_with_unknown_finish_reason_stays_parseable() {
2959 let candidate: ContentCandidate = serde_json::from_value(json!({
2963 "content": {
2964 "parts": [{"text": "done"}],
2965 "role": "model"
2966 },
2967 "finishReason": "FINISH_REASON_FUTURE"
2968 }))
2969 .expect("unknown finish reason should not fail the chunk");
2970
2971 let reason = candidate.finish_reason.expect("finish reason present");
2972 assert_eq!(
2973 map_finish_reason(&reason),
2974 Some(crate::completion::FinishReason::Other(
2975 "FINISH_REASON_FUTURE".to_string()
2976 ))
2977 );
2978 }
2979
2980 #[test]
2981 fn test_completion_response_carries_normalized_metadata() {
2982 let response: GenerateContentResponse = serde_json::from_value(json!({
2983 "responseId": "resp-meta",
2984 "modelVersion": "gemini-2.0-flash-001",
2985 "candidates": [{
2986 "content": {
2987 "parts": [{"text": "hi"}],
2988 "role": "model"
2989 },
2990 "finishReason": "MAX_TOKENS"
2991 }]
2992 }))
2993 .expect("response should deserialize");
2994
2995 let converted: crate::completion::CompletionResponse =
2996 response.try_into().expect("convert response");
2997
2998 assert_eq!(converted.provider, PROVIDER_NAME);
2999 assert_eq!(converted.model.as_deref(), Some("gemini-2.0-flash-001"));
3000 assert_eq!(converted.response_id.as_deref(), Some("resp-meta"));
3001 assert_eq!(converted.message_id, None);
3002 assert_eq!(
3003 converted.finish_reason(),
3004 Some(crate::completion::FinishReason::Length)
3005 );
3006 }
3007
3008 #[test]
3009 fn test_completion_response_upgrades_stop_to_tool_calls() {
3010 let response: GenerateContentResponse = serde_json::from_value(json!({
3013 "responseId": "resp-tool",
3014 "candidates": [{
3015 "content": {
3016 "parts": [{
3017 "functionCall": {
3018 "name": "get_weather",
3019 "args": {"city": "Paris"}
3020 }
3021 }],
3022 "role": "model"
3023 },
3024 "finishReason": "STOP"
3025 }]
3026 }))
3027 .expect("response should deserialize");
3028
3029 let converted: crate::completion::CompletionResponse =
3030 response.try_into().expect("convert response");
3031
3032 assert_eq!(
3033 converted.finish_reason(),
3034 Some(crate::completion::FinishReason::ToolCalls)
3035 );
3036 assert_eq!(converted.model, None);
3037 }
3038
3039 #[test]
3040 fn test_reasoning_signature_is_emitted_in_gemini_part() {
3041 let msg = message::Message::Assistant {
3042 id: None,
3043 content: vec![message::AssistantContent::Reasoning(
3044 message::Reasoning::new_with_signature(
3045 "structured thought",
3046 Some("reuse_sig_456".to_string()),
3047 ),
3048 )],
3049 };
3050
3051 let converted: Content = msg.try_into().expect("convert message");
3052 let first = converted.parts.first().expect("reasoning part");
3053 assert_eq!(first.thought, Some(true));
3054 assert_eq!(first.thought_signature.as_deref(), Some("reuse_sig_456"));
3055 assert!(matches!(
3056 &first.part,
3057 PartKind::Text(text) if text == "structured thought"
3058 ));
3059 }
3060
3061 #[test]
3062 fn test_message_conversion_tool_call() {
3063 let tool_call = message::ToolCall::from_wire(
3064 "call-123",
3065 message::ToolFunction {
3066 name: "test_function".to_string(),
3067 arguments: json!({"arg1": "value1"}),
3068 },
3069 );
3070
3071 let msg = message::Message::Assistant {
3072 id: None,
3073 content: vec![message::AssistantContent::ToolCall(tool_call)],
3074 };
3075
3076 let content: Content = msg.try_into().unwrap();
3077 assert_eq!(content.role, Some(Role::Model));
3078 assert_eq!(content.parts.len(), 1);
3079 if let Some(Part {
3080 part: PartKind::FunctionCall(function_call),
3081 ..
3082 }) = content.parts.first()
3083 {
3084 assert_eq!(function_call.name, "test_function");
3085 assert_eq!(
3086 function_call.args.as_object().unwrap().get("arg1").unwrap(),
3087 "value1"
3088 );
3089 assert_eq!(function_call.id.as_deref(), Some("call-123"));
3090 } else {
3091 panic!("Expected function call part");
3092 }
3093 }
3094
3095 #[test]
3096 fn test_response_function_call_preserves_correlation_id() {
3097 let response: GenerateContentResponse = serde_json::from_value(json!({
3098 "responseId": "response-123",
3099 "candidates": [{
3100 "content": {
3101 "parts": [{
3102 "functionCall": {
3103 "name": "test_function",
3104 "args": {"arg1": "value1"},
3105 "id": "call-123"
3106 }
3107 }],
3108 "role": "model"
3109 },
3110 "finishReason": "STOP"
3111 }]
3112 }))
3113 .expect("response should deserialize");
3114
3115 let converted: crate::completion::CompletionResponse =
3116 response.try_into().expect("response should convert");
3117 let Some(message::AssistantContent::ToolCall(tool_call)) = converted.choice.first() else {
3118 panic!("expected a tool call");
3119 };
3120 assert_eq!(tool_call.id, "call-123");
3121 assert_eq!(
3122 tool_call.provider.as_ref().expect("wire id").call_id,
3123 "call-123"
3124 );
3125 }
3126
3127 #[test]
3128 fn test_vec_schema_conversion() {
3129 let schema_with_ref = json!({
3130 "type": "array",
3131 "items": {
3132 "$ref": "#/$defs/Person"
3133 },
3134 "$defs": {
3135 "Person": {
3136 "type": "object",
3137 "properties": {
3138 "first_name": {
3139 "type": ["string", "null"],
3140 "description": "The person's first name, if provided (null otherwise)"
3141 },
3142 "last_name": {
3143 "type": ["string", "null"],
3144 "description": "The person's last name, if provided (null otherwise)"
3145 },
3146 "job": {
3147 "type": ["string", "null"],
3148 "description": "The person's job, if provided (null otherwise)"
3149 }
3150 },
3151 "required": []
3152 }
3153 }
3154 });
3155
3156 let result: Result<Schema, _> = schema_with_ref.try_into();
3157
3158 match result {
3159 Ok(schema) => {
3160 assert_eq!(schema.r#type, "array");
3161
3162 if let Some(items) = schema.items {
3163 println!("item types: {}", items.r#type);
3164
3165 assert_ne!(items.r#type, "", "Items type should not be empty string!");
3166 assert_eq!(items.r#type, "object", "Items should be object type");
3167 } else {
3168 panic!("Schema should have items field for array type");
3169 }
3170 }
3171 Err(e) => println!("Schema conversion failed: {:?}", e),
3172 }
3173 }
3174
3175 #[test]
3176 fn test_object_schema() {
3177 let simple_schema = json!({
3178 "type": "object",
3179 "properties": {
3180 "name": {
3181 "type": "string"
3182 }
3183 }
3184 });
3185
3186 let schema: Schema = simple_schema.try_into().unwrap();
3187 assert_eq!(schema.r#type, "object");
3188 assert!(schema.properties.is_some());
3189 }
3190
3191 #[test]
3192 fn test_array_with_inline_items() {
3193 let inline_schema = json!({
3194 "type": "array",
3195 "items": {
3196 "type": "object",
3197 "properties": {
3198 "name": {
3199 "type": "string"
3200 }
3201 }
3202 }
3203 });
3204
3205 let schema: Schema = inline_schema.try_into().unwrap();
3206 assert_eq!(schema.r#type, "array");
3207
3208 if let Some(items) = schema.items {
3209 assert_eq!(items.r#type, "object");
3210 assert!(items.properties.is_some());
3211 } else {
3212 panic!("Schema should have items field");
3213 }
3214 }
3215 #[test]
3216 fn test_flattened_schema() {
3217 let ref_schema = json!({
3218 "type": "array",
3219 "items": {
3220 "$ref": "#/$defs/Person"
3221 },
3222 "$defs": {
3223 "Person": {
3224 "type": "object",
3225 "properties": {
3226 "name": { "type": "string" }
3227 }
3228 }
3229 }
3230 });
3231
3232 let flattened = flatten_schema(ref_schema).unwrap();
3233 let schema: Schema = flattened.try_into().unwrap();
3234
3235 assert_eq!(schema.r#type, "array");
3236
3237 if let Some(items) = schema.items {
3238 println!("Flattened items type: '{}'", items.r#type);
3239
3240 assert_eq!(items.r#type, "object");
3241 assert!(items.properties.is_some());
3242 }
3243 }
3244
3245 #[test]
3246 fn test_array_without_items_gets_default() {
3247 let schema_json = json!({
3248 "type": "object",
3249 "properties": {
3250 "service_ids": {
3251 "type": "array",
3252 "description": "A list of service IDs"
3253 }
3254 }
3255 });
3256
3257 let schema: Schema = schema_json.try_into().unwrap();
3258 let props = schema.properties.unwrap();
3259 let service_ids = props.get("service_ids").unwrap();
3260 assert_eq!(service_ids.r#type, "array");
3261 let items = service_ids
3262 .items
3263 .as_ref()
3264 .expect("array schema missing items should get a default");
3265 assert_eq!(items.r#type, "string");
3266 }
3267
3268 #[test]
3269 fn test_tool_parameters_to_schema_maps_no_arg_tool_to_none() {
3270 let schema = tool_parameters_to_schema(json!({"type": "object", "properties": {}}))
3271 .expect("schema conversion");
3272
3273 assert!(schema.is_none());
3274 }
3275
3276 #[test]
3277 fn test_tool_parameters_to_schema_resolves_defs_ref() {
3278 let schema_json = json!({
3279 "type": "object",
3280 "properties": {
3281 "destination": { "$ref": "#/$defs/Destination" }
3282 },
3283 "required": ["destination"],
3284 "$defs": {
3285 "Destination": {
3286 "type": "object",
3287 "properties": {
3288 "city": { "type": "string" }
3289 },
3290 "required": ["city"]
3291 }
3292 }
3293 });
3294
3295 let schema = tool_parameters_to_schema(schema_json)
3296 .expect("schema conversion")
3297 .expect("schema");
3298 let props = schema.properties.expect("properties");
3299 let destination = props.get("destination").expect("destination prop");
3300
3301 assert_eq!(destination.r#type, "object");
3302 assert_eq!(destination.required, Some(vec!["city".to_string()]));
3303 }
3304
3305 #[test]
3306 fn test_tool_parameters_to_schema_handles_nullable_type_arrays() {
3307 let schema_json = json!({
3308 "type": "object",
3309 "properties": {
3310 "nickname": { "type": ["null", "string"] }
3311 }
3312 });
3313
3314 let schema = tool_parameters_to_schema(schema_json)
3315 .expect("schema conversion")
3316 .expect("schema");
3317 let props = schema.properties.expect("properties");
3318 let nickname = props.get("nickname").expect("nickname prop");
3319
3320 assert_eq!(nickname.r#type, "string");
3321 assert_eq!(nickname.nullable, Some(true));
3322 }
3323
3324 #[test]
3325 fn test_txt_document_conversion_to_text_part() {
3326 use crate::message::{DocumentMediaType, UserContent};
3328
3329 let doc = UserContent::document(
3330 "Note: test.md\nPath: /test.md\nContent: Hello World!",
3331 Some(DocumentMediaType::TXT),
3332 );
3333
3334 let content: Content = message::Message::User { content: vec![doc] }
3335 .try_into()
3336 .unwrap();
3337
3338 if let Part {
3339 part: PartKind::Text(text),
3340 ..
3341 } = &content.parts[0]
3342 {
3343 assert!(text.contains("Note: test.md"));
3344 assert!(text.contains("Hello World!"));
3345 } else {
3346 panic!(
3347 "Expected text part for TXT document, got: {:?}",
3348 content.parts[0]
3349 );
3350 }
3351 }
3352
3353 #[test]
3354 fn test_tool_result_with_image_content() {
3355 use crate::message::{
3357 DocumentSourceKind, Image, ImageMediaType, ToolResult, ToolResultContent,
3358 };
3359
3360 let tool_result = ToolResult {
3362 call: message::ToolCallId::new_or_mint("call-123"),
3363 provider: message::ProviderCallId::new("call-123"),
3364 name: "test_tool".to_string(),
3365 content: vec![
3366 ToolResultContent::Text(message::Text::new(r#"{"status": "success"}"#.to_string())),
3367 ToolResultContent::Image(Image {
3368 data: DocumentSourceKind::Base64("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==".to_string()),
3369 media_type: Some(ImageMediaType::PNG),
3370 detail: None,
3371 additional_params: None,
3372 }),
3373 ],
3374 };
3375
3376 let user_content = message::UserContent::ToolResult(tool_result);
3377 let msg = message::Message::User {
3378 content: vec![user_content],
3379 };
3380
3381 let content: Content = msg.try_into().expect("Should convert to Gemini Content");
3383 assert_eq!(content.role, Some(Role::User));
3384 assert_eq!(content.parts.len(), 1);
3385
3386 if let Some(Part {
3388 part: PartKind::FunctionResponse(function_response),
3389 ..
3390 }) = content.parts.first()
3391 {
3392 assert_eq!(function_response.name, "test_tool");
3393 assert_eq!(function_response.id.as_deref(), Some("call-123"));
3394
3395 assert!(function_response.response.is_some());
3397 let response = function_response.response.as_ref().unwrap();
3398 assert_eq!(
3399 response,
3400 &json!({
3401 "result": r#"{"status": "success"}"#
3402 })
3403 );
3404
3405 assert!(function_response.parts.is_some());
3407 let parts = function_response.parts.as_ref().unwrap();
3408 assert_eq!(parts.len(), 1);
3409
3410 let image_part = &parts[0];
3411 assert!(image_part.inline_data.is_some());
3412 let inline_data = image_part.inline_data.as_ref().unwrap();
3413 assert_eq!(inline_data.mime_type, "image/png");
3414 assert!(!inline_data.data.is_empty());
3415 assert_eq!(inline_data.display_name, None);
3416 } else {
3417 panic!("Expected FunctionResponse part");
3418 }
3419 }
3420
3421 #[test]
3422 fn mixed_inline_images_and_text_keep_text_response_and_ordered_parts() {
3423 use crate::message::{ImageMediaType, ToolResult, ToolResultContent};
3424
3425 let message = message::Message::User {
3426 content: vec![message::UserContent::ToolResult(ToolResult {
3427 call: message::ToolCallId::mint(),
3428 provider: None,
3429 name: "ordered_tool".to_string(),
3430 content: vec![
3431 ToolResultContent::image_base64("first-image", Some(ImageMediaType::PNG), None),
3432 ToolResultContent::text("between-images"),
3433 ToolResultContent::image_base64(
3434 "second-image",
3435 Some(ImageMediaType::JPEG),
3436 None,
3437 ),
3438 ],
3439 })],
3440 };
3441
3442 let content: Content = message.try_into().expect("tool result should convert");
3443 let PartKind::FunctionResponse(response) = &content.parts[0].part else {
3444 panic!("expected a function response");
3445 };
3446
3447 assert_eq!(
3448 response.response,
3449 Some(json!({ "result": "between-images" }))
3450 );
3451
3452 let parts = response
3453 .parts
3454 .as_ref()
3455 .expect("images should be inline parts");
3456 assert_eq!(parts.len(), 2);
3457 let first = parts[0].inline_data.as_ref().expect("first inline image");
3458 assert_eq!(first.mime_type, "image/png");
3459 assert_eq!(first.data, "first-image");
3460 assert_eq!(first.display_name, None);
3461 let second = parts[1].inline_data.as_ref().expect("second inline image");
3462 assert_eq!(second.mime_type, "image/jpeg");
3463 assert_eq!(second.data, "second-image");
3464 assert_eq!(second.display_name, None);
3465 }
3466
3467 #[test]
3468 fn mixed_inline_image_and_json_keep_structured_value_and_media_part() {
3469 use crate::message::{ImageMediaType, ToolResult, ToolResultContent};
3470
3471 let message = message::Message::User {
3472 content: vec![message::UserContent::ToolResult(ToolResult {
3473 call: message::ToolCallId::mint(),
3474 provider: None,
3475 name: "ordered_tool".to_string(),
3476 content: vec![
3477 ToolResultContent::json(json!({ "status": "ok" })),
3478 ToolResultContent::image_base64("image-data", Some(ImageMediaType::PNG), None),
3479 ],
3480 })],
3481 };
3482
3483 let content: Content = message.try_into().expect("tool result should convert");
3484 let PartKind::FunctionResponse(response) = &content.parts[0].part else {
3485 panic!("expected a function response");
3486 };
3487
3488 assert_eq!(
3489 response.response,
3490 Some(json!({ "result": { "status": "ok" } }))
3491 );
3492 let parts = response
3493 .parts
3494 .as_ref()
3495 .expect("image should be an inline part");
3496 assert_eq!(parts.len(), 1);
3497 let inline_data = parts[0].inline_data.as_ref().expect("inline image data");
3498 assert_eq!(inline_data.data, "image-data");
3499 assert_eq!(inline_data.display_name, None);
3500 }
3501
3502 #[test]
3503 fn mixed_url_image_and_response_value_is_rejected() {
3504 use crate::message::{DocumentSourceKind, Image, ImageMediaType, ToolResultContent};
3505
3506 let tool_result = message::Message::User {
3507 content: vec![message::UserContent::ToolResult(message::ToolResult {
3508 call: message::ToolCallId::mint(),
3509 provider: None,
3510 name: "url_tool".to_string(),
3511 content: vec![
3512 ToolResultContent::Image(Image {
3513 data: DocumentSourceKind::Url("https://example.com/image.png".to_string()),
3514 media_type: Some(ImageMediaType::PNG),
3515 detail: None,
3516 additional_params: None,
3517 }),
3518 ToolResultContent::text("after-image"),
3519 ],
3520 })],
3521 };
3522
3523 let error = Content::try_from(tool_result)
3524 .expect_err("URL-backed tool result images should be rejected");
3525 assert!(
3526 error
3527 .to_string()
3528 .contains("URL-backed images are not supported"),
3529 "unexpected error: {error}"
3530 );
3531 }
3532
3533 #[test]
3534 fn tool_result_rejects_unsupported_image_media_types() {
3535 use crate::message::{ImageMediaType, ToolResult, ToolResultContent};
3536
3537 for media_type in [
3538 ImageMediaType::GIF,
3539 ImageMediaType::HEIC,
3540 ImageMediaType::HEIF,
3541 ImageMediaType::SVG,
3542 ] {
3543 let message = message::Message::User {
3544 content: vec![message::UserContent::ToolResult(ToolResult {
3545 call: message::ToolCallId::mint(),
3546 provider: None,
3547 name: "image_tool".to_string(),
3548 content: vec![ToolResultContent::image_base64(
3549 "image-data",
3550 Some(media_type),
3551 None,
3552 )],
3553 })],
3554 };
3555
3556 let error = Content::try_from(message)
3557 .expect_err("unsupported tool result image type should be rejected");
3558 assert!(
3559 error
3560 .to_string()
3561 .contains("supported types are JPEG, PNG, and WEBP"),
3562 "unexpected error: {error}"
3563 );
3564 }
3565 }
3566
3567 #[test]
3568 fn structured_json_refs_remain_literal_with_unreferenced_image_parts() {
3569 use crate::message::{ImageMediaType, ToolResult, ToolResultContent};
3570
3571 let message = message::Message::User {
3572 content: vec![message::UserContent::ToolResult(ToolResult {
3573 call: message::ToolCallId::mint(),
3574 provider: None,
3575 name: "collision_tool".to_string(),
3576 content: vec![
3577 ToolResultContent::json(json!({
3578 "literal": {
3579 "$ref": "tool_result_image_0"
3580 }
3581 })),
3582 ToolResultContent::image_base64("image-data", Some(ImageMediaType::PNG), None),
3583 ],
3584 })],
3585 };
3586
3587 let content: Content = message.try_into().expect("tool result should convert");
3588 let PartKind::FunctionResponse(response) = &content.parts[0].part else {
3589 panic!("expected a function response");
3590 };
3591
3592 assert_eq!(
3593 response.response,
3594 Some(json!({
3595 "result": {
3596 "literal": {
3597 "$ref": "tool_result_image_0"
3598 }
3599 }
3600 }))
3601 );
3602 assert_eq!(
3603 response.parts.as_ref().and_then(|parts| {
3604 parts
3605 .first()
3606 .and_then(|part| part.inline_data.as_ref())
3607 .and_then(|part| part.display_name.as_deref())
3608 }),
3609 None
3610 );
3611 }
3612
3613 #[test]
3614 fn tool_result_literal_text_and_structured_json_remain_distinct() {
3615 use crate::message::{ToolResult, ToolResultContent};
3616
3617 let cases = [
3618 (
3619 ToolResultContent::text(r#"{"status":"ok"}"#),
3620 json!({ "result": "{\"status\":\"ok\"}" }),
3621 ),
3622 (
3623 ToolResultContent::json(json!({ "status": "ok" })),
3624 json!({ "result": { "status": "ok" } }),
3625 ),
3626 ];
3627
3628 for (tool_content, expected) in cases {
3629 let message = message::Message::User {
3630 content: vec![message::UserContent::ToolResult(ToolResult {
3631 call: message::ToolCallId::mint(),
3632 provider: None,
3633 name: "test_tool".to_string(),
3634 content: vec![tool_content],
3635 })],
3636 };
3637 let content: Content = message.try_into().expect("tool result should convert");
3638
3639 let PartKind::FunctionResponse(response) = &content.parts[0].part else {
3640 panic!("expected a function response");
3641 };
3642 assert_eq!(response.response.as_ref(), Some(&expected));
3643 }
3644 }
3645
3646 #[test]
3651 fn echoed_minted_handle_never_reaches_the_function_response_id() {
3652 use crate::message::{ToolCall, ToolCallId, ToolFunction, ToolResultContent};
3653
3654 let call = ToolCall::new(
3656 ToolCallId::mint(),
3657 ToolFunction {
3658 name: "lookup".to_string(),
3659 arguments: json!({}),
3660 },
3661 );
3662
3663 let message = message::Message::User {
3664 content: vec![message::UserContent::tool_result(
3665 call.id.as_str(),
3666 "lookup",
3667 vec![ToolResultContent::text("out")],
3668 )],
3669 };
3670 let content: Content = message.try_into().expect("tool result should convert");
3671 let PartKind::FunctionResponse(response) = &content.parts[0].part else {
3672 panic!("expected a function response");
3673 };
3674 assert_eq!(response.id, None);
3675 }
3676
3677 #[test]
3682 fn ingested_nameless_results_resolve_their_name_at_request_assembly() {
3683 use crate::completion::request::CompletionRequest;
3684 use crate::message::{AssistantContent, ToolCall, ToolFunction, ToolResultContent};
3685
3686 let request = CompletionRequest {
3687 preamble: None,
3688 chat_history: vec![
3689 message::Message::user("weather?"),
3690 message::Message::Assistant {
3691 id: None,
3692 content: vec![AssistantContent::ToolCall(ToolCall::from_wire(
3693 "toolu_abc",
3694 ToolFunction {
3695 name: "get_weather".to_owned(),
3696 arguments: json!({"city": "Paris"}),
3697 },
3698 ))],
3699 },
3700 message::Message::User {
3701 content: vec![message::UserContent::tool_result_from_wire(
3702 "toolu_abc",
3703 "",
3704 vec![ToolResultContent::text("sunny")],
3705 )],
3706 },
3707 ],
3708 documents: vec![],
3709 tools: vec![],
3710 temperature: None,
3711 model: None,
3712 output_schema: None,
3713 record_telemetry_content: false,
3714 max_tokens: None,
3715 tool_choice: None,
3716 additional_params: None,
3717 };
3718
3719 let body = create_request_body(request).expect("request should build");
3720 let response_names: Vec<_> = body
3721 .contents
3722 .iter()
3723 .flat_map(|content| &content.parts)
3724 .filter_map(|part| match &part.part {
3725 PartKind::FunctionResponse(response) => Some(response.name.clone()),
3726 _ => None,
3727 })
3728 .collect();
3729 assert_eq!(response_names, ["get_weather"]);
3730 }
3731
3732 #[test]
3734 fn wire_derived_tool_result_keeps_the_provider_id_on_the_wire() {
3735 use crate::message::ToolResultContent;
3736
3737 let message = message::Message::User {
3738 content: vec![message::UserContent::tool_result_from_wire(
3739 "gemini-issued-id",
3740 "lookup",
3741 vec![ToolResultContent::text("out")],
3742 )],
3743 };
3744 let content: Content = message.try_into().expect("tool result should convert");
3745 let PartKind::FunctionResponse(response) = &content.parts[0].part else {
3746 panic!("expected a function response");
3747 };
3748 assert_eq!(response.id.as_deref(), Some("gemini-issued-id"));
3749 }
3750
3751 #[test]
3752 fn test_markdown_document_conversion_to_text_part() {
3753 use crate::message::{DocumentMediaType, UserContent};
3755
3756 let doc = UserContent::document(
3757 "# Heading\n\n* List item",
3758 Some(DocumentMediaType::MARKDOWN),
3759 );
3760
3761 let content: Content = message::Message::User { content: vec![doc] }
3762 .try_into()
3763 .unwrap();
3764
3765 if let Part {
3766 part: PartKind::Text(text),
3767 ..
3768 } = &content.parts[0]
3769 {
3770 assert_eq!(text, "# Heading\n\n* List item");
3771 } else {
3772 panic!(
3773 "Expected text part for MARKDOWN document, got: {:?}",
3774 content.parts[0]
3775 );
3776 }
3777 }
3778
3779 #[test]
3780 fn test_markdown_url_document_conversion_to_file_data_part() {
3781 use crate::message::{DocumentMediaType, DocumentSourceKind, UserContent};
3783
3784 let doc = UserContent::Document(message::Document {
3785 data: DocumentSourceKind::Url(
3786 "https://generativelanguage.googleapis.com/v1beta/files/test-markdown".to_string(),
3787 ),
3788 media_type: Some(DocumentMediaType::MARKDOWN),
3789 additional_params: None,
3790 });
3791
3792 let content: Content = message::Message::User { content: vec![doc] }
3793 .try_into()
3794 .unwrap();
3795
3796 if let Part {
3797 part: PartKind::FileData(file_data),
3798 ..
3799 } = &content.parts[0]
3800 {
3801 assert_eq!(
3802 file_data.file_uri,
3803 "https://generativelanguage.googleapis.com/v1beta/files/test-markdown"
3804 );
3805 assert_eq!(file_data.mime_type.as_deref(), Some("text/markdown"));
3806 } else {
3807 panic!(
3808 "Expected file_data part for URL MARKDOWN document, got: {:?}",
3809 content.parts[0]
3810 );
3811 }
3812 }
3813
3814 #[test]
3815 fn test_tool_result_with_url_image_is_rejected() {
3816 use crate::message::{
3817 DocumentSourceKind, Image, ImageMediaType, ToolResult, ToolResultContent,
3818 };
3819
3820 let tool_result = ToolResult {
3821 call: message::ToolCallId::mint(),
3822 provider: None,
3823 name: "screenshot_tool".to_string(),
3824 content: vec![ToolResultContent::Image(Image {
3825 data: DocumentSourceKind::Url("https://example.com/image.png".to_string()),
3826 media_type: Some(ImageMediaType::PNG),
3827 detail: None,
3828 additional_params: None,
3829 })],
3830 };
3831
3832 let user_content = message::UserContent::ToolResult(tool_result);
3833 let msg = message::Message::User {
3834 content: vec![user_content],
3835 };
3836
3837 let error =
3838 Content::try_from(msg).expect_err("URL-backed tool result images should be rejected");
3839 assert!(
3840 error
3841 .to_string()
3842 .contains("URL-backed images are not supported"),
3843 "unexpected error: {error}"
3844 );
3845 }
3846
3847 #[test]
3848 fn test_create_request_body_with_documents() {
3849 use crate::completion::request::{CompletionRequest, Document};
3851 use crate::message::Message;
3852
3853 let documents = vec![
3854 Document {
3855 id: "doc1".to_string(),
3856 text: "Note: first.md\nContent: First note".to_string(),
3857 additional_props: std::collections::HashMap::new(),
3858 },
3859 Document {
3860 id: "doc2".to_string(),
3861 text: "Note: second.md\nContent: Second note".to_string(),
3862 additional_props: std::collections::HashMap::new(),
3863 },
3864 ];
3865
3866 let documents_message = CompletionRequest {
3867 preamble: None,
3868 chat_history: vec![Message::user("placeholder")],
3869 documents,
3870 tools: vec![],
3871 temperature: None,
3872 model: None,
3873 output_schema: None,
3874 record_telemetry_content: false,
3875 max_tokens: None,
3876 tool_choice: None,
3877 additional_params: None,
3878 }
3879 .normalized_documents()
3880 .unwrap();
3881
3882 let completion_request = CompletionRequest {
3883 preamble: Some("You are a helpful assistant".to_string()),
3884 chat_history: vec![documents_message, Message::user("What are my notes about?")],
3885 documents: vec![],
3886 tools: vec![],
3887 temperature: None,
3888 model: None,
3889 output_schema: None,
3890 record_telemetry_content: false,
3891 max_tokens: None,
3892 tool_choice: None,
3893 additional_params: None,
3894 };
3895
3896 let request = create_request_body(completion_request).unwrap();
3897
3898 assert_eq!(
3900 request.contents.len(),
3901 2,
3902 "Expected 2 contents (documents + user message)"
3903 );
3904
3905 assert_eq!(request.contents[0].role, Some(Role::User));
3907 assert_eq!(
3908 request.contents[0].parts.len(),
3909 2,
3910 "Expected 2 document parts"
3911 );
3912
3913 for part in &request.contents[0].parts {
3915 if let Part {
3916 part: PartKind::Text(text),
3917 ..
3918 } = part
3919 {
3920 assert!(
3921 text.contains("Note:") && text.contains("Content:"),
3922 "Document should contain note metadata"
3923 );
3924 } else {
3925 panic!("Document parts should be text, not {:?}", part);
3926 }
3927 }
3928
3929 assert_eq!(request.contents[1].role, Some(Role::User));
3931 if let Part {
3932 part: PartKind::Text(text),
3933 ..
3934 } = &request.contents[1].parts[0]
3935 {
3936 assert_eq!(text, "What are my notes about?");
3937 } else {
3938 panic!("Expected user message to be text");
3939 }
3940 }
3941
3942 #[test]
3943 fn test_create_request_body_without_documents() {
3944 use crate::completion::request::CompletionRequest;
3946 use crate::message::Message;
3947
3948 let completion_request = CompletionRequest {
3949 preamble: Some("You are a helpful assistant".to_string()),
3950 chat_history: vec![Message::user("Hello")],
3951 documents: vec![], tools: vec![],
3953 temperature: None,
3954 max_tokens: None,
3955 tool_choice: None,
3956 model: None,
3957 output_schema: None,
3958 record_telemetry_content: false,
3959 additional_params: None,
3960 };
3961
3962 let request = create_request_body(completion_request).unwrap();
3963
3964 assert_eq!(request.contents.len(), 1, "Expected only user message");
3966 assert_eq!(request.contents[0].role, Some(Role::User));
3967
3968 if let Part {
3969 part: PartKind::Text(text),
3970 ..
3971 } = &request.contents[0].parts[0]
3972 {
3973 assert_eq!(text, "Hello");
3974 } else {
3975 panic!("Expected user message to be text");
3976 }
3977 }
3978
3979 #[tokio::test]
3980 async fn completion_non_success_preserves_status_and_body() {
3981 use crate::client::completion::CompletionClient;
3982 use crate::completion::CompletionModel as _;
3983 use crate::providers::gemini::Client;
3984 use crate::test_utils::RecordingHttpClient;
3985
3986 let body = r#"{"error":{"code":503,"message":"boom","status":"UNAVAILABLE"}}"#;
3987 let http_client =
3988 RecordingHttpClient::with_error_response(http::StatusCode::SERVICE_UNAVAILABLE, body);
3989 let client = Client::builder()
3990 .api_key("test-key")
3991 .http_client(http_client)
3992 .build()
3993 .expect("build client");
3994 let model = client.completion_model(super::GEMINI_3_FLASH_PREVIEW);
3995 let request = model.completion_request("hello").build();
3996
3997 let error = model
3998 .completion(request)
3999 .await
4000 .expect_err("should fail with non-success status");
4001
4002 assert!(matches!(error, CompletionError::HttpError(_)));
4003 assert_eq!(
4004 error.provider_response_status(),
4005 Some(http::StatusCode::SERVICE_UNAVAILABLE)
4006 );
4007 assert_eq!(error.provider_response_body(), Some(body));
4008 }
4009}