tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

fix(google): use parametersJsonSchema for full JSON Schema support (#6555)

Browse Source 
Signed-off-by: rabi <ramishra@redhat.com>
This commit is contained in:
Rabi Mishra
2026-01-28 22:37:26 +05:30
committed by GitHub
parent c2f5889ae5
commit b4cf0fdf51
2 changed files with 81 additions and 329 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/providers/formats/databricks.rs
+34 -13
View File
@@ -1,6 +1,5 @@
use crate::conversation::message::{Message, MessageContent};
use crate::model::ModelConfig;
use crate::providers::formats::google as gemini_schema;
use crate::providers::utils::{
convert_image, detect_image_path, is_valid_function_name, load_image_file, safely_parse_json,
sanitize_function_name, ImageFormat,
@@ -232,19 +231,35 @@ pub fn format_tools(tools: &[Tool], model_name: &str) -> anyhow::Result<Vec<Valu
return Err(anyhow!("Duplicate tool name: {}", tool.name));
}
let parameters = if is_gemini {
gemini_schema::process_map(tool.input_schema.as_ref(), None)
let has_properties = tool
.input_schema
.get("properties")
.and_then(|v| v.as_object())
.is_some_and(|p| !p.is_empty());
let function_def = if is_gemini {
let mut def = json!({
"name": tool.name,
"description": tool.description,
});
if has_properties {
def["parametersJsonSchema"] = json!(tool.input_schema);
}
def
} else {
json!(tool.input_schema)
let mut def = json!({
"name": tool.name,
"description": tool.description,
});
if has_properties {
def["parameters"] = json!(tool.input_schema);
}
def
};
result.push(json!({
"type": "function",
"function": {
"name": tool.name,
"description": tool.description,
"parameters": parameters,
}
"function": function_def,
}));
}
@@ -719,12 +734,18 @@ mod tests {
);
let spec = format_tools(std::slice::from_ref(&tool), "gemini-2-5-flash")?;
assert!(spec[0]["function"]["parameters"].get("$schema").is_none());
assert_eq!(spec[0]["function"]["parameters"]["type"], "object");
assert!(spec[0]["function"].get("parametersJsonSchema").is_some());
assert_eq!(
spec[0]["function"]["parametersJsonSchema"]["type"],
"object"
);
let spec = format_tools(&[tool], "databricks-gemini-3-pro")?;
assert!(spec[0]["function"]["parameters"].get("$schema").is_none());
assert_eq!(spec[0]["function"]["parameters"]["type"], "object");
assert!(spec[0]["function"].get("parametersJsonSchema").is_some());
assert_eq!(
spec[0]["function"]["parametersJsonSchema"]["type"],
"object"
);
Ok(())
}
crates/goose/src/providers/formats/google.rs
+47 -316
View File
@@ -227,117 +227,21 @@ pub fn format_tools(tools: &[Tool]) -> Vec<Value> {
let mut parameters = Map::new();
parameters.insert("name".to_string(), json!(tool.name));
parameters.insert("description".to_string(), json!(tool.description));
let tool_input_schema = &tool.input_schema;
if tool_input_schema
// Use parametersJsonSchema which supports full JSON Schema including $ref/$defs
if tool
.input_schema
.get("properties")
.and_then(|v| v.as_object())
.is_some_and(|p| !p.is_empty())
{
parameters.insert(
"parameters".to_string(),
process_map(tool_input_schema, None),
);
parameters.insert("parametersJsonSchema".to_string(), json!(tool.input_schema));
}
json!(parameters)
})
.collect()
}
pub fn get_accepted_keys(parent_key: Option<&str>) -> Vec<&str> {
match parent_key {
Some("properties") => vec![
"anyOf",
"allOf",
"type",
"description",
"nullable",
"enum",
"properties",
"required",
"items",
],
Some("items") => vec!["type", "properties", "items", "required"],
_ => vec!["type", "properties", "required", "anyOf", "allOf"],
}
}
pub fn process_value(value: &Value, parent_key: Option<&str>) -> Value {
match value {
Value::Object(map) => process_map(map, parent_key),
Value::Array(arr) if parent_key == Some("type") => arr
.iter()
.find(|v| v.as_str() != Some("null"))
.cloned()
.unwrap_or_else(|| json!("string")),
_ => value.clone(),
}
}
/// Process a JSON map to filter out unsupported attributes, mirroring the logic
/// from the official Google Gemini CLI.
/// See: https://github.com/google-gemini/gemini-cli/blob/8a6509ffeba271a8e7ccb83066a9a31a5d72a647/packages/core/src/tools/tool-registry.ts#L356
pub fn process_map(map: &Map<String, Value>, parent_key: Option<&str>) -> Value {
let accepted_keys = get_accepted_keys(parent_key);
let filtered_map: Map<String, Value> = map
.iter()
.filter_map(|(key, value)| {
if !accepted_keys.contains(&key.as_str()) {
return None;
}
let processed_value = match key.as_str() {
"properties" => {
if let Some(nested_map) = value.as_object() {
let processed_properties: Map<String, Value> = nested_map
.iter()
.map(|(prop_key, prop_value)| {
if let Some(prop_obj) = prop_value.as_object() {
(prop_key.clone(), process_map(prop_obj, Some("properties")))
} else {
(prop_key.clone(), prop_value.clone())
}
})
.collect();
Value::Object(processed_properties)
} else {
value.clone()
}
}
"items" => {
if let Some(items_map) = value.as_object() {
process_map(items_map, Some("items"))
} else {
value.clone()
}
}
"anyOf" | "allOf" => {
if let Some(arr) = value.as_array() {
let processed_arr: Vec<Value> = arr
.iter()
.map(|item| {
item.as_object().map_or_else(
|| item.clone(),
|obj| process_map(obj, parent_key),
)
})
.collect();
Value::Array(processed_arr)
} else {
value.clone()
}
}
_ => process_value(value, Some(key.as_str())),
};
Some((key.clone(), processed_value))
})
.collect();
Value::Object(filtered_map)
}
#[derive(Clone, Copy)]
enum SignedTextHandling {
SignedTextAsThinking,
@@ -843,211 +747,23 @@ mod tests {
#[test]
fn test_tools_to_google_spec_with_valid_tools() {
let params1 = object!({
let params = object!({
"properties": {
"param1": {
"type": "string",
"description": "A parameter",
"field_does_not_accept": ["value1", "value2"]
"description": "A parameter"
}
}
});
let params2 = object!({
"properties": {
"param2": {
"type": "string",
"description": "B parameter",
}
}
});
let params3 = object!({
"properties": {
"body": {
"description": "Review comment text",
"type": "string"
},
"comments": {
"description": "Line-specific comments array of objects to place comments on pull request changes. Requires path and body. For line comments use line or position. For multi-line comments use start_line and line with optional side parameters.",
"type": "array",
"items": {
"additionalProperties": false,
"properties": {
"body": {
"description": "comment body",
"type": "string"
},
"line": {
"anyOf": [
{ "type": "number" },
{ "type": "null" }
],
"description": "line number in the file to comment on. For multi-line comments, the end of the line range"
},
"path": {
"description": "path to the file",
"type": "string"
},
"position": {
"anyOf": [
{ "type": "number" },
{ "type": "null" }
],
"description": "position of the comment in the diff"
},
"side": {
"anyOf": [
{ "type": "string" },
{ "type": "null" }
],
"description": "The side of the diff on which the line resides. For multi-line comments, this is the side for the end of the line range. (LEFT or RIGHT)"
},
"start_line": {
"anyOf": [
{ "type": "number" },
{ "type": "null" }
],
"description": "The first line of the range to which the comment refers. Required for multi-line comments."
},
"start_side": {
"anyOf": [
{ "type": "string" },
{ "type": "null" }
],
"description": "The side of the diff on which the start line resides for multi-line comments. (LEFT or RIGHT)"
}
},
"required": ["path", "body", "position", "line", "side", "start_line", "start_side"],
"type": "object"
}
},
"commitId": {
"description": "SHA of commit to review",
"type": "string"
},
"event": {
"description": "Review action to perform",
"enum": ["APPROVE", "REQUEST_CHANGES", "COMMENT"],
"type": "string"
},
"owner": {
"description": "Repository owner",
"type": "string"
},
"pullNumber": {
"description": "Pull request number",
"type": "number"
}
}
});
let tools = vec![
Tool::new("tool1", "description1", params1),
Tool::new("tool2", "description2", params2),
Tool::new("tool3", "description3", params3),
];
let tools = vec![Tool::new("tool1", "description1", params.clone())];
let result = format_tools(&tools);
assert_eq!(result.len(), 3);
assert_eq!(result.len(), 1);
assert_eq!(result[0]["name"], "tool1");
assert_eq!(result[0]["description"], "description1");
assert_eq!(
result[0]["parameters"]["properties"],
json!({"param1": json!({
"type": "string",
"description": "A parameter"
})})
);
assert_eq!(result[1]["name"], "tool2");
assert_eq!(result[1]["description"], "description2");
assert_eq!(
result[1]["parameters"]["properties"],
json!({"param2": json!({
"type": "string",
"description": "B parameter"
})})
);
assert_eq!(result[2]["name"], "tool3");
assert_eq!(
result[2]["parameters"]["properties"],
json!(
{
"body": {
"description": "Review comment text",
"type": "string"
},
"comments": {
"description": "Line-specific comments array of objects to place comments on pull request changes. Requires path and body. For line comments use line or position. For multi-line comments use start_line and line with optional side parameters.",
"type": "array",
"items": {
"properties": {
"body": {
"description": "comment body",
"type": "string"
},
"line": {
"anyOf": [
{ "type": "number" },
{ "type": "null" }
],
"description": "line number in the file to comment on. For multi-line comments, the end of the line range"
},
"path": {
"description": "path to the file",
"type": "string"
},
"position": {
"anyOf": [
{ "type": "number" },
{ "type": "null" }
],
"description": "position of the comment in the diff"
},
"side": {
"anyOf": [
{ "type": "string" },
{ "type": "null" }
],
"description": "The side of the diff on which the line resides. For multi-line comments, this is the side for the end of the line range. (LEFT or RIGHT)"
},
"start_line": {
"anyOf": [
{ "type": "number" },
{ "type": "null" }
],
"description": "The first line of the range to which the comment refers. Required for multi-line comments."
},
"start_side": {
"anyOf": [
{ "type": "string" },
{ "type": "null" }
],
"description": "The side of the diff on which the start line resides for multi-line comments. (LEFT or RIGHT)"
}
},
"required": ["path", "body", "position", "line", "side", "start_line", "start_side"],
"type": "object"
}
},
"commitId": {
"description": "SHA of commit to review",
"type": "string"
},
"event": {
"description": "Review action to perform",
"enum": ["APPROVE", "REQUEST_CHANGES", "COMMENT"],
"type": "string"
},
"owner": {
"description": "Repository owner",
"type": "string"
},
"pullNumber": {
"description": "Pull request number",
"type": "number"
}
}
)
);
assert!(result[0].get("parametersJsonSchema").is_some());
assert!(result[0].get("parameters").is_none());
assert_eq!(result[0]["parametersJsonSchema"], json!(params));
}
#[test]
@@ -1063,7 +779,7 @@ mod tests {
assert_eq!(result.len(), 1);
assert_eq!(result[0]["name"], "tool1");
assert_eq!(result[0]["description"], "description1");
assert!(result[0]["parameters"].get("properties").is_none());
assert!(result[0].get("parametersJsonSchema").is_none());
}
#[test]
@@ -1187,35 +903,22 @@ mod tests {
}
#[test]
fn test_tools_with_nullable_types_converted_to_single_type() {
// Test that type arrays like ["string", "null"] are converted to single types
fn test_tools_uses_parameters_json_schema() {
let params = object!({
"properties": {
"nullable_field": {
"field": {
"type": ["string", "null"],
"description": "A nullable string field"
},
"regular_field": {
"type": "number",
"description": "A regular number field"
"description": "A field"
}
}
});
let tools = vec![Tool::new("test_tool", "test description", params)];
let tools = vec![Tool::new("test_tool", "test description", params.clone())];
let result = format_tools(&tools);
assert_eq!(result.len(), 1);
assert_eq!(result[0]["name"], "test_tool");
// Verify that the type array was converted to a single string type
let nullable_field = &result[0]["parameters"]["properties"]["nullable_field"];
assert_eq!(nullable_field["type"], "string");
assert_eq!(nullable_field["description"], "A nullable string field");
// Verify that regular types are unchanged
let regular_field = &result[0]["parameters"]["properties"]["regular_field"];
assert_eq!(regular_field["type"], "number");
assert_eq!(regular_field["description"], "A regular number field");
assert!(result[0].get("parametersJsonSchema").is_some());
assert_eq!(result[0]["parametersJsonSchema"], json!(params));
}
fn google_response(parts: Vec<Value>) -> Value {
@@ -1512,4 +1215,32 @@ data: [DONE]"#;
// Only "Complete" should be captured, stream should stop at [DONE]
assert_eq!(text_parts, vec!["Complete"]);
}
#[test]
fn test_format_tools_uses_parameters_json_schema() {
let tool = Tool::new(
"test_tool",
"Test tool with $ref",
object!({
"type": "object",
"$defs": {
"MyType": { "type": "string", "description": "A custom type" }
},
"properties": {
"field": { "$ref": "#/$defs/MyType" }
}
}),
);
let result = format_tools(&[tool]);
assert_eq!(result.len(), 1);
assert_eq!(result[0]["name"], "test_tool");
assert!(result[0].get("parametersJsonSchema").is_some());
assert!(result[0].get("parameters").is_none());
let schema = &result[0]["parametersJsonSchema"];
assert_eq!(schema["properties"]["field"]["$ref"], "#/$defs/MyType");
assert!(schema.get("$defs").is_some());
}
}