litellm-rs 0.1.1

A high-performance AI Gateway written in Rust, providing OpenAI-compatible APIs with intelligent routing, load balancing, and enterprise features
//! Utility modules for the Gateway
//!
//! This module contains various utility functions and types used throughout the gateway.

pub mod async_utils;
pub mod config;
pub mod cost;
pub mod crypto;
pub mod dependency_injection;
pub mod error;
pub mod error_recovery;
pub mod http_client;
pub mod logging;
pub mod memory_pool;
pub mod rate_limiter;
pub mod result_ext;
pub mod shared_state;
pub mod string_pool;
pub mod structured_logging;
pub mod token_cache;
pub mod token_counter;
pub mod type_utils;
pub mod types;
pub mod validation;

// Re-export commonly used types
pub use token_counter::TokenCounter;

use std::time::{SystemTime, UNIX_EPOCH};
use uuid::Uuid;

/// Generate a unique request ID
#[allow(dead_code)]
pub fn generate_request_id() -> String {
    Uuid::new_v4().to_string()
}

/// Get current timestamp in seconds
#[allow(dead_code)]
pub fn current_timestamp() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
}

/// Get current timestamp in milliseconds
#[allow(dead_code)]
pub fn current_timestamp_millis() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis() as u64
}

/// Format bytes as human readable string
#[allow(dead_code)]
pub fn format_bytes(bytes: u64) -> String {
    const UNITS: &[&str] = &["B", "KB", "MB", "GB", "TB"];
    const THRESHOLD: u64 = 1024;

    if bytes < THRESHOLD {
        return format!("{} B", bytes);
    }

    let mut size = bytes as f64;
    let mut unit_index = 0;

    while size >= THRESHOLD as f64 && unit_index < UNITS.len() - 1 {
        size /= THRESHOLD as f64;
        unit_index += 1;
    }

    format!("{:.1} {}", size, UNITS[unit_index])
}

/// Format duration as human readable string
#[allow(dead_code)]
pub fn format_duration(duration_ms: u64) -> String {
    if duration_ms < 1000 {
        format!("{}ms", duration_ms)
    } else if duration_ms < 60_000 {
        format!("{:.1}s", duration_ms as f64 / 1000.0)
    } else if duration_ms < 3_600_000 {
        format!("{:.1}m", duration_ms as f64 / 60_000.0)
    } else {
        format!("{:.1}h", duration_ms as f64 / 3_600_000.0)
    }
}

/// Sanitize string for logging (remove sensitive information)
#[allow(dead_code)]
pub fn sanitize_for_logging(input: &str) -> String {
    // Replace potential API keys, tokens, etc.
    let patterns = [
        (
            r#"(?i)api[_-]?key["']?\s*[:=]\s*["']?([a-zA-Z0-9\-_]{20,})"#,
            "api_key: [REDACTED]",
        ),
        (
            r#"(?i)token["']?\s*[:=]\s*["']?([a-zA-Z0-9\-_\.]{20,})"#,
            "token: [REDACTED]",
        ),
        (
            r#"(?i)password["']?\s*[:=]\s*["']?([^\s"']{8,})"#,
            "password: [REDACTED]",
        ),
        (
            r#"(?i)secret["']?\s*[:=]\s*["']?([a-zA-Z0-9\-_]{16,})"#,
            "secret: [REDACTED]",
        ),
    ];

    let mut result = input.to_string();
    for (pattern, replacement) in &patterns {
        if let Ok(re) = regex::Regex::new(pattern) {
            result = re.replace_all(&result, *replacement).to_string();
        }
    }

    result
}

/// Truncate string to specified length with ellipsis
#[allow(dead_code)]
pub fn truncate_string(s: &str, max_len: usize) -> String {
    if s.len() <= max_len {
        s.to_string()
    } else {
        format!("{}...", &s[..max_len.saturating_sub(3)])
    }
}

/// Check if a string is a valid URL
#[allow(dead_code)]
pub fn is_valid_url(url: &str) -> bool {
    url::Url::parse(url).is_ok()
}

/// Check if a string is a valid email
#[allow(dead_code)]
pub fn is_valid_email(email: &str) -> bool {
    // Simple email validation regex
    let email_regex =
        regex::Regex::new(r#"^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$"#).unwrap();
    email_regex.is_match(email)
}

/// Normalize model name (remove provider prefix if present)
#[allow(dead_code)]
pub fn normalize_model_name(model: &str) -> String {
    // Remove common provider prefixes
    let prefixes = ["openai/", "anthropic/", "azure/", "google/", "bedrock/"];

    for prefix in &prefixes {
        if let Some(stripped) = model.strip_prefix(prefix) {
            return stripped.to_string();
        }
    }

    model.to_string()
}

/// Extract provider from model name
#[allow(dead_code)]
pub fn extract_provider_from_model(model: &str) -> Option<String> {
    model
        .find('/')
        .map(|slash_pos| model[..slash_pos].to_string())
}

/// Merge two JSON values
#[allow(dead_code)]
pub fn merge_json_values(base: &mut serde_json::Value, overlay: &serde_json::Value) {
    match (base, overlay) {
        (serde_json::Value::Object(base_map), serde_json::Value::Object(overlay_map)) => {
            for (key, value) in overlay_map {
                match base_map.get_mut(key) {
                    Some(base_value) => merge_json_values(base_value, value),
                    None => {
                        base_map.insert(key.clone(), value.clone());
                    }
                }
            }
        }
        (base_val, overlay_val) => *base_val = overlay_val.clone(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_format_bytes() {
        assert_eq!(format_bytes(512), "512 B");
        assert_eq!(format_bytes(1024), "1.0 KB");
        assert_eq!(format_bytes(1536), "1.5 KB");
        assert_eq!(format_bytes(1048576), "1.0 MB");
    }

    #[test]
    fn test_format_duration() {
        assert_eq!(format_duration(500), "500ms");
        assert_eq!(format_duration(1500), "1.5s");
        assert_eq!(format_duration(90000), "1.5m");
        assert_eq!(format_duration(7200000), "2.0h");
    }

    #[test]
    fn test_normalize_model_name() {
        assert_eq!(normalize_model_name("openai/gpt-4"), "gpt-4");
        assert_eq!(normalize_model_name("anthropic/claude-3"), "claude-3");
        assert_eq!(normalize_model_name("gpt-4"), "gpt-4");
    }

    #[test]
    fn test_extract_provider_from_model() {
        assert_eq!(
            extract_provider_from_model("openai/gpt-4"),
            Some("openai".to_string())
        );
        assert_eq!(extract_provider_from_model("gpt-4"), None);
    }

    #[test]
    fn test_truncate_string() {
        assert_eq!(truncate_string("hello", 10), "hello");
        assert_eq!(truncate_string("hello world", 8), "hello...");
    }

    #[test]
    fn test_is_valid_url() {
        assert!(is_valid_url("https://api.openai.com/v1"));
        assert!(is_valid_url("http://localhost:8080"));
        assert!(!is_valid_url("not-a-url"));
    }

    #[test]
    fn test_is_valid_email() {
        assert!(is_valid_email("user@example.com"));
        assert!(is_valid_email("test.email+tag@domain.co.uk"));
        assert!(!is_valid_email("invalid-email"));
        assert!(!is_valid_email("@domain.com"));
    }
}