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