109 lines
2.9 KiB
Rust
109 lines
2.9 KiB
Rust
use std::collections::HashMap;
|
|
|
|
use anyhow::{Context, Result};
|
|
use jsonschema::{JSONSchema, ValidationError};
|
|
use serde_json::{Value, json};
|
|
|
|
pub struct SchemaValidator {
|
|
schemas: HashMap<String, JSONSchema>,
|
|
}
|
|
|
|
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<String> = 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<String, Value> {
|
|
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
|
|
}
|