1use crate::attachment;
2use crate::retry::{RetryConfig, execute_with_retry, is_retryable_model_error};
3use adk_core::{
4 CacheCapable, CitationMetadata, CitationSource, Content, ErrorCategory, ErrorComponent,
5 FinishReason, Llm, LlmRequest, LlmResponse, LlmResponseStream, Part, Result, SchemaAdapter,
6 SchemaCache, UsageMetadata,
7};
8use adk_gemini::Gemini;
9use adk_gemini::schema_adapter::GeminiSchemaAdapter;
10use async_trait::async_trait;
11use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64_STANDARD};
12use futures::TryStreamExt;
13
14#[cfg(feature = "gemini-interactions")]
15use super::interactions_target::InteractionTarget;
16
17#[cfg(feature = "gemini-interactions")]
27#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
28pub enum GeminiTransport {
29 #[default]
31 GenerateContent,
32 Interactions,
34}
35
36#[cfg(feature = "gemini-interactions")]
45#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
46pub enum BackgroundMode {
47 #[default]
51 AgentTargetsOnly,
52 Always,
54 Never,
56}
57
58#[cfg(feature = "gemini-interactions")]
69#[derive(Debug, Clone)]
70pub struct InteractionOptions {
71 pub store: bool,
73 pub stateful: bool,
76 pub background: BackgroundMode,
78 pub poll_interval: std::time::Duration,
80}
81
82#[cfg(feature = "gemini-interactions")]
83impl Default for InteractionOptions {
84 fn default() -> Self {
85 Self {
86 store: true,
87 stateful: true,
88 background: BackgroundMode::AgentTargetsOnly,
89 poll_interval: std::time::Duration::from_secs(1),
90 }
91 }
92}
93
94pub struct GeminiModel {
96 client: Gemini,
97 model_name: String,
98 retry_config: RetryConfig,
99 thinking_config: Option<adk_gemini::ThinkingConfig>,
105 #[cfg(feature = "gemini-interactions")]
109 transport: GeminiTransport,
110 #[cfg(feature = "gemini-interactions")]
113 interaction_target: Option<InteractionTarget>,
114 #[cfg(feature = "gemini-interactions")]
116 interaction_options: InteractionOptions,
117}
118
119fn gemini_error_to_adk(e: &adk_gemini::ClientError) -> adk_core::AdkError {
121 fn format_error_chain(e: &dyn std::error::Error) -> String {
122 let mut msg = e.to_string();
123 let mut source = e.source();
124 while let Some(s) = source {
125 msg.push_str(": ");
126 msg.push_str(&s.to_string());
127 source = s.source();
128 }
129 msg
130 }
131
132 let message = format_error_chain(e);
133
134 let (category, code, status_code) = if message.contains("code 429")
137 || message.contains("RESOURCE_EXHAUSTED")
138 || message.contains("rate limit")
139 {
140 (ErrorCategory::RateLimited, "model.gemini.rate_limited", Some(429u16))
141 } else if message.contains("code 503") || message.contains("UNAVAILABLE") {
142 (ErrorCategory::Unavailable, "model.gemini.unavailable", Some(503))
143 } else if message.contains("code 529") || message.contains("OVERLOADED") {
144 (ErrorCategory::Unavailable, "model.gemini.overloaded", Some(529))
145 } else if message.contains("code 408")
146 || message.contains("DEADLINE_EXCEEDED")
147 || message.contains("TIMEOUT")
148 {
149 (ErrorCategory::Timeout, "model.gemini.timeout", Some(408))
150 } else if message.contains("code 401") || message.contains("Invalid API key") {
151 (ErrorCategory::Unauthorized, "model.gemini.unauthorized", Some(401))
152 } else if message.contains("code 400") {
153 (ErrorCategory::InvalidInput, "model.gemini.bad_request", Some(400))
154 } else if message.contains("code 404") {
155 (ErrorCategory::NotFound, "model.gemini.not_found", Some(404))
156 } else if message.contains("invalid generation config") {
157 (ErrorCategory::InvalidInput, "model.gemini.invalid_config", None)
158 } else {
159 (ErrorCategory::Internal, "model.gemini.internal", None)
160 };
161
162 let mut err = adk_core::AdkError::new(ErrorComponent::Model, category, code, message)
163 .with_provider("gemini");
164 if let Some(sc) = status_code {
165 err = err.with_upstream_status(sc);
166 }
167 err
168}
169
170#[cfg(feature = "gemini-interactions")]
181fn interaction_status_to_result(status: adk_gemini::interactions::InteractionStatus) -> Result<()> {
182 use adk_gemini::interactions::InteractionStatus;
183 match status {
184 InteractionStatus::Failed => Err(interaction_terminal_error(
185 "model.gemini.interactions.failed",
186 "the Gemini interaction terminated with status `failed`",
187 )),
188 InteractionStatus::BudgetExceeded => Err(interaction_terminal_error(
189 "model.gemini.interactions.budget_exceeded",
190 "the Gemini interaction terminated with status `budget_exceeded`",
191 )),
192 InteractionStatus::InProgress
193 | InteractionStatus::RequiresAction
194 | InteractionStatus::Completed
195 | InteractionStatus::Cancelled
196 | InteractionStatus::Incomplete => Ok(()),
197 }
198}
199
200#[cfg(feature = "gemini-interactions")]
202fn interaction_terminal_error(code: &'static str, message: &str) -> adk_core::AdkError {
203 adk_core::AdkError::new(
204 ErrorComponent::Model,
205 ErrorCategory::Internal,
206 code,
207 message.to_string(),
208 )
209 .with_provider("gemini")
210}
211
212impl GeminiModel {
213 fn gemini_part_thought_signature(value: &serde_json::Value) -> Option<String> {
214 value.get("thoughtSignature").and_then(serde_json::Value::as_str).map(str::to_string)
215 }
216
217 fn from_client(client: Gemini, model_name: String) -> Self {
223 Self {
224 client,
225 model_name,
226 retry_config: RetryConfig::default(),
227 thinking_config: None,
228 #[cfg(feature = "gemini-interactions")]
229 transport: GeminiTransport::GenerateContent,
230 #[cfg(feature = "gemini-interactions")]
231 interaction_target: None,
232 #[cfg(feature = "gemini-interactions")]
233 interaction_options: InteractionOptions::default(),
234 }
235 }
236
237 pub fn new(api_key: impl Into<String>, model: impl Into<String>) -> Result<Self> {
239 let model_name = model.into();
240 let client = Gemini::with_model(api_key.into(), model_name.clone())
241 .map_err(|e| adk_core::AdkError::model(e.to_string()))?;
242
243 Ok(Self::from_client(client, model_name))
244 }
245
246 #[cfg(feature = "gemini-vertex")]
250 pub fn new_google_cloud(
251 api_key: impl Into<String>,
252 project_id: impl AsRef<str>,
253 location: impl AsRef<str>,
254 model: impl Into<String>,
255 ) -> Result<Self> {
256 let model_name = model.into();
257 let client = Gemini::with_google_cloud_model(
258 api_key.into(),
259 project_id,
260 location,
261 model_name.clone(),
262 )
263 .map_err(|e| adk_core::AdkError::model(e.to_string()))?;
264
265 Ok(Self::from_client(client, model_name))
266 }
267
268 #[cfg(feature = "gemini-vertex")]
272 pub fn new_google_cloud_service_account(
273 service_account_json: &str,
274 project_id: impl AsRef<str>,
275 location: impl AsRef<str>,
276 model: impl Into<String>,
277 ) -> Result<Self> {
278 let model_name = model.into();
279 let client = Gemini::with_google_cloud_service_account_json(
280 service_account_json,
281 project_id.as_ref(),
282 location.as_ref(),
283 model_name.clone(),
284 )
285 .map_err(|e| adk_core::AdkError::model(e.to_string()))?;
286
287 Ok(Self::from_client(client, model_name))
288 }
289
290 #[cfg(feature = "gemini-vertex")]
294 pub fn new_google_cloud_adc(
295 project_id: impl AsRef<str>,
296 location: impl AsRef<str>,
297 model: impl Into<String>,
298 ) -> Result<Self> {
299 let model_name = model.into();
300 let client = Gemini::with_google_cloud_adc_model(
301 project_id.as_ref(),
302 location.as_ref(),
303 model_name.clone(),
304 )
305 .map_err(|e| adk_core::AdkError::model(e.to_string()))?;
306
307 Ok(Self::from_client(client, model_name))
308 }
309
310 #[cfg(feature = "gemini-vertex")]
314 pub fn new_google_cloud_wif(
315 wif_json: &str,
316 project_id: impl AsRef<str>,
317 location: impl AsRef<str>,
318 model: impl Into<String>,
319 ) -> Result<Self> {
320 let model_name = model.into();
321 let client = Gemini::with_google_cloud_wif_json(
322 wif_json,
323 project_id.as_ref(),
324 location.as_ref(),
325 model_name.clone(),
326 )
327 .map_err(|e| adk_core::AdkError::model(e.to_string()))?;
328
329 Ok(Self::from_client(client, model_name))
330 }
331
332 #[must_use]
334 pub fn with_retry_config(mut self, retry_config: RetryConfig) -> Self {
335 self.retry_config = retry_config;
336 self
337 }
338
339 pub fn set_retry_config(&mut self, retry_config: RetryConfig) {
341 self.retry_config = retry_config;
342 }
343
344 pub fn retry_config(&self) -> &RetryConfig {
346 &self.retry_config
347 }
348
349 #[must_use]
373 pub fn with_thinking_config(mut self, thinking_config: adk_gemini::ThinkingConfig) -> Self {
374 self.thinking_config = Some(thinking_config);
375 self
376 }
377
378 pub fn set_thinking_config(&mut self, thinking_config: adk_gemini::ThinkingConfig) {
380 self.thinking_config = Some(thinking_config);
381 }
382
383 pub fn thinking_config(&self) -> Option<&adk_gemini::ThinkingConfig> {
385 self.thinking_config.as_ref()
386 }
387
388 #[cfg(feature = "gemini-interactions")]
412 pub fn use_interactions_api(mut self, enabled: bool) -> Result<Self> {
413 if enabled {
414 let target = InteractionTarget::parse(&self.model_name)?;
415 self.transport = GeminiTransport::Interactions;
416 self.interaction_target = Some(target);
417 } else {
418 self.transport = GeminiTransport::GenerateContent;
419 self.interaction_target = None;
420 }
421 Ok(self)
422 }
423
424 #[cfg(feature = "gemini-interactions")]
431 #[must_use]
432 pub fn interaction_options(mut self, opts: InteractionOptions) -> Self {
433 self.interaction_options = opts;
434 self
435 }
436
437 #[cfg(feature = "gemini-interactions")]
441 pub fn transport(&self) -> GeminiTransport {
442 self.transport
443 }
444
445 #[cfg(feature = "gemini-interactions")]
449 pub fn interaction_options_ref(&self) -> &InteractionOptions {
450 &self.interaction_options
451 }
452
453 #[cfg(feature = "gemini-interactions")]
465 fn resolve_background(&self) -> bool {
466 match self.interaction_options.background {
467 BackgroundMode::Always => true,
468 BackgroundMode::Never => false,
469 BackgroundMode::AgentTargetsOnly => {
470 self.interaction_target.as_ref().is_some_and(InteractionTarget::is_agent)
471 }
472 }
473 }
474
475 fn convert_response(resp: &adk_gemini::GenerationResponse) -> Result<LlmResponse> {
476 let mut converted_parts: Vec<Part> = Vec::new();
477
478 if let Some(parts) = resp.candidates.first().and_then(|c| c.content.parts.as_ref()) {
480 for p in parts {
481 match p {
482 adk_gemini::Part::Text { text, thought, thought_signature } => {
483 if thought == &Some(true) {
484 converted_parts.push(Part::Thinking {
485 thinking: text.clone(),
486 signature: thought_signature.clone(),
487 });
488 } else {
489 converted_parts.push(Part::Text { text: text.clone() });
490 }
491 }
492 adk_gemini::Part::InlineData { inline_data } => {
493 let decoded =
494 BASE64_STANDARD.decode(&inline_data.data).map_err(|error| {
495 adk_core::AdkError::model(format!(
496 "failed to decode inline data from gemini response: {error}"
497 ))
498 })?;
499 converted_parts.push(Part::InlineData {
500 mime_type: inline_data.mime_type.clone(),
501 data: decoded,
502 });
503 }
504 adk_gemini::Part::FunctionCall { function_call, thought_signature } => {
505 converted_parts.push(Part::FunctionCall {
506 name: function_call.name.clone(),
507 args: function_call.args.clone(),
508 id: function_call.id.clone(),
509 thought_signature: thought_signature.clone(),
510 });
511 }
512 adk_gemini::Part::FunctionResponse { function_response, .. } => {
513 converted_parts.push(Part::FunctionResponse {
514 function_response: adk_core::FunctionResponseData::new(
515 function_response.name.clone(),
516 function_response
517 .response
518 .clone()
519 .unwrap_or(serde_json::Value::Null),
520 ),
521 id: None,
522 });
523 }
524 adk_gemini::Part::ToolCall { .. } | adk_gemini::Part::ExecutableCode { .. } => {
525 if let Ok(value) = serde_json::to_value(p) {
526 converted_parts.push(Part::ServerToolCall { server_tool_call: value });
527 }
528 }
529 adk_gemini::Part::ToolResponse { .. }
530 | adk_gemini::Part::CodeExecutionResult { .. } => {
531 let value = serde_json::to_value(p).unwrap_or(serde_json::Value::Null);
532 converted_parts
533 .push(Part::ServerToolResponse { server_tool_response: value });
534 }
535 adk_gemini::Part::FileData { file_data } => {
536 converted_parts.push(Part::FileData {
537 mime_type: file_data.mime_type.clone(),
538 file_uri: file_data.file_uri.clone(),
539 });
540 }
541 }
542 }
543 }
544
545 if let Some(grounding) = resp.candidates.first().and_then(|c| c.grounding_metadata.as_ref())
547 {
548 if let Some(queries) = &grounding.web_search_queries
549 && !queries.is_empty()
550 {
551 let search_info = format!("\n\n🔍 **Searched:** {}", queries.join(", "));
552 converted_parts.push(Part::Text { text: search_info });
553 }
554 if let Some(chunks) = &grounding.grounding_chunks {
555 let sources: Vec<String> = chunks
556 .iter()
557 .filter_map(|c| {
558 c.web.as_ref().and_then(|w| match (&w.title, &w.uri) {
559 (Some(title), Some(uri)) => Some(format!("[{}]({})", title, uri)),
560 (Some(title), None) => Some(title.clone()),
561 (None, Some(uri)) => Some(uri.to_string()),
562 (None, None) => None,
563 })
564 })
565 .collect();
566 if !sources.is_empty() {
567 let sources_info = format!("\n📚 **Sources:** {}", sources.join(" | "));
568 converted_parts.push(Part::Text { text: sources_info });
569 }
570 }
571 }
572
573 let content = if converted_parts.is_empty() {
574 None
575 } else {
576 Some(Content { role: "model".to_string(), parts: converted_parts })
577 };
578
579 let usage_metadata = resp.usage_metadata.as_ref().map(|u| UsageMetadata {
580 prompt_token_count: u.prompt_token_count.unwrap_or(0),
581 candidates_token_count: u.candidates_token_count.unwrap_or(0),
582 total_token_count: u.total_token_count.unwrap_or(0),
583 thinking_token_count: u.thoughts_token_count,
584 cache_read_input_token_count: u.cached_content_token_count,
585 ..Default::default()
586 });
587
588 let finish_reason =
589 resp.candidates.first().and_then(|c| c.finish_reason.as_ref()).map(|fr| match fr {
590 adk_gemini::FinishReason::Stop => FinishReason::Stop,
591 adk_gemini::FinishReason::MaxTokens => FinishReason::MaxTokens,
592 adk_gemini::FinishReason::Safety => FinishReason::Safety,
593 adk_gemini::FinishReason::Recitation => FinishReason::Recitation,
594 _ => FinishReason::Other,
595 });
596
597 let citation_metadata =
598 resp.candidates.first().and_then(|c| c.citation_metadata.as_ref()).map(|meta| {
599 CitationMetadata {
600 citation_sources: meta
601 .citation_sources
602 .iter()
603 .map(|source| CitationSource {
604 uri: source.uri.clone(),
605 title: source.title.clone(),
606 start_index: source.start_index,
607 end_index: source.end_index,
608 license: source.license.clone(),
609 publication_date: source.publication_date.map(|d| d.to_string()),
610 })
611 .collect(),
612 }
613 });
614
615 let provider_metadata = resp
618 .candidates
619 .first()
620 .and_then(|c| c.grounding_metadata.as_ref())
621 .and_then(|g| serde_json::to_value(g).ok());
622
623 Ok(LlmResponse {
624 content,
625 usage_metadata,
626 finish_reason,
627 citation_metadata,
628 partial: false,
629 turn_complete: true,
630 interrupted: false,
631 error_code: None,
632 error_message: None,
633 provider_metadata,
634 interaction_id: None,
635 })
636 }
637
638 fn gemini_function_response_payload(response: serde_json::Value) -> serde_json::Value {
639 match response {
640 serde_json::Value::Object(_) => response,
642 other => serde_json::json!({ "result": other }),
643 }
644 }
645
646 fn merge_object_value(
647 target: &mut serde_json::Map<String, serde_json::Value>,
648 value: serde_json::Value,
649 ) {
650 if let serde_json::Value::Object(object) = value {
651 for (key, value) in object {
652 target.insert(key, value);
653 }
654 }
655 }
656
657 fn build_gemini_tools(
658 tools: &std::collections::HashMap<String, serde_json::Value>,
659 adapter: &dyn SchemaAdapter,
660 cache: &SchemaCache,
661 ) -> Result<(Vec<adk_gemini::Tool>, adk_gemini::ToolConfig)> {
662 let mut gemini_tools = Vec::new();
663 let mut function_declarations = Vec::new();
664 let mut has_provider_native_tools = false;
665 let mut tool_config_json = serde_json::Map::new();
666
667 for (name, tool_decl) in tools {
668 if let Some(provider_tool) = tool_decl.get("x-adk-gemini-tool") {
669 let tool = serde_json::from_value::<adk_gemini::Tool>(provider_tool.clone())
670 .map_err(|error| {
671 adk_core::AdkError::model(format!(
672 "failed to deserialize Gemini native tool '{name}': {error}"
673 ))
674 })?;
675 has_provider_native_tools = true;
676 gemini_tools.push(tool);
677 } else {
678 let normalized_name = adapter.normalize_tool_name(name);
680
681 let schema =
684 tool_decl.get("parameters").cloned().unwrap_or_else(|| adapter.empty_schema());
685 let normalized_schema = cache.get_or_normalize(&schema, adapter);
686
687 let description =
689 tool_decl.get("description").and_then(|v| v.as_str()).unwrap_or("").to_string();
690
691 let mut func_decl_json = serde_json::json!({
692 "name": normalized_name.as_ref(),
693 "description": description,
694 "parameters": normalized_schema,
695 });
696
697 if let Some(response) = tool_decl.get("response") {
699 func_decl_json["response"] = cache.get_or_normalize(response, adapter);
700 }
701
702 if let Some(behavior) = tool_decl.get("behavior") {
704 func_decl_json["behavior"] = behavior.clone();
705 }
706
707 let func_decl =
708 serde_json::from_value::<adk_gemini::FunctionDeclaration>(func_decl_json)
709 .map_err(|error| {
710 adk_core::AdkError::model(format!(
711 "failed to build Gemini function declaration for '{name}': {error}"
712 ))
713 })?;
714 function_declarations.push(func_decl);
715 }
716
717 if let Some(tool_config) = tool_decl.get("x-adk-gemini-tool-config") {
718 Self::merge_object_value(&mut tool_config_json, tool_config.clone());
719 }
720 }
721
722 let has_function_declarations = !function_declarations.is_empty();
723 if has_function_declarations {
724 gemini_tools.push(adk_gemini::Tool::with_functions(function_declarations));
725 }
726
727 if has_provider_native_tools {
728 tool_config_json.insert(
729 "includeServerSideToolInvocations".to_string(),
730 serde_json::Value::Bool(true),
731 );
732 }
733
734 let tool_config = if tool_config_json.is_empty() {
735 adk_gemini::ToolConfig::default()
736 } else {
737 serde_json::from_value::<adk_gemini::ToolConfig>(serde_json::Value::Object(
738 tool_config_json,
739 ))
740 .map_err(|error| {
741 adk_core::AdkError::model(format!(
742 "failed to deserialize Gemini tool configuration: {error}"
743 ))
744 })?
745 };
746
747 Ok((gemini_tools, tool_config))
748 }
749
750 fn stream_chunks_from_response(
751 mut response: LlmResponse,
752 saw_partial_chunk: bool,
753 ) -> (Vec<LlmResponse>, bool) {
754 let is_final = response.finish_reason.is_some();
755
756 if !is_final {
757 response.partial = true;
758 response.turn_complete = false;
759 return (vec![response], true);
760 }
761
762 response.partial = false;
763 response.turn_complete = true;
764
765 if saw_partial_chunk {
766 return (vec![response], true);
767 }
768
769 let synthetic_partial = LlmResponse {
770 content: None,
771 usage_metadata: None,
772 finish_reason: None,
773 citation_metadata: None,
774 partial: true,
775 turn_complete: false,
776 interrupted: false,
777 error_code: None,
778 error_message: None,
779 provider_metadata: None,
780 interaction_id: None,
781 };
782
783 (vec![synthetic_partial, response], true)
784 }
785
786 async fn generate_content_internal(
787 &self,
788 req: LlmRequest,
789 stream: bool,
790 ) -> Result<LlmResponseStream> {
791 let mut builder = self.client.generate_content();
792
793 let mut fn_call_signatures: std::collections::HashMap<String, String> =
798 std::collections::HashMap::new();
799 for content in &req.contents {
800 if content.role == "model" {
801 for part in &content.parts {
802 if let Part::FunctionCall { name, thought_signature: Some(sig), .. } = part {
803 fn_call_signatures.insert(name.clone(), sig.clone());
804 }
805 }
806 }
807 }
808
809 for content in &req.contents {
811 match content.role.as_str() {
812 "user" => {
813 let mut gemini_parts = Vec::new();
815 for part in &content.parts {
816 match part {
817 Part::Text { text } => {
818 gemini_parts.push(adk_gemini::Part::Text {
819 text: text.clone(),
820 thought: None,
821 thought_signature: None,
822 });
823 }
824 Part::Thinking { thinking, signature } => {
825 gemini_parts.push(adk_gemini::Part::Text {
826 text: thinking.clone(),
827 thought: Some(true),
828 thought_signature: signature.clone(),
829 });
830 }
831 Part::InlineData { data, mime_type } => {
832 let encoded = attachment::encode_base64(data);
833 gemini_parts.push(adk_gemini::Part::InlineData {
834 inline_data: adk_gemini::Blob {
835 mime_type: mime_type.clone(),
836 data: encoded,
837 },
838 });
839 }
840 Part::FileData { mime_type, file_uri } => {
841 gemini_parts.push(adk_gemini::Part::Text {
842 text: attachment::file_attachment_to_text(mime_type, file_uri),
843 thought: None,
844 thought_signature: None,
845 });
846 }
847 _ => {}
848 }
849 }
850 if !gemini_parts.is_empty() {
851 let user_content = adk_gemini::Content {
852 role: Some(adk_gemini::Role::User),
853 parts: Some(gemini_parts),
854 };
855 builder = builder.with_message(adk_gemini::Message {
856 content: user_content,
857 role: adk_gemini::Role::User,
858 });
859 }
860 }
861 "model" => {
862 let mut gemini_parts = Vec::new();
864 for part in &content.parts {
865 match part {
866 Part::Text { text } => {
867 gemini_parts.push(adk_gemini::Part::Text {
868 text: text.clone(),
869 thought: None,
870 thought_signature: None,
871 });
872 }
873 Part::Thinking { thinking, signature } => {
874 gemini_parts.push(adk_gemini::Part::Text {
875 text: thinking.clone(),
876 thought: Some(true),
877 thought_signature: signature.clone(),
878 });
879 }
880 Part::FunctionCall { name, args, thought_signature, id } => {
881 gemini_parts.push(adk_gemini::Part::FunctionCall {
882 function_call: adk_gemini::FunctionCall {
883 name: name.clone(),
884 args: args.clone(),
885 id: id.clone(),
886 thought_signature: None,
887 },
888 thought_signature: thought_signature.clone(),
889 });
890 }
891 Part::ServerToolCall { server_tool_call } => {
892 if let Ok(native_part) = serde_json::from_value::<adk_gemini::Part>(
893 server_tool_call.clone(),
894 ) {
895 match native_part {
896 adk_gemini::Part::ToolCall { .. }
897 | adk_gemini::Part::ExecutableCode { .. } => {
898 gemini_parts.push(native_part);
899 continue;
900 }
901 _ => {}
902 }
903 }
904
905 gemini_parts.push(adk_gemini::Part::ToolCall {
906 tool_call: server_tool_call.clone(),
907 thought_signature: Self::gemini_part_thought_signature(
908 server_tool_call,
909 ),
910 });
911 }
912 Part::ServerToolResponse { server_tool_response } => {
913 if let Ok(native_part) = serde_json::from_value::<adk_gemini::Part>(
914 server_tool_response.clone(),
915 ) {
916 match native_part {
917 adk_gemini::Part::ToolResponse { .. }
918 | adk_gemini::Part::CodeExecutionResult { .. } => {
919 gemini_parts.push(native_part);
920 continue;
921 }
922 _ => {}
923 }
924 }
925
926 gemini_parts.push(adk_gemini::Part::ToolResponse {
927 tool_response: server_tool_response.clone(),
928 thought_signature: Self::gemini_part_thought_signature(
929 server_tool_response,
930 ),
931 });
932 }
933 _ => {}
934 }
935 }
936 if !gemini_parts.is_empty() {
937 let model_content = adk_gemini::Content {
938 role: Some(adk_gemini::Role::Model),
939 parts: Some(gemini_parts),
940 };
941 builder = builder.with_message(adk_gemini::Message {
942 content: model_content,
943 role: adk_gemini::Role::Model,
944 });
945 }
946 }
947 "function" => {
948 let mut gemini_parts = Vec::new();
951 for part in &content.parts {
952 if let Part::FunctionResponse { function_response, id } = part {
953 let sig = fn_call_signatures.get(&function_response.name).cloned();
954
955 let mut fr_parts = Vec::new();
957 for inline in &function_response.inline_data {
958 let encoded = attachment::encode_base64(&inline.data);
959 fr_parts.push(adk_gemini::FunctionResponsePart::InlineData {
960 inline_data: adk_gemini::Blob {
961 mime_type: inline.mime_type.clone(),
962 data: encoded,
963 },
964 });
965 }
966 for file in &function_response.file_data {
967 fr_parts.push(adk_gemini::FunctionResponsePart::FileData {
968 file_data: adk_gemini::FileDataRef {
969 mime_type: file.mime_type.clone(),
970 file_uri: file.file_uri.clone(),
971 },
972 });
973 }
974
975 let mut gemini_fr = adk_gemini::tools::FunctionResponse::new(
976 &function_response.name,
977 Self::gemini_function_response_payload(
978 function_response.response.clone(),
979 ),
980 );
981 gemini_fr.parts = fr_parts;
982 gemini_fr.id = id.clone();
985
986 gemini_parts.push(adk_gemini::Part::FunctionResponse {
987 function_response: gemini_fr,
988 thought_signature: sig,
989 });
990 }
991 }
992 if !gemini_parts.is_empty() {
993 let fn_content = adk_gemini::Content {
994 role: Some(adk_gemini::Role::User),
995 parts: Some(gemini_parts),
996 };
997 builder = builder.with_message(adk_gemini::Message {
998 content: fn_content,
999 role: adk_gemini::Role::User,
1000 });
1001 }
1002 }
1003 _ => {}
1004 }
1005 }
1006
1007 if let Some(config) = req.config {
1009 let has_schema = config.response_schema.is_some();
1010 let gen_config = adk_gemini::GenerationConfig {
1011 temperature: config.temperature,
1012 top_p: config.top_p,
1013 top_k: config.top_k,
1014 max_output_tokens: config.max_output_tokens,
1015 response_schema: config.response_schema,
1016 response_mime_type: if has_schema {
1017 Some("application/json".to_string())
1018 } else {
1019 None
1020 },
1021 thinking_config: self.thinking_config.clone(),
1022 ..Default::default()
1023 };
1024 builder = builder.with_generation_config(gen_config);
1025
1026 if let Some(ref name) = config.cached_content {
1028 let handle = self.client.get_cached_content(name);
1029 builder = builder.with_cached_content(&handle);
1030 }
1031 } else if self.thinking_config.is_some() {
1032 let gen_config = adk_gemini::GenerationConfig {
1035 thinking_config: self.thinking_config.clone(),
1036 ..Default::default()
1037 };
1038 builder = builder.with_generation_config(gen_config);
1039 }
1040
1041 if !req.tools.is_empty() {
1043 let adapter = self.schema_adapter();
1044 use std::sync::LazyLock;
1045 static SCHEMA_CACHE: LazyLock<SchemaCache> = LazyLock::new(SchemaCache::new);
1046 let (gemini_tools, tool_config) =
1047 Self::build_gemini_tools(&req.tools, adapter, &SCHEMA_CACHE)?;
1048 for tool in gemini_tools {
1049 builder = builder.with_tool(tool);
1050 }
1051 if tool_config != adk_gemini::ToolConfig::default() {
1052 builder = builder.with_tool_config(tool_config);
1053 }
1054 }
1055
1056 if stream {
1057 adk_telemetry::debug!("Executing streaming request");
1058 let response_stream = builder.execute_stream().await.map_err(|e| {
1059 adk_telemetry::error!(error = %e, "Model request failed");
1060 gemini_error_to_adk(&e)
1061 })?;
1062
1063 let mapped_stream = async_stream::stream! {
1064 let mut stream = response_stream;
1065 let mut saw_partial_chunk = false;
1066 while let Some(result) = stream.try_next().await.transpose() {
1067 match result {
1068 Ok(resp) => {
1069 match Self::convert_response(&resp) {
1070 Ok(llm_resp) => {
1071 let (chunks, next_saw_partial) =
1072 Self::stream_chunks_from_response(llm_resp, saw_partial_chunk);
1073 saw_partial_chunk = next_saw_partial;
1074 for chunk in chunks {
1075 yield Ok(chunk);
1076 }
1077 }
1078 Err(e) => {
1079 adk_telemetry::error!(error = %e, "Failed to convert response");
1080 yield Err(e);
1081 }
1082 }
1083 }
1084 Err(e) => {
1085 adk_telemetry::error!(error = %e, "Stream error");
1086 yield Err(gemini_error_to_adk(&e));
1087 }
1088 }
1089 }
1090 };
1091
1092 Ok(Box::pin(mapped_stream))
1093 } else {
1094 adk_telemetry::debug!("Executing blocking request");
1095 let response = builder.execute().await.map_err(|e| {
1096 adk_telemetry::error!(error = %e, "Model request failed");
1097 gemini_error_to_adk(&e)
1098 })?;
1099
1100 let llm_response = Self::convert_response(&response)?;
1101
1102 let stream = async_stream::stream! {
1103 yield Ok(llm_response);
1104 };
1105
1106 Ok(Box::pin(stream))
1107 }
1108 }
1109
1110 pub async fn create_cached_content(
1115 &self,
1116 system_instruction: &str,
1117 tools: &std::collections::HashMap<String, serde_json::Value>,
1118 ttl_seconds: u32,
1119 ) -> Result<String> {
1120 let mut cache_builder = self
1121 .client
1122 .create_cache()
1123 .with_system_instruction(system_instruction)
1124 .with_ttl(std::time::Duration::from_secs(u64::from(ttl_seconds)));
1125
1126 let adapter = self.schema_adapter();
1127 use std::sync::LazyLock;
1128 static SCHEMA_CACHE: LazyLock<SchemaCache> = LazyLock::new(SchemaCache::new);
1129 let (gemini_tools, tool_config) = Self::build_gemini_tools(tools, adapter, &SCHEMA_CACHE)?;
1130 if !gemini_tools.is_empty() {
1131 cache_builder = cache_builder.with_tools(gemini_tools);
1132 }
1133 if tool_config != adk_gemini::ToolConfig::default() {
1134 cache_builder = cache_builder.with_tool_config(tool_config);
1135 }
1136
1137 let handle = cache_builder
1138 .execute()
1139 .await
1140 .map_err(|e| adk_core::AdkError::model(format!("cache creation failed: {e}")))?;
1141
1142 Ok(handle.name().to_string())
1143 }
1144
1145 pub async fn delete_cached_content(&self, name: &str) -> Result<()> {
1147 let handle = self.client.get_cached_content(name);
1148 handle
1149 .delete()
1150 .await
1151 .map_err(|(_, e)| adk_core::AdkError::model(format!("cache deletion failed: {e}")))?;
1152 Ok(())
1153 }
1154
1155 #[cfg(feature = "gemini-interactions")]
1183 async fn generate_interactions_once(&self, req: LlmRequest) -> Result<LlmResponse> {
1184 use super::interactions_convert;
1185
1186 let target = self.interaction_target.as_ref().ok_or_else(|| {
1189 adk_core::AdkError::new(
1190 ErrorComponent::Model,
1191 ErrorCategory::InvalidInput,
1192 "model.gemini.interactions.missing_target",
1193 "the Interactions transport is active but no validated target is configured",
1194 )
1195 .with_provider("gemini")
1196 })?;
1197
1198 let thinking_level = self.thinking_config.as_ref().and_then(|c| c.thinking_level);
1201
1202 let stateful = self.interaction_options.stateful;
1203 let store = self.interaction_options.store;
1204 let background = self.resolve_background();
1205
1206 let mut request =
1208 interactions_convert::build_request(&req, target, thinking_level, stateful, store)?;
1209 request.background = Some(background);
1210
1211 let interaction = match self.client.send_interaction(request.clone()).await {
1213 Ok(interaction) => interaction,
1214 Err(error) => {
1215 let mapped = gemini_error_to_adk(&error);
1216 if mapped.category == ErrorCategory::NotFound && req.previous_response_id.is_some()
1220 {
1221 let mut fallback = interactions_convert::build_request(
1222 &req,
1223 target,
1224 thinking_level,
1225 false,
1226 store,
1227 )?;
1228 fallback.background = Some(background);
1229 self.client
1230 .send_interaction(fallback)
1231 .await
1232 .map_err(|e| gemini_error_to_adk(&e))?
1233 } else {
1234 return Err(mapped);
1235 }
1236 }
1237 };
1238
1239 let final_interaction = if background {
1241 self.poll_interaction_to_completion(interaction).await?
1242 } else {
1243 interaction
1244 };
1245
1246 interaction_status_to_result(final_interaction.status)?;
1248
1249 Ok(interactions_convert::to_llm_response(&final_interaction))
1250 }
1251
1252 #[cfg(feature = "gemini-interactions")]
1283 async fn poll_interaction_to_completion(
1284 &self,
1285 interaction: adk_gemini::interactions::Interaction,
1286 ) -> Result<adk_gemini::interactions::Interaction> {
1287 const MAX_POLL_ATTEMPTS: u32 = 600;
1292
1293 let mut current = interaction;
1294 let mut attempts: u32 = 0;
1295 while !current.status.is_terminal() && !current.status.requires_action() {
1296 if attempts >= MAX_POLL_ATTEMPTS {
1297 return Err(adk_core::AdkError::new(
1298 ErrorComponent::Model,
1299 ErrorCategory::Timeout,
1300 "model.gemini.interactions.poll_timeout",
1301 format!(
1302 "the Gemini interaction did not reach a terminal or requires_action \
1303 state after {MAX_POLL_ATTEMPTS} poll attempts"
1304 ),
1305 )
1306 .with_provider("gemini"));
1307 }
1308 attempts += 1;
1309 tokio::time::sleep(self.interaction_options.poll_interval).await;
1310 current = self
1311 .client
1312 .get_interaction(¤t.id, false)
1313 .await
1314 .map_err(|e| gemini_error_to_adk(&e))?;
1315 }
1316 Ok(current)
1317 }
1318
1319 #[cfg(feature = "gemini-interactions")]
1344 async fn generate_interactions_stream(&self, req: LlmRequest) -> Result<LlmResponseStream> {
1345 use super::interactions_convert::{self, SseAccumulator, sse_event_to_chunk};
1346
1347 let target = self.interaction_target.as_ref().ok_or_else(|| {
1348 adk_core::AdkError::new(
1349 ErrorComponent::Model,
1350 ErrorCategory::InvalidInput,
1351 "model.gemini.interactions.missing_target",
1352 "the Interactions transport is active but no validated target is configured",
1353 )
1354 .with_provider("gemini")
1355 })?;
1356
1357 let thinking_level = self.thinking_config.as_ref().and_then(|c| c.thinking_level);
1358 let stateful = self.interaction_options.stateful;
1359 let store = self.interaction_options.store;
1360
1361 let mut request =
1362 interactions_convert::build_request(&req, target, thinking_level, stateful, store)?;
1363 request.background = Some(false);
1368
1369 let sse_stream = self
1370 .client
1371 .send_interaction_stream(request)
1372 .await
1373 .map_err(|e| gemini_error_to_adk(&e))?;
1374
1375 let mapped = async_stream::stream! {
1376 let mut sse_stream = sse_stream;
1377 let mut acc = SseAccumulator::new();
1378 while let Some(result) = sse_stream.try_next().await.transpose() {
1379 match result {
1380 Ok(event) => {
1381 if let Some(chunk) = sse_event_to_chunk(event, &mut acc) {
1382 yield chunk;
1383 }
1384 }
1385 Err(e) => {
1386 adk_telemetry::error!(error = %e, "Interaction stream error");
1387 yield Err(gemini_error_to_adk(&e));
1388 }
1389 }
1390 }
1391 };
1392
1393 Ok(Box::pin(mapped))
1394 }
1395}
1396
1397#[async_trait]
1398impl Llm for GeminiModel {
1399 fn name(&self) -> &str {
1400 &self.model_name
1401 }
1402
1403 fn schema_adapter(&self) -> &dyn SchemaAdapter {
1404 use std::sync::LazyLock;
1405 static ADAPTER: LazyLock<GeminiSchemaAdapter> = LazyLock::new(GeminiSchemaAdapter::new);
1406 &*ADAPTER
1407 }
1408
1409 #[cfg(feature = "gemini-interactions")]
1410 fn uses_interactions_api(&self) -> bool {
1411 self.transport == GeminiTransport::Interactions
1412 }
1413
1414 #[adk_telemetry::instrument(
1415 name = "call_llm",
1416 skip(self, req),
1417 fields(
1418 model.name = %self.model_name,
1419 stream = %stream,
1420 request.contents_count = %req.contents.len(),
1421 request.tools_count = %req.tools.len()
1422 )
1423 )]
1424 async fn generate_content(&self, req: LlmRequest, stream: bool) -> Result<LlmResponseStream> {
1425 adk_telemetry::info!("Generating content");
1426 let usage_span = adk_telemetry::llm_generate_span("gemini", &self.model_name, stream);
1427
1428 #[cfg(feature = "gemini-interactions")]
1436 if self.transport == GeminiTransport::Interactions {
1437 if stream {
1446 let mapped =
1447 execute_with_retry(&self.retry_config, is_retryable_model_error, || {
1448 self.generate_interactions_stream(req.clone())
1449 })
1450 .await?;
1451 return Ok(crate::usage_tracking::with_usage_tracking(mapped, usage_span));
1452 }
1453 let response = execute_with_retry(&self.retry_config, is_retryable_model_error, || {
1454 self.generate_interactions_once(req.clone())
1455 })
1456 .await?;
1457 let single = async_stream::stream! {
1458 yield Ok(response);
1459 };
1460 return Ok(crate::usage_tracking::with_usage_tracking(Box::pin(single), usage_span));
1461 }
1462
1463 let result = execute_with_retry(&self.retry_config, is_retryable_model_error, || {
1464 self.generate_content_internal(req.clone(), stream)
1465 })
1466 .await?;
1467 Ok(crate::usage_tracking::with_usage_tracking(result, usage_span))
1468 }
1469}
1470
1471#[cfg(test)]
1472mod native_tool_tests {
1473 use super::*;
1474
1475 fn test_adapter() -> GeminiSchemaAdapter {
1476 GeminiSchemaAdapter::new()
1477 }
1478
1479 fn test_cache() -> SchemaCache {
1480 SchemaCache::new()
1481 }
1482
1483 #[test]
1484 fn test_build_gemini_tools_supports_native_tool_metadata() {
1485 let mut tools = std::collections::HashMap::new();
1486 tools.insert(
1487 "google_search".to_string(),
1488 serde_json::json!({
1489 "x-adk-gemini-tool": {
1490 "google_search": {}
1491 }
1492 }),
1493 );
1494 tools.insert(
1495 "lookup_weather".to_string(),
1496 serde_json::json!({
1497 "name": "lookup_weather",
1498 "description": "lookup weather",
1499 "parameters": {
1500 "type": "object",
1501 "properties": {
1502 "city": { "type": "string" }
1503 }
1504 }
1505 }),
1506 );
1507
1508 let adapter = test_adapter();
1509 let cache = test_cache();
1510 let (gemini_tools, tool_config) = GeminiModel::build_gemini_tools(&tools, &adapter, &cache)
1511 .expect("tool conversion should succeed");
1512
1513 assert_eq!(gemini_tools.len(), 2);
1514 assert_eq!(tool_config.include_server_side_tool_invocations, Some(true));
1515 }
1516
1517 #[test]
1518 fn test_build_gemini_tools_sets_flag_for_builtin_only() {
1519 let mut tools = std::collections::HashMap::new();
1520 tools.insert(
1521 "google_search".to_string(),
1522 serde_json::json!({
1523 "x-adk-gemini-tool": {
1524 "google_search": {}
1525 }
1526 }),
1527 );
1528
1529 let adapter = test_adapter();
1530 let cache = test_cache();
1531 let (_gemini_tools, tool_config) =
1532 GeminiModel::build_gemini_tools(&tools, &adapter, &cache)
1533 .expect("tool conversion should succeed");
1534
1535 assert_eq!(
1536 tool_config.include_server_side_tool_invocations,
1537 Some(true),
1538 "includeServerSideToolInvocations should be set even with only built-in tools"
1539 );
1540 }
1541
1542 #[test]
1543 fn test_build_gemini_tools_no_flag_for_function_only() {
1544 let mut tools = std::collections::HashMap::new();
1545 tools.insert(
1546 "lookup_weather".to_string(),
1547 serde_json::json!({
1548 "name": "lookup_weather",
1549 "description": "lookup weather",
1550 "parameters": {
1551 "type": "object",
1552 "properties": {
1553 "city": { "type": "string" }
1554 }
1555 }
1556 }),
1557 );
1558
1559 let adapter = test_adapter();
1560 let cache = test_cache();
1561 let (_gemini_tools, tool_config) =
1562 GeminiModel::build_gemini_tools(&tools, &adapter, &cache)
1563 .expect("tool conversion should succeed");
1564
1565 assert_eq!(
1566 tool_config.include_server_side_tool_invocations, None,
1567 "includeServerSideToolInvocations should NOT be set for function-only tools"
1568 );
1569 }
1570
1571 #[test]
1572 fn test_build_gemini_tools_merges_native_tool_config() {
1573 let mut tools = std::collections::HashMap::new();
1574 tools.insert(
1575 "google_maps".to_string(),
1576 serde_json::json!({
1577 "x-adk-gemini-tool": {
1578 "google_maps": {
1579 "enable_widget": true
1580 }
1581 },
1582 "x-adk-gemini-tool-config": {
1583 "retrievalConfig": {
1584 "latLng": {
1585 "latitude": 1.23,
1586 "longitude": 4.56
1587 }
1588 }
1589 }
1590 }),
1591 );
1592
1593 let adapter = test_adapter();
1594 let cache = test_cache();
1595 let (_gemini_tools, tool_config) =
1596 GeminiModel::build_gemini_tools(&tools, &adapter, &cache)
1597 .expect("tool conversion should succeed");
1598
1599 assert_eq!(
1600 tool_config.retrieval_config,
1601 Some(serde_json::json!({
1602 "latLng": {
1603 "latitude": 1.23,
1604 "longitude": 4.56
1605 }
1606 }))
1607 );
1608 }
1609
1610 #[test]
1611 fn test_response_schema_is_normalized_like_parameters() {
1612 let mut tools = std::collections::HashMap::new();
1618 tools.insert(
1619 "read_file".to_string(),
1620 serde_json::json!({
1621 "name": "read_file",
1622 "description": "Read a file from the filesystem",
1623 "parameters": {
1624 "$schema": "http://json-schema.org/draft-07/schema#",
1625 "type": "object",
1626 "properties": {
1627 "path": { "type": "string", "description": "File path" }
1628 },
1629 "required": ["path"],
1630 "additionalProperties": false
1631 },
1632 "response": {
1633 "$schema": "http://json-schema.org/draft-07/schema#",
1634 "type": "object",
1635 "properties": {
1636 "content": { "type": "string", "description": "File contents" }
1637 },
1638 "required": ["content"],
1639 "additionalProperties": false
1640 }
1641 }),
1642 );
1643
1644 let adapter = test_adapter();
1645 let cache = test_cache();
1646 let (gemini_tools, _) = GeminiModel::build_gemini_tools(&tools, &adapter, &cache)
1647 .expect("tool conversion should succeed");
1648
1649 let func_tool = gemini_tools
1651 .iter()
1652 .find(|t| matches!(t, adk_gemini::Tool::Function { .. }))
1653 .expect("should have function declarations");
1654
1655 let decls = match func_tool {
1656 adk_gemini::Tool::Function { function_declarations } => function_declarations,
1657 _ => panic!("expected Function tool"),
1658 };
1659
1660 let decl = &decls[0];
1661 let decl_json = serde_json::to_value(decl).unwrap();
1662
1663 let params = &decl_json["parameters"];
1665 assert!(params.get("$schema").is_none(), "parameters.$schema should be stripped");
1666 assert!(
1667 params.get("additionalProperties").is_none(),
1668 "parameters.additionalProperties should be stripped"
1669 );
1670
1671 let response = &decl_json["response"];
1673 assert!(
1674 response.get("$schema").is_none(),
1675 "response.$schema should be stripped (was the bug: copied raw)"
1676 );
1677 assert!(
1678 response.get("additionalProperties").is_none(),
1679 "response.additionalProperties should be stripped (was the bug: copied raw)"
1680 );
1681 }
1682}
1683
1684#[async_trait]
1685impl CacheCapable for GeminiModel {
1686 async fn create_cache(
1687 &self,
1688 system_instruction: &str,
1689 tools: &std::collections::HashMap<String, serde_json::Value>,
1690 ttl_seconds: u32,
1691 ) -> Result<String> {
1692 self.create_cached_content(system_instruction, tools, ttl_seconds).await
1693 }
1694
1695 async fn delete_cache(&self, name: &str) -> Result<()> {
1696 self.delete_cached_content(name).await
1697 }
1698}
1699
1700#[cfg(test)]
1701mod tests {
1702 use super::*;
1703 use adk_core::AdkError;
1704 use std::{
1705 sync::{
1706 Arc,
1707 atomic::{AtomicU32, Ordering},
1708 },
1709 time::Duration,
1710 };
1711
1712 #[test]
1713 fn constructor_is_backward_compatible_and_sync() {
1714 fn accepts_sync_constructor<F>(_f: F)
1715 where
1716 F: Fn(&str, &str) -> Result<GeminiModel>,
1717 {
1718 }
1719
1720 accepts_sync_constructor(|api_key, model| GeminiModel::new(api_key, model));
1721 }
1722
1723 #[test]
1724 fn stream_chunks_from_response_injects_partial_before_lone_final_chunk() {
1725 let response = LlmResponse {
1726 content: Some(Content::new("model").with_text("hello")),
1727 usage_metadata: None,
1728 finish_reason: Some(FinishReason::Stop),
1729 citation_metadata: None,
1730 partial: false,
1731 turn_complete: true,
1732 interrupted: false,
1733 error_code: None,
1734 error_message: None,
1735 provider_metadata: None,
1736 interaction_id: None,
1737 };
1738
1739 let (chunks, saw_partial) = GeminiModel::stream_chunks_from_response(response, false);
1740 assert!(saw_partial);
1741 assert_eq!(chunks.len(), 2);
1742 assert!(chunks[0].partial);
1743 assert!(!chunks[0].turn_complete);
1744 assert!(chunks[0].content.is_none());
1745 assert!(!chunks[1].partial);
1746 assert!(chunks[1].turn_complete);
1747 }
1748
1749 #[test]
1750 fn stream_chunks_from_response_keeps_final_only_when_partial_already_seen() {
1751 let response = LlmResponse {
1752 content: Some(Content::new("model").with_text("done")),
1753 usage_metadata: None,
1754 finish_reason: Some(FinishReason::Stop),
1755 citation_metadata: None,
1756 partial: false,
1757 turn_complete: true,
1758 interrupted: false,
1759 error_code: None,
1760 error_message: None,
1761 provider_metadata: None,
1762 interaction_id: None,
1763 };
1764
1765 let (chunks, saw_partial) = GeminiModel::stream_chunks_from_response(response, true);
1766 assert!(saw_partial);
1767 assert_eq!(chunks.len(), 1);
1768 assert!(!chunks[0].partial);
1769 assert!(chunks[0].turn_complete);
1770 }
1771
1772 #[tokio::test]
1773 async fn execute_with_retry_retries_retryable_errors() {
1774 let retry_config = RetryConfig::default()
1775 .with_max_retries(2)
1776 .with_initial_delay(Duration::from_millis(0))
1777 .with_max_delay(Duration::from_millis(0));
1778 let attempts = Arc::new(AtomicU32::new(0));
1779
1780 let result = execute_with_retry(&retry_config, is_retryable_model_error, || {
1781 let attempts = Arc::clone(&attempts);
1782 async move {
1783 let attempt = attempts.fetch_add(1, Ordering::SeqCst);
1784 if attempt < 2 {
1785 return Err(AdkError::model("code 429 RESOURCE_EXHAUSTED"));
1786 }
1787 Ok("ok")
1788 }
1789 })
1790 .await
1791 .expect("retry should eventually succeed");
1792
1793 assert_eq!(result, "ok");
1794 assert_eq!(attempts.load(Ordering::SeqCst), 3);
1795 }
1796
1797 #[tokio::test]
1798 async fn execute_with_retry_does_not_retry_non_retryable_errors() {
1799 let retry_config = RetryConfig::default()
1800 .with_max_retries(3)
1801 .with_initial_delay(Duration::from_millis(0))
1802 .with_max_delay(Duration::from_millis(0));
1803 let attempts = Arc::new(AtomicU32::new(0));
1804
1805 let error = execute_with_retry(&retry_config, is_retryable_model_error, || {
1806 let attempts = Arc::clone(&attempts);
1807 async move {
1808 attempts.fetch_add(1, Ordering::SeqCst);
1809 Err::<(), _>(AdkError::model("code 400 invalid request"))
1810 }
1811 })
1812 .await
1813 .expect_err("non-retryable error should return immediately");
1814
1815 assert!(error.is_model());
1816 assert_eq!(attempts.load(Ordering::SeqCst), 1);
1817 }
1818
1819 #[tokio::test]
1820 async fn execute_with_retry_respects_disabled_config() {
1821 let retry_config = RetryConfig::disabled().with_max_retries(10);
1822 let attempts = Arc::new(AtomicU32::new(0));
1823
1824 let error = execute_with_retry(&retry_config, is_retryable_model_error, || {
1825 let attempts = Arc::clone(&attempts);
1826 async move {
1827 attempts.fetch_add(1, Ordering::SeqCst);
1828 Err::<(), _>(AdkError::model("code 429 RESOURCE_EXHAUSTED"))
1829 }
1830 })
1831 .await
1832 .expect_err("disabled retries should return first error");
1833
1834 assert!(error.is_model());
1835 assert_eq!(attempts.load(Ordering::SeqCst), 1);
1836 }
1837
1838 #[test]
1839 fn convert_response_preserves_citation_metadata() {
1840 let response = adk_gemini::GenerationResponse {
1841 candidates: vec![adk_gemini::Candidate {
1842 content: adk_gemini::Content {
1843 role: Some(adk_gemini::Role::Model),
1844 parts: Some(vec![adk_gemini::Part::Text {
1845 text: "hello world".to_string(),
1846 thought: None,
1847 thought_signature: None,
1848 }]),
1849 },
1850 safety_ratings: None,
1851 citation_metadata: Some(adk_gemini::CitationMetadata {
1852 citation_sources: vec![adk_gemini::CitationSource {
1853 uri: Some("https://example.com".to_string()),
1854 title: Some("Example".to_string()),
1855 start_index: Some(0),
1856 end_index: Some(5),
1857 license: Some("CC-BY".to_string()),
1858 publication_date: None,
1859 }],
1860 }),
1861 grounding_metadata: None,
1862 finish_reason: Some(adk_gemini::FinishReason::Stop),
1863 index: Some(0),
1864 }],
1865 prompt_feedback: None,
1866 usage_metadata: None,
1867 model_version: None,
1868 response_id: None,
1869 };
1870
1871 let converted =
1872 GeminiModel::convert_response(&response).expect("conversion should succeed");
1873 let metadata = converted.citation_metadata.expect("citation metadata should be mapped");
1874 assert_eq!(metadata.citation_sources.len(), 1);
1875 assert_eq!(metadata.citation_sources[0].uri.as_deref(), Some("https://example.com"));
1876 assert_eq!(metadata.citation_sources[0].start_index, Some(0));
1877 assert_eq!(metadata.citation_sources[0].end_index, Some(5));
1878 }
1879
1880 #[test]
1881 fn convert_response_handles_inline_data_from_model() {
1882 let image_bytes = vec![0x89, 0x50, 0x4E, 0x47];
1883 let encoded = crate::attachment::encode_base64(&image_bytes);
1884
1885 let response = adk_gemini::GenerationResponse {
1886 candidates: vec![adk_gemini::Candidate {
1887 content: adk_gemini::Content {
1888 role: Some(adk_gemini::Role::Model),
1889 parts: Some(vec![
1890 adk_gemini::Part::Text {
1891 text: "Here is the image".to_string(),
1892 thought: None,
1893 thought_signature: None,
1894 },
1895 adk_gemini::Part::InlineData {
1896 inline_data: adk_gemini::Blob {
1897 mime_type: "image/png".to_string(),
1898 data: encoded,
1899 },
1900 },
1901 ]),
1902 },
1903 safety_ratings: None,
1904 citation_metadata: None,
1905 grounding_metadata: None,
1906 finish_reason: Some(adk_gemini::FinishReason::Stop),
1907 index: Some(0),
1908 }],
1909 prompt_feedback: None,
1910 usage_metadata: None,
1911 model_version: None,
1912 response_id: None,
1913 };
1914
1915 let converted =
1916 GeminiModel::convert_response(&response).expect("conversion should succeed");
1917 let content = converted.content.expect("should have content");
1918 assert!(
1919 content
1920 .parts
1921 .iter()
1922 .any(|part| matches!(part, Part::Text { text } if text == "Here is the image"))
1923 );
1924 assert!(content.parts.iter().any(|part| {
1925 matches!(
1926 part,
1927 Part::InlineData { mime_type, data }
1928 if mime_type == "image/png" && data.as_slice() == image_bytes.as_slice()
1929 )
1930 }));
1931 }
1932
1933 #[test]
1934 fn gemini_function_response_payload_preserves_objects() {
1935 let value = serde_json::json!({
1936 "documents": [
1937 { "id": "pricing", "score": 0.91 }
1938 ]
1939 });
1940
1941 let payload = GeminiModel::gemini_function_response_payload(value.clone());
1942
1943 assert_eq!(payload, value);
1944 }
1945
1946 #[test]
1947 fn gemini_function_response_payload_wraps_arrays() {
1948 let payload =
1949 GeminiModel::gemini_function_response_payload(serde_json::json!([{ "id": "pricing" }]));
1950
1951 assert_eq!(payload, serde_json::json!({ "result": [{ "id": "pricing" }] }));
1952 }
1953
1954 fn convert_function_response_to_gemini_fr(
1959 frd: &adk_core::FunctionResponseData,
1960 ) -> adk_gemini::tools::FunctionResponse {
1961 let mut fr_parts = Vec::new();
1962
1963 for inline in &frd.inline_data {
1964 let encoded = crate::attachment::encode_base64(&inline.data);
1965 fr_parts.push(adk_gemini::FunctionResponsePart::InlineData {
1966 inline_data: adk_gemini::Blob {
1967 mime_type: inline.mime_type.clone(),
1968 data: encoded,
1969 },
1970 });
1971 }
1972
1973 for file in &frd.file_data {
1974 fr_parts.push(adk_gemini::FunctionResponsePart::FileData {
1975 file_data: adk_gemini::FileDataRef {
1976 mime_type: file.mime_type.clone(),
1977 file_uri: file.file_uri.clone(),
1978 },
1979 });
1980 }
1981
1982 let mut gemini_fr = adk_gemini::tools::FunctionResponse::new(
1983 &frd.name,
1984 GeminiModel::gemini_function_response_payload(frd.response.clone()),
1985 );
1986 gemini_fr.parts = fr_parts;
1987 gemini_fr
1988 }
1989
1990 #[test]
1991 fn json_only_function_response_has_no_nested_parts() {
1992 let frd = adk_core::FunctionResponseData::new("tool", serde_json::json!({"ok": true}));
1993 let gemini_fr = convert_function_response_to_gemini_fr(&frd);
1994 assert!(gemini_fr.parts.is_empty());
1995 let json = serde_json::to_string(&gemini_fr).unwrap();
1997 assert!(!json.contains("\"parts\""));
1998 }
1999
2000 #[test]
2001 fn function_response_with_inline_data_has_nested_parts() {
2002 let frd = adk_core::FunctionResponseData::with_inline_data(
2003 "chart",
2004 serde_json::json!({"status": "ok"}),
2005 vec![adk_core::InlineDataPart {
2006 mime_type: "image/png".to_string(),
2007 data: vec![0x89, 0x50, 0x4E, 0x47],
2008 }],
2009 );
2010 let gemini_fr = convert_function_response_to_gemini_fr(&frd);
2011 assert_eq!(gemini_fr.parts.len(), 1);
2012 match &gemini_fr.parts[0] {
2013 adk_gemini::FunctionResponsePart::InlineData { inline_data } => {
2014 assert_eq!(inline_data.mime_type, "image/png");
2015 let decoded = BASE64_STANDARD.decode(&inline_data.data).unwrap();
2016 assert_eq!(decoded, vec![0x89, 0x50, 0x4E, 0x47]);
2017 }
2018 other => panic!("expected InlineData, got {other:?}"),
2019 }
2020 }
2021
2022 #[test]
2023 fn function_response_with_file_data_has_nested_parts() {
2024 let frd = adk_core::FunctionResponseData::with_file_data(
2025 "doc",
2026 serde_json::json!({"ok": true}),
2027 vec![adk_core::FileDataPart {
2028 mime_type: "application/pdf".to_string(),
2029 file_uri: "gs://bucket/report.pdf".to_string(),
2030 }],
2031 );
2032 let gemini_fr = convert_function_response_to_gemini_fr(&frd);
2033 assert_eq!(gemini_fr.parts.len(), 1);
2034 match &gemini_fr.parts[0] {
2035 adk_gemini::FunctionResponsePart::FileData { file_data } => {
2036 assert_eq!(file_data.mime_type, "application/pdf");
2037 assert_eq!(file_data.file_uri, "gs://bucket/report.pdf");
2038 }
2039 other => panic!("expected FileData, got {other:?}"),
2040 }
2041 }
2042
2043 #[test]
2044 fn function_response_with_both_inline_and_file_data_ordering() {
2045 let frd = adk_core::FunctionResponseData::with_multimodal(
2046 "multi",
2047 serde_json::json!({}),
2048 vec![
2049 adk_core::InlineDataPart { mime_type: "image/png".to_string(), data: vec![1, 2] },
2050 adk_core::InlineDataPart { mime_type: "image/jpeg".to_string(), data: vec![3, 4] },
2051 ],
2052 vec![adk_core::FileDataPart {
2053 mime_type: "application/pdf".to_string(),
2054 file_uri: "gs://b/f.pdf".to_string(),
2055 }],
2056 );
2057 let gemini_fr = convert_function_response_to_gemini_fr(&frd);
2058 assert_eq!(gemini_fr.parts.len(), 3);
2060 assert!(matches!(&gemini_fr.parts[0], adk_gemini::FunctionResponsePart::InlineData { .. }));
2061 assert!(matches!(&gemini_fr.parts[1], adk_gemini::FunctionResponsePart::InlineData { .. }));
2062 assert!(matches!(&gemini_fr.parts[2], adk_gemini::FunctionResponsePart::FileData { .. }));
2063 }
2064}
2065
2066#[cfg(all(test, feature = "gemini-interactions"))]
2067mod interactions_transport_tests {
2068 use super::*;
2069
2070 #[test]
2075 fn default_interaction_options_match_api_posture() {
2076 let opts = InteractionOptions::default();
2077 assert!(opts.store, "store should default to true");
2078 assert!(opts.stateful, "stateful should default to true");
2079 assert_eq!(opts.background, BackgroundMode::AgentTargetsOnly);
2080 assert_eq!(opts.poll_interval, std::time::Duration::from_secs(1));
2081 }
2082
2083 #[test]
2084 fn new_model_defaults_to_generate_content_transport() {
2085 let model = GeminiModel::new("test-key", "gemini-2.5-flash")
2086 .expect("constructing a Gemini model should not require network");
2087 assert_eq!(model.transport(), GeminiTransport::GenerateContent);
2088 }
2089
2090 #[test]
2091 fn use_interactions_api_enables_transport_for_allowlisted_model() {
2092 let model = GeminiModel::new("test-key", "gemini-2.5-flash")
2093 .expect("construct model")
2094 .use_interactions_api(true)
2095 .expect("allowlisted model should enable the Interactions transport");
2096
2097 assert_eq!(model.transport(), GeminiTransport::Interactions);
2098 assert_eq!(
2099 model.interaction_target,
2100 Some(InteractionTarget::Model("gemini-2.5-flash".to_string()))
2101 );
2102 }
2103
2104 #[test]
2105 fn use_interactions_api_rejects_unsupported_model_with_invalid_input() {
2106 let result = GeminiModel::new("test-key", "gemini-2.0-flash")
2107 .expect("construct model")
2108 .use_interactions_api(true);
2109
2110 let err = match result {
2111 Ok(_) => panic!("unsupported model should be rejected"),
2112 Err(err) => err,
2113 };
2114 assert_eq!(err.category, adk_core::ErrorCategory::InvalidInput);
2115 assert_eq!(err.details.provider.as_deref(), Some("gemini"));
2116 }
2117
2118 #[test]
2119 fn use_interactions_api_false_reverts_to_generate_content() {
2120 let model = GeminiModel::new("test-key", "gemini-2.5-flash")
2121 .expect("construct model")
2122 .use_interactions_api(true)
2123 .expect("enable interactions")
2124 .use_interactions_api(false)
2125 .expect("disabling should always succeed");
2126
2127 assert_eq!(model.transport(), GeminiTransport::GenerateContent);
2128 assert_eq!(model.interaction_target, None);
2129 }
2130
2131 #[test]
2132 fn interaction_options_override_is_stored() {
2133 let opts = InteractionOptions {
2134 store: false,
2135 stateful: false,
2136 background: BackgroundMode::Always,
2137 poll_interval: std::time::Duration::from_millis(250),
2138 };
2139 let model = GeminiModel::new("test-key", "gemini-2.5-flash")
2140 .expect("construct model")
2141 .interaction_options(opts.clone());
2142
2143 let stored = model.interaction_options_ref();
2144 assert_eq!(stored.store, opts.store);
2145 assert_eq!(stored.stateful, opts.stateful);
2146 assert_eq!(stored.background, opts.background);
2147 assert_eq!(stored.poll_interval, opts.poll_interval);
2148 }
2149
2150 #[test]
2151 fn resolve_background_honors_background_mode() {
2152 let always = GeminiModel::new("test-key", "gemini-2.5-flash")
2154 .expect("construct model")
2155 .use_interactions_api(true)
2156 .expect("enable")
2157 .interaction_options(InteractionOptions {
2158 background: BackgroundMode::Always,
2159 ..InteractionOptions::default()
2160 });
2161 assert!(always.resolve_background());
2162
2163 let never = GeminiModel::new("test-key", "gemini-2.5-flash")
2165 .expect("construct model")
2166 .use_interactions_api(true)
2167 .expect("enable")
2168 .interaction_options(InteractionOptions {
2169 background: BackgroundMode::Never,
2170 ..InteractionOptions::default()
2171 });
2172 assert!(!never.resolve_background());
2173
2174 let model_target = GeminiModel::new("test-key", "gemini-2.5-flash")
2176 .expect("construct model")
2177 .use_interactions_api(true)
2178 .expect("enable");
2179 assert!(!model_target.resolve_background());
2180
2181 let agent_target = GeminiModel::new("test-key", "deep-research-preview-04-2026")
2183 .expect("construct model")
2184 .use_interactions_api(true)
2185 .expect("enable");
2186 assert!(agent_target.resolve_background());
2187 }
2188
2189 #[test]
2192 fn terminal_failed_status_maps_to_error() {
2193 use adk_gemini::interactions::InteractionStatus;
2194
2195 let err = interaction_status_to_result(InteractionStatus::Failed)
2196 .expect_err("a failed interaction must surface an error");
2197 assert_eq!(err.category, adk_core::ErrorCategory::Internal);
2198 assert_eq!(err.details.provider.as_deref(), Some("gemini"));
2199 assert_eq!(err.code, "model.gemini.interactions.failed");
2200 }
2201
2202 #[test]
2205 fn terminal_budget_exceeded_status_maps_to_error() {
2206 use adk_gemini::interactions::InteractionStatus;
2207
2208 let err = interaction_status_to_result(InteractionStatus::BudgetExceeded)
2209 .expect_err("a budget_exceeded interaction must surface an error");
2210 assert_eq!(err.category, adk_core::ErrorCategory::Internal);
2211 assert_eq!(err.details.provider.as_deref(), Some("gemini"));
2212 assert_eq!(err.code, "model.gemini.interactions.budget_exceeded");
2213 }
2214
2215 #[test]
2218 fn non_failure_statuses_are_ok() {
2219 use adk_gemini::interactions::InteractionStatus;
2220
2221 for status in [
2222 InteractionStatus::InProgress,
2223 InteractionStatus::RequiresAction,
2224 InteractionStatus::Completed,
2225 InteractionStatus::Cancelled,
2226 InteractionStatus::Incomplete,
2227 ] {
2228 assert!(
2229 interaction_status_to_result(status).is_ok(),
2230 "status {status:?} should not map to an error"
2231 );
2232 }
2233 }
2234
2235 #[test]
2238 fn failed_interaction_resource_surfaces_error() {
2239 use adk_gemini::interactions::{Interaction, InteractionStatus};
2240
2241 let interaction = Interaction {
2242 id: "v1_failed".to_string(),
2243 model: Some("gemini-2.5-flash".to_string()),
2244 agent: None,
2245 status: InteractionStatus::Failed,
2246 steps: Vec::new(),
2247 usage: None,
2248 created: None,
2249 updated: None,
2250 };
2251
2252 let err = interaction_status_to_result(interaction.status)
2253 .expect_err("failed interaction must surface an error");
2254 assert_eq!(err.category, adk_core::ErrorCategory::Internal);
2255 assert_eq!(err.details.provider.as_deref(), Some("gemini"));
2256 }
2257}