Merge branch 'fawney19:main' into main

crates/aether-ai-formats/src/formats/gemini/generate_content/request.rs

@@ -15,7 +15,7 @@ use crate::{
         apply_gemini_request_extensions, canonical_extension_object_mut,
         canonical_openai_reasoning_effort, extract_gemini_model_from_path,
         gemini_contents_to_canonical_messages, gemini_extensions, gemini_generation_config,
-        gemini_generation_config_extra, gemini_openai_extra_body,
+        gemini_generation_config_extra, gemini_google_search_grounding, gemini_openai_extra_body,
         gemini_response_format_to_canonical, gemini_system_to_canonical_instructions,
         gemini_thinking_to_canonical, gemini_tool_choice_to_canonical, gemini_tools_to_canonical,
         gemini_value_by_case, CanonicalContentBlock, CanonicalMessage, CanonicalRequest,
@@ -81,7 +81,7 @@ pub fn from_raw(body_json: &Value, request_path: &str) -> Option<CanonicalReques
             .get("generationConfig")
             .or_else(|| request.get("generation_config")),
     );
-    let (tools, builtin_tools, web_search_options, raw_tools) =
+    let (tools, builtin_tools, web_search_options, raw_tools, google_search_grounding) =
         gemini_tools_to_canonical(request.get("tools"))?;
     canonical.tools = tools;
     canonical.tool_choice = gemini_tool_choice_to_canonical(
@@ -158,6 +158,13 @@ pub fn from_raw(body_json: &Value, request_path: &str) -> Option<CanonicalReques
         canonical_extension_object_mut(&mut canonical.extensions, "gemini")
             .insert("builtin_tools".to_string(), Value::Array(builtin_tools));
     }
+    if let Some(google_search_grounding) = google_search_grounding {
+        let gemini_extension = canonical_extension_object_mut(&mut canonical.extensions, "gemini");
+        gemini_extension.insert(
+            "grounding".to_string(),
+            json!({ "google_search": google_search_grounding }),
+        );
+    }
     if let Some(tool_config) = request
         .get("toolConfig")
         .or_else(|| request.get("tool_config"))
@@ -492,6 +499,10 @@ fn canonical_tools_to_gemini(canonical: &CanonicalRequest) -> Option<Value> {
         .get("openai")
         .and_then(Value::as_object)
         .is_some_and(|value| value.contains_key("web_search_options"));
+    let mut google_search_payload = canonical_google_search_output_payload(canonical);
+    if google_search_payload.is_some() {
+        google_search = true;
+    }
     let mut code_execution = false;
     let mut url_context = false;
 
@@ -528,18 +539,9 @@ fn canonical_tools_to_gemini(canonical: &CanonicalRequest) -> Option<Value> {
         }
         declarations.push(canonical_tool_to_gemini_declaration(tool));
     }
-    if code_execution {
-        tools.push(json!({ "codeExecution": {} }));
-    }
-    if google_search {
-        tools.push(json!({ "googleSearch": {} }));
-    }
-    if url_context {
-        tools.push(json!({ "urlContext": {} }));
-    }
-    if !declarations.is_empty() {
-        tools.push(json!({ "functionDeclarations": declarations }));
-    }
+    let mut emitted_google_search = false;
+    let mut emitted_code_execution = false;
+    let mut emitted_url_context = false;
     if let Some(builtin_tools) = canonical
         .extensions
         .get("gemini")
