litellm_rs/core/providers/replicate/
error.rs

1//! Replicate Error Handling
2//!
3//! Error handling for the Replicate provider
4
5use crate::core::providers::shared::parse_retry_after_from_body;
6use crate::core::providers::unified_provider::ProviderError;
7use crate::core::traits::error_mapper::trait_def::ErrorMapper;
8
9// Standard error helper methods
10crate::impl_provider_error_helpers!("replicate", replicate);
11
12/// Replicate-specific error constructors (non-standard)
13impl ProviderError {
14    /// Create Replicate prediction failed error
15    pub fn replicate_prediction_failed(message: impl Into<String>) -> Self {
16        Self::api_error("replicate", 422, message.into())
17    }
18
19    /// Create Replicate prediction timeout error
20    pub fn replicate_prediction_timeout(message: impl Into<String>) -> Self {
21        Self::timeout("replicate", message)
22    }
23
24    /// Create Replicate prediction canceled error
25    pub fn replicate_prediction_canceled(message: impl Into<String>) -> Self {
26        Self::cancelled("replicate", "prediction", Some(message.into()))
27    }
28}
29
30/// Error mapper for Replicate API responses
31#[derive(Debug)]
32pub struct ReplicateErrorMapper;
33
34impl ErrorMapper<ProviderError> for ReplicateErrorMapper {
35    fn map_http_error(&self, status_code: u16, response_body: &str) -> ProviderError {
36        match status_code {
37            401 => ProviderError::replicate_authentication("Invalid API token"),
38            // Upstream 403 is a permission failure; keep the status.
39            403 => ProviderError::api_error("replicate", status_code, "Permission denied"),
40            404 => {
41                // Try to extract model name from error
42                if response_body.contains("model") || response_body.contains("version") {
43                    ProviderError::replicate_model_not_found(response_body)
44                } else {
45                    ProviderError::replicate_api_error(404, response_body)
46                }
47            }
48            422 => {
49                // Unprocessable entity - usually a prediction failed
50                ProviderError::replicate_prediction_failed(response_body)
51            }
52            429 => {
53                // Rate limit - try to parse retry-after
54                let retry_after = parse_retry_after_from_body(response_body);
55                ProviderError::replicate_rate_limit(retry_after)
56            }
57            500..=599 => ProviderError::provider_unavailable("replicate", response_body),
58            _ => ProviderError::replicate_api_error(status_code, response_body),
59        }
60    }
61
62    fn map_json_error(&self, error_response: &serde_json::Value) -> ProviderError {
63        // Replicate error format:
64        // {"detail": "error message"} or
65        // {"error": {"message": "...", "code": "..."}}
66        let error_msg = error_response
67            .get("detail")
68            .and_then(|d| d.as_str())
69            .or_else(|| {
70                error_response
71                    .get("error")
72                    .and_then(|e| e.get("message"))
73                    .and_then(|m| m.as_str())
74            })
75            .unwrap_or("Unknown error");
76
77        let error_type = error_response
78            .get("type")
79            .or_else(|| error_response.get("error").and_then(|e| e.get("type")))
80            .and_then(|t| t.as_str())
81            .unwrap_or("unknown");
82
83        match error_type {
84            "authentication_error" | "invalid_token" => {
85                ProviderError::replicate_authentication(error_msg)
86            }
87            "rate_limit_exceeded" => ProviderError::replicate_rate_limit(None),
88            "model_not_found" | "version_not_found" => {
89                ProviderError::replicate_model_not_found(error_msg)
90            }
91            "prediction_failed" => ProviderError::replicate_prediction_failed(error_msg),
92            _ => ProviderError::other("replicate", error_msg),
93        }
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    #[test]
102    fn test_replicate_authentication_error() {
103        let err = ProviderError::replicate_authentication("Invalid API token");
104        assert!(err.is_replicate_error());
105        assert!(matches!(err, ProviderError::Authentication { .. }));
106    }
107
108    #[test]
109    fn test_replicate_rate_limit_error() {
110        let err = ProviderError::replicate_rate_limit(Some(60));
111        assert!(err.is_replicate_error());
112        assert!(matches!(err, ProviderError::RateLimit { .. }));
113    }
114
115    #[test]
116    fn test_replicate_model_not_found_error() {
117        let err = ProviderError::replicate_model_not_found("unknown/model");
118        assert!(err.is_replicate_error());
119        assert!(matches!(err, ProviderError::ModelNotFound { .. }));
120    }
121
122    #[test]
123    fn test_replicate_prediction_failed_error() {
124        let err = ProviderError::replicate_prediction_failed("Prediction failed");
125        assert!(err.is_replicate_error());
126        assert!(matches!(err, ProviderError::ApiError { status: 422, .. }));
127    }
128
129    #[test]
130    fn test_replicate_prediction_timeout_error() {
131        let err = ProviderError::replicate_prediction_timeout("Timeout waiting for result");
132        assert!(err.is_replicate_error());
133        assert!(matches!(err, ProviderError::Timeout { .. }));
134    }
135
136    #[test]
137    fn test_error_mapper_401() {
138        let mapper = ReplicateErrorMapper;
139        let err = mapper.map_http_error(401, "Unauthorized");
140        assert!(matches!(err, ProviderError::Authentication { .. }));
141    }
142
143    #[test]
144    fn test_error_mapper_403() {
145        let mapper = ReplicateErrorMapper;
146        // 403 keeps its upstream status as a permission failure.
147        let err = mapper.map_http_error(403, "Forbidden");
148        assert!(matches!(err, ProviderError::ApiError { status: 403, .. }));
149    }
150
151    #[test]
152    fn test_error_mapper_404_model() {
153        let mapper = ReplicateErrorMapper;
154        let err = mapper.map_http_error(404, "model not found");
155        assert!(matches!(err, ProviderError::ModelNotFound { .. }));
156    }
157
158    #[test]
159    fn test_error_mapper_404_other() {
160        let mapper = ReplicateErrorMapper;
161        let err = mapper.map_http_error(404, "resource not found");
162        assert!(matches!(err, ProviderError::ApiError { status: 404, .. }));
163    }
164
165    #[test]
166    fn test_error_mapper_422() {
167        let mapper = ReplicateErrorMapper;
168        let err = mapper.map_http_error(422, "Prediction failed");
169        assert!(matches!(err, ProviderError::ApiError { status: 422, .. }));
170    }
171
172    #[test]
173    fn test_error_mapper_429() {
174        let mapper = ReplicateErrorMapper;
175        let err = mapper.map_http_error(429, "rate limit exceeded");
176        assert!(matches!(err, ProviderError::RateLimit { .. }));
177    }
178
179    #[test]
180    fn test_error_mapper_500() {
181        let mapper = ReplicateErrorMapper;
182        let err = mapper.map_http_error(500, "Internal server error");
183        assert!(matches!(err, ProviderError::ProviderUnavailable { .. }));
184    }
185
186    #[test]
187    fn test_parse_retry_after_with_rate_limit() {
188        let result = parse_retry_after_from_body("rate limit exceeded");
189        assert_eq!(result, Some(60));
190    }
191
192    #[test]
193    fn test_parse_retry_after_with_too_many_requests() {
194        let result = parse_retry_after_from_body("too many requests");
195        assert_eq!(result, Some(60));
196    }
197
198    #[test]
199    fn test_parse_retry_after_no_match() {
200        let result = parse_retry_after_from_body("some other error");
201        assert_eq!(result, None);
202    }
203
204    #[test]
205    fn test_error_mapper_json_authentication() {
206        let mapper = ReplicateErrorMapper;
207        let json = serde_json::json!({
208            "type": "authentication_error",
209            "detail": "Invalid token"
210        });
211        let err = mapper.map_json_error(&json);
212        assert!(matches!(err, ProviderError::Authentication { .. }));
213    }
214
215    #[test]
216    fn test_error_mapper_json_model_not_found() {
217        let mapper = ReplicateErrorMapper;
218        let json = serde_json::json!({
219            "type": "model_not_found",
220            "detail": "Model not found"
221        });
222        let err = mapper.map_json_error(&json);
223        assert!(matches!(err, ProviderError::ModelNotFound { .. }));
224    }
225
226    #[test]
227    fn test_error_mapper_json_rate_limit() {
228        let mapper = ReplicateErrorMapper;
229        let json = serde_json::json!({
230            "type": "rate_limit_exceeded",
231            "detail": "Too many requests"
232        });
233        let err = mapper.map_json_error(&json);
234        assert!(matches!(err, ProviderError::RateLimit { .. }));
235    }
236
237    #[test]
238    fn test_error_mapper_json_unknown() {
239        let mapper = ReplicateErrorMapper;
240        let json = serde_json::json!({
241            "type": "unknown_error",
242            "detail": "Something went wrong"
243        });
244        let err = mapper.map_json_error(&json);
245        assert!(matches!(err, ProviderError::Other { .. }));
246    }
247
248    #[test]
249    fn test_replicate_prediction_canceled() {
250        let err = ProviderError::replicate_prediction_canceled("User cancelled");
251        assert!(err.is_replicate_error());
252        assert!(matches!(err, ProviderError::Cancelled { .. }));
253    }
254
255    #[test]
256    fn test_replicate_configuration_error() {
257        let err = ProviderError::replicate_configuration("Missing API token");
258        assert!(err.is_replicate_error());
259        assert!(matches!(err, ProviderError::Configuration { .. }));
260    }
261}
litellm_rs/core/providers/replicate/error.rs

litellm_rs/core/providers/replicate/
error.rs
