use crate::tools::Tool;
use serde_json::Value as JsonValue;
use std::collections::{HashMap, HashSet};
#[derive(Debug, Clone, Default)]
pub struct ToolArgShape {
pub string_args: HashSet<String>,
pub properties: HashMap<String, JsonValue>,
pub parameter_order: Vec<String>,
}
pub fn build_arg_shapes(tools: &[Tool]) -> HashMap<String, ToolArgShape> {
let mut shapes = HashMap::new();
for tool in tools {
let props = resolve_properties(&tool.parameters);
let mut string_args = HashSet::new();
let mut parameter_order = Vec::new();
if let Some(obj) = props {
for (key, schema) in obj {
parameter_order.push(key.clone());
if is_string_only_schema(&schema) {
string_args.insert(key);
}
}
}
shapes.insert(
tool.name.clone(),
ToolArgShape {
string_args,
properties: HashMap::new(),
parameter_order,
},
);
}
shapes
}
fn resolve_properties(parameters: &JsonValue) -> Option<serde_json::Map<String, JsonValue>> {
let obj = parameters.as_object()?;
obj.get("properties")?.as_object().cloned()
}
pub fn is_string_only_schema(schema: &JsonValue) -> bool {
let mut types = collect_schema_types(schema, 0);
types.remove("null");
types.len() == 1 && types.contains("string")
}
fn collect_schema_types(schema: &JsonValue, depth: usize) -> HashSet<String> {
let mut out = HashSet::new();
if depth > 8 {
return out;
}
let Some(node) = schema.as_object() else {
return out;
};
match node.get("type") {
Some(JsonValue::String(t)) => {
out.insert(t.clone());
}
Some(JsonValue::Array(arr)) => {
for t in arr {
if let JsonValue::String(s) = t {
out.insert(s.clone());
}
}
}
_ => {}
}
if !node.contains_key("type") {
if let Some(JsonValue::Array(en)) = node.get("enum") {
for v in en {
out.insert(json_type_of(v).to_string());
}
}
if let Some(c) = node.get("const") {
out.insert(json_type_of(c).to_string());
}
}
for key in ["anyOf", "oneOf", "allOf"] {
if let Some(JsonValue::Array(branches)) = node.get(key) {
for sub in branches {
out.extend(collect_schema_types(sub, depth + 1));
}
}
}
out
}
fn json_type_of(value: &JsonValue) -> &'static str {
match value {
JsonValue::Null => "null",
JsonValue::Bool(_) => "boolean",
JsonValue::Number(_) => "number",
JsonValue::String(_) => "string",
JsonValue::Array(_) | JsonValue::Object(_) => "object",
}
}
pub fn decode_value(raw: &str) -> JsonValue {
let trimmed = raw.trim();
if trimmed.is_empty() {
return JsonValue::String(trimmed.to_string());
}
match serde_json::from_str::<JsonValue>(trimmed) {
Ok(v) => v,
Err(_) => JsonValue::String(raw.to_string()),
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn tool(name: &str, params: JsonValue) -> Tool {
Tool {
name: name.to_string(),
description: String::new(),
parameters: params,
}
}
#[test]
fn detects_string_only_args() {
let t = tool(
"write",
json!({
"type": "object",
"properties": {
"path": {"type": "string"},
"content": {"type": "string"},
"lines": {"type": "integer"},
"force": {"type": ["boolean", "null"]},
}
}),
);
let shapes = build_arg_shapes(&[t]);
let shape = &shapes["write"];
assert!(shape.string_args.contains("path"));
assert!(shape.string_args.contains("content"));
assert!(!shape.string_args.contains("lines"));
assert!(!shape.string_args.contains("force"));
}
#[test]
fn nullable_string_is_string_only() {
assert!(is_string_only_schema(&json!({"type": ["string", "null"]})));
assert!(!is_string_only_schema(
&json!({"type": ["string", "number"]})
));
}
#[test]
fn enum_infers_types() {
assert!(is_string_only_schema(&json!({"enum": ["a", "b"]})));
assert!(!is_string_only_schema(&json!({"enum": [1, 2]})));
}
#[test]
fn anyof_union_collects_types() {
let schema = json!({"anyOf": [{"type": "string"}, {"type": "null"}]});
assert!(is_string_only_schema(&schema));
}
#[test]
fn decode_value_prefers_json() {
assert_eq!(decode_value("42"), json!(42));
assert_eq!(decode_value("true"), json!(true));
assert_eq!(decode_value(r#""hi""#), json!("hi"));
assert_eq!(decode_value("hello world"), json!("hello world"));
assert_eq!(decode_value(" "), json!(""));
}
}