@@ -547,11 +549,124 @@ fn canonical_tools_to_gemini(canonical: &CanonicalRequest) -> Option<Value> {
         .and_then(|value| value.get("builtin_tools"))
         .and_then(Value::as_array)
     {
-        tools.extend(builtin_tools.iter().cloned());
+        for builtin_tool in builtin_tools {
+            let Some(tool_object) = builtin_tool.as_object() else {
+                tools.push(builtin_tool.clone());
+                continue;
+            };
+            let mut emitted_builtin_portion = false;
+            if let Some(grounding) = gemini_google_search_grounding(tool_object) {
+                google_search = true;
+                if google_search_payload.is_none() {
+                    google_search_payload = Some(grounding.output_payload);
+                }
+                if !emitted_google_search {
+                    tools.push(json!({
+                        "googleSearch": google_search_payload.clone().unwrap_or_else(|| json!({}))
+                    }));
+                    emitted_google_search = true;
+                }
+                emitted_builtin_portion = true;
+            }
+            if let Some(tool) =
+                gemini_builtin_tool_by_case(tool_object, "codeExecution", "code_execution")
+            {
+                if !emitted_code_execution {
+                    tools.push(tool);
+                    emitted_code_execution = true;
+                }
+                emitted_builtin_portion = true;
+            }
+            if let Some(tool) =
+                gemini_builtin_tool_by_case(tool_object, "urlContext", "url_context")
+            {
+                if !emitted_url_context {
+                    tools.push(tool);
+                    emitted_url_context = true;
+                }
+                emitted_builtin_portion = true;
+            }
+            if let Some(tool) = gemini_unhandled_builtin_tool_portion(tool_object) {
+                tools.push(tool);
+            } else if !emitted_builtin_portion {
+                tools.push(builtin_tool.clone());
+            }
+        }
+    }
+    if code_execution && !emitted_code_execution {
+        tools.push(json!({ "codeExecution": {} }));
+    }
+    if google_search && !emitted_google_search {
+        tools.push(json!({
+            "googleSearch": google_search_payload.unwrap_or_else(|| json!({}))
+        }));
+    }
+    if url_context && !emitted_url_context {
+        tools.push(json!({ "urlContext": {} }));
+    }
+    if !declarations.is_empty() {
+        tools.push(json!({ "functionDeclarations": declarations }));
     }
     (!tools.is_empty()).then_some(Value::Array(tools))
 }
 
+fn canonical_google_search_output_payload(canonical: &CanonicalRequest) -> Option<Value> {
+    let google_search = canonical
+        .extensions
+        .get("gemini")
+        .and_then(Value::as_object)
+        .and_then(|value| value.get("grounding"))
+        .and_then(Value::as_object)
+        .and_then(|value| value.get("google_search"))
+        .and_then(Value::as_object)?;
+    google_search
+        .get("legacy")
+        .and_then(Value::as_bool)
+        .filter(|legacy| *legacy)
+        .map(|_| json!({}))
+        .or_else(|| google_search.get("payload").cloned())
+}
+
+fn gemini_builtin_tool_by_case(
+    tool_object: &Map<String, Value>,
+    camel: &str,
+    snake: &str,
+) -> Option<Value> {
+    let payload = tool_object
+        .get(camel)
+        .or_else(|| tool_object.get(snake))
+        .map(gemini_builtin_tool_payload)?;
+    Some(json!({ camel: payload }))
+}
+
+fn gemini_builtin_tool_payload(payload: &Value) -> Value {
+    match payload {
+        Value::Null => json!({}),
+        value => value.clone(),
+    }
+}
+
+fn gemini_unhandled_builtin_tool_portion(tool_object: &Map<String, Value>) -> Option<Value> {
+    let builtin = tool_object
+        .iter()
+        .filter(|(key, _)| {
+            !matches!(
+                key.as_str(),
+                "googleSearch"
+                    | "google_search"
+                    | "googleSearchRetrieval"
+                    | "google_search_retrieval"
+                    | "codeExecution"
+                    | "code_execution"
+                    | "urlContext"
+                    | "url_context"
+            )
+        })
+        .map(|(key, value)| (key.clone(), value.clone()))
+        .collect::<Map<_, _>>();
+    (!builtin.is_empty()).then_some(Value::Object(builtin))
+}
+
 fn canonical_tool_to_gemini_declaration(tool: &CanonicalToolDefinition) -> Value {
     let mut declaration = Map::new();
     declaration.insert("name".to_string(), Value::String(tool.name.clone()));

crates/aether-ai-formats/src/formats/gemini/generate_content/response.rs

@@ -63,7 +63,10 @@ pub fn from_raw(body_json: &Value) -> Option<CanonicalResponse> {
             role: CanonicalRole::Assistant,
             content,
             stop_reason,
-            extensions: Default::default(),
+            extensions: gemini_extensions(
+                candidate_object,
+                &["index", "content", "finishReason", "finish_reason"],
+            ),
         });
     }
     outputs.retain(gemini_response_output_has_visible_content);
