agent2: Fix tool schemas for Gemini (#36507)

Bennet Bo Fenner and Agus Zubiaga created 3 months ago

Release Notes:

- N/A

---------

Co-authored-by: Agus Zubiaga <[e-mail]>

Change summary

crates/agent2/src/agent2.rs       |  1 
crates/agent2/src/thread.rs       |  6 ++--
crates/agent2/src/tool_schema.rs  | 43 +++++++++++++++++++++++++++++++++
crates/google_ai/src/google_ai.rs |  2 
4 files changed, 48 insertions(+), 4 deletions(-)

Detailed changes

crates/agent2/src/agent2.rs

@@ -2,6 +2,7 @@ mod agent;
 mod native_agent_server;
 mod templates;
 mod thread;
+mod tool_schema;
 mod tools;
 
 #[cfg(test)]

crates/agent2/src/thread.rs

@@ -1732,8 +1732,8 @@ where
     fn initial_title(&self, input: Result<Self::Input, serde_json::Value>) -> SharedString;
 
     /// Returns the JSON schema that describes the tool's input.
-    fn input_schema(&self) -> Schema {
-        schemars::schema_for!(Self::Input)
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Schema {
+        crate::tool_schema::root_schema_for::<Self::Input>(format)
     }
 
     /// Some tools rely on a provider for the underlying billing or other reasons.
@@ -1819,7 +1819,7 @@ where
     }
 
     fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
-        let mut json = serde_json::to_value(self.0.input_schema())?;
+        let mut json = serde_json::to_value(self.0.input_schema(format))?;
         adapt_schema_to_format(&mut json, format)?;
         Ok(json)
     }

crates/agent2/src/tool_schema.rs

@@ -0,0 +1,43 @@
+use language_model::LanguageModelToolSchemaFormat;
+use schemars::{
+    JsonSchema, Schema,
+    generate::SchemaSettings,
+    transform::{Transform, transform_subschemas},
+};
+
+pub(crate) fn root_schema_for<T: JsonSchema>(format: LanguageModelToolSchemaFormat) -> Schema {
+    let mut generator = match format {
+        LanguageModelToolSchemaFormat::JsonSchema => SchemaSettings::draft07().into_generator(),
+        LanguageModelToolSchemaFormat::JsonSchemaSubset => SchemaSettings::openapi3()
+            .with(|settings| {
+                settings.meta_schema = None;
+                settings.inline_subschemas = true;
+            })
+            .with_transform(ToJsonSchemaSubsetTransform)
+            .into_generator(),
+    };
+    generator.root_schema_for::<T>()
+}
+
+#[derive(Debug, Clone)]
+struct ToJsonSchemaSubsetTransform;
+
+impl Transform for ToJsonSchemaSubsetTransform {
+    fn transform(&mut self, schema: &mut Schema) {
+        // Ensure that the type field is not an array, this happens when we use
+        // Option<T>, the type will be [T, "null"].
+        if let Some(type_field) = schema.get_mut("type")
+            && let Some(types) = type_field.as_array()
+            && let Some(first_type) = types.first()
+        {
+            *type_field = first_type.clone();
+        }
+
+        // oneOf is not supported, use anyOf instead
+        if let Some(one_of) = schema.remove("oneOf") {
+            schema.insert("anyOf".to_string(), one_of);
+        }
+
+        transform_subschemas(self, schema);
+    }
+}

crates/google_ai/src/google_ai.rs

@@ -266,7 +266,7 @@ pub struct CitationMetadata {
 pub struct PromptFeedback {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub block_reason: Option<String>,
-    pub safety_ratings: Vec<SafetyRating>,
+    pub safety_ratings: Option<Vec<SafetyRating>>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub block_reason_message: Option<String>,
 }