rustytorch/crates/models/rtx-nlg/src/error.rs

//! Error types for RTX-NLG operations

use std::fmt;

/// Result type for RTX-NLG operations
pub type Result<T> = std::result::Result<T, NlgError>;

/// Errors that can occur during NLG operations
#[derive(thiserror::Error, Debug)]
pub enum NlgError {
    /// Model-related errors
    #[error("Model not found: {0}")]
    ModelNotFound(String),

    #[error("Model loading failed: {0}")]
    ModelLoadError(String),

    #[error("Model inference failed: {0}")]
    InferenceError(String),

    /// Generation errors
    #[error("Generation failed: {0}")]
    GenerationError(String),

    #[error("Invalid generation configuration: {0}")]
    InvalidConfig(String),

    #[error("Generation timeout after {timeout_ms}ms")]
    GenerationTimeout { timeout_ms: u64 },

    /// Tokenization errors
    #[error("Tokenization failed: {0}")]
    TokenizationError(String),

    #[error("Invalid token sequence")]
    InvalidTokenSequence,

    /// Quality control errors
    #[error("Content filtered by quality control: {reason}")]
    ContentFiltered { reason: String },

    #[error("Toxicity detected: {score}")]
    ToxicityDetected { score: f32 },

    #[error("Factuality check failed: {0}")]
    FactualityError(String),

    /// Resource errors
    #[error("Out of memory during generation")]
    OutOfMemory,

    #[error("Device error: {0}")]
    DeviceError(String),

    #[error("Cache error: {0}")]
    CacheError(String),

