use std::collections::HashMap; use anyhow::{Context, Result}; use jsonschema::{JSONSchema, ValidationError}; use serde_json::{json, Value}; pub struct SchemaValidator { schemas: HashMap, } impl Default for SchemaValidator { fn default() -> Self { Self::new() } } impl SchemaValidator { pub fn new() -> Self { Self { schemas: HashMap::new(), } } pub fn register_schema(&mut self, tool_name: &str, schema: Value) -> Result<()> { let compiled = JSONSchema::compile(&schema) .map_err(|e| anyhow::anyhow!("Invalid schema for {}: {}", tool_name, e))?; self.schemas.insert(tool_name.to_string(), compiled); Ok(()) } pub fn validate(&self, tool_name: &str, input: &Value) -> Result<()> { let schema = self .schemas .get(tool_name) .with_context(|| format!("No schema registered for tool: {}", tool_name))?; if let Err(errors) = schema.validate(input) { let error_messages: Vec = errors.map(format_validation_error).collect(); return Err(anyhow::anyhow!( "Input validation failed for {}: {}", tool_name, error_messages.join(", ") )); } Ok(()) } } fn format_validation_error(error: ValidationError) -> String { format!("Validation error at {}: {}", error.instance_path, error) } pub fn get_builtin_schemas() -> HashMap { let mut schemas = HashMap::new(); schemas.insert( "web_search".to_string(), json!({ "type": "object", "properties": { "query": { "type": "string", "minLength": 1, "maxLength": 500 }, "max_results": { "type": "integer", "minimum": 1, "maximum": 10, "default": 5 } }, "required": ["query"], "additionalProperties": false }), ); schemas.insert( "code_exec".to_string(), json!({ "type": "object", "properties": { "language": { "type": "string", "enum": ["python", "javascript", "bash", "rust"] }, "code": { "type": "string", "minLength": 1, "maxLength": 10000 }, "timeout": { "type": "integer", "minimum": 1, "maximum": 300, "default": 30 } }, "required": ["language", "code"], "additionalProperties": false }), ); schemas }