litellm/litellm-rust/crates/python-bridge/src/errors.rs

use litellm_core::error::Error;
use pyo3::exceptions::{PyRuntimeError, PyValueError};
use pyo3::prelude::*;

pyo3::create_exception!(
    _native,
    RustBridgeDeclined,
    pyo3::exceptions::PyException,
    "The route declined before calling the provider, so the host may retry on its own path."
);

pyo3::create_exception!(
    _native,
    RustUpstreamError,
    pyo3::exceptions::PyException,
    "The provider call was already issued and failed. Args are (status, message); status is 0 when there was no HTTP response."
);

pub(crate) fn core_error_to_pyerr(err: Error) -> PyErr {
    match err {
        Error::Auth(message) => PyValueError::new_err(message),
        Error::InvalidProvider(_)
        | Error::InvalidRequest(_)
        | Error::InvalidType { .. }
        | Error::MissingField(_) => PyValueError::new_err(err.to_string()),
        other => PyRuntimeError::new_err(other.to_string()),
    }
}

/// Map a core error for a route whose host keeps a Python implementation.
///
/// The distinction the host needs is whether the provider was already called.
/// Everything raised before the request goes out is safe for the host to retry
/// on its own path; anything after it is not, because the provider has already
/// done the work and billed for it.
pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr {
    match err {
        Error::Unsupported(_)
        | Error::Auth(_)
        | Error::InvalidProvider(_)
        | Error::InvalidRequest(_)
        | Error::InvalidType { .. }
        | Error::MissingField(_)
        | Error::MissingApiKey { .. }
        | Error::MissingAzureAiCredentials
        | Error::MissingAzureAiCredentialsOrAdToken
        | Error::Routing(_)
        // Nothing reached the provider, so serving it on Python cannot double
        // bill and is the only way the caller gets an answer at all.
        | Error::Connect(_) => RustBridgeDeclined::new_err(err.to_string()),
        Error::Http { status, body } => {
            RustUpstreamError::new_err((status, format!("{status}: {body}")))
        }
        Error::Network(message) | Error::InvalidResponse(message) => {
            RustUpstreamError::new_err((0u16, message))
        }
    }
}

pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
    let py = module.py();
    module.add("RustBridgeDeclined", py.get_type::<RustBridgeDeclined>())?;
    module.add("RustUpstreamError", py.get_type::<RustUpstreamError>())
}

pub(crate) fn ocr_error_to_pyerr(err: Error) -> PyErr {
    match err {
        Error::MissingField("document_url" | "image_url") => {
            PyValueError::new_err("Document URL is required")
        }
        Error::Http { status, body } => RustUpstreamError::new_err((status, body)),
        other => core_error_to_pyerr(other),
    }
}

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

    #[test]
    fn ocr_errors_preserve_python_validation_and_provider_details() {
        Python::initialize();
        Python::attach(|py| {
            for field in ["document_url", "image_url"] {
                let mapped = ocr_error_to_pyerr(Error::MissingField(field));
                assert!(mapped.is_instance_of::<PyValueError>(py));
                assert_eq!(mapped.value(py).to_string(), "Document URL is required");
            }
            let mapped = ocr_error_to_pyerr(Error::Http {
                status: 429,
                body: r#"{"message":"rate limited"}"#.to_string(),
            });
            assert!(mapped.is_instance_of::<RustUpstreamError>(py));
            let args: (u16, String) = mapped
                .value(py)
                .getattr("args")
                .and_then(|args| args.extract())
                .expect("OCR failures retain status and unprefixed provider message");
            assert_eq!(args, (429, r#"{"message":"rate limited"}"#.to_string()));
        });
    }
}