    /// I/O and serialization errors
    #[error("I/O error: {0}")]
    IoError(#[from] std::io::Error),

    #[error("Serialization error: {0}")]
    SerializationError(#[from] serde_json::Error),

    /// Tensor operations
    #[error("Tensor operation failed: {0}")]
    TensorError(String),

    /// Configuration errors
    #[error("Invalid parameter: {param} = {value}")]
    InvalidParameter { param: String, value: String },

    #[error("Missing required parameter: {0}")]
    MissingParameter(String),

    /// Network and serving errors
    #[error("Network error: {0}")]
    NetworkError(String),

    #[error("Server overloaded: {current_load}/{max_load}")]
    ServerOverloaded {
        current_load: usize,
        max_load: usize,
    },

    #[error("Request timeout")]
    RequestTimeout,

    /// Translation-specific errors
    #[error("Translation failed: {0}")]
    TranslationError(String),

    #[error("Unsupported language pair: {source_lang} to {target_lang}")]
    UnsupportedLanguagePair {
        source_lang: String,
        target_lang: String,
    },

    /// Summarization-specific errors
    #[error("Summarization failed: {0}")]
    SummarizationError(String),

    #[error("Document too short for summarization")]
    DocumentTooShort,

    #[error("Document too long: {length} > {max_length}")]
    DocumentTooLong { length: usize, max_length: usize },

    /// Question answering errors
    #[error("Question answering failed: {0}")]
    QuestionAnsweringError(String),

    #[error("No answer found for question")]
    NoAnswerFound,

    #[error("Context too short for answering")]
    ContextTooShort,

    /// Dialogue system errors
    #[error("Dialogue error: {0}")]
    DialogueError(String),

    #[error("Conversation context too long")]
    ConversationTooLong,

    #[error("Invalid conversation state")]
    InvalidConversationState,

    /// Optimization errors
    #[error("KV cache error: {0}")]
    KvCacheError(String),

    #[error("Speculative decoding failed: {0}")]
    SpeculativeDecodingError(String),

    #[error("Model parallelism error: {0}")]
    ModelParallelismError(String),

    /// Template and batch processing errors
    #[error("Template validation failed for field '{field}': {message}")]
    ValidationFailed { field: String, message: String },

    #[error("Invalid pattern '{pattern}': {error}")]
    InvalidPattern { pattern: String, error: String },

    #[error("Batch queue is full")]
    BatchQueueFull,

    #[error("Batch request timeout")]
    BatchTimeout,

    #[error("Batch request was cancelled")]
    BatchRequestCancelled,

    /// Generic error for wrapping other error types
    #[error("Internal error: {0}")]
    Internal(#[from] anyhow::Error),
}

impl NlgError {
    /// Create a model not found error
    pub fn model_not_found(model_id: impl Into<String>) -> Self {
        Self::ModelNotFound(model_id.into())
    }

    /// Create a generation error
    pub fn generation_error(msg: impl Into<String>) -> Self {
        Self::GenerationError(msg.into())
    }

    /// Create an invalid configuration error
    pub fn invalid_config(msg: impl Into<String>) -> Self {
        Self::InvalidConfig(msg.into())
    }

    /// Create a tokenization error
    pub fn tokenization_error(msg: impl Into<String>) -> Self {
        Self::TokenizationError(msg.into())
    }

    /// Create a content filtered error
    pub fn content_filtered(reason: impl Into<String>) -> Self {
        Self::ContentFiltered {
            reason: reason.into(),
        }
    }

    /// Create a toxicity detected error
    pub fn toxicity_detected(score: f32) -> Self {
        Self::ToxicityDetected { score }
    }

    /// Create a tensor error
    pub fn tensor_error(msg: impl Into<String>) -> Self {
        Self::TensorError(msg.into())
    }

    /// Create an invalid parameter error
    pub fn invalid_parameter(param: impl Into<String>, value: impl Into<String>) -> Self {
        Self::InvalidParameter {
            param: param.into(),
            value: value.into(),
        }
    }

    /// Create a missing parameter error
    pub fn missing_parameter(param: impl Into<String>) -> Self {
        Self::MissingParameter(param.into())
    }

    /// Create a server overloaded error
    pub fn server_overloaded(current_load: usize, max_load: usize) -> Self {
        Self::ServerOverloaded {
            current_load,
            max_load,
        }
    }

    /// Create a translation error
    pub fn translation_error(msg: impl Into<String>) -> Self {
        Self::TranslationError(msg.into())
    }

    /// Create an unsupported language pair error
    pub fn unsupported_language_pair(source: impl Into<String>, target: impl Into<String>) -> Self {
        Self::UnsupportedLanguagePair {
            source_lang: source.into(),
            target_lang: target.into(),
        }
    }

    /// Create a summarization error
    pub fn summarization_error(msg: impl Into<String>) -> Self {
        Self::SummarizationError(msg.into())
    }

    /// Create a document too long error
    pub fn document_too_long(length: usize, max_length: usize) -> Self {
        Self::DocumentTooLong { length, max_length }
    }

    /// Create a question answering error
    pub fn question_answering_error(msg: impl Into<String>) -> Self {
        Self::QuestionAnsweringError(msg.into())
    }

    /// Create a dialogue error
    pub fn dialogue_error(msg: impl Into<String>) -> Self {
        Self::DialogueError(msg.into())
    }

    /// Check if error is recoverable (can retry)
    pub fn is_recoverable(&self) -> bool {
        matches!(
            self,
            Self::GenerationTimeout { .. }
                | Self::OutOfMemory
                | Self::NetworkError(_)
                | Self::ServerOverloaded { .. }
                | Self::RequestTimeout
        )
    }

    /// Check if error is related to content safety
    pub fn is_safety_error(&self) -> bool {
        matches!(
            self,
            Self::ContentFiltered { .. } | Self::ToxicityDetected { .. } | Self::FactualityError(_)
        )
    }

    /// Check if error is a configuration issue
    pub fn is_config_error(&self) -> bool {
        matches!(
            self,
            Self::InvalidConfig(_) | Self::InvalidParameter { .. } | Self::MissingParameter(_)
        )
    }

    /// Get error category for metrics and logging
    pub fn category(&self) -> ErrorCategory {
        match self {
            Self::ModelNotFound(_) | Self::ModelLoadError(_) | Self::InferenceError(_) => {
                ErrorCategory::Model
            }

            Self::GenerationError(_) | Self::GenerationTimeout { .. } => ErrorCategory::Generation,

            Self::TokenizationError(_) | Self::InvalidTokenSequence => ErrorCategory::Tokenization,

            Self::ContentFiltered { .. }
            | Self::ToxicityDetected { .. }
            | Self::FactualityError(_) => ErrorCategory::Safety,

            Self::InvalidConfig(_) | Self::InvalidParameter { .. } | Self::MissingParameter(_) => {
                ErrorCategory::Configuration
            }

            Self::OutOfMemory | Self::DeviceError(_) | Self::CacheError(_) => {
                ErrorCategory::Resource
            }

            Self::NetworkError(_) | Self::ServerOverloaded { .. } | Self::RequestTimeout => {
                ErrorCategory::Network
            }

            Self::TranslationError(_) | Self::UnsupportedLanguagePair { .. } => {
                ErrorCategory::Translation
            }

            Self::SummarizationError(_) | Self::DocumentTooShort | Self::DocumentTooLong { .. } => {
                ErrorCategory::Summarization
            }

            Self::QuestionAnsweringError(_) | Self::NoAnswerFound | Self::ContextTooShort => {
                ErrorCategory::QuestionAnswering
            }

            Self::DialogueError(_) | Self::ConversationTooLong | Self::InvalidConversationState => {
                ErrorCategory::Dialogue
            }

            Self::KvCacheError(_)
            | Self::SpeculativeDecodingError(_)
            | Self::ModelParallelismError(_) => ErrorCategory::Optimization,

            Self::ValidationFailed { .. } | Self::InvalidPattern { .. } => {
                ErrorCategory::Configuration
            }

            Self::BatchQueueFull | Self::BatchTimeout | Self::BatchRequestCancelled => {
                ErrorCategory::Network
            }

            _ => ErrorCategory::Internal,
        }
    }
}

/// Error categories for metrics and monitoring
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ErrorCategory {
    Model,
    Generation,
    Tokenization,
    Safety,
    Configuration,
    Resource,
    Network,
    Translation,
    Summarization,
    QuestionAnswering,
    Dialogue,
    Optimization,
    Internal,
}

impl fmt::Display for ErrorCategory {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let name = match self {
            Self::Model => "model",
            Self::Generation => "generation",
            Self::Tokenization => "tokenization",
            Self::Safety => "safety",
            Self::Configuration => "configuration",
            Self::Resource => "resource",
            Self::Network => "network",
            Self::Translation => "translation",
            Self::Summarization => "summarization",
            Self::QuestionAnswering => "question_answering",
            Self::Dialogue => "dialogue",
            Self::Optimization => "optimization",
            Self::Internal => "internal",
        };
        write!(f, "{name}")
    }
}

/// Error context for debugging and monitoring
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ErrorContext {
    pub category: String,
    pub recoverable: bool,
    pub safety_related: bool,
    pub timestamp: chrono::DateTime<chrono::Utc>,
    pub request_id: Option<String>,
    pub model_id: Option<String>,
    pub user_id: Option<String>,
}

impl ErrorContext {
    /// Create error context from NlgError
    pub fn from_error(error: &NlgError) -> Self {
        Self {
            category: error.category().to_string(),
            recoverable: error.is_recoverable(),
            safety_related: error.is_safety_error(),
            timestamp: chrono::Utc::now(),
            request_id: None,
            model_id: None,
            user_id: None,
        }
    }

    /// Add request ID to context
    pub fn with_request_id(mut self, request_id: impl Into<String>) -> Self {
        self.request_id = Some(request_id.into());
        self
    }

    /// Add model ID to context
    pub fn with_model_id(mut self, model_id: impl Into<String>) -> Self {
        self.model_id = Some(model_id.into());
        self
    }

    /// Add user ID to context
    pub fn with_user_id(mut self, user_id: impl Into<String>) -> Self {
        self.user_id = Some(user_id.into());
        self
    }
}

// Implement From for rtx-tensor errors to allow ? operator
impl From<rtx_tensor::TensorError> for NlgError {
    fn from(err: rtx_tensor::TensorError) -> Self {
        Self::TensorError(err.to_string())
    }
}

// Implement From for String to allow easier error creation
impl From<String> for NlgError {
    fn from(s: String) -> Self {
        Self::GenerationError(s)
    }
}

// Implement From for &str to allow easier error creation
impl From<&str> for NlgError {
    fn from(s: &str) -> Self {
        Self::GenerationError(s.to_string())
    }
}

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

    #[test]
    fn test_error_creation() {
        let error = NlgError::model_not_found("test_model");
        assert!(matches!(error, NlgError::ModelNotFound(_)));
        assert_eq!(error.category(), ErrorCategory::Model);
        assert!(!error.is_recoverable());
        assert!(!error.is_safety_error());
    }

    #[test]
    fn test_safety_error() {
        let error = NlgError::toxicity_detected(0.9);
        assert!(matches!(error, NlgError::ToxicityDetected { .. }));
        assert_eq!(error.category(), ErrorCategory::Safety);
        assert!(error.is_safety_error());
    }

    #[test]
    fn test_recoverable_error() {
        let error = NlgError::GenerationTimeout { timeout_ms: 5000 };
        assert!(error.is_recoverable());
        assert_eq!(error.category(), ErrorCategory::Generation);
    }

    #[test]
    fn test_error_context() {
        let error = NlgError::generation_error("test error");
        let context = ErrorContext::from_error(&error)
            .with_request_id("req_123")
            .with_model_id("model_456");

        assert_eq!(context.category, "generation");
        assert_eq!(context.request_id, Some("req_123".to_string()));
        assert_eq!(context.model_id, Some("model_456".to_string()));
    }

    #[test]
    fn test_error_display() {
        let error = NlgError::server_overloaded(100, 50);
        let error_str = format!("{}", error);
        assert!(error_str.contains("Server overloaded"));
        assert!(error_str.contains("100"));
        assert!(error_str.contains("50"));
    }
}