mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-17 00:17:46 +08:00
fix(formats): sanitize Gemini tool schemas
This commit is contained in:
ZheFox
@@ -1,4 +1,4 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
@@ -920,24 +920,194 @@ fn insert_f64(output: &mut Map<String, Value>, key: &str, value: Option<f64>) {
|
||||
}
|
||||
|
||||
fn clean_gemini_schema(value: &mut Value) {
|
||||
match value {
|
||||
Value::Object(object) => {
|
||||
for inner in object.values_mut() {
|
||||
clean_gemini_schema(inner);
|
||||
}
|
||||
if object.get("type").and_then(Value::as_str) == Some("object")
|
||||
&& !object.contains_key("properties")
|
||||
{
|
||||
object.insert("properties".to_string(), Value::Object(Map::new()));
|
||||
let root = value.clone();
|
||||
*value = json_schema_to_gemini_schema(&root, &root, &mut BTreeSet::new());
|
||||
}
|
||||
|
||||
fn json_schema_to_gemini_schema(
|
||||
value: &Value,
|
||||
root: &Value,
|
||||
resolving_refs: &mut BTreeSet<String>,
|
||||
) -> Value {
|
||||
let Some(object) = value.as_object() else {
|
||||
return json!({});
|
||||
};
|
||||
|
||||
if let Some(reference) = object.get("$ref").and_then(Value::as_str) {
|
||||
if let Some(pointer) = reference.strip_prefix('#') {
|
||||
if resolving_refs.insert(reference.to_string()) {
|
||||
if let Some(resolved) = root.pointer(pointer).and_then(Value::as_object) {
|
||||
let mut merged = resolved.clone();
|
||||
for (key, value) in object {
|
||||
if key != "$ref" {
|
||||
merged.insert(key.clone(), value.clone());
|
||||
}
|
||||
}
|
||||
let schema = clean_gemini_schema_object(&merged, root, resolving_refs);
|
||||
resolving_refs.remove(reference);
|
||||
return Value::Object(schema);
|
||||
}
|
||||
resolving_refs.remove(reference);
|
||||
}
|
||||
}
|
||||
Value::Array(items) => {
|
||||
for item in items {
|
||||
clean_gemini_schema(item);
|
||||
}
|
||||
|
||||
Value::Object(clean_gemini_schema_object(object, root, resolving_refs))
|
||||
}
|
||||
|
||||
fn clean_gemini_schema_object(
|
||||
object: &Map<String, Value>,
|
||||
root: &Value,
|
||||
resolving_refs: &mut BTreeSet<String>,
|
||||
) -> Map<String, Value> {
|
||||
let mut schema = Map::new();
|
||||
|
||||
for key in ["title", "description", "format", "pattern"] {
|
||||
if let Some(value) = object.get(key).filter(|value| value.is_string()) {
|
||||
schema.insert(key.to_string(), value.clone());
|
||||
}
|
||||
}
|
||||
for key in ["default", "example"] {
|
||||
if let Some(value) = object.get(key) {
|
||||
schema.insert(key.to_string(), value.clone());
|
||||
}
|
||||
}
|
||||
for key in ["minimum", "maximum"] {
|
||||
if let Some(value) = object.get(key).filter(|value| value.is_number()) {
|
||||
schema.insert(key.to_string(), value.clone());
|
||||
}
|
||||
}
|
||||
for key in [
|
||||
"minItems",
|
||||
"maxItems",
|
||||
"minLength",
|
||||
"maxLength",
|
||||
"minProperties",
|
||||
"maxProperties",
|
||||
] {
|
||||
if let Some(value) = object.get(key).and_then(gemini_int64_string) {
|
||||
schema.insert(key.to_string(), Value::String(value));
|
||||
}
|
||||
}
|
||||
if let Some(value) = object.get("nullable").filter(|value| value.is_boolean()) {
|
||||
schema.insert("nullable".to_string(), value.clone());
|
||||
}
|
||||
for key in ["required", "propertyOrdering"] {
|
||||
if let Some(values) = object.get(key).and_then(gemini_string_array) {
|
||||
schema.insert(key.to_string(), values);
|
||||
}
|
||||
}
|
||||
if let Some(values) = object.get("enum").and_then(Value::as_array) {
|
||||
let values = values
|
||||
.iter()
|
||||
.filter(|value| value.is_string())
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
if !values.is_empty() {
|
||||
schema.insert("enum".to_string(), Value::Array(values));
|
||||
}
|
||||
}
|
||||
if let Some(properties) = object.get("properties").and_then(Value::as_object) {
|
||||
schema.insert(
|
||||
"properties".to_string(),
|
||||
Value::Object(
|
||||
properties
|
||||
.iter()
|
||||
.map(|(name, value)| {
|
||||
(
|
||||
name.clone(),
|
||||
json_schema_to_gemini_schema(value, root, resolving_refs),
|
||||
)
|
||||
})
|
||||
.collect(),
|
||||
),
|
||||
);
|
||||
}
|
||||
if let Some(items) = object.get("items") {
|
||||
schema.insert(
|
||||
"items".to_string(),
|
||||
json_schema_to_gemini_schema(items, root, resolving_refs),
|
||||
);
|
||||
}
|
||||
|
||||
let explicit_any_of = object
|
||||
.get("anyOf")
|
||||
.or_else(|| object.get("oneOf"))
|
||||
.and_then(Value::as_array)
|
||||
.map(|items| {
|
||||
Value::Array(
|
||||
items
|
||||
.iter()
|
||||
.map(|item| json_schema_to_gemini_schema(item, root, resolving_refs))
|
||||
.collect(),
|
||||
)
|
||||
});
|
||||
if let Some(any_of) = explicit_any_of {
|
||||
schema.insert("anyOf".to_string(), any_of);
|
||||
}
|
||||
|
||||
match object.get("type") {
|
||||
Some(Value::String(schema_type)) => {
|
||||
schema.insert("type".to_string(), Value::String(schema_type.clone()));
|
||||
}
|
||||
Some(Value::Array(types)) => {
|
||||
let mut non_null_types = types
|
||||
.iter()
|
||||
.filter_map(Value::as_str)
|
||||
.filter(|schema_type| *schema_type != "null")
|
||||
.map(ToOwned::to_owned)
|
||||
.collect::<Vec<_>>();
|
||||
let mut seen_types = BTreeSet::new();
|
||||
non_null_types.retain(|schema_type| seen_types.insert(schema_type.clone()));
|
||||
let nullable = types.iter().any(|value| value.as_str() == Some("null"));
|
||||
|
||||
match non_null_types.as_slice() {
|
||||
[schema_type] => {
|
||||
schema.insert("type".to_string(), Value::String(schema_type.clone()));
|
||||
}
|
||||
[] if nullable => {
|
||||
schema.insert("type".to_string(), Value::String("null".to_string()));
|
||||
}
|
||||
[] => {}
|
||||
_ if !schema.contains_key("anyOf") => {
|
||||
schema.insert(
|
||||
"anyOf".to_string(),
|
||||
Value::Array(
|
||||
non_null_types
|
||||
.iter()
|
||||
.map(|schema_type| json!({ "type": schema_type }))
|
||||
.collect(),
|
||||
),
|
||||
);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
if nullable && !non_null_types.is_empty() {
|
||||
schema.insert("nullable".to_string(), Value::Bool(true));
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
|
||||
if schema.get("type").and_then(Value::as_str) == Some("object")
|
||||
&& !schema.contains_key("properties")
|
||||
{
|
||||
schema.insert("properties".to_string(), Value::Object(Map::new()));
|
||||
}
|
||||
schema
|
||||
}
|
||||
|
||||
fn gemini_int64_string(value: &Value) -> Option<String> {
|
||||
match value {
|
||||
Value::String(value) => Some(value.clone()),
|
||||
Value::Number(value) => Some(value.to_string()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn gemini_string_array(value: &Value) -> Option<Value> {
|
||||
let values = value.as_array()?;
|
||||
values.iter().all(Value::is_string).then(|| value.clone())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -945,6 +1115,79 @@ mod tests {
|
||||
use super::*;
|
||||
use crate::CanonicalContentBlock;
|
||||
|
||||
#[test]
|
||||
fn canonical_tool_declaration_sanitizes_json_schema_for_gemini() {
|
||||
let declaration = canonical_tool_to_gemini_declaration(&CanonicalToolDefinition {
|
||||
name: "inspect".to_string(),
|
||||
description: None,
|
||||
parameters: Some(json!({
|
||||
"$defs": {
|
||||
"Target": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"secret": {
|
||||
"type": "string",
|
||||
"encrypted": true
|
||||
}
|
||||
},
|
||||
"required": ["secret"],
|
||||
"additionalProperties": false
|
||||
}
|
||||
},
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"target": {
|
||||
"oneOf": [
|
||||
{"$ref": "#/$defs/Target"},
|
||||
{"type": "null"}
|
||||
]
|
||||
},
|
||||
"mode": {
|
||||
"type": ["string", "null"],
|
||||
"enum": [1, "fast"]
|
||||
},
|
||||
"value": {
|
||||
"type": ["string", "integer"]
|
||||
}
|
||||
}
|
||||
})),
|
||||
strict: None,
|
||||
extensions: BTreeMap::new(),
|
||||
});
|
||||
|
||||
assert_eq!(
|
||||
declaration["parameters"],
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"target": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"secret": {"type": "string"}
|
||||
},
|
||||
"required": ["secret"]
|
||||
},
|
||||
{"type": "null"}
|
||||
]
|
||||
},
|
||||
"mode": {
|
||||
"type": "string",
|
||||
"nullable": true,
|
||||
"enum": ["fast"]
|
||||
},
|
||||
"value": {
|
||||
"anyOf": [
|
||||
{"type": "string"},
|
||||
{"type": "integer"}
|
||||
]
|
||||
}
|
||||
}
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn canonical_tool_result_to_gemini_request_preserves_function_response_id() {
|
||||
let mut tool_name_by_id = BTreeMap::new();
|
||||
|
||||
Reference in New Issue
Block a user