tiny_agent_rs/schemas/
schema.rs1use 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#[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
60pub 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
107pub fn schema_type_name<T>() -> &'static str {
109 type_name::<T>()
110}