From cbce50f591c78108f3dfb1279ef857009198ef68 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 23 Oct 2026 14:38:14 -0700 Subject: [PATCH] refactor(rust-error): fix the remaining sweep findings in gateway, llms, auth and cache crates Auth rejections hide their backend cause on every inference renderer, path parameters reject in the surface envelope, Textract's unsupported document and Azure DI's unparseable poll URL get typed variants, the credential chain exposes its last failure as the source, gateway-mcp asks predicates over rmcp errors, and the S3 and Azure blob transports classify through litellm_http::Error. Co-Authored-By: Claude Opus 5.5 --- litellm-rust/Cargo.lock | 5 + litellm-rust/crates/auth-aws/src/aws.rs | 8 +- litellm-rust/crates/auth-azure/src/resolve.rs | 10 +- litellm-rust/crates/auth-gcp/src/auth.rs | 29 ++- litellm-rust/crates/auth-types/src/error.rs | 69 ++++-- .../crates/cache-azure-blob/src/transport.rs | 3 +- litellm-rust/crates/cache-s3/src/transport.rs | 5 +- .../gateway-inference/src/chat_completions.rs | 7 +- .../crates/gateway-inference/src/error.rs | 55 ++++- .../crates/gateway-inference/src/request.rs | 13 ++ litellm-rust/crates/gateway-mcp/src/error.rs | 120 ++++++++++- litellm-rust/crates/gateway-mcp/src/native.rs | 15 +- .../crates/gateway-mcp/src/native/catalog.rs | 11 +- litellm-rust/crates/http/src/error.rs | 37 ++++ litellm-rust/crates/http/src/outbound.rs | 22 +- litellm-rust/crates/inference-ocr/Cargo.toml | 1 + .../crates/inference-ocr/src/document.rs | 19 +- litellm-rust/crates/inference-ocr/src/wire.rs | 19 +- .../inference-ocr/tests/ocr/aws_textract.rs | 15 +- .../tests/ocr/azure_document_intelligence.rs | 18 +- .../llms/src/aws_textract/ocr/common_utils.rs | 83 ++++---- .../document_intelligence/transformation.rs | 12 +- .../crates/llms/src/base_llm/ocr/error.rs | 30 +++ .../llms/src/reducto/ocr/transformation.rs | 19 +- .../crates/python-compat/src/error.rs | 2 +- .../crates/python-compat/src/literal.rs | 34 +-- .../crates/python-compat/src/pickle.rs | 4 +- .../token-counter-tiktoken/src/ranks.rs | 197 +++++++++++------- 28 files changed, 609 insertions(+), 253 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 7131cbb2bdd..9592091eb60 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3649,6 +3649,7 @@ name = "litellm-callbacks-legacy-python" version = "0.1.0" dependencies = [ "litellm-auth", + "litellm-error", "litellm-host", "litellm-host-python", "proptest", @@ -3656,6 +3657,7 @@ dependencies = [ "rstest", "serde_json", "strum", + "thiserror 2.0.19", ] [[package]] @@ -4118,6 +4120,7 @@ dependencies = [ "time", "tokio", "tracing", + "url", "wiremock", ] @@ -4476,7 +4479,9 @@ version = "0.1.0" dependencies = [ "litellm-core-utils", "litellm-error", + "litellm-http", "litellm-secrets-types", + "reqwest", "rstest", "rustify", "rustify_derive", diff --git a/litellm-rust/crates/auth-aws/src/aws.rs b/litellm-rust/crates/auth-aws/src/aws.rs index 1ef736e6b55..81aa4b068bc 100644 --- a/litellm-rust/crates/auth-aws/src/aws.rs +++ b/litellm-rust/crates/auth-aws/src/aws.rs @@ -1123,10 +1123,12 @@ mod tests { .. }) )); - assert_eq!( - litellm_error::Report::new(&error).to_string(), - "provider authentication failed: could not load AWS profile credentials: profile not found" + let report = litellm_error::Report::new(&error).to_string(); + assert!( + report.contains("could not load AWS profile credentials"), + "{report}" ); + assert_eq!(report.matches("profile not found").count(), 1, "{report}"); } #[ignore] diff --git a/litellm-rust/crates/auth-azure/src/resolve.rs b/litellm-rust/crates/auth-azure/src/resolve.rs index 44101998dd5..25b71741342 100644 --- a/litellm-rust/crates/auth-azure/src/resolve.rs +++ b/litellm-rust/crates/auth-azure/src/resolve.rs @@ -140,7 +140,13 @@ impl AzureAuthService { Err(error) => failures.push(error), } } - Err(Error::CredentialChain(failures)) + match failures.pop() { + Some(last) => Err(Error::CredentialChain { + last: Box::new(last), + earlier: failures, + }), + None => Ok(None), + } } AzureCredentialPlan::Missing => Ok(None), } @@ -723,7 +729,7 @@ mod tests { .await .unwrap_err(); - assert!(matches!(&error, Error::CredentialChain(errors) if errors.len() == 2)); + assert!(matches!(&error, Error::CredentialChain { earlier, .. } if earlier.len() == 1)); let report = litellm_error::Report::new(&error).to_string(); assert!( report.contains("managed-identity did not resolve"), diff --git a/litellm-rust/crates/auth-gcp/src/auth.rs b/litellm-rust/crates/auth-gcp/src/auth.rs index d3a9e7f3d2c..8ac1e0f54ba 100644 --- a/litellm-rust/crates/auth-gcp/src/auth.rs +++ b/litellm-rust/crates/auth-gcp/src/auth.rs @@ -73,13 +73,17 @@ impl VertexAuth { .or_else(|| non_empty_env(env_lookup, VERTEXAI_API_KEY_ENV)); let project_id = get_vertex_ai_project(config, env_lookup); - if !has_authorization && static_token.is_none() { - let access = self.get_access_token(config, env_lookup).await?; - return Ok(VertexEnvironment { - headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?, - project_id: project_id.unwrap_or(access.project_id), - }); - } + let static_token = match (has_authorization, static_token) { + (true, _) => None, + (false, Some(token)) => Some(token), + (false, None) => { + let access = self.get_access_token(config, env_lookup).await?; + return Ok(VertexEnvironment { + headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?, + project_id: project_id.unwrap_or(access.project_id), + }); + } + }; let project_id = match project_id { Some(project_id) => project_id, @@ -90,14 +94,9 @@ impl VertexAuth { .await? } }; - let headers = if has_authorization { - headers - } else { - apply_credential( - headers, - static_token.as_deref().expect("static token was checked"), - CredentialPlacement::Bearer, - )? + let headers = match static_token { + Some(token) => apply_credential(headers, &token, CredentialPlacement::Bearer)?, + None => headers, }; Ok(VertexEnvironment { headers, diff --git a/litellm-rust/crates/auth-types/src/error.rs b/litellm-rust/crates/auth-types/src/error.rs index b0fc43a91b7..a8ba7b64e03 100644 --- a/litellm-rust/crates/auth-types/src/error.rs +++ b/litellm-rust/crates/auth-types/src/error.rs @@ -12,8 +12,12 @@ pub enum Error { EmptyCallerCredential { caller: &'static str }, #[error("provider authentication failed")] Provider(#[source] Source), - #[error("no credential in the chain succeeded: {}", reports(.0))] - CredentialChain(Vec), + #[error("no credential in the chain succeeded{}", earlier_attempts(earlier))] + CredentialChain { + #[source] + last: Box, + earlier: Vec, + }, #[error( "Missing {provider} API Key - Set `api_key` or the {environment_variable} environment variable" )] @@ -35,12 +39,16 @@ pub enum Error { InvalidHeader, } -fn reports(failures: &[Error]) -> String { - failures +fn earlier_attempts(earlier: &[Error]) -> String { + if earlier.is_empty() { + return String::new(); + } + let attempts = earlier .iter() .map(|failure| Report::new(failure).to_string()) .collect::>() - .join("; ") + .join("; "); + format!(" (earlier attempts: {attempts})") } impl Error { @@ -59,7 +67,7 @@ impl Error { | Self::CredentialAcquisition(_) | Self::EmptyCallerCredential { .. } | Self::Provider(_) - | Self::CredentialChain(_) + | Self::CredentialChain { .. } | Self::MissingApiBase { .. } | Self::InvalidHeader => false, } @@ -70,7 +78,9 @@ impl Error { Self::InvalidConfiguration(_) | Self::EmptyCallerCredential { .. } | Self::InvalidHeader => true, - Self::CredentialAcquisition(_) | Self::Provider(_) | Self::CredentialChain(_) => false, + Self::CredentialAcquisition(_) | Self::Provider(_) | Self::CredentialChain { .. } => { + false + } Self::MissingParam { spec, .. } => match spec.kind { ParamKind::Setting => true, ParamKind::Credential => false, @@ -178,14 +188,6 @@ mod tests { Error::EmptyCallerCredential { caller: "token provider" }, "token provider returned an empty credential" )] - #[case::chain( - Error::CredentialChain(vec![ - Error::CredentialAcquisition(ErrorDetail::Unresolved { subject: "OIDC reference" }), - Error::EmptyCallerCredential { caller: "token provider" }, - ]), - "no credential in the chain succeeded: credential acquisition failed: OIDC reference did not resolve \ - to a value; token provider returned an empty credential" - )] fn the_report_names_every_layer_in_order(#[case] error: Error, #[case] expected: &str) { assert!( Report::new(&error).to_string().contains(expected), @@ -238,9 +240,42 @@ mod tests { litellm_error::find::(&error).map(io::Error::kind), Some(io::ErrorKind::PermissionDenied) ); + assert!(!error.to_string().contains("access denied"), "{error}"); + let report = Report::new(&error).to_string(); + assert_eq!(report.matches("access denied").count(), 1, "{report}"); + } + + #[test] + fn a_chain_keeps_its_last_failure_as_the_source_and_renders_every_attempt_once() { + let error = Error::CredentialChain { + last: Box::new(Error::provider(denied())), + earlier: vec![ + Error::CredentialAcquisition(ErrorDetail::Unresolved { + subject: "OIDC reference", + }), + Error::EmptyCallerCredential { + caller: "token provider", + }, + ], + }; + assert_eq!( - Report::new(&error).to_string(), - "provider authentication failed: access denied" + litellm_error::find::(&error).map(io::Error::kind), + Some(io::ErrorKind::PermissionDenied) + ); + let report = Report::new(&error).to_string(); + assert_eq!(report.matches("access denied").count(), 1, "{report}"); + assert_eq!( + report.matches("OIDC reference did not resolve").count(), + 1, + "{report}" + ); + assert_eq!( + report + .matches("token provider returned an empty credential") + .count(), + 1, + "{report}" ); } diff --git a/litellm-rust/crates/cache-azure-blob/src/transport.rs b/litellm-rust/crates/cache-azure-blob/src/transport.rs index e3f2cce8700..62feee8fae0 100644 --- a/litellm-rust/crates/cache-azure-blob/src/transport.rs +++ b/litellm-rust/crates/cache-azure-blob/src/transport.rs @@ -24,12 +24,13 @@ impl HttpClient for ReqwestTransport { Body::SeekableStream(stream) => outgoing.body(reqwest::Body::wrap_stream(stream)), }; let response = outgoing.send().await.map_err(|error| { + let error = litellm_http::Error::from(error); let kind = if error.is_connect() { ErrorKind::Connection } else { ErrorKind::Io }; - azure_core::Error::new(kind, litellm_http::Error::from(error)) + azure_core::Error::new(kind, error) })?; let status = response.status().as_u16().into(); let mut headers = Headers::new(); diff --git a/litellm-rust/crates/cache-s3/src/transport.rs b/litellm-rust/crates/cache-s3/src/transport.rs index 5b6e867ede5..bf1711276f2 100644 --- a/litellm-rust/crates/cache-s3/src/transport.rs +++ b/litellm-rust/crates/cache-s3/src/transport.rs @@ -29,12 +29,11 @@ impl HttpConnector for ReqwestHttpClient { .try_into_http1x() .map_err(|error| ConnectorError::other(error.into(), None))? .map(reqwest::Body::wrap); - let request = reqwest::Request::try_from(request) + let request = litellm_http::outbound::from_http(request) .map_err(|error| ConnectorError::other(error.into(), None))?; let response = client.execute(request).await.map_err(|error| { - let timeout = error.is_timeout(); let error = litellm_http::Error::from(error); - if timeout { + if error.is_timeout() { ConnectorError::timeout(error.into()) } else { ConnectorError::io(error.into()) diff --git a/litellm-rust/crates/gateway-inference/src/chat_completions.rs b/litellm-rust/crates/gateway-inference/src/chat_completions.rs index 2ab70961171..fb8a6f4609d 100644 --- a/litellm-rust/crates/gateway-inference/src/chat_completions.rs +++ b/litellm-rust/crates/gateway-inference/src/chat_completions.rs @@ -3,13 +3,16 @@ use std::sync::Arc; use axum::{ Json, - extract::{Path, State}, + extract::State, response::{IntoResponse, Response}, }; use litellm_inference_chat::types::ChatCompletionsCall; use serde_json::{Map, Value}; -use crate::{Error, Gateway, JsonObject, request}; +use crate::{ + Error, Gateway, JsonObject, + request::{self, Path}, +}; pub(crate) async fn create( State(gateway): State>, diff --git a/litellm-rust/crates/gateway-inference/src/error.rs b/litellm-rust/crates/gateway-inference/src/error.rs index 117c5f5eda7..525ccee825f 100644 --- a/litellm-rust/crates/gateway-inference/src/error.rs +++ b/litellm-rust/crates/gateway-inference/src/error.rs @@ -13,6 +13,8 @@ pub enum Error { Auth(#[from] litellm_gateway_auth::Error), #[error("invalid request body")] InvalidBody(#[from] BodyError), + #[error("invalid path parameter")] + InvalidPath(#[source] axum::extract::rejection::PathRejection), #[error( "{route}: Invalid model name passed in model={model}. Call `/v1/models` to view available models for your key." )] @@ -70,6 +72,7 @@ impl Error { pub fn status(&self) -> StatusCode { match self { Self::Auth(error) => error.status(), + Self::InvalidPath(rejection) => rejection.status(), Self::Unsupported(_) => StatusCode::NOT_IMPLEMENTED, Self::BodyTooLarge => StatusCode::PAYLOAD_TOO_LARGE, Self::InvalidBody(_) | Self::UnknownModel { .. } => StatusCode::BAD_REQUEST, @@ -89,7 +92,7 @@ impl Error { let event = json!({ "type": "error", "code": self.status().as_u16().to_string(), - "message": litellm_error::Report::new(self).to_string(), + "message": self.message(), "param": null, }); format!("event: error\ndata: {event}\n\n") @@ -98,7 +101,7 @@ impl Error { fn openai_body(&self) -> Value { let status = self.status(); json!({"error": { - "message": litellm_error::Report::new(self).to_string(), + "message": self.message(), "type": error_type(status), "param": null, "code": status.as_u16(), @@ -112,7 +115,7 @@ impl Error { Self::Route(error) if let Some((_, body)) = error.upstream_response() => { body.to_owned() } - other => litellm_error::Report::new(other).to_string(), + other => other.message(), }; let parsed = serde_json::from_str::(&raw).ok(); let envelope = match parsed { @@ -129,6 +132,20 @@ impl Error { pub fn sse_frame(&self) -> String { format!("event: error\ndata: {}\n\n", self.body(None)) } + + fn message(&self) -> String { + match self { + Self::Auth(error) => error.to_string(), + Self::InvalidBody(_) + | Self::InvalidPath(_) + | Self::UnknownModel { .. } + | Self::Route(_) + | Self::Ocr(_) + | Self::Unsupported(_) + | Self::BodyTooLarge + | Self::Internal => litellm_error::Report::new(self).to_string(), + } + } } fn route_status(error: &InferenceError) -> StatusCode { @@ -177,13 +194,13 @@ fn error_type(status: StatusCode) -> &'static str { } fn envelope(status: StatusCode, message: &str) -> Map { - let Value::Object(envelope) = json!({ - "type": "error", - "error": {"type": error_type(status), "message": message}, - }) else { - unreachable!("a json object literal is an object") - }; - envelope + Map::from_iter([ + ("type".to_owned(), Value::from("error")), + ( + "error".to_owned(), + json!({"type": error_type(status), "message": message}), + ), + ]) } fn is_anthropic_error(object: &Map) -> bool { @@ -350,6 +367,24 @@ mod tests { assert_eq!(body["error"]["code"], 400); } + fn auth_outage() -> Error { + Error::Auth(litellm_gateway_auth::Error::Unavailable( + litellm_error::Source::new(std::io::Error::other("private-backend-detail")), + )) + } + + #[rstest] + #[case::openai(auth_outage().openai_body()["error"]["message"].clone())] + #[case::anthropic(auth_outage().body(None)["error"]["message"].clone())] + #[case::responses_sse(Value::from(auth_outage().responses_sse_frame()))] + #[case::anthropic_sse(Value::from(auth_outage().sse_frame()))] + fn an_auth_failure_hides_its_backend_cause(#[case] rendered: Value) { + let top = auth_outage().to_string(); + let rendered = rendered.to_string(); + assert!(rendered.contains(&top), "{rendered}"); + assert!(!rendered.contains("private-backend-detail"), "{rendered}"); + } + #[rstest] #[case::provider_answer(upstream(429, "slow down"))] #[case::caller_fault(invalid_json())] diff --git a/litellm-rust/crates/gateway-inference/src/request.rs b/litellm-rust/crates/gateway-inference/src/request.rs index 27a9cfbc6a0..d41794242ff 100644 --- a/litellm-rust/crates/gateway-inference/src/request.rs +++ b/litellm-rust/crates/gateway-inference/src/request.rs @@ -21,6 +21,19 @@ pub struct JsonObject(pub Map); pub struct RequestId(pub Option); +pub(crate) struct Path(pub T); + +impl FromRequestParts for Path { + type Rejection = Error; + + async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { + axum::extract::Path::::from_request_parts(parts, state) + .await + .map(|axum::extract::Path(value)| Self(value)) + .map_err(Error::InvalidPath) + } +} + impl FromRequestParts for RequestId { type Rejection = std::convert::Infallible; diff --git a/litellm-rust/crates/gateway-mcp/src/error.rs b/litellm-rust/crates/gateway-mcp/src/error.rs index adc602f7c07..9f08abaf26c 100644 --- a/litellm-rust/crates/gateway-mcp/src/error.rs +++ b/litellm-rust/crates/gateway-mcp/src/error.rs @@ -128,9 +128,78 @@ impl Error { } } + fn upstream(&self) -> Option<&ServiceError> { + match self { + Self::Upstream(error) => Some(error), + Self::MissingField(_) + | Self::UpstreamCursor + | Self::AmbiguousResourceServer + | Self::InvalidServersHeader(_) + | Self::Json(_) + | Self::Query(_) + | Self::Path(_) + | Self::BodyTooLarge(_) + | Self::Forbidden + | Self::HostNotAllowed + | Self::OriginNotAllowed + | Self::ForeignSession + | Self::UnknownSession + | Self::SessionBusy + | Self::InvalidPagination + | Self::UnexpectedResult + | Self::AuthorizationUnavailable(_) + | Self::Cancelled => None, + } + } + + pub(crate) fn is_upstream(&self) -> bool { + self.upstream().is_some() + } + + pub(crate) fn is_timeout(&self) -> bool { + self.upstream().is_some_and(|error| match error { + ServiceError::Timeout { .. } => true, + ServiceError::McpError(_) + | ServiceError::TransportSend(_) + | ServiceError::TransportClosed + | ServiceError::UnexpectedResponse + | ServiceError::SubscriptionLagged { .. } + | ServiceError::Cancelled { .. } + | ServiceError::InputRequiredRoundsExceeded { .. } => false, + _ => false, + }) + } + + pub(crate) fn is_unreachable(&self) -> bool { + self.upstream().is_some_and(|error| match error { + ServiceError::TransportSend(_) | ServiceError::TransportClosed => true, + ServiceError::McpError(_) + | ServiceError::Timeout { .. } + | ServiceError::UnexpectedResponse + | ServiceError::SubscriptionLagged { .. } + | ServiceError::Cancelled { .. } + | ServiceError::InputRequiredRoundsExceeded { .. } => false, + _ => false, + }) + } + + pub(crate) fn mcp_error(&self) -> Option<&ErrorData> { + self.upstream().and_then(|error| match error { + ServiceError::McpError(error) => Some(error), + ServiceError::TransportSend(_) + | ServiceError::TransportClosed + | ServiceError::Timeout { .. } + | ServiceError::UnexpectedResponse + | ServiceError::SubscriptionLagged { .. } + | ServiceError::Cancelled { .. } + | ServiceError::InputRequiredRoundsExceeded { .. } => None, + _ => None, + }) + } + pub fn into_mcp(self) -> ErrorData { - if let Self::Upstream(ServiceError::McpError(error)) = self { - return error; + if let Some(error) = self.mcp_error() { + return error.clone(); } let message = self.detail(); if self.status() == StatusCode::FORBIDDEN { @@ -249,9 +318,12 @@ mod tests { fault: ConfigFault::ZeroTimeout, }; assert!(format!("{error:?}").contains("ZeroTimeout")); + let report = Report::new(&error).to_string(); + assert!(report.contains("docs"), "{report}"); assert_eq!( - Report::new(&error).to_string(), - "MCP server docs is misconfigured: timeout must be positive" + report.matches("timeout must be positive").count(), + 1, + "{report}" ); } @@ -272,6 +344,46 @@ mod tests { assert_eq!(error.status(), status); } + fn mcp_failure() -> Error { + Error::Upstream(ServiceError::McpError(ErrorData::invalid_params( + "bad", None, + ))) + } + + #[rstest] + #[case::timeout( + Error::Upstream(ServiceError::Timeout { timeout: std::time::Duration::from_secs(1) }), + true, + false, + false + )] + #[case::closed(Error::Upstream(ServiceError::TransportClosed), false, true, false)] + #[case::mcp_error(mcp_failure(), false, false, true)] + #[case::unexpected_response( + Error::Upstream(ServiceError::UnexpectedResponse), + false, + false, + false + )] + #[case::gateway_cancelled(Error::Cancelled, false, false, false)] + fn upstream_failures_answer_what_went_wrong( + #[case] error: Error, + #[case] timeout: bool, + #[case] unreachable: bool, + #[case] mcp_error: bool, + ) { + assert_eq!(error.is_timeout(), timeout); + assert_eq!(error.is_unreachable(), unreachable); + assert_eq!(error.mcp_error().is_some(), mcp_error); + } + + #[rstest] + fn an_upstream_mcp_error_passes_through_unchanged() { + let mapped = mcp_failure().into_mcp(); + assert_eq!(mapped.code, ErrorCode::INVALID_PARAMS); + assert_eq!(mapped.message, "bad"); + } + #[rstest] #[case::request(Error::MissingField("name"), ErrorCode::INVALID_PARAMS)] #[case::forbidden(Error::Forbidden, ErrorCode(-32003))] diff --git a/litellm-rust/crates/gateway-mcp/src/native.rs b/litellm-rust/crates/gateway-mcp/src/native.rs index 69306c7dbd2..7a83f1316e9 100644 --- a/litellm-rust/crates/gateway-mcp/src/native.rs +++ b/litellm-rust/crates/gateway-mcp/src/native.rs @@ -178,14 +178,15 @@ impl NativeGateway { .ok_or(Error::Forbidden)?; let mut upstream = request; upstream.name = tool.tool.name.clone(); - match self + let result = self .send(server, CallToolRequest::new(upstream).into(), &context) - .await - { - Err(Error::Upstream(rmcp::service::ServiceError::McpError(error))) => { - Ok(CallToolResult::error(vec![ContentBlock::text(error.message)]).into()) - } - result => result, + .await; + match result.as_ref().err().and_then(Error::mcp_error) { + Some(error) => Ok(CallToolResult::error(vec![ContentBlock::text( + error.message.clone(), + )]) + .into()), + None => result, } } Operation::ListPrompts(params) => { diff --git a/litellm-rust/crates/gateway-mcp/src/native/catalog.rs b/litellm-rust/crates/gateway-mcp/src/native/catalog.rs index 78ff7842b2d..6dd16d1fe63 100644 --- a/litellm-rust/crates/gateway-mcp/src/native/catalog.rs +++ b/litellm-rust/crates/gateway-mcp/src/native/catalog.rs @@ -93,14 +93,9 @@ impl NativeGateway { Ok(tools) => ServerOutcome::Ok { tool_count: tools.len(), }, - Err(Error::Upstream(rmcp::service::ServiceError::Timeout { .. })) => { - ServerOutcome::Timeout - } - Err(Error::Upstream( - rmcp::service::ServiceError::TransportClosed - | rmcp::service::ServiceError::TransportSend(_), - )) => ServerOutcome::Unreachable, - Err(Error::Upstream(_)) => ServerOutcome::UpstreamError, + Err(error) if error.is_timeout() => ServerOutcome::Timeout, + Err(error) if error.is_unreachable() => ServerOutcome::Unreachable, + Err(error) if error.is_upstream() => ServerOutcome::UpstreamError, Err(_) => ServerOutcome::Internal, }; (server.prefix(), outcome) diff --git a/litellm-rust/crates/http/src/error.rs b/litellm-rust/crates/http/src/error.rs index 6f025876b21..acf8cbf445d 100644 --- a/litellm-rust/crates/http/src/error.rs +++ b/litellm-rust/crates/http/src/error.rs @@ -347,6 +347,41 @@ impl Error { } } + pub fn is_timeout(&self) -> bool { + match &*self.0 { + Repr::Timeout(_) | Repr::HandshakeTimeout(_) | Repr::DownloadTimeout(_) => true, + Repr::Read { .. } + | Repr::InvalidPem { .. } + | Repr::Client(_) + | Repr::TlsConfig(_) + | Repr::RequestBody(_) + | Repr::Request(_) + | Repr::RequestUrl(_) + | Repr::ComputedHeader(_) + | Repr::UnsignableBody + | Repr::UnsignableHeader(_) + | Repr::Signature(_) + | Repr::Connect(_) + | Repr::Network(_) + | Repr::Status { .. } + | Repr::WebSocketUrl(_) + | Repr::WebSocketScheme(_) + | Repr::WebSocketSchemeRewrite(_) + | Repr::Handshake(_) + | Repr::WebSocket(_) + | Repr::WebSocketClosed + | Repr::BlockedUrl + | Repr::Resolve(_) + | Repr::DownloadDisabled + | Repr::DownloadTooLarge + | Repr::TooManyRedirects + | Repr::MissingRedirectLocation + | Repr::InvalidRedirect(_) + | Repr::DownloadStatus(_) + | Repr::Download(_) => false, + } + } + pub fn upstream_response(&self) -> Option<(u16, &str)> { match &*self.0 { Repr::Status { status, body } => Some((*status, body)), @@ -552,6 +587,7 @@ mod tests { assert_eq!(report.matches(&root_cause).count(), 1, "{report}"); assert!(!error.to_string().contains(&root_cause)); assert!(error.is_connect()); + assert!(!error.is_timeout()); assert!(!report.contains("secret"), "{report}"); assert!(!report.contains("private"), "{report}"); } @@ -581,6 +617,7 @@ mod tests { assert!(error.is_timeout()); let error = crate::Error::from(error); assert!(error.dispatched()); + assert!(error.is_timeout()); assert!(!error.is_connect()); assert_eq!(error.status(), Some(408)); } diff --git a/litellm-rust/crates/http/src/outbound.rs b/litellm-rust/crates/http/src/outbound.rs index 0cf7212608e..55e046245b5 100644 --- a/litellm-rust/crates/http/src/outbound.rs +++ b/litellm-rust/crates/http/src/outbound.rs @@ -41,13 +41,15 @@ pub fn json( } else { with_headers.header(CONTENT_TYPE, "application/json") }; - let request = builder.body(body).map_err(Repr::Request)?; - let mut request = - Request::try_from(request).map_err(|error| Repr::RequestUrl(error.without_url()))?; + let mut request = from_http(builder.body(body).map_err(Repr::Request)?)?; *request.timeout_mut() = timeout; Ok(request) } +pub fn from_http>(request: http::Request) -> Result { + Request::try_from(request).map_err(|error| Repr::RequestUrl(error.without_url()).into()) +} + pub fn signed(mut request: Request, signer: &dyn RequestSigner) -> Result { let signature = signer.sign(&request)?; request.headers_mut().extend(signature); @@ -171,4 +173,18 @@ mod tests { Err(error) if error.is_request() )); } + + #[test] + fn a_relative_http_request_is_the_callers_fault_and_never_reached_the_provider() { + let request = http::Request::get("/private?api_key=secret") + .body(Vec::::new()) + .unwrap(); + + let error = from_http(request).unwrap_err(); + + assert!(error.is_request()); + assert!(!error.dispatched()); + let report = litellm_error::Report::new(&error).to_string(); + assert!(!report.contains("secret"), "{report}"); + } } diff --git a/litellm-rust/crates/inference-ocr/Cargo.toml b/litellm-rust/crates/inference-ocr/Cargo.toml index 8d18b1f3aa6..bbbdaddb050 100644 --- a/litellm-rust/crates/inference-ocr/Cargo.toml +++ b/litellm-rust/crates/inference-ocr/Cargo.toml @@ -35,4 +35,5 @@ litellm-llms = { workspace = true, features = ["test-support"] } rand.workspace = true rstest.workspace = true time.workspace = true +url.workspace = true wiremock.workspace = true diff --git a/litellm-rust/crates/inference-ocr/src/document.rs b/litellm-rust/crates/inference-ocr/src/document.rs index a68639ce45d..44c1049102a 100644 --- a/litellm-rust/crates/inference-ocr/src/document.rs +++ b/litellm-rust/crates/inference-ocr/src/document.rs @@ -243,15 +243,16 @@ mod tests { } #[test] - fn file_encoding_rejects_empty_bytes_and_invalid_explicit_mime() { + fn file_encoding_rejects_empty_bytes() { assert!(encode_file_document(b"", None, None).is_err()); - for mime in [ - "text/plain;bad", - "text/plain/extra", - " text/plain", - "text/plain\n", - ] { - assert!(encode_file_document(b"abc", None, Some(mime)).is_err()); - } + } + + #[rstest] + #[case::bad_parameter("text/plain;bad")] + #[case::extra_segment("text/plain/extra")] + #[case::leading_space(" text/plain")] + #[case::trailing_newline("text/plain\n")] + fn file_encoding_rejects_an_invalid_explicit_mime(#[case] mime: &str) { + assert!(encode_file_document(b"abc", None, Some(mime)).is_err()); } } diff --git a/litellm-rust/crates/inference-ocr/src/wire.rs b/litellm-rust/crates/inference-ocr/src/wire.rs index ae9a7563d43..56d22be6523 100644 --- a/litellm-rust/crates/inference-ocr/src/wire.rs +++ b/litellm-rust/crates/inference-ocr/src/wire.rs @@ -190,16 +190,13 @@ mod tests { assert!(is_supported_request("deepseek-ocr", Some("vertex_ai"))); } - #[test] - fn missing_document_source_has_a_typed_public_error() { - for document in [ - serde_json::json!({"type": "document_url"}), - serde_json::json!({"type": "image_url"}), - ] { - assert!(matches!( - decode_document(document), - Err(Error::MissingDocumentUrl) - )); - } + #[rstest] + #[case::document_url(serde_json::json!({"type": "document_url"}))] + #[case::image_url(serde_json::json!({"type": "image_url"}))] + fn missing_document_source_has_a_typed_public_error(#[case] document: serde_json::Value) { + assert!(matches!( + decode_document(document), + Err(Error::MissingDocumentUrl) + )); } } diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs b/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs index c4096767188..f0e72b19fc1 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs @@ -191,12 +191,15 @@ async fn a_multi_page_rejection_reaches_the_caller_with_the_single_page_limit(#[ .await .unwrap_err(); - let Error::Provider { status, body, .. } = error else { - panic!("expected a provider error, got {error:?}"); - }; - assert_eq!(status, 400); + let message = litellm_error::Report::new(&error).to_string(); + assert!(error.is_request()); + assert_eq!(error.http_status_code(), Some(400)); assert!( - body.contains("multi-page documents are not supported"), - "{body}" + message.contains("Request has unsupported document format"), + "{message}" + ); + assert!( + message.contains("multi-page documents are not supported"), + "{message}" ); } diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/azure_document_intelligence.rs b/litellm-rust/crates/inference-ocr/tests/ocr/azure_document_intelligence.rs index b700c4fd128..98875cc777c 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/azure_document_intelligence.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/azure_document_intelligence.rs @@ -393,12 +393,15 @@ async fn malformed_provider_pages_report_the_response_path( } #[rstest] -#[case::missing(None)] -#[case::relative(Some("/relative"))] -#[case::cross_origin(Some("http://example.com/operation"))] -#[case::with_userinfo(Some("http://user:password@127.0.0.1/operation"))] +#[case::missing(None, false)] +#[case::relative(Some("/relative"), true)] +#[case::cross_origin(Some("http://example.com/operation"), false)] +#[case::with_userinfo(Some("http://user:password@127.0.0.1/operation"), false)] #[tokio::test] -async fn an_unusable_operation_location_is_rejected(#[case] location: Option<&str>) { +async fn an_unusable_operation_location_is_rejected( + #[case] location: Option<&str>, + #[case] unparseable: bool, +) { let response = location .into_iter() .fold(ResponseTemplate::new(202), |response, location| { @@ -411,6 +414,11 @@ async fn an_unusable_operation_location_is_rejected(#[case] location: Option<&st .unwrap_err(); assert!(error.to_string().contains("operation-location"), "{error}"); + assert!(error.is_response()); + assert_eq!( + litellm_error::find::(&error).is_some(), + unparseable + ); assert_eq!(received(&upstream).await.len(), 1); } diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs index 112e9cf299b..08df14c988d 100644 --- a/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs @@ -17,6 +17,7 @@ const AWS_JSON_CONTENT_TYPE: &str = "application/x-amz-json-1.1"; const TARGET_HEADER: &str = "X-Amz-Target"; const CONTENT_TYPE_HEADER: &str = "Content-Type"; const UNSUPPORTED_DOCUMENT: &str = "UnsupportedDocumentException"; +const SYNC_API_GUIDANCE: &str = "aws_textract uses Textract's synchronous API, which reads a JPEG, PNG, or a single-page PDF or TIFF; other formats and multi-page documents are not supported"; const SYNC_DOCUMENT_MAX_BYTES: usize = 10 * 1024 * 1024; const HEALTH_CHECK_IMAGE_DATA_URI: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4//8/AAX+Av4N70a4AAAAAElFTkSuQmCC"; @@ -322,19 +323,17 @@ struct AwsError { /// with a bare "unsupported document format", which reads like a corrupt file. /// Say what the synchronous API accepts. pub(super) fn error_class(body: String, status: u16, headers: reqwest::header::HeaderMap) -> Error { - let unsupported = serde_json::from_str::(&body) - .ok() - .filter(|error| error.kind.ends_with(UNSUPPORTED_DOCUMENT)); - Error::Provider { - status, - body: match unsupported { - Some(error) => format!( - "{UNSUPPORTED_DOCUMENT}: {}. aws_textract uses Textract's synchronous API, which reads a JPEG, PNG, or a single-page PDF or TIFF; other formats and multi-page documents are not supported", - error.message - ), - None => body, + match serde_json::from_str::(&body) { + Ok(error) if error.kind.ends_with(UNSUPPORTED_DOCUMENT) => Error::RejectedDocument { + provider: "AWS Textract", + message: error.message, + guidance: SYNC_API_GUIDANCE, + }, + Ok(_) | Err(_) => Error::Provider { + status, + body, + headers, }, - headers, } } @@ -408,11 +407,11 @@ mod tests { fn a_model_outside_the_operations_is_refused_with_the_supported_names(#[case] model: &str) { let error = TextractOperation::from_model(model).unwrap_err(); - assert_eq!( - error.to_string(), - format!( - "invalid model: aws_textract has no model {model:?} - use one of: detect-document-text, analyze-document" - ) + let message = error.to_string(); + assert!(message.contains(&format!("{model:?}")), "{message}"); + assert!( + message.contains("detect-document-text, analyze-document"), + "{message}" ); assert_eq!(error.http_status_code(), Some(400)); } @@ -542,46 +541,46 @@ mod tests { #[rstest] #[case::bare_type( r#"{"__type":"UnsupportedDocumentException","Message":"Request has unsupported document format"}"#, - Some("Request has unsupported document format") + "Request has unsupported document format" )] #[case::namespaced_type( r#"{"__type":"com.amazonaws.textract#UnsupportedDocumentException","Message":"bad"}"#, - Some("bad") + "bad" )] #[case::lowercase_message( r#"{"__type":"UnsupportedDocumentException","message":"bad"}"#, - Some("bad") + "bad" )] - #[case::other_exception(r#"{"__type":"AccessDeniedException","Message":"no"}"#, None)] - #[case::json_without_a_type(r#"{"Message":"no"}"#, None)] - #[case::not_json("bad gateway", None)] - fn only_an_unsupported_document_gains_the_sync_api_hint( + fn an_unsupported_document_is_the_callers_fault_with_the_sync_api_hint( #[case] body: &str, - #[case] hinted_message: Option<&str>, + #[case] provider_message: &str, ) { + let error = error_class(body.into(), 400, reqwest::header::HeaderMap::new()); + + let message = error.to_string(); + assert!(error.is_request()); + assert_eq!(error.http_status_code(), Some(400)); + assert!(error.upstream_response().is_none()); + assert!(message.contains(provider_message), "{message}"); + assert!(message.contains(HINT), "{message}"); + } + + #[rstest] + #[case::other_exception(r#"{"__type":"AccessDeniedException","Message":"no"}"#)] + #[case::json_without_a_type(r#"{"Message":"no"}"#)] + #[case::not_json("bad gateway")] + fn any_other_rejection_passes_through_untouched(#[case] body: &str) { let response_headers = reqwest::header::HeaderMap::from_iter([( reqwest::header::HeaderName::from_static("x-amzn-requestid"), reqwest::header::HeaderValue::from_static("abc"), )]); - let Error::Provider { - status, - body: reported, - headers, - } = error_class(body.into(), 400, response_headers.clone()) - else { - panic!("expected a provider error"); - }; + let error = error_class(body.into(), 403, response_headers.clone()); - assert_eq!(status, 400); - assert_eq!(headers, response_headers); - match hinted_message { - Some(message) => { - assert!(reported.contains(message), "{reported}"); - assert!(reported.contains(HINT), "{reported}"); - } - None => assert_eq!(reported, body), - } + assert_eq!( + error.upstream_response(), + Some((403, body, &response_headers)) + ); } #[rstest] diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs index 06906407124..d7233db0a76 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs @@ -13,7 +13,7 @@ use tokio::time::Instant; use crate::base_llm::ocr::{ document::InlineDocument, - error::{Error, PagesFault}, + error::{Error, PagesFault, PollUrl}, handler::{CallHooks, OcrClient, read_json_response}, settings::OcrSettings, transformation::{ @@ -457,8 +457,14 @@ async fn read_operation_response( .and_then(|value| value.to_str().ok()) .ok_or(Error::PollLocation)? .to_string(); - let original = Url::parse(original_url).map_err(|_| Error::PollOrigin)?; - let operation = Url::parse(&location).map_err(|_| Error::PollOrigin)?; + let original = Url::parse(original_url).map_err(|source| Error::InvalidPollUrl { + url: PollUrl::Submission, + source, + })?; + let operation = Url::parse(&location).map_err(|source| Error::InvalidPollUrl { + url: PollUrl::OperationLocation, + source, + })?; if original.origin() != operation.origin() || !operation.username().is_empty() || operation.password().is_some() diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs index 40302251320..249c05f4dbc 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs @@ -31,6 +31,12 @@ pub enum Error { provider: &'static str, guidance: &'static str, }, + #[error("{provider} cannot read this document: {message} - {guidance}")] + RejectedDocument { + provider: &'static str, + message: String, + guidance: &'static str, + }, #[error("Invalid `req_format`. Expected 'native' or 'litellm'.")] RequestFormat, #[error("invalid OCR request field: {path}")] @@ -71,6 +77,12 @@ pub enum Error { NumericRange(&'static str), #[error("OCR accepted response is missing a valid operation-location")] PollLocation, + #[error("OCR {url} is not a valid URL")] + InvalidPollUrl { + url: PollUrl, + #[source] + source: url::ParseError, + }, #[error("OCR operation-location must use the submission origin without credentials")] PollOrigin, #[error("OCR polling timed out")] @@ -120,6 +132,14 @@ pub enum PagesFault { Range, } +#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] +pub enum PollUrl { + #[error("submission URL")] + Submission, + #[error("operation-location")] + OperationLocation, +} + impl From for Error { fn from(fault: litellm_host::machine::MachineFault) -> Self { Self::Host(fault) @@ -162,6 +182,7 @@ impl Error { | Self::DocumentTask(_) | Self::InvalidMimeType(_) | Self::UnsupportedDocument { .. } + | Self::RejectedDocument { .. } | Self::RequestFormat | Self::RequestField { .. } | Self::Arguments(_) @@ -178,6 +199,7 @@ impl Error { | Self::OperationStatus(_) | Self::NumericRange(_) | Self::PollLocation + | Self::InvalidPollUrl { .. } | Self::PollOrigin | Self::PollTimeout | Self::Unsupported(_) @@ -203,6 +225,7 @@ impl Error { | Self::DocumentTask(_) | Self::InvalidMimeType(_) | Self::UnsupportedDocument { .. } + | Self::RejectedDocument { .. } | Self::RequestFormat | Self::RequestField { .. } | Self::Arguments(_) @@ -219,6 +242,7 @@ impl Error { | Self::OperationStatus(_) | Self::NumericRange(_) | Self::PollLocation + | Self::InvalidPollUrl { .. } | Self::PollOrigin | Self::PollTimeout | Self::Unsupported(_) @@ -242,6 +266,7 @@ impl Error { | Self::DocumentTask(_) | Self::InvalidMimeType(_) | Self::UnsupportedDocument { .. } + | Self::RejectedDocument { .. } | Self::RequestFormat | Self::RequestField { .. } | Self::Arguments(_) @@ -258,6 +283,7 @@ impl Error { | Self::OperationStatus(_) | Self::NumericRange(_) | Self::PollLocation + | Self::InvalidPollUrl { .. } | Self::PollOrigin | Self::PollTimeout | Self::InvalidProvider(_) @@ -290,6 +316,7 @@ impl Error { Self::EmptyFile | Self::InvalidMimeType(_) | Self::UnsupportedDocument { .. } + | Self::RejectedDocument { .. } | Self::RequestFormat | Self::RequestField { .. } | Self::Arguments(_) @@ -315,6 +342,7 @@ impl Error { | Self::OperationStatus(_) | Self::NumericRange(_) | Self::PollLocation + | Self::InvalidPollUrl { .. } | Self::PollOrigin | Self::PollTimeout | Self::Unsupported(_) @@ -333,6 +361,7 @@ impl Error { | Self::OperationStatus(_) | Self::NumericRange(_) | Self::PollLocation + | Self::InvalidPollUrl { .. } | Self::PollOrigin | Self::PollTimeout => true, Self::Http(error) => error.is_response(), @@ -342,6 +371,7 @@ impl Error { | Self::DocumentTask(_) | Self::InvalidMimeType(_) | Self::UnsupportedDocument { .. } + | Self::RejectedDocument { .. } | Self::RequestFormat | Self::RequestField { .. } | Self::Arguments(_) diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index 9173468bba4..bd86bab2782 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -563,6 +563,8 @@ async fn upload_bytes_async( #[cfg(test)] mod tests { + use rstest::rstest; + use super::*; #[tokio::test] @@ -635,15 +637,16 @@ mod tests { ); } + #[rstest] + #[case::fractional_pages(json!({"num_pages":1.5}))] + #[case::non_numeric_pages(json!({"num_pages":[]}))] + #[case::non_numeric_credits(json!({"credits":{}}))] + fn invalid_usage_is_rejected(#[case] usage: Value) { + assert!(serde_json::from_value::(json!({"usage":usage})).is_err()); + } + #[test] - fn usage_uses_shared_validation_while_block_page_numbers_are_best_effort() { - for usage in [ - json!({"num_pages":1.5}), - json!({"num_pages":[]}), - json!({"credits":{}}), - ] { - assert!(serde_json::from_value::(json!({"usage":usage})).is_err()); - } + fn block_page_numbers_are_best_effort() { let response = serde_json::from_value(json!({"result":{"chunks":[{"blocks":[ {"content":"ignored", "bbox":{"page":"invalid"}}, {"content":"kept", "bbox":{"page":2.5}, "extra":null} diff --git a/litellm-rust/crates/python-compat/src/error.rs b/litellm-rust/crates/python-compat/src/error.rs index 328c54aa934..6547907056e 100644 --- a/litellm-rust/crates/python-compat/src/error.rs +++ b/litellm-rust/crates/python-compat/src/error.rs @@ -20,7 +20,7 @@ pub enum Error { InvalidJsonKey(&'static str), #[error("Out of range float values are not JSON compliant")] NonFiniteFloat, - #[error("integer does not fit in the JSON number range")] + #[error("integer does not fit in 64 bits")] IntegerOutOfRange, #[error("Object of type {0} cannot be pickled as plain data")] NotPicklable(&'static str), diff --git a/litellm-rust/crates/python-compat/src/literal.rs b/litellm-rust/crates/python-compat/src/literal.rs index 12e9036572f..336c6dfcb2e 100644 --- a/litellm-rust/crates/python-compat/src/literal.rs +++ b/litellm-rust/crates/python-compat/src/literal.rs @@ -58,13 +58,6 @@ impl Term { kind: Kind::Other, } } - - fn is_number(&self) -> bool { - matches!( - self.value, - Value::Int(_) | Value::Float(_) | Value::Complex { .. } - ) - } } struct Parser<'a> { @@ -204,18 +197,25 @@ impl Parser<'_> { self.enter()?; let operand = self.unary()?; self.depth -= 1; - if operand.kind != Kind::Constant || !operand.is_number() { + if operand.kind != Kind::Constant { return Err(Error::InvalidLiteral(self.offset + at)); } - let value = if !negative { - operand.value - } else { - match operand.value { - Value::Int(value) => Value::Int(-value), - Value::Float(value) => Value::Float(-value), - Value::Complex { re, im } => Value::Complex { re: -re, im: -im }, - _ => unreachable!("checked numeric above"), - } + let value = match (operand.value, negative) { + (value @ (Value::Int(_) | Value::Float(_) | Value::Complex { .. }), false) => value, + (Value::Int(value), true) => Value::Int(-value), + (Value::Float(value), true) => Value::Float(-value), + (Value::Complex { re, im }, true) => Value::Complex { re: -re, im: -im }, + ( + Value::None + | Value::Bool(_) + | Value::Str(_) + | Value::Bytes(_) + | Value::List(_) + | Value::Tuple(_) + | Value::Dict(_) + | Value::Set(_), + _, + ) => return Err(Error::InvalidLiteral(self.offset + at)), }; Ok(Term { value, diff --git a/litellm-rust/crates/python-compat/src/pickle.rs b/litellm-rust/crates/python-compat/src/pickle.rs index 78161d33a73..dd7ac6a0e4f 100644 --- a/litellm-rust/crates/python-compat/src/pickle.rs +++ b/litellm-rust/crates/python-compat/src/pickle.rs @@ -63,7 +63,7 @@ impl Serialize for Pickled<'_> { Value::Bool(value) => serializer.serialize_bool(*value), Value::Int(value) => { let value = i64::try_from(value) - .map_err(|_| serde::ser::Error::custom("integer out of i64 range"))?; + .map_err(|_| serde::ser::Error::custom(Error::IntegerOutOfRange))?; serializer.serialize_i64(value) } Value::Float(value) => serializer.serialize_f64(*value), @@ -91,7 +91,7 @@ impl Serialize for Pickled<'_> { map.end() } value @ (Value::Set(_) | Value::Complex { .. }) => Err(serde::ser::Error::custom( - format!("{} cannot be pickled as plain data", value.type_name()), + Error::NotPicklable(value.type_name()), )), } } diff --git a/litellm-rust/crates/token-counter-tiktoken/src/ranks.rs b/litellm-rust/crates/token-counter-tiktoken/src/ranks.rs index 3a9cf7d2d3f..36c30e711c1 100644 --- a/litellm-rust/crates/token-counter-tiktoken/src/ranks.rs +++ b/litellm-rust/crates/token-counter-tiktoken/src/ranks.rs @@ -83,7 +83,23 @@ pub(super) fn load( name: &str, load_file: impl FnOnce(&str) -> std::io::Result, ) -> Result<(&'static Loaded, &'static str), Error> { - let (requested, canonical, file, cache) = match name { + let (requested, canonical, file, cache) = encoding(name)?; + let loaded = cache.get_or_try_init(|| { + let ranks = load_file(file).map_err(|source| Error::Read { file, source })?; + build(canonical, &ranks) + })?; + Ok((loaded, requested)) +} + +type Encoding = ( + &'static str, + &'static str, + &'static str, + &'static OnceCell, +); + +fn encoding(name: &str) -> Result { + Ok(match name { "cl100k_base" => ("cl100k_base", "cl100k_base", CL100K, &CL100K_ENCODER), "o200k_base" => ("o200k_base", "o200k_base", O200K, &O200K_ENCODER), "o200k_harmony" => ("o200k_harmony", "o200k_harmony", O200K, &HARMONY_ENCODER), @@ -92,12 +108,7 @@ pub(super) fn load( "r50k_base" => ("r50k_base", "r50k_base", P50K, &R50K_ENCODER), "gpt2" => ("gpt2", "r50k_base", P50K, &R50K_ENCODER), _ => return Err(Error::UnsupportedTokenizer(name.to_owned())), - }; - let loaded = cache.get_or_try_init(|| { - let ranks = load_file(file).map_err(|source| Error::Read { file, source })?; - build(canonical, &ranks) - })?; - Ok((loaded, requested)) + }) } fn build(name: &str, ranks: &str) -> Result { @@ -219,79 +230,117 @@ mod tests { ) } - #[test] - fn packaged_encodings_match_embedded_encodings_and_reuse_successful_loads() { - for name in [ + #[rstest] + #[case::cl100k_base("cl100k_base", None)] + #[case::o200k_base("o200k_base", None)] + #[case::o200k_harmony("o200k_harmony", None)] + #[case::p50k_base("p50k_base", None)] + #[case::p50k_edit("p50k_edit", None)] + #[case::r50k_base_shared_with_gpt2("r50k_base", Some("gpt2"))] + fn packaged_encodings_load_once_and_match_embedded_encodings( + #[case] name: &str, + #[case] alias: Option<&str>, + ) { + assert!( + TiktokenTokenizer::from_cached_ranks(name, |_| { + Err(std::io::Error::other("unreadable vocabulary")) + }) + .is_err() + ); + let loads = std::sync::atomic::AtomicUsize::new(0); + let barrier = std::sync::Barrier::new(4); + let encoders = std::thread::scope(|scope| { + let tasks: Vec<_> = (0..4) + .map(|_| { + scope.spawn(|| { + barrier.wait(); + TiktokenTokenizer::from_cached_ranks(name, |file| { + loads.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + read_packaged_ranks(file) + }) + .unwrap() + }) + }) + .collect(); + tasks + .into_iter() + .map(|task| task.join().unwrap()) + .collect::>() + }); + assert_eq!(loads.into_inner(), 1); + let actual = &encoders[0]; + let expected = TiktokenTokenizer::from_name(name).unwrap(); + assert_eq!(actual.special_tokens(), expected.special_tokens()); + let specials: Vec<_> = expected.special_tokens().into_iter().collect(); + let special_text = specials.join(" "); + assert_eq!( + actual.encode_special(&special_text, &specials).unwrap(), + expected.encode_special(&special_text, &specials).unwrap() + ); + assert_eq!(actual.encode(&special_text), expected.encode(&special_text)); + let cached = + TiktokenTokenizer::from_cached_ranks(name, |_| panic!("reloaded cached vocabulary")) + .unwrap(); + assert_eq!(cached.encode("cached"), expected.encode("cached")); + assert_eq!(cached.name(), name); + assert!(expected.vocabulary().is_none()); + assert_vocabulary_lookups(name, actual); + if let Some(alias) = alias { + let shared = TiktokenTokenizer::from_cached_ranks(alias, |_| { + panic!("reloaded the shared vocabulary") + }) + .unwrap(); + assert_eq!(shared.name(), alias); + assert_vocabulary_lookups(alias, &shared); + } + } + + fn uncached(name: &str) -> TiktokenTokenizer { + static BUILT: std::sync::Mutex> = + std::sync::Mutex::new(Vec::new()); + let (requested, canonical, file, _) = encoding(name).unwrap(); + let mut built = BUILT.lock().unwrap(); + let loaded = match built.iter().find(|(built, _)| *built == requested) { + Some((_, loaded)) => *loaded, + None => { + let ranks = read_packaged_ranks(file).unwrap(); + let loaded: &'static Loaded = + Box::leak(Box::new(build(canonical, &ranks).unwrap())); + built.push((requested, loaded)); + loaded + } + }; + TiktokenTokenizer { + encoder: &loaded.bpe, + vocabulary: Some(&loaded.vocabulary), + name: requested, + } + } + + #[rstest] + fn packaged_ranks_encode_like_embedded_encodings( + #[values( "cl100k_base", "o200k_base", "o200k_harmony", "p50k_base", "p50k_edit", "r50k_base", - "gpt2", - ] { - if name != "gpt2" { - assert!( - TiktokenTokenizer::from_cached_ranks(name, |_| { - Err(std::io::Error::other("unreadable vocabulary")) - }) - .is_err() - ); - } - let loads = std::sync::atomic::AtomicUsize::new(0); - let barrier = std::sync::Barrier::new(4); - let encoders = std::thread::scope(|scope| { - let tasks: Vec<_> = (0..4) - .map(|_| { - scope.spawn(|| { - barrier.wait(); - TiktokenTokenizer::from_cached_ranks(name, |file| { - loads.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - read_packaged_ranks(file) - }) - .unwrap() - }) - }) - .collect(); - tasks - .into_iter() - .map(|task| task.join().unwrap()) - .collect::>() - }); - assert_eq!(loads.into_inner(), usize::from(name != "gpt2")); - let actual = &encoders[0]; - let expected = TiktokenTokenizer::from_name(name).unwrap(); - assert_eq!(actual.special_tokens(), expected.special_tokens()); - let specials: Vec<_> = expected.special_tokens().into_iter().collect(); - let special_text = specials.join(" "); - assert_eq!( - actual.encode_special(&special_text, &specials).unwrap(), - expected.encode_special(&special_text, &specials).unwrap() - ); - for text in [ - "", - "café 漢字 ع 🙂", - "a\r\nb\t ", - " hello 123456789", - &special_text, - ] { - let ids = expected.encode(text); - assert_eq!(actual.encode(text), ids, "{name}: {text:?}"); - assert_eq!(actual.count_tokens(text), ids.len(), "{name}: {text:?}"); - assert_eq!( - actual.decode_bytes(&ids).unwrap(), - expected.decode_bytes(&ids).unwrap() - ); - } - let cached = TiktokenTokenizer::from_cached_ranks(name, |_| { - panic!("reloaded cached vocabulary") - }) - .unwrap(); - assert_eq!(cached.encode("cached"), expected.encode("cached")); - assert_eq!(cached.name(), name); - assert!(expected.vocabulary().is_none()); - assert_vocabulary_lookups(name, actual); - } + "gpt2" + )] + name: &str, + #[values("", "café 漢字 ع 🙂", "a\r\nb\t ", " hello 123456789")] text: &str, + ) { + let actual = uncached(name); + let expected = TiktokenTokenizer::from_name(name).unwrap(); + + let ids = expected.encode(text); + assert_eq!(actual.encode(text), ids); + assert_eq!(actual.count_tokens(text), ids.len()); + assert_eq!( + actual.decode_bytes(&ids).unwrap(), + expected.decode_bytes(&ids).unwrap() + ); } /// The token-level lookups tiktoken's Python `Encoding` exposes, checked against the