Skip to main content

tiny_agent_rs/schemas/
schema.rs

1use schemars::schema::{ObjectValidation, RootSchema, Schema, SchemaObject};
2use serde::de::DeserializeOwned;
3use serde_json::Value;
4use std::{
5    any::{type_name, TypeId},
6    sync::Arc,
7};
8
9/// Cached JSON schema handle associated with a response type.
10#[derive(Clone, Debug)]
11pub struct SchemaHandle {
12    schema_name: &'static str,
13    type_name: &'static str,
14    type_id: TypeId,
15    schema_json: Arc<Value>,
16}
17
18impl SchemaHandle {
19    pub fn from_root_schema<T: 'static>(
20        schema_name: &'static str,
21        type_name: &'static str,
22        root: RootSchema,
23    ) -> Self {
24        let schema_json = serde_json::to_value(root)
25            .unwrap_or_else(|err| panic!("failed to serialize schema for {}: {}", type_name, err));
26
27        Self {
28            schema_name,
29            type_name,
30            type_id: TypeId::of::<T>(),
31            schema_json: Arc::new(schema_json),
32        }
33    }
34
35    pub fn schema_name(&self) -> &'static str {
36        self.schema_name
37    }
38
39    pub fn type_name(&self) -> &'static str {
40        self.type_name
41    }
42
43    pub fn type_id(&self) -> TypeId {
44        self.type_id
45    }
46
47    pub fn schema_json(&self) -> &Value {
48        self.schema_json.as_ref()
49    }
50
51    pub fn schema_json_arc(&self) -> Arc<Value> {
52        Arc::clone(&self.schema_json)
53    }
54}
55
56pub trait CompletionSchema: DeserializeOwned + Send + Sync + 'static {
57    fn schema() -> &'static SchemaHandle;
58}
59
60/// Apply doc comments captured by the procedural macro to the generated schema metadata.
61pub fn apply_doc_comments(
62    root: &mut RootSchema,
63    title: &'static str,
64    description: Option<&'static str>,
65    field_docs: &[(&'static str, &'static str)],
66) {
67    let schema_object = &mut root.schema;
68    apply_struct_metadata(schema_object, title, description);
69
70    if let Some(object_validation) = schema_object.object.as_mut() {
71        apply_field_metadata(object_validation.as_mut(), field_docs);
72    }
73}
74
75fn apply_struct_metadata(
76    schema_object: &mut SchemaObject,
77    title: &'static str,
78    description: Option<&'static str>,
79) {
80    let metadata = schema_object.metadata();
81
82    if metadata.title.is_none() {
83        metadata.title = Some(title.to_string());
84    }
85
86    if let Some(description) = description {
87        if metadata.description.is_none() {
88            metadata.description = Some(description.to_string());
89        }
90    }
91}
92
93fn apply_field_metadata(
94    object_validation: &mut ObjectValidation,
95    field_docs: &[(&'static str, &'static str)],
96) {
97    for (field, doc) in field_docs {
98        if let Some(Schema::Object(field_object)) = object_validation.properties.get_mut(*field) {
99            let metadata = field_object.metadata();
100            if metadata.description.is_none() {
101                metadata.description = Some((*doc).to_string());
102            }
103        }
104    }
105}
106
107/// Helper so callers can retrieve the Rust type name of a schema provider.
108pub fn schema_type_name<T>() -> &'static str {
109    type_name::<T>()
110}