feat: add comfyui sdk
This commit is contained in:
216
cargos/comfyui-sdk/utils/error_handling.rs
Normal file
216
cargos/comfyui-sdk/utils/error_handling.rs
Normal file
@@ -0,0 +1,216 @@
|
||||
//! Error handling utilities
|
||||
|
||||
use crate::error::{ComfyUIError, Result};
|
||||
use std::time::Duration;
|
||||
use tokio::time::sleep;
|
||||
|
||||
/// Retry configuration
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RetryConfig {
|
||||
pub max_attempts: u32,
|
||||
pub initial_delay: Duration,
|
||||
pub max_delay: Duration,
|
||||
pub backoff_multiplier: f64,
|
||||
}
|
||||
|
||||
impl Default for RetryConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_attempts: 3,
|
||||
initial_delay: Duration::from_millis(1000),
|
||||
max_delay: Duration::from_secs(30),
|
||||
backoff_multiplier: 2.0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Retries an async operation with exponential backoff
|
||||
pub async fn retry_with_backoff<F, Fut, T>(
|
||||
operation: F,
|
||||
config: RetryConfig,
|
||||
) -> Result<T>
|
||||
where
|
||||
F: Fn() -> Fut,
|
||||
Fut: std::future::Future<Output = Result<T>>,
|
||||
{
|
||||
let mut last_error = None;
|
||||
let mut delay = config.initial_delay;
|
||||
|
||||
for attempt in 1..=config.max_attempts {
|
||||
match operation().await {
|
||||
Ok(result) => return Ok(result),
|
||||
Err(error) => {
|
||||
last_error = Some(error);
|
||||
|
||||
if attempt < config.max_attempts {
|
||||
log::warn!("Attempt {} failed, retrying in {:?}", attempt, delay);
|
||||
sleep(delay).await;
|
||||
|
||||
// Exponential backoff
|
||||
delay = std::cmp::min(
|
||||
Duration::from_millis(
|
||||
(delay.as_millis() as f64 * config.backoff_multiplier) as u64
|
||||
),
|
||||
config.max_delay,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Err(last_error.unwrap_or_else(|| ComfyUIError::new("Retry failed with no error")))
|
||||
}
|
||||
|
||||
/// Checks if an error is retryable
|
||||
pub fn is_retryable_error(error: &ComfyUIError) -> bool {
|
||||
match error {
|
||||
ComfyUIError::Http(reqwest_error) => {
|
||||
// Retry on network errors, timeouts, and 5xx status codes
|
||||
reqwest_error.is_timeout() ||
|
||||
reqwest_error.is_connect() ||
|
||||
reqwest_error.status().map_or(false, |status| status.is_server_error())
|
||||
}
|
||||
ComfyUIError::WebSocket(_) => true, // Most WebSocket errors are retryable
|
||||
ComfyUIError::Connection(_) => true,
|
||||
ComfyUIError::Timeout(_) => true,
|
||||
ComfyUIError::Io(_) => true,
|
||||
_ => false, // Don't retry validation errors, etc.
|
||||
}
|
||||
}
|
||||
|
||||
/// Retries an operation only if the error is retryable
|
||||
pub async fn retry_if_retryable<F, Fut, T>(
|
||||
operation: F,
|
||||
config: RetryConfig,
|
||||
) -> Result<T>
|
||||
where
|
||||
F: Fn() -> Fut,
|
||||
Fut: std::future::Future<Output = Result<T>>,
|
||||
{
|
||||
let mut last_error = None;
|
||||
let mut delay = config.initial_delay;
|
||||
|
||||
for attempt in 1..=config.max_attempts {
|
||||
match operation().await {
|
||||
Ok(result) => return Ok(result),
|
||||
Err(error) => {
|
||||
if !is_retryable_error(&error) {
|
||||
return Err(error);
|
||||
}
|
||||
|
||||
last_error = Some(error);
|
||||
|
||||
if attempt < config.max_attempts {
|
||||
log::warn!("Attempt {} failed with retryable error, retrying in {:?}", attempt, delay);
|
||||
sleep(delay).await;
|
||||
|
||||
// Exponential backoff
|
||||
delay = std::cmp::min(
|
||||
Duration::from_millis(
|
||||
(delay.as_millis() as f64 * config.backoff_multiplier) as u64
|
||||
),
|
||||
config.max_delay,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Err(last_error.unwrap_or_else(|| ComfyUIError::new("Retry failed with no error")))
|
||||
}
|
||||
|
||||
/// Wraps an operation with timeout
|
||||
pub async fn with_timeout<F, T>(
|
||||
operation: F,
|
||||
timeout: Duration,
|
||||
) -> Result<T>
|
||||
where
|
||||
F: std::future::Future<Output = Result<T>>,
|
||||
{
|
||||
match tokio::time::timeout(timeout, operation).await {
|
||||
Ok(result) => result,
|
||||
Err(_) => Err(ComfyUIError::timeout(format!(
|
||||
"Operation timed out after {:?}",
|
||||
timeout
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
/// Error context helper
|
||||
pub trait ErrorContext<T> {
|
||||
fn with_context(self, context: &str) -> Result<T>;
|
||||
}
|
||||
|
||||
impl<T> ErrorContext<T> for Result<T> {
|
||||
fn with_context(self, context: &str) -> Result<T> {
|
||||
self.map_err(|error| {
|
||||
ComfyUIError::new(format!("{}: {}", context, error))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::sync::atomic::{AtomicU32, Ordering};
|
||||
use std::sync::Arc;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_retry_success_on_second_attempt() {
|
||||
let counter = Arc::new(AtomicU32::new(0));
|
||||
let counter_clone = counter.clone();
|
||||
|
||||
let config = RetryConfig {
|
||||
max_attempts: 3,
|
||||
initial_delay: Duration::from_millis(10),
|
||||
max_delay: Duration::from_millis(100),
|
||||
backoff_multiplier: 2.0,
|
||||
};
|
||||
|
||||
let result = retry_with_backoff(
|
||||
|| {
|
||||
let counter = counter_clone.clone();
|
||||
async move {
|
||||
let count = counter.fetch_add(1, Ordering::SeqCst);
|
||||
if count == 0 {
|
||||
Err(ComfyUIError::connection("First attempt fails"))
|
||||
} else {
|
||||
Ok("Success")
|
||||
}
|
||||
}
|
||||
},
|
||||
config,
|
||||
).await;
|
||||
|
||||
assert!(result.is_ok());
|
||||
assert_eq!(result.unwrap(), "Success");
|
||||
assert_eq!(counter.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_retry_exhausted() {
|
||||
let counter = Arc::new(AtomicU32::new(0));
|
||||
let counter_clone = counter.clone();
|
||||
|
||||
let config = RetryConfig {
|
||||
max_attempts: 2,
|
||||
initial_delay: Duration::from_millis(10),
|
||||
max_delay: Duration::from_millis(100),
|
||||
backoff_multiplier: 2.0,
|
||||
};
|
||||
|
||||
let result: Result<&str> = retry_with_backoff(
|
||||
|| {
|
||||
let counter = counter_clone.clone();
|
||||
async move {
|
||||
counter.fetch_add(1, Ordering::SeqCst);
|
||||
Err(ComfyUIError::connection("Always fails"))
|
||||
}
|
||||
},
|
||||
config,
|
||||
).await;
|
||||
|
||||
assert!(result.is_err());
|
||||
assert_eq!(counter.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
}
|
||||
171
cargos/comfyui-sdk/utils/event_emitter.rs
Normal file
171
cargos/comfyui-sdk/utils/event_emitter.rs
Normal file
@@ -0,0 +1,171 @@
|
||||
//! Event emitter utilities
|
||||
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
use std::collections::HashMap;
|
||||
use crate::types::{ExecutionProgress, ExecutionResult, ExecutionError, ExecutionCallbacks};
|
||||
|
||||
/// Event emitter for execution events
|
||||
#[derive(Clone)]
|
||||
pub struct EventEmitter {
|
||||
callbacks: Arc<RwLock<HashMap<String, Arc<dyn ExecutionCallbacks>>>>,
|
||||
}
|
||||
|
||||
impl EventEmitter {
|
||||
/// Creates a new event emitter
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
callbacks: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Registers callbacks with a unique ID
|
||||
pub async fn register_callbacks(&self, id: String, callbacks: Arc<dyn ExecutionCallbacks>) {
|
||||
let mut callbacks_map = self.callbacks.write().await;
|
||||
callbacks_map.insert(id, callbacks);
|
||||
}
|
||||
|
||||
/// Unregisters callbacks by ID
|
||||
pub async fn unregister_callbacks(&self, id: &str) {
|
||||
let mut callbacks_map = self.callbacks.write().await;
|
||||
callbacks_map.remove(id);
|
||||
}
|
||||
|
||||
/// Emits a progress event
|
||||
pub async fn emit_progress(&self, progress: ExecutionProgress) {
|
||||
let callbacks_map = self.callbacks.read().await;
|
||||
for callback in callbacks_map.values() {
|
||||
callback.on_progress(progress.clone());
|
||||
}
|
||||
}
|
||||
|
||||
/// Emits an executing event
|
||||
pub async fn emit_executing(&self, node_id: String) {
|
||||
let callbacks_map = self.callbacks.read().await;
|
||||
for callback in callbacks_map.values() {
|
||||
callback.on_executing(node_id.clone());
|
||||
}
|
||||
}
|
||||
|
||||
/// Emits an executed event
|
||||
pub async fn emit_executed(&self, result: ExecutionResult) {
|
||||
let callbacks_map = self.callbacks.read().await;
|
||||
for callback in callbacks_map.values() {
|
||||
callback.on_executed(result.clone());
|
||||
}
|
||||
}
|
||||
|
||||
/// Emits an error event
|
||||
pub async fn emit_error(&self, error: ExecutionError) {
|
||||
let callbacks_map = self.callbacks.read().await;
|
||||
for callback in callbacks_map.values() {
|
||||
callback.on_error(error.clone());
|
||||
}
|
||||
}
|
||||
|
||||
/// Gets the number of registered callbacks
|
||||
pub async fn callback_count(&self) -> usize {
|
||||
let callbacks_map = self.callbacks.read().await;
|
||||
callbacks_map.len()
|
||||
}
|
||||
|
||||
/// Clears all callbacks
|
||||
pub async fn clear(&self) {
|
||||
let mut callbacks_map = self.callbacks.write().await;
|
||||
callbacks_map.clear();
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for EventEmitter {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Simple callback implementation for testing
|
||||
pub struct SimpleCallbacks {
|
||||
pub on_progress_fn: Option<Box<dyn Fn(ExecutionProgress) + Send + Sync>>,
|
||||
pub on_executing_fn: Option<Box<dyn Fn(String) + Send + Sync>>,
|
||||
pub on_executed_fn: Option<Box<dyn Fn(ExecutionResult) + Send + Sync>>,
|
||||
pub on_error_fn: Option<Box<dyn Fn(ExecutionError) + Send + Sync>>,
|
||||
}
|
||||
|
||||
impl SimpleCallbacks {
|
||||
/// Creates new simple callbacks
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
on_progress_fn: None,
|
||||
on_executing_fn: None,
|
||||
on_executed_fn: None,
|
||||
on_error_fn: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Sets the progress callback
|
||||
pub fn with_progress<F>(mut self, f: F) -> Self
|
||||
where
|
||||
F: Fn(ExecutionProgress) + Send + Sync + 'static,
|
||||
{
|
||||
self.on_progress_fn = Some(Box::new(f));
|
||||
self
|
||||
}
|
||||
|
||||
/// Sets the executing callback
|
||||
pub fn with_executing<F>(mut self, f: F) -> Self
|
||||
where
|
||||
F: Fn(String) + Send + Sync + 'static,
|
||||
{
|
||||
self.on_executing_fn = Some(Box::new(f));
|
||||
self
|
||||
}
|
||||
|
||||
/// Sets the executed callback
|
||||
pub fn with_executed<F>(mut self, f: F) -> Self
|
||||
where
|
||||
F: Fn(ExecutionResult) + Send + Sync + 'static,
|
||||
{
|
||||
self.on_executed_fn = Some(Box::new(f));
|
||||
self
|
||||
}
|
||||
|
||||
/// Sets the error callback
|
||||
pub fn with_error<F>(mut self, f: F) -> Self
|
||||
where
|
||||
F: Fn(ExecutionError) + Send + Sync + 'static,
|
||||
{
|
||||
self.on_error_fn = Some(Box::new(f));
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl ExecutionCallbacks for SimpleCallbacks {
|
||||
fn on_progress(&self, progress: ExecutionProgress) {
|
||||
if let Some(ref f) = self.on_progress_fn {
|
||||
f(progress);
|
||||
}
|
||||
}
|
||||
|
||||
fn on_executing(&self, node_id: String) {
|
||||
if let Some(ref f) = self.on_executing_fn {
|
||||
f(node_id);
|
||||
}
|
||||
}
|
||||
|
||||
fn on_executed(&self, result: ExecutionResult) {
|
||||
if let Some(ref f) = self.on_executed_fn {
|
||||
f(result);
|
||||
}
|
||||
}
|
||||
|
||||
fn on_error(&self, error: ExecutionError) {
|
||||
if let Some(ref f) = self.on_error_fn {
|
||||
f(error);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for SimpleCallbacks {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
12
cargos/comfyui-sdk/utils/mod.rs
Normal file
12
cargos/comfyui-sdk/utils/mod.rs
Normal file
@@ -0,0 +1,12 @@
|
||||
//! Utility functions for ComfyUI SDK
|
||||
|
||||
pub mod validation;
|
||||
pub mod template_parser;
|
||||
pub mod event_emitter;
|
||||
pub mod error_handling;
|
||||
|
||||
// Re-export main utilities
|
||||
pub use validation::*;
|
||||
pub use template_parser::*;
|
||||
pub use event_emitter::*;
|
||||
pub use error_handling::*;
|
||||
256
cargos/comfyui-sdk/utils/template_parser.rs
Normal file
256
cargos/comfyui-sdk/utils/template_parser.rs
Normal file
@@ -0,0 +1,256 @@
|
||||
//! Template parsing utilities for parameter substitution
|
||||
|
||||
use std::collections::HashMap;
|
||||
use regex::Regex;
|
||||
use crate::types::{ComfyUIWorkflow, ParameterValues, ComfyUINode, VariableSyntax};
|
||||
use crate::error::{ComfyUIError, Result};
|
||||
|
||||
/// Applies parameters to a workflow template
|
||||
pub fn apply_parameters(
|
||||
workflow: &ComfyUIWorkflow,
|
||||
parameters: &ParameterValues,
|
||||
) -> Result<ComfyUIWorkflow> {
|
||||
let mut resolved_workflow = HashMap::new();
|
||||
|
||||
for (node_id, node) in workflow {
|
||||
let resolved_node = apply_parameters_to_node(node, parameters)?;
|
||||
resolved_workflow.insert(node_id.clone(), resolved_node);
|
||||
}
|
||||
|
||||
Ok(resolved_workflow)
|
||||
}
|
||||
|
||||
/// Applies parameters to a single node
|
||||
fn apply_parameters_to_node(
|
||||
node: &ComfyUINode,
|
||||
parameters: &ParameterValues,
|
||||
) -> Result<ComfyUINode> {
|
||||
let mut resolved_inputs = HashMap::new();
|
||||
|
||||
for (input_name, input_value) in &node.inputs {
|
||||
let resolved_value = apply_parameters_to_value(input_value, parameters)?;
|
||||
resolved_inputs.insert(input_name.clone(), resolved_value);
|
||||
}
|
||||
|
||||
Ok(ComfyUINode {
|
||||
class_type: node.class_type.clone(),
|
||||
inputs: resolved_inputs,
|
||||
_meta: node._meta.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Applies parameters to a JSON value
|
||||
fn apply_parameters_to_value(
|
||||
value: &serde_json::Value,
|
||||
parameters: &ParameterValues,
|
||||
) -> Result<serde_json::Value> {
|
||||
match value {
|
||||
serde_json::Value::String(s) => {
|
||||
let resolved_string = substitute_variables(s, parameters, VariableSyntax::DoubleBrace)?;
|
||||
Ok(serde_json::Value::String(resolved_string))
|
||||
}
|
||||
serde_json::Value::Array(arr) => {
|
||||
let mut resolved_array = Vec::new();
|
||||
for item in arr {
|
||||
resolved_array.push(apply_parameters_to_value(item, parameters)?);
|
||||
}
|
||||
Ok(serde_json::Value::Array(resolved_array))
|
||||
}
|
||||
serde_json::Value::Object(obj) => {
|
||||
let mut resolved_object = serde_json::Map::new();
|
||||
for (key, val) in obj {
|
||||
let resolved_key = substitute_variables(key, parameters, VariableSyntax::DoubleBrace)?;
|
||||
let resolved_val = apply_parameters_to_value(val, parameters)?;
|
||||
resolved_object.insert(resolved_key, resolved_val);
|
||||
}
|
||||
Ok(serde_json::Value::Object(resolved_object))
|
||||
}
|
||||
_ => Ok(value.clone()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Substitutes variables in a string using the specified syntax
|
||||
pub fn substitute_variables(
|
||||
template: &str,
|
||||
parameters: &ParameterValues,
|
||||
syntax: VariableSyntax,
|
||||
) -> Result<String> {
|
||||
let pattern = match syntax {
|
||||
VariableSyntax::DoubleBrace => r"\{\{([^}]+)\}\}",
|
||||
VariableSyntax::DollarBrace => r"\$\{([^}]+)\}",
|
||||
VariableSyntax::AtBrace => r"@\{([^}]+)\}",
|
||||
};
|
||||
|
||||
let regex = Regex::new(pattern)
|
||||
.map_err(|e| ComfyUIError::new(format!("Invalid regex pattern: {}", e)))?;
|
||||
|
||||
let mut result = template.to_string();
|
||||
let mut offset = 0i32;
|
||||
|
||||
for captures in regex.captures_iter(template) {
|
||||
let full_match = captures.get(0).unwrap();
|
||||
let var_name = captures.get(1).unwrap().as_str().trim();
|
||||
|
||||
// Get parameter value
|
||||
let replacement = match parameters.get(var_name) {
|
||||
Some(value) => value_to_string(value)?,
|
||||
None => {
|
||||
return Err(ComfyUIError::template_validation(
|
||||
format!("Parameter '{}' not found", var_name)
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
// Calculate positions with offset
|
||||
let start = (full_match.start() as i32 + offset) as usize;
|
||||
let end = (full_match.end() as i32 + offset) as usize;
|
||||
|
||||
// Replace the variable
|
||||
result.replace_range(start..end, &replacement);
|
||||
|
||||
// Update offset
|
||||
offset += replacement.len() as i32 - full_match.len() as i32;
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Converts a JSON value to string for substitution
|
||||
fn value_to_string(value: &serde_json::Value) -> Result<String> {
|
||||
match value {
|
||||
serde_json::Value::String(s) => Ok(s.clone()),
|
||||
serde_json::Value::Number(n) => Ok(n.to_string()),
|
||||
serde_json::Value::Bool(b) => Ok(b.to_string()),
|
||||
serde_json::Value::Null => Ok("null".to_string()),
|
||||
_ => {
|
||||
// For complex types, serialize to JSON
|
||||
serde_json::to_string(value)
|
||||
.map_err(ComfyUIError::from)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Extracts variable names from a template string
|
||||
pub fn extract_variables(template: &str, syntax: VariableSyntax) -> Result<Vec<String>> {
|
||||
let pattern = match syntax {
|
||||
VariableSyntax::DoubleBrace => r"\{\{([^}]+)\}\}",
|
||||
VariableSyntax::DollarBrace => r"\$\{([^}]+)\}",
|
||||
VariableSyntax::AtBrace => r"@\{([^}]+)\}",
|
||||
};
|
||||
|
||||
let regex = Regex::new(pattern)
|
||||
.map_err(|e| ComfyUIError::new(format!("Invalid regex pattern: {}", e)))?;
|
||||
|
||||
let mut variables = Vec::new();
|
||||
for captures in regex.captures_iter(template) {
|
||||
if let Some(var_match) = captures.get(1) {
|
||||
let var_name = var_match.as_str().trim().to_string();
|
||||
if !variables.contains(&var_name) {
|
||||
variables.push(var_name);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(variables)
|
||||
}
|
||||
|
||||
/// Extracts all variables from a workflow
|
||||
pub fn extract_workflow_variables(workflow: &ComfyUIWorkflow) -> Result<Vec<String>> {
|
||||
let mut all_variables = Vec::new();
|
||||
|
||||
for node in workflow.values() {
|
||||
let node_variables = extract_node_variables(node)?;
|
||||
for var in node_variables {
|
||||
if !all_variables.contains(&var) {
|
||||
all_variables.push(var);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(all_variables)
|
||||
}
|
||||
|
||||
/// Extracts variables from a single node
|
||||
fn extract_node_variables(node: &ComfyUINode) -> Result<Vec<String>> {
|
||||
let mut variables = Vec::new();
|
||||
|
||||
for input_value in node.inputs.values() {
|
||||
let value_variables = extract_value_variables(input_value)?;
|
||||
for var in value_variables {
|
||||
if !variables.contains(&var) {
|
||||
variables.push(var);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(variables)
|
||||
}
|
||||
|
||||
/// Extracts variables from a JSON value
|
||||
fn extract_value_variables(value: &serde_json::Value) -> Result<Vec<String>> {
|
||||
let mut variables = Vec::new();
|
||||
|
||||
match value {
|
||||
serde_json::Value::String(s) => {
|
||||
let string_vars = extract_variables(s, VariableSyntax::DoubleBrace)?;
|
||||
variables.extend(string_vars);
|
||||
}
|
||||
serde_json::Value::Array(arr) => {
|
||||
for item in arr {
|
||||
let item_vars = extract_value_variables(item)?;
|
||||
variables.extend(item_vars);
|
||||
}
|
||||
}
|
||||
serde_json::Value::Object(obj) => {
|
||||
for (key, val) in obj {
|
||||
let key_vars = extract_variables(key, VariableSyntax::DoubleBrace)?;
|
||||
variables.extend(key_vars);
|
||||
|
||||
let val_vars = extract_value_variables(val)?;
|
||||
variables.extend(val_vars);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
|
||||
// Remove duplicates
|
||||
variables.sort();
|
||||
variables.dedup();
|
||||
|
||||
Ok(variables)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn test_substitute_variables() {
|
||||
let mut parameters = HashMap::new();
|
||||
parameters.insert("name".to_string(), json!("test"));
|
||||
parameters.insert("value".to_string(), json!(42));
|
||||
|
||||
let template = "Hello {{name}}, value is {{value}}";
|
||||
let result = substitute_variables(template, ¶meters, VariableSyntax::DoubleBrace).unwrap();
|
||||
|
||||
assert_eq!(result, "Hello test, value is 42");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_variables() {
|
||||
let template = "{{var1}} and {{var2}} and {{var1}} again";
|
||||
let variables = extract_variables(template, VariableSyntax::DoubleBrace).unwrap();
|
||||
|
||||
assert_eq!(variables, vec!["var1", "var2"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_missing_parameter() {
|
||||
let parameters = HashMap::new();
|
||||
let template = "Hello {{missing}}";
|
||||
|
||||
let result = substitute_variables(template, ¶meters, VariableSyntax::DoubleBrace);
|
||||
assert!(result.is_err());
|
||||
}
|
||||
}
|
||||
276
cargos/comfyui-sdk/utils/validation.rs
Normal file
276
cargos/comfyui-sdk/utils/validation.rs
Normal file
@@ -0,0 +1,276 @@
|
||||
//! Parameter validation utilities
|
||||
|
||||
use std::collections::HashMap;
|
||||
use crate::types::{ParameterSchema, ParameterValues, ValidationResult, ValidationError, ParameterType};
|
||||
use regex::Regex;
|
||||
|
||||
/// Validates parameters against their schemas
|
||||
pub fn validate_parameters(
|
||||
parameters: &ParameterValues,
|
||||
schemas: &HashMap<String, ParameterSchema>,
|
||||
) -> ValidationResult {
|
||||
let mut errors = Vec::new();
|
||||
|
||||
// Check required parameters
|
||||
for (name, schema) in schemas {
|
||||
if schema.required.unwrap_or(false) {
|
||||
if !parameters.contains_key(name) {
|
||||
errors.push(ValidationError::new(
|
||||
name,
|
||||
"Required parameter is missing",
|
||||
));
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
// Validate parameter if present
|
||||
if let Some(value) = parameters.get(name) {
|
||||
if let Err(error) = validate_parameter_value(name, value, schema) {
|
||||
errors.push(error);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check for unknown parameters
|
||||
for name in parameters.keys() {
|
||||
if !schemas.contains_key(name) {
|
||||
errors.push(ValidationError::new(
|
||||
name,
|
||||
"Unknown parameter",
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
if errors.is_empty() {
|
||||
ValidationResult::success()
|
||||
} else {
|
||||
ValidationResult::failure(errors)
|
||||
}
|
||||
}
|
||||
|
||||
/// Validates a single parameter value against its schema
|
||||
pub fn validate_parameter_value(
|
||||
name: &str,
|
||||
value: &serde_json::Value,
|
||||
schema: &ParameterSchema,
|
||||
) -> Result<(), ValidationError> {
|
||||
// Check type
|
||||
match (&schema.param_type, value) {
|
||||
(ParameterType::String, serde_json::Value::String(s)) => {
|
||||
validate_string_constraints(name, s, schema)?;
|
||||
}
|
||||
(ParameterType::Number, serde_json::Value::Number(n)) => {
|
||||
validate_number_constraints(name, n, schema)?;
|
||||
}
|
||||
(ParameterType::Boolean, serde_json::Value::Bool(_)) => {
|
||||
// Boolean values are always valid
|
||||
}
|
||||
(ParameterType::Array, serde_json::Value::Array(arr)) => {
|
||||
validate_array_constraints(name, arr, schema)?;
|
||||
}
|
||||
(ParameterType::Object, serde_json::Value::Object(_)) => {
|
||||
// Object validation could be more complex, but for now just check type
|
||||
}
|
||||
_ => {
|
||||
return Err(ValidationError::with_value(
|
||||
name,
|
||||
format!("Expected type {:?}, got {:?}", schema.param_type, get_value_type(value)),
|
||||
value.clone(),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
// Check enum constraints
|
||||
if let Some(enum_values) = &schema.r#enum {
|
||||
if !enum_values.contains(value) {
|
||||
return Err(ValidationError::with_value(
|
||||
name,
|
||||
format!("Value must be one of: {:?}", enum_values),
|
||||
value.clone(),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Validates string constraints
|
||||
fn validate_string_constraints(
|
||||
name: &str,
|
||||
value: &str,
|
||||
schema: &ParameterSchema,
|
||||
) -> Result<(), ValidationError> {
|
||||
// Check pattern
|
||||
if let Some(pattern) = &schema.pattern {
|
||||
let regex = Regex::new(pattern).map_err(|_| {
|
||||
ValidationError::new(name, "Invalid regex pattern in schema")
|
||||
})?;
|
||||
|
||||
if !regex.is_match(value) {
|
||||
return Err(ValidationError::with_value(
|
||||
name,
|
||||
format!("String does not match pattern: {}", pattern),
|
||||
serde_json::Value::String(value.to_string()),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
// Check length constraints (using min/max as length bounds for strings)
|
||||
if let Some(min) = schema.min {
|
||||
if (value.len() as f64) < min {
|
||||
return Err(ValidationError::with_value(
|
||||
name,
|
||||
format!("String length must be at least {}", min),
|
||||
serde_json::Value::String(value.to_string()),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(max) = schema.max {
|
||||
if (value.len() as f64) > max {
|
||||
return Err(ValidationError::with_value(
|
||||
name,
|
||||
format!("String length must be at most {}", max),
|
||||
serde_json::Value::String(value.to_string()),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Validates number constraints
|
||||
fn validate_number_constraints(
|
||||
name: &str,
|
||||
value: &serde_json::Number,
|
||||
schema: &ParameterSchema,
|
||||
) -> Result<(), ValidationError> {
|
||||
let num_value = value.as_f64().unwrap_or(0.0);
|
||||
|
||||
if let Some(min) = schema.min {
|
||||
if num_value < min {
|
||||
return Err(ValidationError::with_value(
|
||||
name,
|
||||
format!("Number must be at least {}", min),
|
||||
serde_json::Value::Number(value.clone()),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(max) = schema.max {
|
||||
if num_value > max {
|
||||
return Err(ValidationError::with_value(
|
||||
name,
|
||||
format!("Number must be at most {}", max),
|
||||
serde_json::Value::Number(value.clone()),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Validates array constraints
|
||||
fn validate_array_constraints(
|
||||
name: &str,
|
||||
value: &[serde_json::Value],
|
||||
schema: &ParameterSchema,
|
||||
) -> Result<(), ValidationError> {
|
||||
// Check length constraints
|
||||
if let Some(min) = schema.min {
|
||||
if (value.len() as f64) < min {
|
||||
return Err(ValidationError::with_value(
|
||||
name,
|
||||
format!("Array length must be at least {}", min),
|
||||
serde_json::Value::Array(value.to_vec()),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(max) = schema.max {
|
||||
if (value.len() as f64) > max {
|
||||
return Err(ValidationError::with_value(
|
||||
name,
|
||||
format!("Array length must be at most {}", max),
|
||||
serde_json::Value::Array(value.to_vec()),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
// Validate items if schema is provided
|
||||
if let Some(items_schema) = &schema.items {
|
||||
for (index, item) in value.iter().enumerate() {
|
||||
let item_name = format!("{}[{}]", name, index);
|
||||
if let Err(error) = validate_parameter_value(&item_name, item, items_schema) {
|
||||
return Err(error);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Gets the type of a JSON value
|
||||
fn get_value_type(value: &serde_json::Value) -> &'static str {
|
||||
match value {
|
||||
serde_json::Value::Null => "null",
|
||||
serde_json::Value::Bool(_) => "boolean",
|
||||
serde_json::Value::Number(_) => "number",
|
||||
serde_json::Value::String(_) => "string",
|
||||
serde_json::Value::Array(_) => "array",
|
||||
serde_json::Value::Object(_) => "object",
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn test_validate_required_parameter() {
|
||||
let mut schemas = HashMap::new();
|
||||
schemas.insert("test".to_string(), ParameterSchema {
|
||||
param_type: ParameterType::String,
|
||||
required: Some(true),
|
||||
default: None,
|
||||
description: None,
|
||||
r#enum: None,
|
||||
min: None,
|
||||
max: None,
|
||||
pattern: None,
|
||||
items: None,
|
||||
properties: None,
|
||||
});
|
||||
|
||||
let parameters = HashMap::new();
|
||||
let result = validate_parameters(¶meters, &schemas);
|
||||
|
||||
assert!(!result.valid);
|
||||
assert_eq!(result.errors.len(), 1);
|
||||
assert_eq!(result.errors[0].path, "test");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_validate_string_parameter() {
|
||||
let mut schemas = HashMap::new();
|
||||
schemas.insert("test".to_string(), ParameterSchema {
|
||||
param_type: ParameterType::String,
|
||||
required: Some(true),
|
||||
default: None,
|
||||
description: None,
|
||||
r#enum: None,
|
||||
min: Some(3.0),
|
||||
max: Some(10.0),
|
||||
pattern: None,
|
||||
items: None,
|
||||
properties: None,
|
||||
});
|
||||
|
||||
let mut parameters = HashMap::new();
|
||||
parameters.insert("test".to_string(), json!("hello"));
|
||||
|
||||
let result = validate_parameters(¶meters, &schemas);
|
||||
assert!(result.valid);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user