fix(ai): preserve combined Gemini builtin tools

2026-10-04 16:37:46 +08:00 · 2026-05-26 11:14:24 +08:00
parent b75953bf4c
commit 12ff87949d
2 changed files with 100 additions and 4 deletions
@@ -554,6 +554,7 @@ fn canonical_tools_to_gemini(canonical: &CanonicalRequest) -> Option<Value> {
                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() {
@@ -565,7 +566,7 @@ fn canonical_tools_to_gemini(canonical: &CanonicalRequest) -> Option<Value> {
                    }));
                    emitted_google_search = true;
                }
-                continue;
+                emitted_builtin_portion = true;
            }
            if let Some(tool) =
                gemini_builtin_tool_by_case(tool_object, "codeExecution", "code_execution")
@@ -574,7 +575,7 @@ fn canonical_tools_to_gemini(canonical: &CanonicalRequest) -> Option<Value> {
                    tools.push(tool);
                    emitted_code_execution = true;
                }
-                continue;
+                emitted_builtin_portion = true;
            }
            if let Some(tool) =
                gemini_builtin_tool_by_case(tool_object, "urlContext", "url_context")
@@ -583,9 +584,13 @@ fn canonical_tools_to_gemini(canonical: &CanonicalRequest) -> Option<Value> {
                    tools.push(tool);
                    emitted_url_context = true;
                }
-                continue;
+                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());
            }
-            tools.push(builtin_tool.clone());
        }
    }
    if code_execution && !emitted_code_execution {
@@ -641,6 +646,27 @@ fn gemini_builtin_tool_payload(payload: &Value) -> Value {
    }
 }

+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()));
@@ -6293,6 +6293,76 @@ mod tests {
            .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!({