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 uri: None,
503 annotations: None,
504 });
505 }
506 adk_gemini::Part::FunctionCall { function_call, thought_signature } => {
507 converted_parts.push(Part::FunctionCall {
508 name: function_call.name.clone(),
509 args: function_call.args.clone(),
510 id: function_call.id.clone(),
511 thought_signature: thought_signature.clone(),
512 });
513 }
514 adk_gemini::Part::FunctionResponse { function_response, .. } => {
515 converted_parts.push(Part::FunctionResponse {
516 function_response: adk_core::FunctionResponseData::new(
517 function_response.name.clone(),
518 function_response
519 .response
520 .clone()
521 .unwrap_or(serde_json::Value::Null),
522 ),
523 id: None,
524 annotations: None,
525 });
526 }
527 adk_gemini::Part::ToolCall { .. } | adk_gemini::Part::ExecutableCode { .. } => {
528 if let Ok(value) = serde_json::to_value(p) {
529 converted_parts.push(Part::ServerToolCall { server_tool_call: value });
530 }
531 }
532 adk_gemini::Part::ToolResponse { .. }
533 | adk_gemini::Part::CodeExecutionResult { .. } => {
534 let value = serde_json::to_value(p).unwrap_or(serde_json::Value::Null);
535 converted_parts
536 .push(Part::ServerToolResponse { server_tool_response: value });
537 }
538 adk_gemini::Part::FileData { file_data } => {
539 converted_parts.push(Part::FileData {
540 mime_type: file_data.mime_type.clone(),
541 file_uri: file_data.file_uri.clone(),
542 annotations: None,
543 });
544 }
545 }
546 }
547 }
548
549 if let Some(grounding) = resp.candidates.first().and_then(|c| c.grounding_metadata.as_ref())
551 {
552 if let Some(queries) = &grounding.web_search_queries
553 && !queries.is_empty()
554 {
555 let search_info = format!("\n\n🔍 **Searched:** {}", queries.join(", "));
556 converted_parts.push(Part::Text { text: search_info });
557 }
558 if let Some(chunks) = &grounding.grounding_chunks {
559 let sources: Vec<String> = chunks
560 .iter()
561 .filter_map(|c| {
562 c.web.as_ref().and_then(|w| match (&w.title, &w.uri) {
563 (Some(title), Some(uri)) => Some(format!("[{}]({})", title, uri)),
564 (Some(title), None) => Some(title.clone()),
565 (None, Some(uri)) => Some(uri.to_string()),
566 (None, None) => None,
567 })
568 })
569 .collect();
570 if !sources.is_empty() {
571 let sources_info = format!("\n📚 **Sources:** {}", sources.join(" | "));
572 converted_parts.push(Part::Text { text: sources_info });
573 }
574 }
575 }
576
577 let content = if converted_parts.is_empty() {
578 None
579 } else {
580 Some(Content { role: "model".to_string(), parts: converted_parts })
581 };
582
583 let usage_metadata = resp.usage_metadata.as_ref().map(|u| UsageMetadata {
584 prompt_token_count: u.prompt_token_count.unwrap_or(0),
585 candidates_token_count: u.candidates_token_count.unwrap_or(0),
586 total_token_count: u.total_token_count.unwrap_or(0),
587 thinking_token_count: u.thoughts_token_count,
588 cache_read_input_token_count: u.cached_content_token_count,
589 ..Default::default()
590 });
591
592 let finish_reason =
593 resp.candidates.first().and_then(|c| c.finish_reason.as_ref()).map(|fr| match fr {
594 adk_gemini::FinishReason::Stop => FinishReason::Stop,
595 adk_gemini::FinishReason::MaxTokens => FinishReason::MaxTokens,
596 adk_gemini::FinishReason::Safety => FinishReason::Safety,
597 adk_gemini::FinishReason::Recitation => FinishReason::Recitation,
598 _ => FinishReason::Other,
599 });
600
601 let citation_metadata =
602 resp.candidates.first().and_then(|c| c.citation_metadata.as_ref()).map(|meta| {
603 CitationMetadata {
604 citation_sources: meta
605 .citation_sources
606 .iter()
607 .map(|source| CitationSource {
608 uri: source.uri.clone(),
609 title: source.title.clone(),
610 start_index: source.start_index,
611 end_index: source.end_index,
612 license: source.license.clone(),
613 publication_date: source.publication_date.map(|d| d.to_string()),
614 })
615 .collect(),
616 }
617 });
618
619 let provider_metadata = resp
622 .candidates
623 .first()
624 .and_then(|c| c.grounding_metadata.as_ref())
625 .and_then(|g| serde_json::to_value(g).ok());
626
627 Ok(LlmResponse {
628 content,
629 usage_metadata,
630 finish_reason,
631 citation_metadata,
632 partial: false,
633 turn_complete: true,
634 interrupted: false,
635 error_code: None,
636 error_message: None,
637 provider_metadata,
638 interaction_id: None,
639 })
640 }
641
642 fn gemini_function_response_payload(response: serde_json::Value) -> serde_json::Value {
643 match response {
644 serde_json::Value::Object(_) => response,
646 other => serde_json::json!({ "result": other }),
647 }
648 }
649
650 fn merge_object_value(
651 target: &mut serde_json::Map<String, serde_json::Value>,
652 value: serde_json::Value,
653 ) {
654 if let serde_json::Value::Object(object) = value {
655 for (key, value) in object {
656 target.insert(key, value);
657 }
658 }
659 }
660
661 fn build_gemini_tools(
662 tools: &std::collections::HashMap<String, serde_json::Value>,
663 adapter: &dyn SchemaAdapter,
664 cache: &SchemaCache,
665 ) -> Result<(Vec<adk_gemini::Tool>, adk_gemini::ToolConfig)> {
666 let mut gemini_tools = Vec::new();
667 let mut function_declarations = Vec::new();
668 let mut has_provider_native_tools = false;
669 let mut tool_config_json = serde_json::Map::new();
670
671 for (name, tool_decl) in tools {
672 if let Some(provider_tool) = tool_decl.get("x-adk-gemini-tool") {
673 let tool = serde_json::from_value::<adk_gemini::Tool>(provider_tool.clone())
674 .map_err(|error| {
675 adk_core::AdkError::model(format!(
676 "failed to deserialize Gemini native tool '{name}': {error}"
677 ))
678 })?;
679 has_provider_native_tools = true;
680 gemini_tools.push(tool);
681 } else {
682 let normalized_name = adapter.normalize_tool_name(name);
684
685 let schema =
688 tool_decl.get("parameters").cloned().unwrap_or_else(|| adapter.empty_schema());
689 let normalized_schema = cache.normalize(&schema);
690
691 let description =
693 tool_decl.get("description").and_then(|v| v.as_str()).unwrap_or("").to_string();
694
695 let mut func_decl_json = serde_json::json!({
696 "name": normalized_name.as_ref(),
697 "description": description,
698 "parameters": normalized_schema,
699 });
700
701 if let Some(response) = tool_decl.get("response") {
703 func_decl_json["response"] = cache.normalize(response);
704 }
705
706 if let Some(behavior) = tool_decl.get("behavior") {
708 func_decl_json["behavior"] = behavior.clone();
709 }
710
711 let func_decl =
712 serde_json::from_value::<adk_gemini::FunctionDeclaration>(func_decl_json)
713 .map_err(|error| {
714 adk_core::AdkError::model(format!(
715 "failed to build Gemini function declaration for '{name}': {error}"
716 ))
717 })?;
718 function_declarations.push(func_decl);
719 }
720
721 if let Some(tool_config) = tool_decl.get("x-adk-gemini-tool-config") {
722 Self::merge_object_value(&mut tool_config_json, tool_config.clone());
723 }
724 }
725
726 let has_function_declarations = !function_declarations.is_empty();
727 if has_function_declarations {
728 gemini_tools.push(adk_gemini::Tool::with_functions(function_declarations));
729 }
730
731 if has_provider_native_tools {
732 tool_config_json.insert(
733 "includeServerSideToolInvocations".to_string(),
734 serde_json::Value::Bool(true),
735 );
736 }
737
738 let tool_config = if tool_config_json.is_empty() {
739 adk_gemini::ToolConfig::default()
740 } else {
741 serde_json::from_value::<adk_gemini::ToolConfig>(serde_json::Value::Object(
742 tool_config_json,
743 ))
744 .map_err(|error| {
745 adk_core::AdkError::model(format!(
746 "failed to deserialize Gemini tool configuration: {error}"
747 ))
748 })?
749 };
750
751 Ok((gemini_tools, tool_config))
752 }
753
754 fn stream_chunks_from_response(
755 mut response: LlmResponse,
756 saw_partial_chunk: bool,
757 ) -> (Vec<LlmResponse>, bool) {
758 let is_final = response.finish_reason.is_some();
759
760 if !is_final {
761 response.partial = true;
762 response.turn_complete = false;
763 return (vec![response], true);
764 }
765
766 response.partial = false;
767 response.turn_complete = true;
768
769 if saw_partial_chunk {
770 return (vec![response], true);
771 }
772
773 let synthetic_partial = LlmResponse {
774 content: None,
775 usage_metadata: None,
776 finish_reason: None,
777 citation_metadata: None,
778 partial: true,
779 turn_complete: false,
780 interrupted: false,
781 error_code: None,
782 error_message: None,
783 provider_metadata: None,
784 interaction_id: None,
785 };
786
787 (vec![synthetic_partial, response], true)
788 }
789
790 async fn generate_content_internal(
791 &self,
792 req: LlmRequest,
793 stream: bool,
794 ) -> Result<LlmResponseStream> {
795 let mut builder = self.client.generate_content();
796
797 let mut fn_call_signatures: std::collections::HashMap<String, String> =
802 std::collections::HashMap::new();
803 for content in &req.contents {
804 if content.role == "model" {
805 for part in &content.parts {
806 if let Part::FunctionCall { name, thought_signature: Some(sig), .. } = part {
807 fn_call_signatures.insert(name.clone(), sig.clone());
808 }
809 }
810 }
811 }
812
813 for content in &req.contents {
815 match content.role.as_str() {
816 "user" => {
817 let mut gemini_parts = Vec::new();
819 for part in &content.parts {
820 match part {
821 Part::Text { text } => {
822 gemini_parts.push(adk_gemini::Part::Text {
823 text: text.clone(),
824 thought: None,
825 thought_signature: None,
826 });
827 }
828 Part::Thinking { thinking, signature } => {
829 gemini_parts.push(adk_gemini::Part::Text {
830 text: thinking.clone(),
831 thought: Some(true),
832 thought_signature: signature.clone(),
833 });
834 }
835 Part::InlineData { data, mime_type, .. } => {
836 let encoded = attachment::encode_base64(data);
837 gemini_parts.push(adk_gemini::Part::InlineData {
838 inline_data: adk_gemini::Blob {
839 mime_type: mime_type.clone(),
840 data: encoded,
841 },
842 });
843 }
844 Part::FileData { mime_type, file_uri, .. } => {
845 gemini_parts.push(adk_gemini::Part::Text {
846 text: attachment::file_attachment_to_text(mime_type, file_uri),
847 thought: None,
848 thought_signature: None,
849 });
850 }
851 _ => {}
852 }
853 }
854 if !gemini_parts.is_empty() {
855 let user_content = adk_gemini::Content {
856 role: Some(adk_gemini::Role::User),
857 parts: Some(gemini_parts),
858 };
859 builder = builder.with_message(adk_gemini::Message {
860 content: user_content,
861 role: adk_gemini::Role::User,
862 });
863 }
864 }
865 "model" => {
866 let mut gemini_parts = Vec::new();
868 for part in &content.parts {
869 match part {
870 Part::Text { text } => {
871 gemini_parts.push(adk_gemini::Part::Text {
872 text: text.clone(),
873 thought: None,
874 thought_signature: None,
875 });
876 }
877 Part::Thinking { thinking, signature } => {
878 gemini_parts.push(adk_gemini::Part::Text {
879 text: thinking.clone(),
880 thought: Some(true),
881 thought_signature: signature.clone(),
882 });
883 }
884 Part::FunctionCall { name, args, thought_signature, id } => {
885 gemini_parts.push(adk_gemini::Part::FunctionCall {
886 function_call: adk_gemini::FunctionCall {
887 name: name.clone(),
888 args: args.clone(),
889 id: id.clone(),
890 thought_signature: None,
891 },
892 thought_signature: thought_signature.clone(),
893 });
894 }
895 Part::ServerToolCall { server_tool_call } => {
896 if let Ok(native_part) = serde_json::from_value::<adk_gemini::Part>(
897 server_tool_call.clone(),
898 ) {
899 match native_part {
900 adk_gemini::Part::ToolCall { .. }
901 | adk_gemini::Part::ExecutableCode { .. } => {
902 gemini_parts.push(native_part);
903 continue;
904 }
905 _ => {}
906 }
907 }
908
909 gemini_parts.push(adk_gemini::Part::ToolCall {
910 tool_call: server_tool_call.clone(),
911 thought_signature: Self::gemini_part_thought_signature(
912 server_tool_call,
913 ),
914 });
915 }
916 Part::ServerToolResponse { server_tool_response } => {
917 if let Ok(native_part) = serde_json::from_value::<adk_gemini::Part>(
918 server_tool_response.clone(),
919 ) {
920 match native_part {
921 adk_gemini::Part::ToolResponse { .. }
922 | adk_gemini::Part::CodeExecutionResult { .. } => {
923 gemini_parts.push(native_part);
924 continue;
925 }
926 _ => {}
927 }
928 }
929
930 gemini_parts.push(adk_gemini::Part::ToolResponse {
931 tool_response: server_tool_response.clone(),
932 thought_signature: Self::gemini_part_thought_signature(
933 server_tool_response,
934 ),
935 });
936 }
937 _ => {}
938 }
939 }
940 if !gemini_parts.is_empty() {
941 let model_content = adk_gemini::Content {
942 role: Some(adk_gemini::Role::Model),
943 parts: Some(gemini_parts),
944 };
945 builder = builder.with_message(adk_gemini::Message {
946 content: model_content,
947 role: adk_gemini::Role::Model,
948 });
949 }
950 }
951 "function" => {
952 let mut gemini_parts = Vec::new();
955 for part in &content.parts {
956 if let Part::FunctionResponse { function_response, id, .. } = part {
957 let sig = fn_call_signatures.get(&function_response.name).cloned();
958
959 let mut fr_parts = Vec::new();
961 for inline in &function_response.inline_data {
962 let encoded = attachment::encode_base64(&inline.data);
963 fr_parts.push(adk_gemini::FunctionResponsePart::InlineData {
964 inline_data: adk_gemini::Blob {
965 mime_type: inline.mime_type.clone(),
966 data: encoded,
967 },
968 });
969 }
970 for file in &function_response.file_data {
971 fr_parts.push(adk_gemini::FunctionResponsePart::FileData {
972 file_data: adk_gemini::FileDataRef {
973 mime_type: file.mime_type.clone(),
974 file_uri: file.file_uri.clone(),
975 },
976 });
977 }
978
979 let mut gemini_fr = adk_gemini::tools::FunctionResponse::new(
980 &function_response.name,
981 Self::gemini_function_response_payload(
982 function_response.response.clone(),
983 ),
984 );
985 gemini_fr.parts = fr_parts;
986 gemini_fr.id = id.clone();
989
990 gemini_parts.push(adk_gemini::Part::FunctionResponse {
991 function_response: gemini_fr,
992 thought_signature: sig,
993 });
994 }
995 }
996 if !gemini_parts.is_empty() {
997 let fn_content = adk_gemini::Content {
998 role: Some(adk_gemini::Role::User),
999 parts: Some(gemini_parts),
1000 };
1001 builder = builder.with_message(adk_gemini::Message {
1002 content: fn_content,
1003 role: adk_gemini::Role::User,
1004 });
1005 }
1006 }
1007 _ => {}
1008 }
1009 }
1010
1011 if let Some(config) = req.config {
1013 let has_schema = config.response_schema.is_some();
1014 let gen_config = adk_gemini::GenerationConfig {
1015 temperature: config.temperature,
1016 top_p: config.top_p,
1017 top_k: config.top_k,
1018 max_output_tokens: config.max_output_tokens,
1019 response_schema: config.response_schema,
1020 response_mime_type: if has_schema {
1021 Some("application/json".to_string())
1022 } else {
1023 None
1024 },
1025 thinking_config: self.thinking_config.clone(),
1026 ..Default::default()
1027 };
1028 builder = builder.with_generation_config(gen_config);
1029
1030 if let Some(ref name) = config.cached_content {
1032 let handle = self.client.get_cached_content(name);
1033 builder = builder.with_cached_content(&handle);
1034 }
1035 } else if self.thinking_config.is_some() {
1036 let gen_config = adk_gemini::GenerationConfig {
1039 thinking_config: self.thinking_config.clone(),
1040 ..Default::default()
1041 };
1042 builder = builder.with_generation_config(gen_config);
1043 }
1044
1045 if !req.tools.is_empty() {
1047 let adapter = self.schema_adapter();
1048 use std::sync::LazyLock;
1049 static SCHEMA_CACHE: LazyLock<SchemaCache> = LazyLock::new(|| {
1050 SchemaCache::for_adapter(std::sync::Arc::new(GeminiSchemaAdapter::new()))
1051 });
1052 let (gemini_tools, tool_config) =
1053 Self::build_gemini_tools(&req.tools, adapter, &SCHEMA_CACHE)?;
1054 for tool in gemini_tools {
1055 builder = builder.with_tool(tool);
1056 }
1057 if tool_config != adk_gemini::ToolConfig::default() {
1058 builder = builder.with_tool_config(tool_config);
1059 }
1060 }
1061
1062 if stream {
1063 adk_telemetry::debug!("Executing streaming request");
1064 let response_stream = builder.execute_stream().await.map_err(|e| {
1065 adk_telemetry::error!(error = %e, "Model request failed");
1066 gemini_error_to_adk(&e)
1067 })?;
1068
1069 let mapped_stream = async_stream::stream! {
1070 let mut stream = response_stream;
1071 let mut saw_partial_chunk = false;
1072 while let Some(result) = stream.try_next().await.transpose() {
1073 match result {
1074 Ok(resp) => {
1075 match Self::convert_response(&resp) {
1076 Ok(llm_resp) => {
1077 let (chunks, next_saw_partial) =
1078 Self::stream_chunks_from_response(llm_resp, saw_partial_chunk);
1079 saw_partial_chunk = next_saw_partial;
1080 for chunk in chunks {
1081 yield Ok(chunk);
1082 }
1083 }
1084 Err(e) => {
1085 adk_telemetry::error!(error = %e, "Failed to convert response");
1086 yield Err(e);
1087 }
1088 }
1089 }
1090 Err(e) => {
1091 adk_telemetry::error!(error = %e, "Stream error");
1092 yield Err(gemini_error_to_adk(&e));
1093 }
1094 }
1095 }
1096 };
1097
1098 Ok(Box::pin(mapped_stream))
1099 } else {
1100 adk_telemetry::debug!("Executing blocking request");
1101 let response = builder.execute().await.map_err(|e| {
1102 adk_telemetry::error!(error = %e, "Model request failed");
1103 gemini_error_to_adk(&e)
1104 })?;
1105
1106 let llm_response = Self::convert_response(&response)?;
1107
1108 let stream = async_stream::stream! {
1109 yield Ok(llm_response);
1110 };
1111
1112 Ok(Box::pin(stream))
1113 }
1114 }
1115
1116 pub async fn create_cached_content(
1121 &self,
1122 system_instruction: &str,
1123 tools: &std::collections::HashMap<String, serde_json::Value>,
1124 ttl_seconds: u32,
1125 ) -> Result<String> {
1126 let mut cache_builder = self
1127 .client
1128 .create_cache()
1129 .with_system_instruction(system_instruction)
1130 .with_ttl(std::time::Duration::from_secs(u64::from(ttl_seconds)));
1131
1132 let adapter = self.schema_adapter();
1133 use std::sync::LazyLock;
1134 static SCHEMA_CACHE: LazyLock<SchemaCache> = LazyLock::new(|| {
1135 SchemaCache::for_adapter(std::sync::Arc::new(GeminiSchemaAdapter::new()))
1136 });
1137 let (gemini_tools, tool_config) = Self::build_gemini_tools(tools, adapter, &SCHEMA_CACHE)?;
1138 if !gemini_tools.is_empty() {
1139 cache_builder = cache_builder.with_tools(gemini_tools);
1140 }
1141 if tool_config != adk_gemini::ToolConfig::default() {
1142 cache_builder = cache_builder.with_tool_config(tool_config);
1143 }
1144
1145 let handle = cache_builder
1146 .execute()
1147 .await
1148 .map_err(|e| adk_core::AdkError::model(format!("cache creation failed: {e}")))?;
1149
1150 Ok(handle.name().to_string())
1151 }
1152
1153 pub async fn delete_cached_content(&self, name: &str) -> Result<()> {
1155 let handle = self.client.get_cached_content(name);
1156 handle
1157 .delete()
1158 .await
1159 .map_err(|(_, e)| adk_core::AdkError::model(format!("cache deletion failed: {e}")))?;
1160 Ok(())
1161 }
1162
1163 #[cfg(feature = "gemini-interactions")]
1191 async fn generate_interactions_once(&self, req: LlmRequest) -> Result<LlmResponse> {
1192 use super::interactions_convert;
1193
1194 let target = self.interaction_target.as_ref().ok_or_else(|| {
1197 adk_core::AdkError::new(
1198 ErrorComponent::Model,
1199 ErrorCategory::InvalidInput,
1200 "model.gemini.interactions.missing_target",
1201 "the Interactions transport is active but no validated target is configured",
1202 )
1203 .with_provider("gemini")
1204 })?;
1205
1206 let thinking_level = self.thinking_config.as_ref().and_then(|c| c.thinking_level);
1209
1210 let stateful = self.interaction_options.stateful;
1211 let store = self.interaction_options.store;
1212 let background = self.resolve_background();
1213
1214 let mut request =
1216 interactions_convert::build_request(&req, target, thinking_level, stateful, store)?;
1217 request.background = Some(background);
1218
1219 let interaction = match self.client.send_interaction(request.clone()).await {
1221 Ok(interaction) => interaction,
1222 Err(error) => {
1223 let mapped = gemini_error_to_adk(&error);
1224 if mapped.category == ErrorCategory::NotFound && req.previous_response_id.is_some()
1228 {
1229 let mut fallback = interactions_convert::build_request(
1230 &req,
1231 target,
1232 thinking_level,
1233 false,
1234 store,
1235 )?;
1236 fallback.background = Some(background);
1237 self.client
1238 .send_interaction(fallback)
1239 .await
1240 .map_err(|e| gemini_error_to_adk(&e))?
1241 } else {
1242 return Err(mapped);
1243 }
1244 }
1245 };
1246
1247 let final_interaction = if background {
1249 self.poll_interaction_to_completion(interaction).await?
1250 } else {
1251 interaction
1252 };
1253
1254 interaction_status_to_result(final_interaction.status)?;
1256
1257 Ok(interactions_convert::to_llm_response(&final_interaction))
1258 }
1259
1260 #[cfg(feature = "gemini-interactions")]
1291 async fn poll_interaction_to_completion(
1292 &self,
1293 interaction: adk_gemini::interactions::Interaction,
1294 ) -> Result<adk_gemini::interactions::Interaction> {
1295 const MAX_POLL_ATTEMPTS: u32 = 600;
1300
1301 let mut current = interaction;
1302 let mut attempts: u32 = 0;
1303 while !current.status.is_terminal() && !current.status.requires_action() {
1304 if attempts >= MAX_POLL_ATTEMPTS {
1305 return Err(adk_core::AdkError::new(
1306 ErrorComponent::Model,
1307 ErrorCategory::Timeout,
1308 "model.gemini.interactions.poll_timeout",
1309 format!(
1310 "the Gemini interaction did not reach a terminal or requires_action \
1311 state after {MAX_POLL_ATTEMPTS} poll attempts"
1312 ),
1313 )
1314 .with_provider("gemini"));
1315 }
1316 attempts += 1;
1317 tokio::time::sleep(self.interaction_options.poll_interval).await;
1318 current = self
1319 .client
1320 .get_interaction(¤t.id, false)
1321 .await
1322 .map_err(|e| gemini_error_to_adk(&e))?;
1323 }
1324 Ok(current)
1325 }
1326
1327 #[cfg(feature = "gemini-interactions")]
1352 async fn generate_interactions_stream(&self, req: LlmRequest) -> Result<LlmResponseStream> {
1353 use super::interactions_convert::{self, SseAccumulator, sse_event_to_chunk};
1354
1355 let target = self.interaction_target.as_ref().ok_or_else(|| {
1356 adk_core::AdkError::new(
1357 ErrorComponent::Model,
1358 ErrorCategory::InvalidInput,
1359 "model.gemini.interactions.missing_target",
1360 "the Interactions transport is active but no validated target is configured",
1361 )
1362 .with_provider("gemini")
1363 })?;
1364
1365 let thinking_level = self.thinking_config.as_ref().and_then(|c| c.thinking_level);
1366 let stateful = self.interaction_options.stateful;
1367 let store = self.interaction_options.store;
1368
1369 let mut request =
1370 interactions_convert::build_request(&req, target, thinking_level, stateful, store)?;
1371 request.background = Some(false);
1376
1377 let sse_stream = self
1378 .client
1379 .send_interaction_stream(request)
1380 .await
1381 .map_err(|e| gemini_error_to_adk(&e))?;
1382
1383 let mapped = async_stream::stream! {
1384 let mut sse_stream = sse_stream;
1385 let mut acc = SseAccumulator::new();
1386 while let Some(result) = sse_stream.try_next().await.transpose() {
1387 match result {
1388 Ok(event) => {
1389 if let Some(chunk) = sse_event_to_chunk(event, &mut acc) {
1390 yield chunk;
1391 }
1392 }
1393 Err(e) => {
1394 adk_telemetry::error!(error = %e, "Interaction stream error");
1395 yield Err(gemini_error_to_adk(&e));
1396 }
1397 }
1398 }
1399 };
1400
1401 Ok(Box::pin(mapped))
1402 }
1403}
1404
1405#[async_trait]
1406impl Llm for GeminiModel {
1407 fn name(&self) -> &str {
1408 &self.model_name
1409 }
1410
1411 fn schema_adapter(&self) -> &dyn SchemaAdapter {
1412 use std::sync::LazyLock;
1413 static ADAPTER: LazyLock<GeminiSchemaAdapter> = LazyLock::new(GeminiSchemaAdapter::new);
1414 &*ADAPTER
1415 }
1416
1417 #[cfg(feature = "gemini-interactions")]
1418 fn uses_interactions_api(&self) -> bool {
1419 self.transport == GeminiTransport::Interactions
1420 }
1421
1422 #[adk_telemetry::instrument(
1427 name = "model.generate_content",
1428 skip(self, req),
1429 fields(
1430 model.name = %self.model_name,
1431 stream = %stream,
1432 request.contents_count = %req.contents.len(),
1433 request.tools_count = %req.tools.len()
1434 )
1435 )]
1436 async fn generate_content(&self, req: LlmRequest, stream: bool) -> Result<LlmResponseStream> {
1437 adk_telemetry::info!("Generating content");
1438 let usage_span = adk_telemetry::llm_generate_span("gemini", &self.model_name, stream);
1439
1440 #[cfg(feature = "gemini-interactions")]
1448 if self.transport == GeminiTransport::Interactions {
1449 if stream {
1458 let mapped =
1459 execute_with_retry(&self.retry_config, is_retryable_model_error, || {
1460 self.generate_interactions_stream(req.clone())
1461 })
1462 .await?;
1463 return Ok(crate::usage_tracking::with_usage_tracking(mapped, usage_span));
1464 }
1465 let response = execute_with_retry(&self.retry_config, is_retryable_model_error, || {
1466 self.generate_interactions_once(req.clone())
1467 })
1468 .await?;
1469 let single = async_stream::stream! {
1470 yield Ok(response);
1471 };
1472 return Ok(crate::usage_tracking::with_usage_tracking(Box::pin(single), usage_span));
1473 }
1474
1475 let result = execute_with_retry(&self.retry_config, is_retryable_model_error, || {
1476 self.generate_content_internal(req.clone(), stream)
1477 })
1478 .await?;
1479 Ok(crate::usage_tracking::with_usage_tracking(result, usage_span))
1480 }
1481}
1482
1483#[cfg(test)]
1484mod native_tool_tests {
1485 use super::*;
1486
1487 fn test_adapter() -> GeminiSchemaAdapter {
1488 GeminiSchemaAdapter::new()
1489 }
1490
1491 fn test_cache() -> SchemaCache {
1492 SchemaCache::for_adapter(std::sync::Arc::new(GeminiSchemaAdapter::new()))
1493 }
1494
1495 #[test]
1496 fn test_build_gemini_tools_supports_native_tool_metadata() {
1497 let mut tools = std::collections::HashMap::new();
1498 tools.insert(
1499 "google_search".to_string(),
1500 serde_json::json!({
1501 "x-adk-gemini-tool": {
1502 "google_search": {}
1503 }
1504 }),
1505 );
1506 tools.insert(
1507 "lookup_weather".to_string(),
1508 serde_json::json!({
1509 "name": "lookup_weather",
1510 "description": "lookup weather",
1511 "parameters": {
1512 "type": "object",
1513 "properties": {
1514 "city": { "type": "string" }
1515 }
1516 }
1517 }),
1518 );
1519
1520 let adapter = test_adapter();
1521 let cache = test_cache();
1522 let (gemini_tools, tool_config) = GeminiModel::build_gemini_tools(&tools, &adapter, &cache)
1523 .expect("tool conversion should succeed");
1524
1525 assert_eq!(gemini_tools.len(), 2);
1526 assert_eq!(tool_config.include_server_side_tool_invocations, Some(true));
1527 }
1528
1529 #[test]
1530 fn test_build_gemini_tools_sets_flag_for_builtin_only() {
1531 let mut tools = std::collections::HashMap::new();
1532 tools.insert(
1533 "google_search".to_string(),
1534 serde_json::json!({
1535 "x-adk-gemini-tool": {
1536 "google_search": {}
1537 }
1538 }),
1539 );
1540
1541 let adapter = test_adapter();
1542 let cache = test_cache();
1543 let (_gemini_tools, tool_config) =
1544 GeminiModel::build_gemini_tools(&tools, &adapter, &cache)
1545 .expect("tool conversion should succeed");
1546
1547 assert_eq!(
1548 tool_config.include_server_side_tool_invocations,
1549 Some(true),
1550 "includeServerSideToolInvocations should be set even with only built-in tools"
1551 );
1552 }
1553
1554 #[test]
1555 fn test_build_gemini_tools_no_flag_for_function_only() {
1556 let mut tools = std::collections::HashMap::new();
1557 tools.insert(
1558 "lookup_weather".to_string(),
1559 serde_json::json!({
1560 "name": "lookup_weather",
1561 "description": "lookup weather",
1562 "parameters": {
1563 "type": "object",
1564 "properties": {
1565 "city": { "type": "string" }
1566 }
1567 }
1568 }),
1569 );
1570
1571 let adapter = test_adapter();
1572 let cache = test_cache();
1573 let (_gemini_tools, tool_config) =
1574 GeminiModel::build_gemini_tools(&tools, &adapter, &cache)
1575 .expect("tool conversion should succeed");
1576
1577 assert_eq!(
1578 tool_config.include_server_side_tool_invocations, None,
1579 "includeServerSideToolInvocations should NOT be set for function-only tools"
1580 );
1581 }
1582
1583 #[test]
1584 fn test_build_gemini_tools_merges_native_tool_config() {
1585 let mut tools = std::collections::HashMap::new();
1586 tools.insert(
1587 "google_maps".to_string(),
1588 serde_json::json!({
1589 "x-adk-gemini-tool": {
1590 "google_maps": {
1591 "enable_widget": true
1592 }
1593 },
1594 "x-adk-gemini-tool-config": {
1595 "retrievalConfig": {
1596 "latLng": {
1597 "latitude": 1.23,
1598 "longitude": 4.56
1599 }
1600 }
1601 }
1602 }),
1603 );
1604
1605 let adapter = test_adapter();
1606 let cache = test_cache();
1607 let (_gemini_tools, tool_config) =
1608 GeminiModel::build_gemini_tools(&tools, &adapter, &cache)
1609 .expect("tool conversion should succeed");
1610
1611 assert_eq!(
1612 tool_config.retrieval_config,
1613 Some(serde_json::json!({
1614 "latLng": {
1615 "latitude": 1.23,
1616 "longitude": 4.56
1617 }
1618 }))
1619 );
1620 }
1621
1622 #[test]
1623 fn test_response_schema_is_normalized_like_parameters() {
1624 let mut tools = std::collections::HashMap::new();
1630 tools.insert(
1631 "read_file".to_string(),
1632 serde_json::json!({
1633 "name": "read_file",
1634 "description": "Read a file from the filesystem",
1635 "parameters": {
1636 "$schema": "http://json-schema.org/draft-07/schema#",
1637 "type": "object",
1638 "properties": {
1639 "path": { "type": "string", "description": "File path" }
1640 },
1641 "required": ["path"],
1642 "additionalProperties": false
1643 },
1644 "response": {
1645 "$schema": "http://json-schema.org/draft-07/schema#",
1646 "type": "object",
1647 "properties": {
1648 "content": { "type": "string", "description": "File contents" }
1649 },
1650 "required": ["content"],
1651 "additionalProperties": false
1652 }
1653 }),
1654 );
1655
1656 let adapter = test_adapter();
1657 let cache = test_cache();
1658 let (gemini_tools, _) = GeminiModel::build_gemini_tools(&tools, &adapter, &cache)
1659 .expect("tool conversion should succeed");
1660
1661 let func_tool = gemini_tools
1663 .iter()
1664 .find(|t| matches!(t, adk_gemini::Tool::Function { .. }))
1665 .expect("should have function declarations");
1666
1667 let decls = match func_tool {
1668 adk_gemini::Tool::Function { function_declarations } => function_declarations,
1669 _ => panic!("expected Function tool"),
1670 };
1671
1672 let decl = &decls[0];
1673 let decl_json = serde_json::to_value(decl).unwrap();
1674
1675 let params = &decl_json["parameters"];
1677 assert!(params.get("$schema").is_none(), "parameters.$schema should be stripped");
1678 assert!(
1679 params.get("additionalProperties").is_none(),
1680 "parameters.additionalProperties should be stripped"
1681 );
1682
1683 let response = &decl_json["response"];
1685 assert!(
1686 response.get("$schema").is_none(),
1687 "response.$schema should be stripped (was the bug: copied raw)"
1688 );
1689 assert!(
1690 response.get("additionalProperties").is_none(),
1691 "response.additionalProperties should be stripped (was the bug: copied raw)"
1692 );
1693 }
1694}
1695
1696#[async_trait]
1697impl CacheCapable for GeminiModel {
1698 async fn create_cache(
1699 &self,
1700 system_instruction: &str,
1701 tools: &std::collections::HashMap<String, serde_json::Value>,
1702 ttl_seconds: u32,
1703 ) -> Result<String> {
1704 self.create_cached_content(system_instruction, tools, ttl_seconds).await
1705 }
1706
1707 async fn delete_cache(&self, name: &str) -> Result<()> {
1708 self.delete_cached_content(name).await
1709 }
1710}
1711
1712#[cfg(test)]
1713mod tests {
1714 use super::*;
1715 use adk_core::AdkError;
1716 use std::{
1717 sync::{
1718 Arc,
1719 atomic::{AtomicU32, Ordering},
1720 },
1721 time::Duration,
1722 };
1723
1724 #[test]
1725 fn constructor_is_backward_compatible_and_sync() {
1726 fn accepts_sync_constructor<F>(_f: F)
1727 where
1728 F: Fn(&str, &str) -> Result<GeminiModel>,
1729 {
1730 }
1731
1732 accepts_sync_constructor(|api_key, model| GeminiModel::new(api_key, model));
1733 }
1734
1735 #[test]
1736 fn stream_chunks_from_response_injects_partial_before_lone_final_chunk() {
1737 let response = LlmResponse {
1738 content: Some(Content::new("model").with_text("hello")),
1739 usage_metadata: None,
1740 finish_reason: Some(FinishReason::Stop),
1741 citation_metadata: None,
1742 partial: false,
1743 turn_complete: true,
1744 interrupted: false,
1745 error_code: None,
1746 error_message: None,
1747 provider_metadata: None,
1748 interaction_id: None,
1749 };
1750
1751 let (chunks, saw_partial) = GeminiModel::stream_chunks_from_response(response, false);
1752 assert!(saw_partial);
1753 assert_eq!(chunks.len(), 2);
1754 assert!(chunks[0].partial);
1755 assert!(!chunks[0].turn_complete);
1756 assert!(chunks[0].content.is_none());
1757 assert!(!chunks[1].partial);
1758 assert!(chunks[1].turn_complete);
1759 }
1760
1761 #[test]
1762 fn stream_chunks_from_response_keeps_final_only_when_partial_already_seen() {
1763 let response = LlmResponse {
1764 content: Some(Content::new("model").with_text("done")),
1765 usage_metadata: None,
1766 finish_reason: Some(FinishReason::Stop),
1767 citation_metadata: None,
1768 partial: false,
1769 turn_complete: true,
1770 interrupted: false,
1771 error_code: None,
1772 error_message: None,
1773 provider_metadata: None,
1774 interaction_id: None,
1775 };
1776
1777 let (chunks, saw_partial) = GeminiModel::stream_chunks_from_response(response, true);
1778 assert!(saw_partial);
1779 assert_eq!(chunks.len(), 1);
1780 assert!(!chunks[0].partial);
1781 assert!(chunks[0].turn_complete);
1782 }
1783
1784 #[tokio::test]
1785 async fn execute_with_retry_retries_retryable_errors() {
1786 let retry_config = RetryConfig::default()
1787 .with_max_retries(2)
1788 .with_initial_delay(Duration::from_millis(0))
1789 .with_max_delay(Duration::from_millis(0));
1790 let attempts = Arc::new(AtomicU32::new(0));
1791
1792 let result = execute_with_retry(&retry_config, is_retryable_model_error, || {
1793 let attempts = Arc::clone(&attempts);
1794 async move {
1795 let attempt = attempts.fetch_add(1, Ordering::SeqCst);
1796 if attempt < 2 {
1797 return Err(AdkError::model("code 429 RESOURCE_EXHAUSTED"));
1798 }
1799 Ok("ok")
1800 }
1801 })
1802 .await
1803 .expect("retry should eventually succeed");
1804
1805 assert_eq!(result, "ok");
1806 assert_eq!(attempts.load(Ordering::SeqCst), 3);
1807 }
1808
1809 #[tokio::test]
1810 async fn execute_with_retry_does_not_retry_non_retryable_errors() {
1811 let retry_config = RetryConfig::default()
1812 .with_max_retries(3)
1813 .with_initial_delay(Duration::from_millis(0))
1814 .with_max_delay(Duration::from_millis(0));
1815 let attempts = Arc::new(AtomicU32::new(0));
1816
1817 let error = execute_with_retry(&retry_config, is_retryable_model_error, || {
1818 let attempts = Arc::clone(&attempts);
1819 async move {
1820 attempts.fetch_add(1, Ordering::SeqCst);
1821 Err::<(), _>(AdkError::model("code 400 invalid request"))
1822 }
1823 })
1824 .await
1825 .expect_err("non-retryable error should return immediately");
1826
1827 assert!(error.is_model());
1828 assert_eq!(attempts.load(Ordering::SeqCst), 1);
1829 }
1830
1831 #[tokio::test]
1832 async fn execute_with_retry_respects_disabled_config() {
1833 let retry_config = RetryConfig::disabled().with_max_retries(10);
1834 let attempts = Arc::new(AtomicU32::new(0));
1835
1836 let error = execute_with_retry(&retry_config, is_retryable_model_error, || {
1837 let attempts = Arc::clone(&attempts);
1838 async move {
1839 attempts.fetch_add(1, Ordering::SeqCst);
1840 Err::<(), _>(AdkError::model("code 429 RESOURCE_EXHAUSTED"))
1841 }
1842 })
1843 .await
1844 .expect_err("disabled retries should return first error");
1845
1846 assert!(error.is_model());
1847 assert_eq!(attempts.load(Ordering::SeqCst), 1);
1848 }
1849
1850 #[test]
1851 fn convert_response_preserves_citation_metadata() {
1852 let response = adk_gemini::GenerationResponse {
1853 candidates: vec![adk_gemini::Candidate {
1854 content: adk_gemini::Content {
1855 role: Some(adk_gemini::Role::Model),
1856 parts: Some(vec![adk_gemini::Part::Text {
1857 text: "hello world".to_string(),
1858 thought: None,
1859 thought_signature: None,
1860 }]),
1861 },
1862 safety_ratings: None,
1863 citation_metadata: Some(adk_gemini::CitationMetadata {
1864 citation_sources: vec![adk_gemini::CitationSource {
1865 uri: Some("https://example.com".to_string()),
1866 title: Some("Example".to_string()),
1867 start_index: Some(0),
1868 end_index: Some(5),
1869 license: Some("CC-BY".to_string()),
1870 publication_date: None,
1871 }],
1872 }),
1873 grounding_metadata: None,
1874 finish_reason: Some(adk_gemini::FinishReason::Stop),
1875 index: Some(0),
1876 }],
1877 prompt_feedback: None,
1878 usage_metadata: None,
1879 model_version: None,
1880 response_id: None,
1881 };
1882
1883 let converted =
1884 GeminiModel::convert_response(&response).expect("conversion should succeed");
1885 let metadata = converted.citation_metadata.expect("citation metadata should be mapped");
1886 assert_eq!(metadata.citation_sources.len(), 1);
1887 assert_eq!(metadata.citation_sources[0].uri.as_deref(), Some("https://example.com"));
1888 assert_eq!(metadata.citation_sources[0].start_index, Some(0));
1889 assert_eq!(metadata.citation_sources[0].end_index, Some(5));
1890 }
1891
1892 #[test]
1893 fn convert_response_handles_inline_data_from_model() {
1894 let image_bytes = vec![0x89, 0x50, 0x4E, 0x47];
1895 let encoded = crate::attachment::encode_base64(&image_bytes);
1896
1897 let response = adk_gemini::GenerationResponse {
1898 candidates: vec![adk_gemini::Candidate {
1899 content: adk_gemini::Content {
1900 role: Some(adk_gemini::Role::Model),
1901 parts: Some(vec![
1902 adk_gemini::Part::Text {
1903 text: "Here is the image".to_string(),
1904 thought: None,
1905 thought_signature: None,
1906 },
1907 adk_gemini::Part::InlineData {
1908 inline_data: adk_gemini::Blob {
1909 mime_type: "image/png".to_string(),
1910 data: encoded,
1911 },
1912 },
1913 ]),
1914 },
1915 safety_ratings: None,
1916 citation_metadata: None,
1917 grounding_metadata: None,
1918 finish_reason: Some(adk_gemini::FinishReason::Stop),
1919 index: Some(0),
1920 }],
1921 prompt_feedback: None,
1922 usage_metadata: None,
1923 model_version: None,
1924 response_id: None,
1925 };
1926
1927 let converted =
1928 GeminiModel::convert_response(&response).expect("conversion should succeed");
1929 let content = converted.content.expect("should have content");
1930 assert!(
1931 content
1932 .parts
1933 .iter()
1934 .any(|part| matches!(part, Part::Text { text } if text == "Here is the image"))
1935 );
1936 assert!(content.parts.iter().any(|part| {
1937 matches!(
1938 part,
1939 Part::InlineData { mime_type, data, .. }
1940 if mime_type == "image/png" && data.as_slice() == image_bytes.as_slice()
1941 )
1942 }));
1943 }
1944
1945 #[test]
1946 fn gemini_function_response_payload_preserves_objects() {
1947 let value = serde_json::json!({
1948 "documents": [
1949 { "id": "pricing", "score": 0.91 }
1950 ]
1951 });
1952
1953 let payload = GeminiModel::gemini_function_response_payload(value.clone());
1954
1955 assert_eq!(payload, value);
1956 }
1957
1958 #[test]
1959 fn gemini_function_response_payload_wraps_arrays() {
1960 let payload =
1961 GeminiModel::gemini_function_response_payload(serde_json::json!([{ "id": "pricing" }]));
1962
1963 assert_eq!(payload, serde_json::json!({ "result": [{ "id": "pricing" }] }));
1964 }
1965
1966 fn convert_function_response_to_gemini_fr(
1971 frd: &adk_core::FunctionResponseData,
1972 ) -> adk_gemini::tools::FunctionResponse {
1973 let mut fr_parts = Vec::new();
1974
1975 for inline in &frd.inline_data {
1976 let encoded = crate::attachment::encode_base64(&inline.data);
1977 fr_parts.push(adk_gemini::FunctionResponsePart::InlineData {
1978 inline_data: adk_gemini::Blob {
1979 mime_type: inline.mime_type.clone(),
1980 data: encoded,
1981 },
1982 });
1983 }
1984
1985 for file in &frd.file_data {
1986 fr_parts.push(adk_gemini::FunctionResponsePart::FileData {
1987 file_data: adk_gemini::FileDataRef {
1988 mime_type: file.mime_type.clone(),
1989 file_uri: file.file_uri.clone(),
1990 },
1991 });
1992 }
1993
1994 let mut gemini_fr = adk_gemini::tools::FunctionResponse::new(
1995 &frd.name,
1996 GeminiModel::gemini_function_response_payload(frd.response.clone()),
1997 );
1998 gemini_fr.parts = fr_parts;
1999 gemini_fr
2000 }
2001
2002 #[test]
2003 fn json_only_function_response_has_no_nested_parts() {
2004 let frd = adk_core::FunctionResponseData::new("tool", serde_json::json!({"ok": true}));
2005 let gemini_fr = convert_function_response_to_gemini_fr(&frd);
2006 assert!(gemini_fr.parts.is_empty());
2007 let json = serde_json::to_string(&gemini_fr).unwrap();
2009 assert!(!json.contains("\"parts\""));
2010 }
2011
2012 #[test]
2013 fn function_response_with_inline_data_has_nested_parts() {
2014 let frd = adk_core::FunctionResponseData::with_inline_data(
2015 "chart",
2016 serde_json::json!({"status": "ok"}),
2017 vec![adk_core::InlineDataPart {
2018 mime_type: "image/png".to_string(),
2019 data: vec![0x89, 0x50, 0x4E, 0x47],
2020 uri: None,
2021 annotations: None,
2022 }],
2023 );
2024 let gemini_fr = convert_function_response_to_gemini_fr(&frd);
2025 assert_eq!(gemini_fr.parts.len(), 1);
2026 match &gemini_fr.parts[0] {
2027 adk_gemini::FunctionResponsePart::InlineData { inline_data } => {
2028 assert_eq!(inline_data.mime_type, "image/png");
2029 let decoded = BASE64_STANDARD.decode(&inline_data.data).unwrap();
2030 assert_eq!(decoded, vec![0x89, 0x50, 0x4E, 0x47]);
2031 }
2032 other => panic!("expected InlineData, got {other:?}"),
2033 }
2034 }
2035
2036 #[test]
2037 fn function_response_with_file_data_has_nested_parts() {
2038 let frd = adk_core::FunctionResponseData::with_file_data(
2039 "doc",
2040 serde_json::json!({"ok": true}),
2041 vec![adk_core::FileDataPart {
2042 mime_type: "application/pdf".to_string(),
2043 file_uri: "gs://bucket/report.pdf".to_string(),
2044 annotations: None,
2045 }],
2046 );
2047 let gemini_fr = convert_function_response_to_gemini_fr(&frd);
2048 assert_eq!(gemini_fr.parts.len(), 1);
2049 match &gemini_fr.parts[0] {
2050 adk_gemini::FunctionResponsePart::FileData { file_data } => {
2051 assert_eq!(file_data.mime_type, "application/pdf");
2052 assert_eq!(file_data.file_uri, "gs://bucket/report.pdf");
2053 }
2054 other => panic!("expected FileData, got {other:?}"),
2055 }
2056 }
2057
2058 #[test]
2059 fn function_response_with_both_inline_and_file_data_ordering() {
2060 let frd = adk_core::FunctionResponseData::with_multimodal(
2061 "multi",
2062 serde_json::json!({}),
2063 vec![
2064 adk_core::InlineDataPart {
2065 mime_type: "image/png".to_string(),
2066 data: vec![1, 2],
2067 uri: None,
2068 annotations: None,
2069 },
2070 adk_core::InlineDataPart {
2071 mime_type: "image/jpeg".to_string(),
2072 data: vec![3, 4],
2073 uri: None,
2074 annotations: None,
2075 },
2076 ],
2077 vec![adk_core::FileDataPart {
2078 mime_type: "application/pdf".to_string(),
2079 file_uri: "gs://b/f.pdf".to_string(),
2080 annotations: None,
2081 }],
2082 );
2083 let gemini_fr = convert_function_response_to_gemini_fr(&frd);
2084 assert_eq!(gemini_fr.parts.len(), 3);
2086 assert!(matches!(&gemini_fr.parts[0], adk_gemini::FunctionResponsePart::InlineData { .. }));
2087 assert!(matches!(&gemini_fr.parts[1], adk_gemini::FunctionResponsePart::InlineData { .. }));
2088 assert!(matches!(&gemini_fr.parts[2], adk_gemini::FunctionResponsePart::FileData { .. }));
2089 }
2090}
2091
2092#[cfg(all(test, feature = "gemini-interactions"))]
2093mod interactions_transport_tests {
2094 use super::*;
2095
2096 #[test]
2101 fn default_interaction_options_match_api_posture() {
2102 let opts = InteractionOptions::default();
2103 assert!(opts.store, "store should default to true");
2104 assert!(opts.stateful, "stateful should default to true");
2105 assert_eq!(opts.background, BackgroundMode::AgentTargetsOnly);
2106 assert_eq!(opts.poll_interval, std::time::Duration::from_secs(1));
2107 }
2108
2109 #[test]
2110 fn new_model_defaults_to_generate_content_transport() {
2111 let model = GeminiModel::new("test-key", "gemini-2.5-flash")
2112 .expect("constructing a Gemini model should not require network");
2113 assert_eq!(model.transport(), GeminiTransport::GenerateContent);
2114 }
2115
2116 #[test]
2117 fn use_interactions_api_enables_transport_for_allowlisted_model() {
2118 let model = GeminiModel::new("test-key", "gemini-2.5-flash")
2119 .expect("construct model")
2120 .use_interactions_api(true)
2121 .expect("allowlisted model should enable the Interactions transport");
2122
2123 assert_eq!(model.transport(), GeminiTransport::Interactions);
2124 assert_eq!(
2125 model.interaction_target,
2126 Some(InteractionTarget::Model("gemini-2.5-flash".to_string()))
2127 );
2128 }
2129
2130 #[test]
2131 fn use_interactions_api_rejects_unsupported_model_with_invalid_input() {
2132 let result = GeminiModel::new("test-key", "gemini-2.0-flash")
2133 .expect("construct model")
2134 .use_interactions_api(true);
2135
2136 let err = match result {
2137 Ok(_) => panic!("unsupported model should be rejected"),
2138 Err(err) => err,
2139 };
2140 assert_eq!(err.category, adk_core::ErrorCategory::InvalidInput);
2141 assert_eq!(err.details.provider.as_deref(), Some("gemini"));
2142 }
2143
2144 #[test]
2145 fn use_interactions_api_false_reverts_to_generate_content() {
2146 let model = GeminiModel::new("test-key", "gemini-2.5-flash")
2147 .expect("construct model")
2148 .use_interactions_api(true)
2149 .expect("enable interactions")
2150 .use_interactions_api(false)
2151 .expect("disabling should always succeed");
2152
2153 assert_eq!(model.transport(), GeminiTransport::GenerateContent);
2154 assert_eq!(model.interaction_target, None);
2155 }
2156
2157 #[test]
2158 fn interaction_options_override_is_stored() {
2159 let opts = InteractionOptions {
2160 store: false,
2161 stateful: false,
2162 background: BackgroundMode::Always,
2163 poll_interval: std::time::Duration::from_millis(250),
2164 };
2165 let model = GeminiModel::new("test-key", "gemini-2.5-flash")
2166 .expect("construct model")
2167 .interaction_options(opts.clone());
2168
2169 let stored = model.interaction_options_ref();
2170 assert_eq!(stored.store, opts.store);
2171 assert_eq!(stored.stateful, opts.stateful);
2172 assert_eq!(stored.background, opts.background);
2173 assert_eq!(stored.poll_interval, opts.poll_interval);
2174 }
2175
2176 #[test]
2177 fn resolve_background_honors_background_mode() {
2178 let always = GeminiModel::new("test-key", "gemini-2.5-flash")
2180 .expect("construct model")
2181 .use_interactions_api(true)
2182 .expect("enable")
2183 .interaction_options(InteractionOptions {
2184 background: BackgroundMode::Always,
2185 ..InteractionOptions::default()
2186 });
2187 assert!(always.resolve_background());
2188
2189 let never = GeminiModel::new("test-key", "gemini-2.5-flash")
2191 .expect("construct model")
2192 .use_interactions_api(true)
2193 .expect("enable")
2194 .interaction_options(InteractionOptions {
2195 background: BackgroundMode::Never,
2196 ..InteractionOptions::default()
2197 });
2198 assert!(!never.resolve_background());
2199
2200 let model_target = GeminiModel::new("test-key", "gemini-2.5-flash")
2202 .expect("construct model")
2203 .use_interactions_api(true)
2204 .expect("enable");
2205 assert!(!model_target.resolve_background());
2206
2207 let agent_target = GeminiModel::new("test-key", "deep-research-preview-04-2026")
2209 .expect("construct model")
2210 .use_interactions_api(true)
2211 .expect("enable");
2212 assert!(agent_target.resolve_background());
2213 }
2214
2215 #[test]
2218 fn terminal_failed_status_maps_to_error() {
2219 use adk_gemini::interactions::InteractionStatus;
2220
2221 let err = interaction_status_to_result(InteractionStatus::Failed)
2222 .expect_err("a failed interaction must surface an error");
2223 assert_eq!(err.category, adk_core::ErrorCategory::Internal);
2224 assert_eq!(err.details.provider.as_deref(), Some("gemini"));
2225 assert_eq!(err.code, "model.gemini.interactions.failed");
2226 }
2227
2228 #[test]
2231 fn terminal_budget_exceeded_status_maps_to_error() {
2232 use adk_gemini::interactions::InteractionStatus;
2233
2234 let err = interaction_status_to_result(InteractionStatus::BudgetExceeded)
2235 .expect_err("a budget_exceeded interaction must surface an error");
2236 assert_eq!(err.category, adk_core::ErrorCategory::Internal);
2237 assert_eq!(err.details.provider.as_deref(), Some("gemini"));
2238 assert_eq!(err.code, "model.gemini.interactions.budget_exceeded");
2239 }
2240
2241 #[test]
2244 fn non_failure_statuses_are_ok() {
2245 use adk_gemini::interactions::InteractionStatus;
2246
2247 for status in [
2248 InteractionStatus::InProgress,
2249 InteractionStatus::RequiresAction,
2250 InteractionStatus::Completed,
2251 InteractionStatus::Cancelled,
2252 InteractionStatus::Incomplete,
2253 ] {
2254 assert!(
2255 interaction_status_to_result(status).is_ok(),
2256 "status {status:?} should not map to an error"
2257 );
2258 }
2259 }
2260
2261 #[test]
2264 fn failed_interaction_resource_surfaces_error() {
2265 use adk_gemini::interactions::{Interaction, InteractionStatus};
2266
2267 let interaction = Interaction {
2268 id: "v1_failed".to_string(),
2269 model: Some("gemini-2.5-flash".to_string()),
2270 agent: None,
2271 status: InteractionStatus::Failed,
2272 steps: Vec::new(),
2273 usage: None,
2274 created: None,
2275 updated: None,
2276 environment_id: None,
2277 };
2278
2279 let err = interaction_status_to_result(interaction.status)
2280 .expect_err("failed interaction must surface an error");
2281 assert_eq!(err.category, adk_core::ErrorCategory::Internal);
2282 assert_eq!(err.details.provider.as_deref(), Some("gemini"));
2283 }
2284}