mirror of
https://github.com/openai/codex.git
synced 2026-10-21 17:48:44 +00:00
Split tool JSON Schema code into focused modules (#42161)
## What changed - Move schema types, traversal helpers, and large-schema compaction into dedicated `json_schema` submodules. - Keep the public schema types re-exported from the existing module and colocate focused unit tests with their implementations. GitOrigin-RevId: 953b96c28ae3880f601b588eb016fa126c438852
This commit is contained in:
Celia Chen
@@ -1,190 +1,26 @@
|
||||
use serde::Deserialize;
|
||||
use serde::Serialize;
|
||||
//! Normalize tool schemas and prune unreachable definitions before applying the size policy.
|
||||
|
||||
mod compaction;
|
||||
mod traversal;
|
||||
mod types;
|
||||
|
||||
pub use types::AdditionalProperties;
|
||||
pub use types::JsonSchema;
|
||||
pub use types::JsonSchemaPrimitiveType;
|
||||
pub use types::JsonSchemaType;
|
||||
|
||||
use compaction::compact_large_tool_schema;
|
||||
use traversal::DefinitionTraversal;
|
||||
use traversal::for_each_schema_child;
|
||||
|
||||
use serde_json::Value as JsonValue;
|
||||
use serde_json::json;
|
||||
use std::collections::BTreeMap;
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
const DEFINITION_TABLE_KEYS: [&str; 2] = ["$defs", "definitions"];
|
||||
const SCHEMA_CHILD_KEYS: [&str; 4] = ["items", "anyOf", "oneOf", "allOf"];
|
||||
const COMPOSITION_SCHEMA_KEYS: [&str; 3] = ["anyOf", "oneOf", "allOf"];
|
||||
|
||||
/// Primitive JSON Schema type names we support in tool definitions.
|
||||
///
|
||||
/// This mirrors the OpenAI Structured Outputs subset for JSON Schema `type`:
|
||||
/// string, number, boolean, integer, object, array, and null.
|
||||
/// Keywords such as `enum`, `const`, `anyOf`, `oneOf`, and `allOf` are modeled
|
||||
/// separately.
|
||||
/// See <https://developers.openai.com/api/docs/guides/structured-outputs#supported-schemas>.
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum JsonSchemaPrimitiveType {
|
||||
String,
|
||||
Number,
|
||||
Boolean,
|
||||
Integer,
|
||||
Object,
|
||||
Array,
|
||||
Null,
|
||||
}
|
||||
|
||||
/// JSON Schema `type` supports either a single type name or a union of names.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(untagged)]
|
||||
pub enum JsonSchemaType {
|
||||
Single(JsonSchemaPrimitiveType),
|
||||
Multiple(Vec<JsonSchemaPrimitiveType>),
|
||||
}
|
||||
|
||||
/// Generic JSON-Schema subset needed for our tool definitions.
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
|
||||
pub struct JsonSchema {
|
||||
#[serde(rename = "$ref", skip_serializing_if = "Option::is_none")]
|
||||
pub schema_ref: Option<String>,
|
||||
#[serde(rename = "type", skip_serializing_if = "Option::is_none")]
|
||||
pub schema_type: Option<JsonSchemaType>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub description: Option<String>,
|
||||
/// Responses-only marker for reviewed encrypted tool parameters.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub encrypted: Option<bool>,
|
||||
#[serde(rename = "enum", skip_serializing_if = "Option::is_none")]
|
||||
pub enum_values: Option<Vec<JsonValue>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub items: Option<Box<JsonSchema>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub properties: Option<BTreeMap<String, JsonSchema>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub required: Option<Vec<String>>,
|
||||
#[serde(
|
||||
rename = "additionalProperties",
|
||||
skip_serializing_if = "Option::is_none"
|
||||
)]
|
||||
pub additional_properties: Option<AdditionalProperties>,
|
||||
#[serde(rename = "anyOf", skip_serializing_if = "Option::is_none")]
|
||||
pub any_of: Option<Vec<JsonSchema>>,
|
||||
#[serde(rename = "oneOf", skip_serializing_if = "Option::is_none")]
|
||||
pub one_of: Option<Vec<JsonSchema>>,
|
||||
#[serde(rename = "allOf", skip_serializing_if = "Option::is_none")]
|
||||
pub all_of: Option<Vec<JsonSchema>>,
|
||||
#[serde(rename = "$defs", skip_serializing_if = "Option::is_none")]
|
||||
pub defs: Option<BTreeMap<String, JsonSchema>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub definitions: Option<BTreeMap<String, JsonSchema>>,
|
||||
}
|
||||
|
||||
impl JsonSchema {
|
||||
/// Construct a scalar/object/array schema with a single JSON Schema type.
|
||||
fn typed(schema_type: JsonSchemaPrimitiveType, description: Option<String>) -> Self {
|
||||
Self {
|
||||
schema_type: Some(JsonSchemaType::Single(schema_type)),
|
||||
description,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn any_of(variants: Vec<JsonSchema>, description: Option<String>) -> Self {
|
||||
Self {
|
||||
description,
|
||||
any_of: Some(variants),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn one_of(variants: Vec<JsonSchema>, description: Option<String>) -> Self {
|
||||
Self {
|
||||
description,
|
||||
one_of: Some(variants),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn all_of(variants: Vec<JsonSchema>, description: Option<String>) -> Self {
|
||||
Self {
|
||||
description,
|
||||
all_of: Some(variants),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn boolean(description: Option<String>) -> Self {
|
||||
Self::typed(JsonSchemaPrimitiveType::Boolean, description)
|
||||
}
|
||||
|
||||
pub fn string(description: Option<String>) -> Self {
|
||||
Self::typed(JsonSchemaPrimitiveType::String, description)
|
||||
}
|
||||
|
||||
pub fn with_encrypted(mut self) -> Self {
|
||||
self.encrypted = Some(true);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn number(description: Option<String>) -> Self {
|
||||
Self::typed(JsonSchemaPrimitiveType::Number, description)
|
||||
}
|
||||
|
||||
pub fn integer(description: Option<String>) -> Self {
|
||||
Self::typed(JsonSchemaPrimitiveType::Integer, description)
|
||||
}
|
||||
|
||||
pub fn null(description: Option<String>) -> Self {
|
||||
Self::typed(JsonSchemaPrimitiveType::Null, description)
|
||||
}
|
||||
|
||||
pub fn string_enum(values: Vec<JsonValue>, description: Option<String>) -> Self {
|
||||
Self {
|
||||
schema_type: Some(JsonSchemaType::Single(JsonSchemaPrimitiveType::String)),
|
||||
description,
|
||||
enum_values: Some(values),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn array(items: JsonSchema, description: Option<String>) -> Self {
|
||||
Self {
|
||||
schema_type: Some(JsonSchemaType::Single(JsonSchemaPrimitiveType::Array)),
|
||||
description,
|
||||
items: Some(Box::new(items)),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn object(
|
||||
properties: BTreeMap<String, JsonSchema>,
|
||||
required: Option<Vec<String>>,
|
||||
additional_properties: Option<AdditionalProperties>,
|
||||
) -> Self {
|
||||
Self {
|
||||
schema_type: Some(JsonSchemaType::Single(JsonSchemaPrimitiveType::Object)),
|
||||
properties: Some(properties),
|
||||
required,
|
||||
additional_properties,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether additional properties are allowed, and if so, any required schema.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
#[serde(untagged)]
|
||||
pub enum AdditionalProperties {
|
||||
Boolean(bool),
|
||||
Schema(Box<JsonSchema>),
|
||||
}
|
||||
|
||||
impl From<bool> for AdditionalProperties {
|
||||
fn from(value: bool) -> Self {
|
||||
Self::Boolean(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<JsonSchema> for AdditionalProperties {
|
||||
fn from(value: JsonSchema) -> Self {
|
||||
Self::Schema(Box::new(value))
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse the tool `input_schema` or return an error for invalid schema.
|
||||
pub fn parse_tool_input_schema(input_schema: &JsonValue) -> Result<JsonSchema, serde_json::Error> {
|
||||
let mut input_schema = prepare_tool_input_schema(input_schema);
|
||||
@@ -217,237 +53,6 @@ fn deserialize_tool_input_schema(input_schema: JsonValue) -> Result<JsonSchema,
|
||||
Ok(schema)
|
||||
}
|
||||
|
||||
// Use compact normalized JSON bytes as a cheap local proxy for the 1k-token
|
||||
// schema budget.
|
||||
const MAX_COMPACT_TOOL_SCHEMA_BYTES: usize = 5_000;
|
||||
const MAX_COMPACT_TOOL_SCHEMA_DEPTH: usize = 3;
|
||||
|
||||
/// Shrink unusually large tool schemas while preserving the top-level argument
|
||||
/// surface. Compaction is best-effort rather than a hard cap: it runs only
|
||||
/// after schema sanitization/pruning and applies increasingly lossy passes
|
||||
/// while the schema remains over budget.
|
||||
fn compact_large_tool_schema(value: &mut JsonValue) {
|
||||
for pass in LARGE_SCHEMA_COMPACTION_PASSES {
|
||||
if compact_schema_fits_budget(value) {
|
||||
break;
|
||||
}
|
||||
pass(value);
|
||||
}
|
||||
}
|
||||
|
||||
type LargeSchemaCompactionPass = fn(&mut JsonValue);
|
||||
|
||||
const LARGE_SCHEMA_COMPACTION_PASSES: &[LargeSchemaCompactionPass] = &[
|
||||
strip_schema_descriptions,
|
||||
drop_schema_definitions,
|
||||
collapse_deep_schema_objects_from_root,
|
||||
prune_schema_compositions,
|
||||
];
|
||||
|
||||
fn collapse_deep_schema_objects_from_root(value: &mut JsonValue) {
|
||||
collapse_deep_schema_objects(value, /*depth*/ 0);
|
||||
}
|
||||
|
||||
fn compact_schema_fits_budget(value: &JsonValue) -> bool {
|
||||
compact_normalized_schema_len(value) <= MAX_COMPACT_TOOL_SCHEMA_BYTES
|
||||
}
|
||||
|
||||
fn compact_normalized_schema_len(value: &JsonValue) -> usize {
|
||||
serde_json::from_value::<JsonSchema>(value.clone())
|
||||
.and_then(|schema| serde_json::to_vec(&schema))
|
||||
.map(|json| json.len())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum DefinitionTraversal {
|
||||
Include,
|
||||
Skip,
|
||||
}
|
||||
|
||||
fn for_each_schema_child(
|
||||
map: &serde_json::Map<String, JsonValue>,
|
||||
definition_traversal: DefinitionTraversal,
|
||||
visitor: &mut impl FnMut(&JsonValue),
|
||||
) {
|
||||
if let Some(properties) = map.get("properties")
|
||||
&& let Some(properties_map) = properties.as_object()
|
||||
{
|
||||
for value in properties_map.values() {
|
||||
visitor(value);
|
||||
}
|
||||
}
|
||||
|
||||
for key in SCHEMA_CHILD_KEYS {
|
||||
if let Some(value) = map.get(key) {
|
||||
visitor(value);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(additional_properties) = map.get("additionalProperties")
|
||||
&& !matches!(additional_properties, JsonValue::Bool(_))
|
||||
{
|
||||
visitor(additional_properties);
|
||||
}
|
||||
|
||||
if definition_traversal == DefinitionTraversal::Include {
|
||||
for key in DEFINITION_TABLE_KEYS {
|
||||
if let Some(definitions) = map.get(key)
|
||||
&& let Some(definitions_map) = definitions.as_object()
|
||||
{
|
||||
for value in definitions_map.values() {
|
||||
visitor(value);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn strip_schema_descriptions(value: &mut JsonValue) {
|
||||
match value {
|
||||
JsonValue::Array(values) => {
|
||||
for value in values {
|
||||
strip_schema_descriptions(value);
|
||||
}
|
||||
}
|
||||
JsonValue::Object(map) => {
|
||||
map.remove("description");
|
||||
for_each_schema_child_mut(map, DefinitionTraversal::Include, &mut |value| {
|
||||
strip_schema_descriptions(value);
|
||||
});
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
fn for_each_schema_child_mut(
|
||||
map: &mut serde_json::Map<String, JsonValue>,
|
||||
definition_traversal: DefinitionTraversal,
|
||||
visitor: &mut impl FnMut(&mut JsonValue),
|
||||
) {
|
||||
if let Some(properties) = map.get_mut("properties")
|
||||
&& let Some(properties_map) = properties.as_object_mut()
|
||||
{
|
||||
for value in properties_map.values_mut() {
|
||||
visitor(value);
|
||||
}
|
||||
}
|
||||
|
||||
for key in SCHEMA_CHILD_KEYS {
|
||||
if let Some(value) = map.get_mut(key) {
|
||||
visitor(value);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(additional_properties) = map.get_mut("additionalProperties")
|
||||
&& !matches!(additional_properties, JsonValue::Bool(_))
|
||||
{
|
||||
visitor(additional_properties);
|
||||
}
|
||||
|
||||
if definition_traversal == DefinitionTraversal::Include {
|
||||
for key in DEFINITION_TABLE_KEYS {
|
||||
if let Some(definitions) = map.get_mut(key)
|
||||
&& let Some(definitions_map) = definitions.as_object_mut()
|
||||
{
|
||||
for value in definitions_map.values_mut() {
|
||||
visitor(value);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Replace local definition refs with empty schemas before dropping root
|
||||
/// definition tables, so downstream behavior does not depend on how a schema
|
||||
/// parser handles refs to missing definitions.
|
||||
fn drop_schema_definitions(value: &mut JsonValue) {
|
||||
rewrite_definition_refs_to_empty_schemas(value);
|
||||
|
||||
let JsonValue::Object(map) = value else {
|
||||
return;
|
||||
};
|
||||
|
||||
for key in DEFINITION_TABLE_KEYS {
|
||||
map.remove(key);
|
||||
}
|
||||
}
|
||||
|
||||
fn rewrite_definition_refs_to_empty_schemas(value: &mut JsonValue) {
|
||||
match value {
|
||||
JsonValue::Array(values) => {
|
||||
for value in values {
|
||||
rewrite_definition_refs_to_empty_schemas(value);
|
||||
}
|
||||
}
|
||||
JsonValue::Object(map) => {
|
||||
if map
|
||||
.get("$ref")
|
||||
.and_then(JsonValue::as_str)
|
||||
.and_then(parse_local_definition_ref)
|
||||
.is_some()
|
||||
{
|
||||
*value = json!({});
|
||||
return;
|
||||
}
|
||||
|
||||
for_each_schema_child_mut(map, DefinitionTraversal::Skip, &mut |value| {
|
||||
rewrite_definition_refs_to_empty_schemas(value);
|
||||
});
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
fn collapse_deep_schema_objects(value: &mut JsonValue, depth: usize) {
|
||||
match value {
|
||||
JsonValue::Array(values) => {
|
||||
for value in values {
|
||||
collapse_deep_schema_objects(value, depth);
|
||||
}
|
||||
}
|
||||
JsonValue::Object(map) => {
|
||||
if depth >= MAX_COMPACT_TOOL_SCHEMA_DEPTH && is_complex_schema_object(map) {
|
||||
*value = json!({});
|
||||
return;
|
||||
}
|
||||
|
||||
for_each_schema_child_mut(map, DefinitionTraversal::Skip, &mut |value| {
|
||||
collapse_deep_schema_objects(value, depth + 1);
|
||||
});
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
fn prune_schema_compositions(value: &mut JsonValue) {
|
||||
match value {
|
||||
JsonValue::Array(values) => {
|
||||
for value in values {
|
||||
prune_schema_compositions(value);
|
||||
}
|
||||
}
|
||||
JsonValue::Object(map) => {
|
||||
if has_composition_keyword(map) {
|
||||
*value = json!({});
|
||||
return;
|
||||
}
|
||||
|
||||
for_each_schema_child_mut(map, DefinitionTraversal::Skip, &mut |value| {
|
||||
prune_schema_compositions(value);
|
||||
});
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
fn is_complex_schema_object(map: &serde_json::Map<String, JsonValue>) -> bool {
|
||||
SCHEMA_CHILD_KEYS.iter().any(|key| map.contains_key(*key))
|
||||
|| map.contains_key("properties")
|
||||
|| map.contains_key("additionalProperties")
|
||||
|| map.contains_key("$ref")
|
||||
}
|
||||
|
||||
fn has_composition_keyword(map: &serde_json::Map<String, JsonValue>) -> bool {
|
||||
COMPOSITION_SCHEMA_KEYS
|
||||
.into_iter()
|
||||
|
||||
@@ -0,0 +1,165 @@
|
||||
//! Apply the legacy best-effort size policy after schema sanitization and pruning.
|
||||
//! Pass order and normalized-byte accounting are compatibility behavior.
|
||||
|
||||
use super::DEFINITION_TABLE_KEYS;
|
||||
use super::JsonSchema;
|
||||
use super::SCHEMA_CHILD_KEYS;
|
||||
use super::has_composition_keyword;
|
||||
use super::parse_local_definition_ref;
|
||||
use super::traversal::DefinitionTraversal;
|
||||
use super::traversal::for_each_schema_child_mut;
|
||||
use serde_json::Value as JsonValue;
|
||||
use serde_json::json;
|
||||
|
||||
// Use compact normalized JSON bytes as a cheap local proxy for the 1k-token
|
||||
// schema budget.
|
||||
const MAX_COMPACT_TOOL_SCHEMA_BYTES: usize = 5_000;
|
||||
const MAX_COMPACT_TOOL_SCHEMA_DEPTH: usize = 3;
|
||||
|
||||
/// Shrink unusually large tool schemas while preserving the top-level argument
|
||||
/// surface. Compaction is best-effort rather than a hard cap: it runs only
|
||||
/// after schema sanitization/pruning and applies increasingly lossy passes
|
||||
/// while the schema remains over budget.
|
||||
pub(super) fn compact_large_tool_schema(value: &mut JsonValue) {
|
||||
for pass in LARGE_SCHEMA_COMPACTION_PASSES {
|
||||
if compact_schema_fits_budget(value) {
|
||||
break;
|
||||
}
|
||||
pass(value);
|
||||
}
|
||||
}
|
||||
|
||||
type LargeSchemaCompactionPass = fn(&mut JsonValue);
|
||||
|
||||
const LARGE_SCHEMA_COMPACTION_PASSES: &[LargeSchemaCompactionPass] = &[
|
||||
strip_schema_descriptions,
|
||||
drop_schema_definitions,
|
||||
collapse_deep_schema_objects_from_root,
|
||||
prune_schema_compositions,
|
||||
];
|
||||
|
||||
fn collapse_deep_schema_objects_from_root(value: &mut JsonValue) {
|
||||
collapse_deep_schema_objects(value, /*depth*/ 0);
|
||||
}
|
||||
|
||||
fn compact_schema_fits_budget(value: &JsonValue) -> bool {
|
||||
compact_normalized_schema_len(value) <= MAX_COMPACT_TOOL_SCHEMA_BYTES
|
||||
}
|
||||
|
||||
fn compact_normalized_schema_len(value: &JsonValue) -> usize {
|
||||
serde_json::from_value::<JsonSchema>(value.clone())
|
||||
.and_then(|schema| serde_json::to_vec(&schema))
|
||||
.map(|json| json.len())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn strip_schema_descriptions(value: &mut JsonValue) {
|
||||
match value {
|
||||
JsonValue::Array(values) => {
|
||||
for value in values {
|
||||
strip_schema_descriptions(value);
|
||||
}
|
||||
}
|
||||
JsonValue::Object(map) => {
|
||||
map.remove("description");
|
||||
for_each_schema_child_mut(map, DefinitionTraversal::Include, &mut |value| {
|
||||
strip_schema_descriptions(value);
|
||||
});
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
/// Replace local definition refs with empty schemas before dropping root
|
||||
/// definition tables, so downstream behavior does not depend on how a schema
|
||||
/// parser handles refs to missing definitions.
|
||||
fn drop_schema_definitions(value: &mut JsonValue) {
|
||||
rewrite_definition_refs_to_empty_schemas(value);
|
||||
|
||||
let JsonValue::Object(map) = value else {
|
||||
return;
|
||||
};
|
||||
|
||||
for key in DEFINITION_TABLE_KEYS {
|
||||
map.remove(key);
|
||||
}
|
||||
}
|
||||
|
||||
fn rewrite_definition_refs_to_empty_schemas(value: &mut JsonValue) {
|
||||
match value {
|
||||
JsonValue::Array(values) => {
|
||||
for value in values {
|
||||
rewrite_definition_refs_to_empty_schemas(value);
|
||||
}
|
||||
}
|
||||
JsonValue::Object(map) => {
|
||||
if map
|
||||
.get("$ref")
|
||||
.and_then(JsonValue::as_str)
|
||||
.and_then(parse_local_definition_ref)
|
||||
.is_some()
|
||||
{
|
||||
*value = json!({});
|
||||
return;
|
||||
}
|
||||
|
||||
for_each_schema_child_mut(map, DefinitionTraversal::Skip, &mut |value| {
|
||||
rewrite_definition_refs_to_empty_schemas(value);
|
||||
});
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
fn collapse_deep_schema_objects(value: &mut JsonValue, depth: usize) {
|
||||
match value {
|
||||
JsonValue::Array(values) => {
|
||||
for value in values {
|
||||
collapse_deep_schema_objects(value, depth);
|
||||
}
|
||||
}
|
||||
JsonValue::Object(map) => {
|
||||
if depth >= MAX_COMPACT_TOOL_SCHEMA_DEPTH && is_complex_schema_object(map) {
|
||||
*value = json!({});
|
||||
return;
|
||||
}
|
||||
|
||||
for_each_schema_child_mut(map, DefinitionTraversal::Skip, &mut |value| {
|
||||
collapse_deep_schema_objects(value, depth + 1);
|
||||
});
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
fn prune_schema_compositions(value: &mut JsonValue) {
|
||||
match value {
|
||||
JsonValue::Array(values) => {
|
||||
for value in values {
|
||||
prune_schema_compositions(value);
|
||||
}
|
||||
}
|
||||
JsonValue::Object(map) => {
|
||||
if has_composition_keyword(map) {
|
||||
*value = json!({});
|
||||
return;
|
||||
}
|
||||
|
||||
for_each_schema_child_mut(map, DefinitionTraversal::Skip, &mut |value| {
|
||||
prune_schema_compositions(value);
|
||||
});
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
fn is_complex_schema_object(map: &serde_json::Map<String, JsonValue>) -> bool {
|
||||
SCHEMA_CHILD_KEYS.iter().any(|key| map.contains_key(*key))
|
||||
|| map.contains_key("properties")
|
||||
|| map.contains_key("additionalProperties")
|
||||
|| map.contains_key("$ref")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "compaction_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,127 @@
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
#[test]
|
||||
fn collapse_deep_schema_objects_traverses_schema_children() {
|
||||
let mut schema = serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"object_parent": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"complex": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"nested": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"leaf": { "type": "string" }
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"scalar": {
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
},
|
||||
"array_parent": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"nested": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"leaf": { "type": "string" }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"map_parent": {
|
||||
"type": "object",
|
||||
"additionalProperties": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"nested": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"leaf": { "type": "string" }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"union_parent": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"nested": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"leaf": { "type": "string" }
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
{ "type": "string" }
|
||||
]
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
super::collapse_deep_schema_objects(&mut schema, /*depth*/ 0);
|
||||
|
||||
assert_eq!(
|
||||
schema,
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"object_parent": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"complex": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"nested": {}
|
||||
}
|
||||
},
|
||||
"scalar": {
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
},
|
||||
"array_parent": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"nested": {}
|
||||
}
|
||||
}
|
||||
},
|
||||
"map_parent": {
|
||||
"type": "object",
|
||||
"additionalProperties": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"nested": {}
|
||||
}
|
||||
}
|
||||
},
|
||||
"union_parent": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"nested": {}
|
||||
}
|
||||
},
|
||||
{ "type": "string" }
|
||||
]
|
||||
}
|
||||
}
|
||||
})
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
//! Visit schema-valued children while leaving enum values and other instance data untouched.
|
||||
|
||||
use super::DEFINITION_TABLE_KEYS;
|
||||
use super::SCHEMA_CHILD_KEYS;
|
||||
use serde_json::Value as JsonValue;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(super) enum DefinitionTraversal {
|
||||
Include,
|
||||
Skip,
|
||||
}
|
||||
|
||||
pub(super) fn for_each_schema_child(
|
||||
map: &serde_json::Map<String, JsonValue>,
|
||||
definition_traversal: DefinitionTraversal,
|
||||
visitor: &mut impl FnMut(&JsonValue),
|
||||
) {
|
||||
if let Some(properties) = map.get("properties")
|
||||
&& let Some(properties_map) = properties.as_object()
|
||||
{
|
||||
for value in properties_map.values() {
|
||||
visitor(value);
|
||||
}
|
||||
}
|
||||
|
||||
for key in SCHEMA_CHILD_KEYS {
|
||||
if let Some(value) = map.get(key) {
|
||||
visitor(value);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(additional_properties) = map.get("additionalProperties")
|
||||
&& !matches!(additional_properties, JsonValue::Bool(_))
|
||||
{
|
||||
visitor(additional_properties);
|
||||
}
|
||||
|
||||
if definition_traversal == DefinitionTraversal::Include {
|
||||
for key in DEFINITION_TABLE_KEYS {
|
||||
if let Some(definitions) = map.get(key)
|
||||
&& let Some(definitions_map) = definitions.as_object()
|
||||
{
|
||||
for value in definitions_map.values() {
|
||||
visitor(value);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn for_each_schema_child_mut(
|
||||
map: &mut serde_json::Map<String, JsonValue>,
|
||||
definition_traversal: DefinitionTraversal,
|
||||
visitor: &mut impl FnMut(&mut JsonValue),
|
||||
) {
|
||||
if let Some(properties) = map.get_mut("properties")
|
||||
&& let Some(properties_map) = properties.as_object_mut()
|
||||
{
|
||||
for value in properties_map.values_mut() {
|
||||
visitor(value);
|
||||
}
|
||||
}
|
||||
|
||||
for key in SCHEMA_CHILD_KEYS {
|
||||
if let Some(value) = map.get_mut(key) {
|
||||
visitor(value);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(additional_properties) = map.get_mut("additionalProperties")
|
||||
&& !matches!(additional_properties, JsonValue::Bool(_))
|
||||
{
|
||||
visitor(additional_properties);
|
||||
}
|
||||
|
||||
if definition_traversal == DefinitionTraversal::Include {
|
||||
for key in DEFINITION_TABLE_KEYS {
|
||||
if let Some(definitions) = map.get_mut(key)
|
||||
&& let Some(definitions_map) = definitions.as_object_mut()
|
||||
{
|
||||
for value in definitions_map.values_mut() {
|
||||
visitor(value);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,186 @@
|
||||
//! Define serializable tool schema types and their constructors.
|
||||
|
||||
use serde::Deserialize;
|
||||
use serde::Serialize;
|
||||
use serde_json::Value as JsonValue;
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
/// Primitive JSON Schema type names we support in tool definitions.
|
||||
///
|
||||
/// This mirrors the OpenAI Structured Outputs subset for JSON Schema `type`:
|
||||
/// string, number, boolean, integer, object, array, and null.
|
||||
/// Keywords such as `enum`, `const`, `anyOf`, `oneOf`, and `allOf` are modeled
|
||||
/// separately.
|
||||
/// See <https://developers.openai.com/api/docs/guides/structured-outputs#supported-schemas>.
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum JsonSchemaPrimitiveType {
|
||||
String,
|
||||
Number,
|
||||
Boolean,
|
||||
Integer,
|
||||
Object,
|
||||
Array,
|
||||
Null,
|
||||
}
|
||||
|
||||
/// JSON Schema `type` supports either a single type name or a union of names.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(untagged)]
|
||||
pub enum JsonSchemaType {
|
||||
Single(JsonSchemaPrimitiveType),
|
||||
Multiple(Vec<JsonSchemaPrimitiveType>),
|
||||
}
|
||||
|
||||
/// Generic JSON-Schema subset needed for our tool definitions.
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
|
||||
pub struct JsonSchema {
|
||||
#[serde(rename = "$ref", skip_serializing_if = "Option::is_none")]
|
||||
pub schema_ref: Option<String>,
|
||||
#[serde(rename = "type", skip_serializing_if = "Option::is_none")]
|
||||
pub schema_type: Option<JsonSchemaType>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub description: Option<String>,
|
||||
/// Responses-only marker for reviewed encrypted tool parameters.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub encrypted: Option<bool>,
|
||||
#[serde(rename = "enum", skip_serializing_if = "Option::is_none")]
|
||||
pub enum_values: Option<Vec<JsonValue>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub items: Option<Box<JsonSchema>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub properties: Option<BTreeMap<String, JsonSchema>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub required: Option<Vec<String>>,
|
||||
#[serde(
|
||||
rename = "additionalProperties",
|
||||
skip_serializing_if = "Option::is_none"
|
||||
)]
|
||||
pub additional_properties: Option<AdditionalProperties>,
|
||||
#[serde(rename = "anyOf", skip_serializing_if = "Option::is_none")]
|
||||
pub any_of: Option<Vec<JsonSchema>>,
|
||||
#[serde(rename = "oneOf", skip_serializing_if = "Option::is_none")]
|
||||
pub one_of: Option<Vec<JsonSchema>>,
|
||||
#[serde(rename = "allOf", skip_serializing_if = "Option::is_none")]
|
||||
pub all_of: Option<Vec<JsonSchema>>,
|
||||
#[serde(rename = "$defs", skip_serializing_if = "Option::is_none")]
|
||||
pub defs: Option<BTreeMap<String, JsonSchema>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub definitions: Option<BTreeMap<String, JsonSchema>>,
|
||||
}
|
||||
|
||||
impl JsonSchema {
|
||||
/// Construct a scalar/object/array schema with a single JSON Schema type.
|
||||
fn typed(schema_type: JsonSchemaPrimitiveType, description: Option<String>) -> Self {
|
||||
Self {
|
||||
schema_type: Some(JsonSchemaType::Single(schema_type)),
|
||||
description,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn any_of(variants: Vec<JsonSchema>, description: Option<String>) -> Self {
|
||||
Self {
|
||||
description,
|
||||
any_of: Some(variants),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn one_of(variants: Vec<JsonSchema>, description: Option<String>) -> Self {
|
||||
Self {
|
||||
description,
|
||||
one_of: Some(variants),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn all_of(variants: Vec<JsonSchema>, description: Option<String>) -> Self {
|
||||
Self {
|
||||
description,
|
||||
all_of: Some(variants),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn boolean(description: Option<String>) -> Self {
|
||||
Self::typed(JsonSchemaPrimitiveType::Boolean, description)
|
||||
}
|
||||
|
||||
pub fn string(description: Option<String>) -> Self {
|
||||
Self::typed(JsonSchemaPrimitiveType::String, description)
|
||||
}
|
||||
|
||||
pub fn with_encrypted(mut self) -> Self {
|
||||
self.encrypted = Some(true);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn number(description: Option<String>) -> Self {
|
||||
Self::typed(JsonSchemaPrimitiveType::Number, description)
|
||||
}
|
||||
|
||||
pub fn integer(description: Option<String>) -> Self {
|
||||
Self::typed(JsonSchemaPrimitiveType::Integer, description)
|
||||
}
|
||||
|
||||
pub fn null(description: Option<String>) -> Self {
|
||||
Self::typed(JsonSchemaPrimitiveType::Null, description)
|
||||
}
|
||||
|
||||
pub fn string_enum(values: Vec<JsonValue>, description: Option<String>) -> Self {
|
||||
Self {
|
||||
schema_type: Some(JsonSchemaType::Single(JsonSchemaPrimitiveType::String)),
|
||||
description,
|
||||
enum_values: Some(values),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn array(items: JsonSchema, description: Option<String>) -> Self {
|
||||
Self {
|
||||
schema_type: Some(JsonSchemaType::Single(JsonSchemaPrimitiveType::Array)),
|
||||
description,
|
||||
items: Some(Box::new(items)),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn object(
|
||||
properties: BTreeMap<String, JsonSchema>,
|
||||
required: Option<Vec<String>>,
|
||||
additional_properties: Option<AdditionalProperties>,
|
||||
) -> Self {
|
||||
Self {
|
||||
schema_type: Some(JsonSchemaType::Single(JsonSchemaPrimitiveType::Object)),
|
||||
properties: Some(properties),
|
||||
required,
|
||||
additional_properties,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether additional properties are allowed, and if so, any required schema.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
#[serde(untagged)]
|
||||
pub enum AdditionalProperties {
|
||||
Boolean(bool),
|
||||
Schema(Box<JsonSchema>),
|
||||
}
|
||||
|
||||
impl From<bool> for AdditionalProperties {
|
||||
fn from(value: bool) -> Self {
|
||||
Self::Boolean(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<JsonSchema> for AdditionalProperties {
|
||||
fn from(value: JsonSchema) -> Self {
|
||||
Self::Schema(Box::new(value))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "types_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,18 @@
|
||||
//! Verify serialization of the supported tool schema representation.
|
||||
|
||||
use super::JsonSchema;
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
#[test]
|
||||
fn json_schema_serializes_encrypted_marker() {
|
||||
let schema = JsonSchema::string(Some("Secret value".to_string())).with_encrypted();
|
||||
|
||||
assert_eq!(
|
||||
serde_json::to_value(schema).expect("serialize schema"),
|
||||
serde_json::json!({
|
||||
"type": "string",
|
||||
"description": "Secret value",
|
||||
"encrypted": true,
|
||||
})
|
||||
);
|
||||
}
|
||||
@@ -24,20 +24,6 @@ fn parse_tool_input_schema_coerces_boolean_schemas() {
|
||||
assert_eq!(schema, JsonSchema::string(/*description*/ None));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn json_schema_serializes_encrypted_marker() {
|
||||
let schema = JsonSchema::string(Some("Secret value".to_string())).with_encrypted();
|
||||
|
||||
assert_eq!(
|
||||
serde_json::to_value(schema).expect("serialize schema"),
|
||||
serde_json::json!({
|
||||
"type": "string",
|
||||
"description": "Secret value",
|
||||
"encrypted": true,
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_tool_input_schema_infers_object_shape_and_defaults_properties() {
|
||||
// Example schema shape:
|
||||
@@ -1289,132 +1275,6 @@ fn parse_large_tool_input_schema_preserves_object_enum_literal_descriptions() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn collapse_deep_schema_objects_traverses_schema_children() {
|
||||
let mut schema = serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"object_parent": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"complex": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"nested": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"leaf": { "type": "string" }
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"scalar": {
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
},
|
||||
"array_parent": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"nested": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"leaf": { "type": "string" }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"map_parent": {
|
||||
"type": "object",
|
||||
"additionalProperties": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"nested": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"leaf": { "type": "string" }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"union_parent": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"nested": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"leaf": { "type": "string" }
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
{ "type": "string" }
|
||||
]
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
super::collapse_deep_schema_objects(&mut schema, /*depth*/ 0);
|
||||
|
||||
assert_eq!(
|
||||
schema,
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"object_parent": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"complex": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"nested": {}
|
||||
}
|
||||
},
|
||||
"scalar": {
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
},
|
||||
"array_parent": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"nested": {}
|
||||
}
|
||||
}
|
||||
},
|
||||
"map_parent": {
|
||||
"type": "object",
|
||||
"additionalProperties": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"nested": {}
|
||||
}
|
||||
}
|
||||
},
|
||||
"union_parent": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"nested": {}
|
||||
}
|
||||
},
|
||||
{ "type": "string" }
|
||||
]
|
||||
}
|
||||
}
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_tool_input_schema_preserves_string_enum_constraints() {
|
||||
// Example schema shape:
|
||||
|
||||
Reference in New Issue
Block a user