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

rig_core/tool/
schema.rs
