rig_core/tool/
schema.rs

1//! A tool's parameters derived from its arguments type, so the schema the
2//! model sees and the type its arguments are parsed into cannot disagree:
3//!
4//! ```
5//! use schemars::JsonSchema;
6//! use serde::Deserialize;
7//!
8//! #[derive(Deserialize, JsonSchema)]
9//! struct Args {
10//!     /// How many lines to read.
11//!     limit: Option<u32>,
12//! }
13//!
14//! let schema = rig_core::tool::args_schema::<Args>();
15//! let limit = serde_json::json!({"type": "integer", "minimum": 0,
16//!     "description": "How many lines to read."});
17//! assert_eq!(schema["properties"]["limit"], limit);
18//! assert_eq!(schema["additionalProperties"], false);
19//! ```
20
21use schemars::generate::SchemaSettings;
22use schemars::transform::RecursiveTransform;
23use schemars::{JsonSchema, Schema};
24use serde_json::Value;
25
26/// The JSON schema of the arguments type `A`, in the plain shape tool
27/// definitions use: inline, without `$schema`, `title`, `default`, `format`
28/// (providers support it unevenly) or the type's own doc comment, and with
29/// `Option` fields optional, not nullable (no `null` type, enum value or
30/// `anyOf` alternative), and each object refusing keys
31/// it does not declare (`additionalProperties: false`, as
32/// `deny_unknown_fields` says; a map keeps its own).
33pub fn args_schema<A: JsonSchema>() -> Value {
34    let mut settings = SchemaSettings::draft2020_12().with_transform(RecursiveTransform(plain));
35    settings.inline_subschemas = true;
36    settings.meta_schema = None;
37    let mut schema = settings.into_generator().into_root_schema_for::<A>();
38    if let Some(root) = schema.as_object_mut() {
39        root.shift_remove("description");
40    }
41    schema.to_value()
42}
43
44/// Drops from one subschema what [`args_schema`] leaves out, keeping the
45/// other keys' order, closes an object to undeclared keys, joins a doc
46/// comment's lines as rustdoc does, and makes an `Option` optional rather
47/// than nullable: no `null` type or enum value, and no `{"type": "null"}`
48/// alternative (the one left is merged in).
49fn plain(schema: &mut Schema) {
50    let Some(schema) = schema.as_object_mut() else {
51        return;
52    };
53    if let Some(Value::Array(alternatives)) = schema.get_mut("anyOf") {
54        let null = serde_json::json!({"type": "null"});
55        let before = alternatives.len();
56        alternatives.retain(|alternative| *alternative != null);
57        if let [Value::Object(only)] = alternatives.as_slice()
58            && before > 1
59        {
60            let only = only.clone();
61            schema.shift_remove("anyOf");
62            for (key, value) in only {
63                schema.entry(key).or_insert(value);
64            }
65        }
66    }
67    if let Some(Value::String(text)) = schema.get_mut("description") {
68        let paragraphs: Vec<String> = text.split("\n\n").map(|p| p.replace('\n', " ")).collect();
69        *text = paragraphs.join("\n\n");
70    }
71    for key in ["title", "default", "format"] {
72        schema.shift_remove(key);
73    }
74    if schema.contains_key("properties") && !schema.contains_key("additionalProperties") {
75        schema.insert("additionalProperties".to_owned(), Value::Bool(false));
76    }
77    if let Some(Value::Array(types)) = schema.get_mut("type")
78        && types.iter().any(|kind| kind == "null")
79    {
80        types.retain(|kind| kind != "null");
81        if let [kind] = types.as_slice() {
82            let kind = kind.clone();
83            schema.insert("type".to_owned(), kind);
84        }
85        if let Some(Value::Array(values)) = schema.get_mut("enum") {
86            values.retain(|value| !value.is_null());
87        }
88    }
89}
rig_core/tool/schema.rs

rig_core/tool/
schema.rs