@@ -161,7 +164,7 @@ fn canonical_to_gemini_response(
     let mut candidates = Vec::new();
     for output in outputs {
         let parts = canonical_blocks_to_gemini_parts(&output.content)?;
-        candidates.push(json!({
+        let mut candidate = json!({
             "index": output.index,
             "content": {
                 "role": "model",
@@ -170,7 +173,15 @@ fn canonical_to_gemini_response(
             "finishReason": canonical_stop_reason_to_gemini(
                 output.stop_reason.as_ref().or(canonical.stop_reason.as_ref())
             ),
-        }));
+        });
+        if let Some(candidate_object) = candidate.as_object_mut() {
+            if let Some(gemini) = output.extensions.get("gemini").and_then(Value::as_object) {
+                for (key, value) in gemini {
+                    candidate_object.entry(key.clone()).or_insert(value.clone());
+                }
+            }
+        }
+        candidates.push(candidate);
     }
 
     let mut response = Map::new();

crates/aether-ai-formats/src/protocol/canonical.rs

@@ -3217,28 +3217,136 @@ pub(crate) type GeminiCanonicalTools = (
     Vec<Value>,
     Option<Value>,
     Option<Value>,
+    Option<Value>,
 );
 
+#[derive(Debug, Clone)]
+pub(crate) struct GeminiGoogleSearchGrounding {
+    pub source_field: &'static str,
+    pub source_dialect: &'static str,
+    pub legacy: bool,
+    pub payload: Value,
+    pub raw_payload: Value,
+    pub output_payload: Value,
+}
+
+pub(crate) fn gemini_google_search_grounding(
+    tool_object: &Map<String, Value>,
+) -> Option<GeminiGoogleSearchGrounding> {
+    for (field, source_dialect, legacy) in [
+        ("googleSearch", "gemini_current", false),
+        ("google_search", "gemini_current", false),
+        ("googleSearchRetrieval", "vertex_legacy", true),
+        ("google_search_retrieval", "vertex_legacy", true),
+    ] {
+        if let Some(raw_payload) = tool_object.get(field) {
+            let raw_payload = normalize_gemini_tool_payload(raw_payload);
+            let payload = lower_camelize_json_object_keys(&raw_payload);
+            let output_payload = if legacy { json!({}) } else { payload.clone() };
+            return Some(GeminiGoogleSearchGrounding {
+                source_field: field,
+                source_dialect,
+                legacy,
+                payload,
+                raw_payload,
+                output_payload,
+            });
+        }
+    }
+    None
+}
+
+pub(crate) fn gemini_google_search_grounding_extension(
+    grounding: &GeminiGoogleSearchGrounding,
+) -> Value {
+    json!({
+        "enabled": true,
+        "source_field": grounding.source_field,
+        "source_dialect": grounding.source_dialect,
+        "payload": grounding.payload,
+        "raw_payload": grounding.raw_payload,
+        "legacy": grounding.legacy,
+    })
+}
+
+fn normalize_gemini_tool_payload(payload: &Value) -> Value {
+    match payload {
+        Value::Null => json!({}),
+        value => value.clone(),
+    }
+}
+
+fn lower_camelize_json_object_keys(value: &Value) -> Value {
+    match value {
+        Value::Object(object) => Value::Object(
+            object
+                .iter()
+                .map(|(key, value)| {
+                    (
+                        snake_to_lower_camel(key),
+                        lower_camelize_json_object_keys(value),
+                    )
+                })
+                .collect(),
+        ),
+        Value::Array(items) => {
+            Value::Array(items.iter().map(lower_camelize_json_object_keys).collect())
+        }
+        other => other.clone(),
+    }
+}
+
+fn snake_to_lower_camel(key: &str) -> String {
+    let mut output = String::with_capacity(key.len());
+    let mut uppercase_next = false;
+    for character in key.chars() {
+        if character == '_' {
+            uppercase_next = true;
+            continue;
+        }
+        if uppercase_next {
+            for uppercase in character.to_uppercase() {
+                output.push(uppercase);
+            }
+            uppercase_next = false;
+        } else {
+            output.push(character);
+        }
+    }
+    output
+}
+
+fn gemini_builtin_tool_portion(tool_object: &Map<String, Value>) -> Option<Value> {
+    let builtin = tool_object
+        .iter()
+        .filter(|(key, _)| {
+            key.as_str() != "functionDeclarations" && key.as_str() != "function_declarations"
+        })
+        .map(|(key, value)| (key.clone(), value.clone()))
+        .collect::<Map<_, _>>();
+    (!builtin.is_empty()).then_some(Value::Object(builtin))
+}
+
 pub(crate) fn gemini_tools_to_canonical(value: Option<&Value>) -> Option<GeminiCanonicalTools> {
     let Some(value) = value else {
-        return Some((Vec::new(), Vec::new(), None, None));
+        return Some((Vec::new(), Vec::new(), None, None, None));
     };
     let tools = value.as_array()?;
     let mut canonical = Vec::new();
     let mut builtin_tools = Vec::new();
     let mut web_search_options = None;
+    let mut google_search_grounding = None;
     for tool in tools {
         let tool_object = tool.as_object()?;
-        if tool_object.get("googleSearch").is_some() || tool_object.get("google_search").is_some() {
+        if let Some(grounding) = gemini_google_search_grounding(tool_object) {
             web_search_options = Some(json!({}));
-            builtin_tools.push(tool.clone());
+            if google_search_grounding.is_none() {
+                google_search_grounding =
+                    Some(gemini_google_search_grounding_extension(&grounding));
+            }
         }
-        if tool_object.get("codeExecution").is_some()
-            || tool_object.get("code_execution").is_some()
-            || tool_object.get("urlContext").is_some()
-            || tool_object.get("url_context").is_some()
-        {
-            builtin_tools.push(tool.clone());
+        if let Some(builtin_tool) = gemini_builtin_tool_portion(tool_object) {
+            builtin_tools.push(builtin_tool);
         }
         let declarations = tool_object
             .get("functionDeclarations")
@@ -3273,6 +3381,7 @@ pub(crate) fn gemini_tools_to_canonical(value: Option<&Value>) -> Option<GeminiC
         builtin_tools,
         web_search_options,
         Some(value.clone()),
+        google_search_grounding,
     ))
 }
 
@@ -4166,7 +4275,13 @@ pub(crate) fn apply_gemini_request_extensions(
         output_object.insert("cachedContent".to_string(), cached_content);
     }
     if let Some(raw_tools) = gemini.get("raw_tools").cloned() {
-        output_object.insert("tools".to_string(), raw_tools);
+        if should_reuse_raw_gemini_tools(gemini) {
+            output_object.insert("tools".to_string(), raw_tools);
+        } else {
+            output_object
+                .entry("tools".to_string())
+                .or_insert(raw_tools);
+        }
     }
     if let Some(raw_tool_config) = gemini.get("raw_tool_config").cloned() {
         output_object.insert("toolConfig".to_string(), raw_tool_config);
@@ -4174,6 +4289,25 @@ pub(crate) fn apply_gemini_request_extensions(
     Some(())
 }
 
+fn should_reuse_raw_gemini_tools(gemini: &Map<String, Value>) -> bool {
+    let Some(google_search) = gemini
+        .get("grounding")
+        .and_then(Value::as_object)
+        .and_then(|grounding| grounding.get("google_search"))
+        .and_then(Value::as_object)
+    else {
+        return true;
+    };
+    google_search
+        .get("legacy")
+        .and_then(Value::as_bool)
+        .is_none_or(|legacy| !legacy)
+        && google_search
+            .get("source_field")
+            .and_then(Value::as_str)
+            .is_some_and(|source_field| source_field == "googleSearch")
+}
+
 pub(crate) fn assistant_image_placeholder(url: Option<&str>, has_data: bool) -> String {
     match (url, has_data) {
         (Some(url), false) if !url.trim().is_empty() => format!("[Image: {url}]"),
@@ -6000,6 +6134,235 @@ mod tests {
         assert_eq!(rebuilt["toolConfig"], request["toolConfig"]);
     }
 
+    #[test]
+    fn gemini_request_adapter_normalizes_google_search_grounding_aliases() {
+        let cases = [
+            (
+                "current_camel",
+                json!({"googleSearch": {"excludeDomains": ["example.com"]}}),
+                "googleSearch",
+                false,
+                json!({"excludeDomains": ["example.com"]}),
+                json!({"excludeDomains": ["example.com"]}),
+            ),
+            (
+                "current_snake",
+                json!({"google_search": {"exclude_domains": ["example.com"]}}),
+                "google_search",
+                false,
+                json!({"excludeDomains": ["example.com"]}),
+                json!({"excludeDomains": ["example.com"]}),
+            ),
+            (
+                "legacy_snake",
+                json!({
+                    "google_search_retrieval": {
+                        "dynamic_retrieval_config": {
+                            "mode": "MODE_DYNAMIC",
+                            "dynamic_threshold": 0.7
+                        }
+                    }
+                }),
+                "google_search_retrieval",
+                true,
+                json!({
+                    "dynamicRetrievalConfig": {
+                        "mode": "MODE_DYNAMIC",
+                        "dynamicThreshold": 0.7
+                    }
+                }),
+                json!({}),
+            ),
+            (
+                "legacy_camel",
+                json!({
+                    "googleSearchRetrieval": {
+                        "dynamicRetrievalConfig": {
+                            "mode": "MODE_DYNAMIC",
+                            "dynamicThreshold": 0.7
+                        }
+                    }
+                }),
+                "googleSearchRetrieval",
+                true,
+                json!({
+                    "dynamicRetrievalConfig": {
+                        "mode": "MODE_DYNAMIC",
+                        "dynamicThreshold": 0.7
+                    }
+                }),
+                json!({}),
+            ),
+        ];
+
+        for (
+            name,
+            tool,
+            source_field,
+            legacy,
+            expected_extension_payload,
+            expected_output_payload,
+        ) in cases
+        {
+            let request = json!({
+                "model": "gemini-2.5-pro",
+                "contents": [{"role": "user", "parts": [{"text": "search"}]}],
+                "tools": [tool]
+            });
+
+            let canonical = from_gemini_to_canonical_request(
+                &request,
+                "/v1beta/models/gemini-2.5-pro:generateContent",
+            )
+            .unwrap_or_else(|| panic!("{name}: canonical request"));
+
+            assert_eq!(
+                canonical
+                    .extensions
+                    .get("openai")
+                    .and_then(|value| value.get("web_search_options")),
+                Some(&json!({})),
+                "{name}: web search option"
+            );
+            let google_search = canonical
+                .extensions
+                .get("gemini")
+                .and_then(|value| value.get("grounding"))
+                .and_then(|value| value.get("google_search"))
+                .unwrap_or_else(|| panic!("{name}: gemini google_search grounding"));
+            assert_eq!(
+                google_search.get("source_field").and_then(Value::as_str),
+                Some(source_field),
+                "{name}: source field"
+            );
+            assert_eq!(
+                google_search.get("legacy").and_then(Value::as_bool),
+                Some(legacy),
+                "{name}: legacy flag"
+            );
+            assert_eq!(
+                google_search.get("payload"),
+                Some(&expected_extension_payload),
+                "{name}: normalized payload"
+            );
+
+            let rebuilt =
+                canonical_to_gemini_request(&canonical, "gemini-upstream", false).unwrap();
+            assert_eq!(
+                rebuilt["tools"],
+                json!([{"googleSearch": expected_output_payload}]),
+                "{name}: canonical output"
+            );
+        }
+    }
+
+    #[test]
+    fn gemini_request_adapter_keeps_agent_search_retrieval_separate_from_google_search() {
+        let request = json!({
+            "model": "gemini-2.5-pro",
+            "contents": [{"role": "user", "parts": [{"text": "private data"}]}],
+            "tools": [{
+                "retrieval": {
+                    "vertexAiSearch": {
+                        "datastore": "projects/p/locations/global/collections/default_collection/dataStores/d"
+                    }
+                }
+            }]
+        });
+
+        let canonical = from_gemini_to_canonical_request(
+            &request,
+            "/v1beta/models/gemini-2.5-pro:generateContent",
+        )
+        .expect("canonical request");
+
+        assert_eq!(
+            canonical
+                .extensions
+                .get("openai")
+                .and_then(|value| value.get("web_search_options")),
+            None
+        );
+
+        let rebuilt = canonical_to_gemini_request(&canonical, "gemini-upstream", false).unwrap();
+        assert_eq!(rebuilt["tools"], request["tools"]);
+        assert!(rebuilt["tools"]
+            .as_array()
+            .unwrap()
+            .iter()
+            .all(|tool| tool.get("googleSearch").is_none()));
+    }
+
+    #[test]
+    fn gemini_request_adapter_preserves_combined_search_builtin_tool_fields() {
+        let cases = [
+            (
+                "current_snake",
+                json!({
+                    "google_search": {},
+                    "code_execution": {},
+                    "url_context": {},
+                    "retrieval": {
+                        "vertexAiSearch": {
+                            "datastore": "projects/p/locations/global/collections/default_collection/dataStores/d"
+                        }
+                    }
+                }),
+            ),
+            (
+                "legacy_snake",
+                json!({
+                    "google_search_retrieval": {
+                        "dynamic_retrieval_config": {
+                            "mode": "MODE_DYNAMIC",
+                            "dynamic_threshold": 0.7
+                        }
+                    },
+                    "code_execution": {},
+                    "url_context": {}
+                }),
+            ),
+        ];
+
+        for (name, tool) in cases {
+            let request = json!({
+                "model": "gemini-2.5-pro",
+                "contents": [{"role": "user", "parts": [{"text": "search with builtins"}]}],
+                "tools": [tool]
+            });
+
+            let canonical = from_gemini_to_canonical_request(
+                &request,
+                "/v1beta/models/gemini-2.5-pro:generateContent",
+            )
+            .unwrap_or_else(|| panic!("{name}: canonical request"));
+
+            let rebuilt =
+                canonical_to_gemini_request(&canonical, "gemini-upstream", false).unwrap();
+            let tools = rebuilt["tools"]
+                .as_array()
+                .unwrap_or_else(|| panic!("{name}: tools array"));
+            assert!(
+                tools.iter().any(|tool| tool.get("googleSearch").is_some()),
+                "{name}: google search should be preserved"
+            );
+            assert!(
+                tools.iter().any(|tool| tool.get("codeExecution").is_some()),
+                "{name}: code execution should be preserved"
+            );
+            assert!(
+                tools.iter().any(|tool| tool.get("urlContext").is_some()),
+                "{name}: URL context should be preserved"
+            );
+            if name == "current_snake" {
+                assert!(
+                    tools.iter().any(|tool| tool.get("retrieval").is_some()),
+                    "{name}: unhandled retrieval should be preserved"
+                );
+            }
+        }
+    }
+
     #[test]
     fn gemini_response_adapter_preserves_thought_signature_tool_and_usage() {
         let response = json!({
@@ -6073,6 +6436,48 @@ mod tests {
         assert_eq!(rebuilt["usageMetadata"]["thoughtsTokenCount"], 2);
     }
 
+    #[test]
+    fn gemini_response_adapter_preserves_grounding_metadata() {
+        let grounding_metadata = json!({
+            "webSearchQueries": ["query"],
+            "searchEntryPoint": {"renderedContent": "<style></style>"},
+            "groundingChunks": [{
+                "web": {
+                    "uri": "https://example.com",
+                    "title": "Example"
+                }
+            }],
+            "groundingSupports": []
+        });
+        let response = json!({
+            "responseId": "resp_grounded",
+            "modelVersion": "gemini-2.5-pro",
+            "candidates": [{
+                "index": 0,
+                "finishReason": "STOP",
+                "groundingMetadata": grounding_metadata,
+                "content": {
+                    "parts": [{"text": "grounded answer"}]
+                }
+            }]
+        });
+
+        let canonical = from_gemini_to_canonical_response(&response).expect("canonical response");
+        assert_eq!(
+            canonical.outputs[0]
+                .extensions
+                .get("gemini")
+                .and_then(|value| value.get("groundingMetadata")),
+            Some(&grounding_metadata)
+        );
+
+        let rebuilt = canonical_to_gemini_response(&canonical, &json!({})).expect("gemini");
+        assert_eq!(
+            rebuilt["candidates"][0]["groundingMetadata"],
+            grounding_metadata
+        );
+    }
+
     #[test]
     fn canonical_response_preserves_openai_choices_and_gemini_candidates() {
         let openai_response = json!({