zai_rs/toolkits/
core.rs

1//! Core traits and types with enhanced type safety
2
3use std::{
4    borrow::Cow,
5    collections::{HashMap, hash_map::DefaultHasher},
6    hash::{Hash, Hasher},
7    sync::Arc,
8};
9
10use async_trait::async_trait;
11use jsonschema;
12use once_cell::sync::Lazy;
13use parking_lot::RwLock;
14use serde::{Deserialize, Serialize};
15use tracing::warn;
16
17use crate::toolkits::error::{ToolResult, error_context};
18
19/// Type-erased tool trait for dynamic dispatch
20#[async_trait]
21pub trait DynTool: Send + Sync {
22    /// Get the tool's metadata
23    fn metadata(&self) -> &ToolMetadata;
24
25    /// Execute with JSON input/output
26    async fn execute_json(&self, input: serde_json::Value) -> ToolResult<serde_json::Value>;
27
28    /// Get input schema
29    fn input_schema(&self) -> serde_json::Value;
30
31    /// Get the tool name
32    fn name(&self) -> &str {
33        &self.metadata().name
34    }
35
36    /// Clone the tool as a boxed trait object
37    fn clone_box(&self) -> Box<dyn DynTool>;
38}
39
40/// Global schema cache for compiled JSON schemas
41static SCHEMA_CACHE: Lazy<RwLock<HashMap<u64, Arc<jsonschema::Validator>>>> =
42    Lazy::new(|| RwLock::new(HashMap::new()));
43
44/// Maximum number of compiled schemas to cache
45const SCHEMA_CACHE_MAX_SIZE: usize = 256;
46
47/// Enhanced tool metadata with better type information and memory optimization
48#[derive(Debug, Clone, Serialize, Deserialize)]
49pub struct ToolMetadata {
50    /// Tool name (must be unique)
51    pub name: Cow<'static, str>,
52
53    /// Tool description
54    pub description: Cow<'static, str>,
55
56    /// Tool version
57    pub version: Cow<'static, str>,
58
59    /// Tool author
60    pub author: Option<Cow<'static, str>>,
61
62    /// Tool tags for categorization
63    pub tags: Vec<Cow<'static, str>>,
64
65    /// Whether the tool is enabled
66    pub enabled: bool,
67
68    /// Additional metadata
69    pub metadata: HashMap<Cow<'static, str>, serde_json::Value>,
70}
71
72impl ToolMetadata {
73    /// Create new metadata with validation
74    pub fn new(name: impl Into<String>, description: impl Into<String>) -> ToolResult<Self> {
75        let name = name.into();
76        let description = description.into();
77
78        // Validate tool name
79        if name.trim().is_empty() {
80            return Err(error_context().invalid_parameters("Tool name cannot be empty"));
81        }
82        if name.contains(|c: char| !c.is_alphanumeric() && c != '_') {
83            return Err(error_context()
84                .invalid_parameters("Tool name must be alphanumeric with underscores only"));
85        }
86
87        Ok(Self {
88            name: Cow::Owned(name),
89            description: Cow::Owned(description),
90            version: Cow::Borrowed("1.0.0"),
91            author: None,
92            tags: Vec::new(),
93            enabled: true,
94            metadata: HashMap::new(),
95        })
96    }
97
98    /// Builder pattern methods
99    pub fn version(mut self, version: impl Into<Cow<'static, str>>) -> Self {
100        self.version = version.into();
101        self
102    }
103
104    pub fn author(mut self, author: impl Into<Cow<'static, str>>) -> Self {
105        self.author = Some(author.into());
106        self
107    }
108
109    pub fn tags<T: Into<Cow<'static, str>>>(mut self, tags: impl IntoIterator<Item = T>) -> Self {
110        self.tags = tags.into_iter().map(Into::into).collect();
111        self
112    }
113
114    pub fn enabled(mut self, enabled: bool) -> Self {
115        self.enabled = enabled;
116        self
117    }
118
119    pub fn with_metadata(
120        mut self,
121        key: impl Into<Cow<'static, str>>,
122        value: serde_json::Value,
123    ) -> Self {
124        self.metadata.insert(key.into(), value);
125        self
126    }
127}
128
129/// Helper functions for type conversions (avoiding orphan rule issues)
130pub mod conversions {
131    use crate::toolkits::error::{ToolResult, error_context};
132
133    /// Convert a value to JSON
134    pub fn to_json<T: serde::Serialize>(value: T) -> ToolResult<serde_json::Value> {
135        serde_json::to_value(value).map_err(|e| error_context().serialization_error(e))
136    }
137
138    /// Extract string from JSON value
139    pub fn from_json_string(value: serde_json::Value) -> ToolResult<String> {
140        match value {
141            serde_json::Value::String(s) => Ok(s),
142            _ => Err(error_context().invalid_parameters("Expected string value")),
143        }
144    }
145
146    /// Extract i32 from JSON value
147    pub fn from_json_i32(value: serde_json::Value) -> ToolResult<i32> {
148        match value {
149            serde_json::Value::Number(n) => n
150                .as_i64()
151                .and_then(|i| i.try_into().ok())
152                .ok_or_else(|| error_context().invalid_parameters("Expected i32 value")),
153            _ => Err(error_context().invalid_parameters("Expected number value")),
154        }
155    }
156
157    /// Extract f64 from JSON value
158    pub fn from_json_f64(value: serde_json::Value) -> ToolResult<f64> {
159        match value {
160            serde_json::Value::Number(n) => n
161                .as_f64()
162                .ok_or_else(|| error_context().invalid_parameters("Expected f64 value")),
163            _ => Err(error_context().invalid_parameters("Expected number value")),
164        }
165    }
166
167    /// Extract bool from JSON value
168    pub fn from_json_bool(value: serde_json::Value) -> ToolResult<bool> {
169        match value {
170            serde_json::Value::Bool(b) => Ok(b),
171            _ => Err(error_context().invalid_parameters("Expected boolean value")),
172        }
173    }
174}
175
176// -----------------------------
177// Single-struct dynamic FunctionTool
178// -----------------------------
179
180/// Type alias for the complex handler type to reduce complexity warnings
181pub(crate) type ToolHandler = std::sync::Arc<
182    dyn Fn(
183            serde_json::Value,
184        ) -> std::pin::Pin<
185            Box<
186                dyn std::future::Future<
187                        Output = crate::toolkits::error::ToolResult<serde_json::Value>,
188                    > + Send,
189            >,
190        > + Send
191        + Sync,
192>;
193
194/// A single-struct tool that carries metadata, JSON schema, and an async
195/// handler
196pub struct FunctionTool {
197    metadata: ToolMetadata,
198    input_schema: serde_json::Value,
199    compiled_schema: Arc<jsonschema::Validator>,
200    handler: ToolHandler,
201}
202
203impl Clone for FunctionTool {
204    fn clone(&self) -> Self {
205        Self {
206            metadata: self.metadata.clone(),
207            input_schema: self.input_schema.clone(),
208            compiled_schema: Arc::clone(&self.compiled_schema),
209            handler: self.handler.clone(),
210        }
211    }
212}
213
214impl FunctionTool {
215    pub fn builder(name: impl Into<String>, description: impl Into<String>) -> FunctionToolBuilder {
216        FunctionToolBuilder::new(name, description)
217    }
218    /// Convenience: build a FunctionTool directly from a full JSON schema and a
219    /// handler
220    pub fn from_schema<F, Fut>(
221        name: impl Into<String>,
222        description: impl Into<String>,
223        schema: serde_json::Value,
224        f: F,
225    ) -> crate::toolkits::error::ToolResult<FunctionTool>
226    where
227        F: Fn(serde_json::Value) -> Fut + Send + Sync + 'static,
228        Fut: std::future::Future<Output = crate::toolkits::error::ToolResult<serde_json::Value>>
229            + Send
230            + 'static,
231    {
232        Self::builder(name, description)
233            .schema(schema)
234            .handler(f)
235            .build()
236    }
237    /// Build a FunctionTool from a full JSON spec (supports two shapes):
238    /// 1) {"name":..., "description":..., "parameters": {...}}
239    /// 2) {"type":"function", "function": {"name":..., "description":...,
240    ///    "parameters": {...}}}
241    pub fn from_function_spec<F, Fut>(
242        spec: serde_json::Value,
243        f: F,
244    ) -> crate::toolkits::error::ToolResult<FunctionTool>
245    where
246        F: Fn(serde_json::Value) -> Fut + Send + Sync + 'static,
247        Fut: std::future::Future<Output = crate::toolkits::error::ToolResult<serde_json::Value>>
248            + Send
249            + 'static,
250    {
251        let (name, description, parameters) = parse_function_spec_details(&spec)?;
252        let mut builder = Self::builder(name, description);
253        if let Some(p) = parameters {
254            builder = builder.schema(p);
255        }
256        builder.handler(f).build()
257    }
258
259    /// Read a JSON function spec from a file and build a FunctionTool.
260    pub fn from_function_spec_file<F, Fut>(
261        path: impl AsRef<std::path::Path>,
262        f: F,
263    ) -> crate::toolkits::error::ToolResult<FunctionTool>
264    where
265        F: Fn(serde_json::Value) -> Fut + Send + Sync + 'static,
266        Fut: std::future::Future<Output = crate::toolkits::error::ToolResult<serde_json::Value>>
267            + Send
268            + 'static,
269    {
270        let content = std::fs::read_to_string(path).map_err(|e| {
271            error_context().invalid_parameters(format!("Failed to read spec file: {}", e))
272        })?;
273        let spec: serde_json::Value = serde_json::from_str(&content)
274            .map_err(|e| error_context().invalid_parameters(format!("Invalid JSON: {}", e)))?;
275        Self::from_function_spec(spec, f)
276    }
277}
278
279/// Compile JSON schema with caching for better performance
280fn compile_schema_cached(schema: &serde_json::Value) -> ToolResult<Arc<jsonschema::Validator>> {
281    let mut hasher = DefaultHasher::new();
282    schema.to_string().hash(&mut hasher);
283    let hash = hasher.finish();
284
285    // Check cache first
286    {
287        let cache = SCHEMA_CACHE.read();
288        if let Some(cached) = cache.get(&hash) {
289            return Ok(Arc::clone(cached));
290        }
291    }
292
293    // Compile and cache
294    let validator = jsonschema::validator_for(schema).map_err(|e| {
295        error_context().schema_validation(format!("Failed to compile schema: {}", e))
296    })?;
297
298    let validator = Arc::new(validator);
299
300    {
301        let mut cache = SCHEMA_CACHE.write();
302        // Evict oldest entries if cache is full
303        if cache.len() >= SCHEMA_CACHE_MAX_SIZE {
304            // Remove approximately 10% of entries (oldest by insertion order)
305            let remove_count = (SCHEMA_CACHE_MAX_SIZE / 10).max(1);
306            let keys: Vec<u64> = cache.keys().take(remove_count).copied().collect();
307            for k in keys {
308                cache.remove(&k);
309            }
310        }
311        cache.insert(hash, Arc::clone(&validator));
312    }
313
314    Ok(validator)
315}
316
317/// (internal) Parses the name, description, and parameters from a JSON function
318/// spec.
319pub(crate) fn parse_function_spec_details(
320    spec: &serde_json::Value,
321) -> crate::toolkits::error::ToolResult<(String, String, Option<serde_json::Value>)> {
322    use serde_json::Value;
323    let obj = match spec {
324        Value::Object(map) => map,
325        _ => return Err(error_context().invalid_parameters("Function spec must be a JSON object")),
326    };
327    // Shape 2 with outer {type:function, function:{...}}
328    let (name, desc, params) = if obj.get("type").and_then(|v| v.as_str()) == Some("function") {
329        let f = obj
330            .get("function")
331            .and_then(|v| v.as_object())
332            .ok_or_else(|| error_context().invalid_parameters("Missing 'function' object"))?;
333        let name = f
334            .get("name")
335            .and_then(|v| v.as_str())
336            .ok_or_else(|| error_context().invalid_parameters("Missing function.name"))?
337            .to_string();
338        let desc = f
339            .get("description")
340            .and_then(|v| v.as_str())
341            .unwrap_or("")
342            .to_string();
343        let params = f.get("parameters").cloned();
344        (name, desc, params)
345    } else {
346        // Shape 1 inner {name, description, parameters}
347        let name = obj
348            .get("name")
349            .and_then(|v| v.as_str())
350            .ok_or_else(|| error_context().invalid_parameters("Missing name"))?
351            .to_string();
352        let desc = obj
353            .get("description")
354            .and_then(|v| v.as_str())
355            .unwrap_or("")
356            .to_string();
357        let params = obj.get("parameters").cloned();
358        (name, desc, params)
359    };
360    Ok((name, desc, params))
361}
362
363/// Builder for FunctionTool
364pub struct FunctionToolBuilder {
365    metadata: ToolMetadata,
366    input_schema: Option<serde_json::Value>,
367    // Optional staged schema pieces for convenience building when schema() is omitted or for
368    // merging
369    staged_properties: Option<serde_json::Map<String, serde_json::Value>>,
370    staged_required: Vec<String>,
371    handler: Option<ToolHandler>,
372}
373
374impl FunctionToolBuilder {
375    pub fn new(name: impl Into<String>, description: impl Into<String>) -> Self {
376        let name_str = name.into();
377        let desc_str = description.into();
378        let metadata = ToolMetadata::new(&name_str, &desc_str).unwrap_or_else(|e| {
379            warn!(
380                "Invalid tool name '{}': {}. Falling back to 'unknown'.",
381                name_str,
382                e
383            );
384            ToolMetadata {
385                name: Cow::Borrowed("unknown"),
386                description: Cow::Owned(desc_str),
387                version: Cow::Borrowed("1.0.0"),
388                author: None,
389                tags: Vec::new(),
390                enabled: true,
391                metadata: HashMap::new(),
392            }
393        });
394        Self {
395            metadata,
396            input_schema: None,
397            staged_properties: None,
398            staged_required: Vec::new(),
399            handler: None,
400        }
401    }
402
403    pub fn schema(mut self, schema: serde_json::Value) -> Self {
404        self.input_schema = Some(schema);
405        self
406    }
407
408    pub fn metadata(mut self, f: impl FnOnce(ToolMetadata) -> ToolMetadata) -> Self {
409        self.metadata = f(self.metadata);
410        self
411    }
412
413    pub fn handler<F, Fut>(mut self, f: F) -> Self
414    where
415        F: Fn(serde_json::Value) -> Fut + Send + Sync + 'static,
416        Fut: std::future::Future<Output = crate::toolkits::error::ToolResult<serde_json::Value>>
417            + Send
418            + 'static,
419    {
420        let wrapped = move |args: serde_json::Value| -> std::pin::Pin<
421            Box<
422                dyn std::future::Future<
423                        Output = crate::toolkits::error::ToolResult<serde_json::Value>,
424                    > + Send,
425            >,
426        > { Box::pin(f(args)) };
427        self.handler = Some(std::sync::Arc::new(wrapped));
428        self
429    }
430
431    /// Chain API: add one property to the schema. If `schema(json!(...))` is
432    /// also provided, the property will be merged into its `properties`
433    /// object.
434    pub fn property(mut self, name: impl Into<String>, schema: serde_json::Value) -> Self {
435        let name = name.into();
436        let entry = self
437            .staged_properties
438            .get_or_insert_with(serde_json::Map::new);
439        entry.insert(name, schema);
440        self
441    }
442
443    /// Chain API: mark a property as required. Will be merged with any provided
444    /// schema's `required`.
445    pub fn required(mut self, name: impl Into<String>) -> Self {
446        self.staged_required.push(name.into());
447        self
448    }
449
450    pub fn build(mut self) -> crate::toolkits::error::ToolResult<FunctionTool> {
451        let handler = self
452            .handler
453            .ok_or_else(|| error_context().invalid_parameters("FunctionTool handler not set"))?;
454        // Start with provided schema or an empty object to fill
455        let mut schema = self
456            .input_schema
457            .take()
458            .unwrap_or_else(|| serde_json::json!({}));
459
460        // If schema is an object, we can augment it; otherwise leave it as-is
461        if let serde_json::Value::Object(ref mut obj) = schema {
462            // Ensure required base shape
463            obj.entry("type")
464                .or_insert(serde_json::Value::String("object".to_string()));
465            obj.entry("additionalProperties")
466                .or_insert(serde_json::Value::Bool(false));
467
468            // Merge staged properties (if any)
469            if let Some(staged) = self.staged_properties.take() {
470                let props = obj
471                    .entry("properties")
472                    .or_insert_with(|| serde_json::Value::Object(serde_json::Map::new()));
473                if let serde_json::Value::Object(props_obj) = props {
474                    for (k, v) in staged {
475                        props_obj.insert(k, v);
476                    }
477                }
478            }
479            // Merge staged required (if any), de-duplicated
480            if !self.staged_required.is_empty() {
481                use std::collections::BTreeSet;
482                let mut set: BTreeSet<String> = obj
483                    .get("required")
484                    .and_then(|v| v.as_array())
485                    .map(|arr| {
486                        arr.iter()
487                            .filter_map(|v| v.as_str().map(|s| s.to_string()))
488                            .collect()
489                    })
490                    .unwrap_or_default();
491                for r in self.staged_required.into_iter() {
492                    set.insert(r);
493                }
494                obj.insert(
495                    "required".to_string(),
496                    serde_json::Value::Array(
497                        set.into_iter().map(serde_json::Value::String).collect(),
498                    ),
499                );
500            }
501        } else {
502            // If schema is not an object and also not provided, enforce default
503            // But since we only hit here when schema is not an object (provided
504            // by user), we leave it.
505        }
506
507        // If user provided nothing and schema is empty object, ensure defaults
508        if let serde_json::Value::Object(ref mut obj) = schema {
509            obj.entry("type")
510                .or_insert(serde_json::Value::String("object".to_string()));
511            obj.entry("additionalProperties")
512                .or_insert(serde_json::Value::Bool(false));
513            // Ensure properties exists when we staged some but merging didn't set (edge
514            // case)
515            if obj.get("properties").is_none() {
516                obj.insert(
517                    "properties".to_string(),
518                    serde_json::Value::Object(serde_json::Map::new()),
519                );
520            }
521        }
522
523        let compiled_schema = compile_schema_cached(&schema).map_err(|e| {
524            error_context()
525                .with_tool(self.metadata.name.clone())
526                .schema_validation(format!("Failed to compile schema: {}", e))
527        })?;
528
529        Ok(FunctionTool {
530            metadata: self.metadata,
531            input_schema: schema,
532            compiled_schema,
533            handler,
534        })
535    }
536}
537
538#[async_trait]
539impl DynTool for FunctionTool {
540    fn metadata(&self) -> &ToolMetadata {
541        &self.metadata
542    }
543
544    async fn execute_json(&self, input: serde_json::Value) -> ToolResult<serde_json::Value> {
545        // Validate the input against the compiled schema
546        if let Err(validation_error) = self.compiled_schema.validate(&input) {
547            return Err(error_context()
548                .with_tool(self.name())
549                .invalid_parameters(format!("Input validation failed: {}", validation_error)));
550        }
551
552        // If validation passes, execute the handler
553        (self.handler)(input).await
554    }
555
556    fn input_schema(&self) -> serde_json::Value {
557        self.input_schema.clone()
558    }
559
560    fn clone_box(&self) -> Box<dyn DynTool> {
561        Box::new(self.clone())
562    }
563}
564
565#[cfg(test)]
566mod tests {
567    use super::*;
568    use crate::toolkits::ToolError;
569
570    #[test]
571    fn test_tool_metadata_new() {
572        let metadata = ToolMetadata::new("test_tool", "A test tool").unwrap();
573        assert_eq!(metadata.name, "test_tool");
574        assert_eq!(metadata.description, "A test tool");
575        assert_eq!(metadata.version, "1.0.0");
576        assert!(metadata.enabled);
577    }
578
579    #[test]
580    fn test_tool_metadata_invalid_name_empty() {
581        let result = ToolMetadata::new("", "A test tool");
582        assert!(result.is_err());
583        match result.unwrap_err() {
584            ToolError::InvalidParameters { .. } => {},
585            _ => panic!("Expected InvalidParameters error"),
586        }
587    }
588
589    #[test]
590    fn test_tool_metadata_invalid_name_special_chars() {
591        let result = ToolMetadata::new("test-tool!", "A test tool");
592        assert!(result.is_err());
593        match result.unwrap_err() {
594            ToolError::InvalidParameters { .. } => {},
595            _ => panic!("Expected InvalidParameters error"),
596        }
597    }
598
599    #[test]
600    fn test_tool_metadata_builder() {
601        let metadata = ToolMetadata::new("test_tool", "A test tool")
602            .unwrap()
603            .version("2.0.0")
604            .author("Test Author")
605            .tags(vec!["tag1", "tag2"])
606            .enabled(false);
607
608        assert_eq!(metadata.version, "2.0.0");
609        assert_eq!(metadata.author, Some(Cow::Borrowed("Test Author")));
610        assert_eq!(metadata.tags.len(), 2);
611        assert!(!metadata.enabled);
612    }
613
614    #[test]
615    fn test_conversions_to_json() {
616        let value = conversions::to_json(42).unwrap();
617        assert_eq!(value, 42);
618    }
619
620    #[test]
621    fn test_conversions_from_json_string() {
622        let value = serde_json::Value::String("hello".to_string());
623        let result = conversions::from_json_string(value).unwrap();
624        assert_eq!(result, "hello");
625    }
626
627    #[test]
628    fn test_conversions_from_json_string_invalid() {
629        let value = serde_json::Value::Number(42.into());
630        let result = conversions::from_json_string(value);
631        assert!(result.is_err());
632    }
633
634    #[test]
635    fn test_conversions_from_json_i32() {
636        let value = serde_json::Value::Number(42.into());
637        let result = conversions::from_json_i32(value).unwrap();
638        assert_eq!(result, 42);
639    }
640
641    #[test]
642    fn test_conversions_from_json_f64() {
643        let value = serde_json::json!(3.5);
644        let result = conversions::from_json_f64(value).unwrap();
645        assert_eq!(result, 3.5);
646    }
647
648    #[test]
649    fn test_conversions_from_json_bool() {
650        let value = serde_json::Value::Bool(true);
651        let result = conversions::from_json_bool(value).unwrap();
652        assert!(result);
653    }
654
655    #[test]
656    fn test_function_tool_builder() {
657        let tool = FunctionTool::builder("test_tool", "A test tool")
658            .property("param1", serde_json::json!({"type": "string"}))
659            .property("param2", serde_json::json!({"type": "number"}))
660            .required("param1")
661            .handler(|_args| async move { Ok(serde_json::json!({"result": "ok"})) })
662            .build();
663
664        assert!(tool.is_ok());
665        let tool = tool.unwrap();
666        assert_eq!(tool.name(), "test_tool");
667    }
668
669    #[test]
670    fn test_function_tool_clone() {
671        let tool1 = FunctionTool::builder("test_tool", "A test tool")
672            .property("param1", serde_json::json!({"type": "string"}))
673            .required("param1")
674            .handler(|_args| async move { Ok(serde_json::json!({"result": "ok"})) })
675            .build()
676            .unwrap();
677
678        let tool2 = tool1.clone();
679        assert_eq!(tool1.name(), tool2.name());
680        assert_eq!(tool1.input_schema(), tool2.input_schema());
681    }
682
683    #[test]
684    fn test_parse_function_spec_shape1() {
685        let spec = serde_json::json!({
686            "name": "test_tool",
687            "description": "A test tool",
688            "parameters": {
689                "type": "object",
690                "properties": {
691                    "param1": {"type": "string"}
692                }
693            }
694        });
695
696        let (name, description, parameters) = parse_function_spec_details(&spec).unwrap();
697        assert_eq!(name, "test_tool");
698        assert_eq!(description, "A test tool");
699        assert!(parameters.is_some());
700    }
701
702    #[test]
703    fn test_parse_function_spec_shape2() {
704        let spec = serde_json::json!({
705            "type": "function",
706            "function": {
707                "name": "test_tool",
708                "description": "A test tool",
709                "parameters": {
710                    "type": "object",
711                    "properties": {
712                        "param1": {"type": "string"}
713                    }
714                }
715            }
716        });
717
718        let (name, description, parameters) = parse_function_spec_details(&spec).unwrap();
719        assert_eq!(name, "test_tool");
720        assert_eq!(description, "A test tool");
721        assert!(parameters.is_some());
722    }
723
724    #[test]
725    fn test_parse_function_spec_invalid() {
726        let spec = serde_json::Value::String("invalid".to_string());
727        let result = parse_function_spec_details(&spec);
728        assert!(result.is_err());
729    }
730}
zai_rs/toolkits/core.rs

zai_rs/toolkits/
core.rs
