scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-30 06:43:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 249

refactor(rust-error): fix the remaining sweep findings in gateway, llms, auth and cache crates

Browse source 
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 <noreply@anthropic.com>
This commit is contained in:
Yujong Lee 2026-10-23 14:38:14 -07:00
parent 22141f2649
commit cbce50f591
28 changed files with 609 additions and 253 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm-rust/Cargo.lock generated
View file

@ -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",

8
litellm-rust/crates/auth-aws/src/aws.rs
View file

@ -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]

10
litellm-rust/crates/auth-azure/src/resolve.rs
View file

@ -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"),

29
litellm-rust/crates/auth-gcp/src/auth.rs
View file

@ -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,

69
litellm-rust/crates/auth-types/src/error.rs
View file

@ -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>),
#[error("no credential in the chain succeeded{}", earlier_attempts(earlier))]
CredentialChain {
#[source]
last: Box<Error>,
earlier: Vec<Error>,
},
#[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::<Vec<_>>()
.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::<io::Error>(&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::<io::Error>(&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}"
);
}

3
litellm-rust/crates/cache-azure-blob/src/transport.rs
View file

@ -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();

5
litellm-rust/crates/cache-s3/src/transport.rs
View file

@ -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())

7
litellm-rust/crates/gateway-inference/src/chat_completions.rs
View file

@ -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<Arc<Gateway>>,

55
litellm-rust/crates/gateway-inference/src/error.rs
View file

@ -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::<Value>(&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<String, Value> {
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<String, Value>) -> 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())]

13
litellm-rust/crates/gateway-inference/src/request.rs
View file

@ -21,6 +21,19 @@ pub struct JsonObject(pub Map<String, Value>);
pub struct RequestId(pub Option<String>);
pub(crate) struct Path<T>(pub T);
impl<T: serde::de::DeserializeOwned + Send, S: Send + Sync> FromRequestParts<S> for Path<T> {
type Rejection = Error;
async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
axum::extract::Path::<T>::from_request_parts(parts, state)
.await
.map(|axum::extract::Path(value)| Self(value))
.map_err(Error::InvalidPath)
}
}
impl<S: Send + Sync> FromRequestParts<S> for RequestId {
type Rejection = std::convert::Infallible;

120
litellm-rust/crates/gateway-mcp/src/error.rs
View file

@ -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))]

15
litellm-rust/crates/gateway-mcp/src/native.rs
View file

@ -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) => {

11
litellm-rust/crates/gateway-mcp/src/native/catalog.rs
View file

@ -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)

37
litellm-rust/crates/http/src/error.rs
View file

@ -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));
}

22
litellm-rust/crates/http/src/outbound.rs
View file

@ -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<B: Into<reqwest::Body>>(request: http::Request<B>) -> Result<Request, Error> {
Request::try_from(request).map_err(|error| Repr::RequestUrl(error.without_url()).into())
}
pub fn signed(mut request: Request, signer: &dyn RequestSigner) -> Result<Request, Error> {
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::<u8>::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}");
}
}

1
litellm-rust/crates/inference-ocr/Cargo.toml
View file

@ -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

19
litellm-rust/crates/inference-ocr/src/document.rs
View file

@ -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());
}
}

19
litellm-rust/crates/inference-ocr/src/wire.rs
View file

@ -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)
));
}
}

15
litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs
View file

@ -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}"
);
}

18
litellm-rust/crates/inference-ocr/tests/ocr/azure_document_intelligence.rs
View file

@ -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::<url::ParseError>(&error).is_some(),
unparseable
);
assert_eq!(received(&upstream).await.len(), 1);
}

83
litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs
View file

@ -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::<AwsError>(&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::<AwsError>(&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("<html>bad gateway</html>", 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("<html>bad gateway</html>")]
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]

12
litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs
View file

@ -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()

30
litellm-rust/crates/llms/src/base_llm/ocr/error.rs
View file

@ -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<litellm_host::machine::MachineFault> 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(_)

19
litellm-rust/crates/llms/src/reducto/ocr/transformation.rs
View file

@ -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::<ReductoResponse>(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::<ReductoResponse>(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}

2
litellm-rust/crates/python-compat/src/error.rs
View file

@ -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),

34
litellm-rust/crates/python-compat/src/literal.rs
View file

@ -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,

4
litellm-rust/crates/python-compat/src/pickle.rs
View file

@ -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()),
)),
}
}

197
litellm-rust/crates/token-counter-tiktoken/src/ranks.rs
View file

@ -83,7 +83,23 @@ pub(super) fn load(
name: &str,
load_file: impl FnOnce(&str) -> std::io::Result<String>,
) -> 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<Loaded>,
);
fn encoding(name: &str) -> Result<Encoding, Error> {
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<Loaded, Error> {
@ -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::<Vec<_>>()
});
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<Vec<(&'static str, &'static Loaded)>> =
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::<Vec<_>>()
});
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