diff --git a/.dockerignore b/.dockerignore index 6b80caeaf9f..f3a80fee3e4 100644 --- a/.dockerignore +++ b/.dockerignore @@ -49,6 +49,8 @@ build/ *.egg-info/ .DS_Store **/node_modules +ui/litellm-dashboard/.next +ui/litellm-dashboard/out litellm-rust/target/ litellm/rust_bridge/_native*.so *.log diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 9658baeb89a..12ad124fa20 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -1,17 +1,17 @@ ## Relevant issues - + ## Linear ticket - + ## Pre-Submission checklist **Please complete all items before asking a LiteLLM maintainer to review your PR** - [ ] I have added meaningful tests -- [ ] My PR passes all unit tests on [`make test-unit`](https://docs.litellm.ai/docs/extras/contributing_code) +- [ ] My PR passes all CI/CD checks (e.g., lint, format, unit tests) - [ ] My PR's scope is as isolated as possible; it only solves 1 specific problem - [ ] I have requested a Greptile review by commenting `@greptileai` and received a **Confidence Score of at least 4/5** before requesting a maintainer review @@ -19,29 +19,13 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slack (#pr-review)](https://join.slack.com/t/litellmossslack/shared_invite/zt-3o7nkuyfr-p_kbNJj8taRfXGgQI1~YyA). -## CI (LiteLLM team) - -> **CI status guideline:** -> -> - 50-55 passing tests: main is stable with minor issues. -> - 45-49 passing tests: acceptable but needs attention -> - <= 40 passing tests: unstable; be careful with your merges and assess the risk. - -- [ ] **Branch creation CI run** - Link: - -- [ ] **CI run for the last commit** - Link: - -- [ ] **Merge / cherry-pick CI run** - Links: - ## Screenshots / Proof of Fix - + ## Type diff --git a/Dockerfile b/Dockerfile index 681681f28cc..b6fef1a21fc 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,12 +1,33 @@ +# syntax=docker/dockerfile:1.7 + # Base image for building ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 # Runtime image ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a +# Pinned by digest like the other base images; bump explicitly on Node upgrades. +ARG UI_BUILD_IMAGE=node:20.18-alpine3.20@sha256:3488b10bf958af7125a176419d2d8a9937d895bf124012aae811651988d2ffe6 FROM $UV_IMAGE AS uvbin +# Admin UI builder. Pinned to the build platform so the architecture-independent +# Next.js static export compiles once natively even in a multi-arch build, +# instead of once per target arch under QEMU. +FROM --platform=$BUILDPLATFORM $UI_BUILD_IMAGE AS ui-builder + +ENV NEXT_TELEMETRY_DISABLED=1 \ + npm_config_fund=false \ + npm_config_audit=false + +WORKDIR /ui + +COPY ui/litellm-dashboard/package.json ui/litellm-dashboard/package-lock.json ./ +RUN --mount=type=cache,target=/root/.npm npm ci --prefer-offline + +COPY ui/litellm-dashboard/ ./ +RUN npm run build + # Builder stage FROM $LITELLM_BUILD_IMAGE AS builder @@ -48,7 +69,13 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr # Copy full source tree COPY . . -# Build Admin UI before final sync +# Replace the committed UI bundle with the one built from this exact source. +# Clearing first drops the committed bundle's content-hashed chunks that COPY +# would otherwise leave behind alongside the fresh ones. +RUN rm -rf litellm/proxy/_experimental/out +COPY --from=ui-builder /ui/out/. litellm/proxy/_experimental/out/ + +# Build Admin UI before final sync (applies the enterprise color override when present) RUN sed -i 's/\r$//' docker/build_admin_ui.sh && chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh # Install project and workspace packages (fast - deps already cached) diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index 50ef55e3261..b3af953511d 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -1,12 +1,33 @@ +# syntax=docker/dockerfile:1.7 + # Base image for building ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 # Runtime image ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a +# Pinned by digest like the other base images; bump explicitly on Node upgrades. +ARG UI_BUILD_IMAGE=node:20.18-alpine3.20@sha256:3488b10bf958af7125a176419d2d8a9937d895bf124012aae811651988d2ffe6 FROM $UV_IMAGE AS uvbin +# Admin UI builder. Pinned to the build platform so the architecture-independent +# Next.js static export compiles once natively even in a multi-arch build, +# instead of once per target arch under QEMU. +FROM --platform=$BUILDPLATFORM $UI_BUILD_IMAGE AS ui-builder + +ENV NEXT_TELEMETRY_DISABLED=1 \ + npm_config_fund=false \ + npm_config_audit=false + +WORKDIR /ui + +COPY ui/litellm-dashboard/package.json ui/litellm-dashboard/package-lock.json ./ +RUN --mount=type=cache,target=/root/.npm npm ci --prefer-offline + +COPY ui/litellm-dashboard/ ./ +RUN npm run build + FROM $LITELLM_BUILD_IMAGE AS builder WORKDIR /app @@ -46,7 +67,13 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr # Copy full source tree COPY . . -# Build Admin UI before final sync +# Replace the committed UI bundle with the one built from this exact source. +# Clearing first drops the committed bundle's content-hashed chunks that COPY +# would otherwise leave behind alongside the fresh ones. +RUN rm -rf litellm/proxy/_experimental/out +COPY --from=ui-builder /ui/out/. litellm/proxy/_experimental/out/ + +# Build Admin UI before final sync (applies the enterprise color override when present) RUN sed -i 's/\r$//' docker/build_admin_ui.sh && chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh # Install project and workspace packages (fast - deps already cached) diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 6bb925aa723..c24cb9008f0 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -1,11 +1,32 @@ +# syntax=docker/dockerfile:1.7 + # Base images ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG PROXY_EXTRAS_SOURCE=published ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a +# Pinned by digest like the other base images; bump explicitly on Node upgrades. +ARG UI_BUILD_IMAGE=node:20.18-alpine3.20@sha256:3488b10bf958af7125a176419d2d8a9937d895bf124012aae811651988d2ffe6 FROM $UV_IMAGE AS uvbin +# Admin UI builder. Pinned to the build platform so the architecture-independent +# Next.js static export compiles once natively even in a multi-arch build, +# instead of once per target arch under QEMU. +FROM --platform=$BUILDPLATFORM $UI_BUILD_IMAGE AS ui-builder + +ENV NEXT_TELEMETRY_DISABLED=1 \ + npm_config_fund=false \ + npm_config_audit=false + +WORKDIR /ui + +COPY ui/litellm-dashboard/package.json ui/litellm-dashboard/package-lock.json ./ +RUN --mount=type=cache,target=/root/.npm npm ci --prefer-offline + +COPY ui/litellm-dashboard/ ./ +RUN npm run build + FROM $LITELLM_BUILD_IMAGE AS builder ARG PROXY_EXTRAS_SOURCE WORKDIR /app @@ -53,6 +74,12 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ # Copy full source tree COPY . . +# Replace the committed UI bundle with the one built from this exact source. +# Clearing first drops the committed bundle's content-hashed chunks that COPY +# would otherwise leave behind alongside the fresh ones. +RUN rm -rf litellm/proxy/_experimental/out +COPY --from=ui-builder /ui/out/. litellm/proxy/_experimental/out/ + # Set non-root flag for build time consistency ENV LITELLM_NON_ROOT=true diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ce86a0ee6ac..9bffe9f9ec6 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -673,6 +673,7 @@ dependencies = [ "serde", "serde_json", "thiserror 2.0.18", + "tokio", ] [[package]] diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml index 2f414159158..4055be36785 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -25,7 +25,7 @@ futures-util.workspace = true serde_json.workspace = true base64.workspace = true axum = { workspace = true, features = ["ws"], optional = true } -serde = { workspace = true, optional = true } +serde.workspace = true subtle = { workspace = true, optional = true } # sha2 hashes the master key into user_api_key_hash (matches the proxy's # SHA-256 hash_token) so the plaintext credential never enters a log payload. @@ -34,7 +34,7 @@ pyo3 = { workspace = true, features = ["auto-initialize"], optional = true } [features] default = [] -server = ["dep:axum", "dep:subtle", "dep:serde", "dep:sha2"] +server = ["dep:axum", "dep:subtle", "dep:sha2"] # Build the gateway's config from the proxy YAML via an embedded Python # interpreter (links libpython; requires `litellm` importable at runtime). python-config = ["dep:pyo3"] diff --git a/litellm-rust/crates/ai-gateway/src/constants.rs b/litellm-rust/crates/ai-gateway/src/constants.rs index 3116a4c9932..109b648f5db 100644 --- a/litellm-rust/crates/ai-gateway/src/constants.rs +++ b/litellm-rust/crates/ai-gateway/src/constants.rs @@ -26,4 +26,5 @@ pub(crate) const DEFAULT_MAX_BATCH_SIZE: usize = 256; pub(crate) const DEFAULT_FLUSH_INTERVAL_MS: u64 = 500; /// Provider attributed to realtime sessions in the logging payload. +#[cfg(feature = "server")] pub(crate) const DEFAULT_PROVIDER: &str = "openai"; diff --git a/litellm-rust/crates/ai-gateway/src/integrations/README.md b/litellm-rust/crates/ai-gateway/src/integrations/README.md new file mode 100644 index 00000000000..16a162dac57 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/README.md @@ -0,0 +1,127 @@ +# LiteLLM Rust integrations + +This directory contains Rust-native equivalents of LiteLLM integration hooks. +The first supported surfaces are terminal custom loggers and pre/during-call +custom guardrails. + +## File layout + +Every integration is a folder: + +- `mod.rs` contains the implementation, trait, runner, or adapter +- `types.rs` contains the integration-local request, response, error, and future + types + +Do not add new flat integration files such as `custom_logger.rs`. Shared wire +contracts that are used by multiple integrations can stay in +`integrations/types.rs`. + +Call ordering and lifecycle timing live in `litellm-core/src/call_lifecycle`. +Call-type modules, such as OCR, adapt their request and response shapes into +that generic lifecycle runner. + +## CustomLogger + +Implement `CustomLogger` when Rust code needs to observe terminal success or +failure events. Method names intentionally match Python `CustomLogger` names. + +```rust +use litellm_ai_gateway::integrations::custom_logger::{ + CallbackTiming, CallbackValue, CustomLogger, LogFuture, ModelCallDetails, +}; + +struct RecordingLogger; + +impl CustomLogger for RecordingLogger { + fn async_log_success_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + response_obj: &'a CallbackValue, + timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + let model = &model_call_details.model; + let provider = &model_call_details.custom_llm_provider; + let call_type = model_call_details.call_type.to_string(); + let request_id = model_call_details.request_id.as_deref(); + let response_object = &response_obj.object; + let duration = timing.end_time - timing.start_time; + let standard_payload = model_call_details.standard_logging_payload.as_ref(); + + Ok(()) + }) + } + + fn async_log_failure_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + response_obj: Option<&'a CallbackValue>, + timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + let error = model_call_details.failure_error.as_ref(); + let response_object = response_obj.map(|value| value.object.as_str()); + let duration = timing.end_time - timing.start_time; + + Ok(()) + }) + } +} +``` + +Use `CustomLoggerRunner` to fan out terminal events to configured loggers. The +runner is a no-op when no loggers are configured, which is the expected fast +path for requests without callbacks. + +## CustomGuardrail + +Implement `CustomGuardrail` when Rust code needs to run pre-call or native +during-call checks. Method names intentionally match Python `CustomGuardrail` +entrypoints inherited from Python `CustomLogger`. + +```rust +use litellm_ai_gateway::integrations::custom_guardrail::{ + CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailEventHook, + GuardrailFuture, GuardrailRequest, +}; + +struct BlocklistedPromptGuardrail; + +impl CustomGuardrail for BlocklistedPromptGuardrail { + fn guardrail_name(&self) -> &str { + "blocklisted-prompt" + } + + fn supported_event_hooks(&self) -> &[GuardrailEventHook] { + &[GuardrailEventHook::PreCall] + } + + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a GuardrailContext, + request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + Box::pin(async move { + if request.data.to_string().contains("blocked phrase") { + return Ok(GuardrailDecision::Block( + litellm_ai_gateway::integrations::custom_guardrail::GuardrailError::blocked( + "blocked phrase detected", + ), + )); + } + Ok(GuardrailDecision::Allow(request)) + }) + } +} +``` + +Use `CustomGuardrailRunner::run_pre_call` for `pre_call` guardrails and +`CustomGuardrailRunner::run_during_call` for `during_call` guardrails. A +`GuardrailDecision::Mask` continues with modified request data. +`GuardrailDecision::Block` short-circuits the provider call. + +## Current boundary + +These are Rust-only primitives. Python callback and guardrail adapters are a +separate layer that should implement these Rust traits instead of changing the +runner interfaces. diff --git a/litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/mod.rs b/litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/mod.rs new file mode 100644 index 00000000000..e5d4ce3a708 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/mod.rs @@ -0,0 +1,468 @@ +//! Rust mirror of Python `CustomGuardrail` entrypoints used by the proxy. +//! +//! This module is intentionally Rust-only: Python/PyO3 adapters are a later +//! layer that should implement this trait rather than changing the runner. + +use std::future::Future; +use std::sync::Arc; + +use crate::integrations::custom_logger::{ + CallbackTiming, CallbackValue, CustomLoggerRunner, LoggingError, ModelCallDetails, +}; + +pub mod types; + +pub use types::{ + GuardrailContext, GuardrailDecision, GuardrailDispatchReport, GuardrailError, + GuardrailEventHook, GuardrailFuture, GuardrailRequest, +}; + +pub trait CustomGuardrail: Send + Sync { + fn guardrail_name(&self) -> &str; + + fn supported_event_hooks(&self) -> &[GuardrailEventHook]; + + /// Python 1:1 name: `async_pre_call_hook(user_api_key_dict, cache, data, call_type)`. + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a GuardrailContext, + request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + Box::pin(async move { Ok(GuardrailDecision::Allow(request)) }) + } + + /// Python 1:1 name: `async_moderation_hook(data, user_api_key_dict, call_type)`. + fn async_moderation_hook<'a>( + &'a self, + _context: &'a GuardrailContext, + request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + Box::pin(async move { Ok(GuardrailDecision::Allow(request)) }) + } +} + +pub struct CustomGuardrailRunner { + guardrails: Vec>, +} + +impl CustomGuardrailRunner { + pub fn new(guardrails: Vec>) -> Self { + Self { guardrails } + } + + pub fn is_empty(&self) -> bool { + self.guardrails.is_empty() + } + + pub async fn run_pre_call( + &self, + context: &GuardrailContext, + request: GuardrailRequest, + ) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> { + self.run_hook(GuardrailEventHook::PreCall, context, request) + .await + } + + pub async fn run_during_call( + &self, + context: &GuardrailContext, + request: GuardrailRequest, + ) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> { + self.run_hook(GuardrailEventHook::DuringCall, context, request) + .await + } + + pub async fn run_before_provider( + &self, + event_hook: GuardrailEventHook, + context: &GuardrailContext, + request: GuardrailRequest, + provider: F, + ) -> Result + where + F: FnOnce(GuardrailRequest) -> Fut, + Fut: Future>, + { + let (request, _) = self.run_hook(event_hook, context, request).await?; + provider(request).await + } + + pub async fn run_pre_call_with_failure_logging( + &self, + context: &GuardrailContext, + request: GuardrailRequest, + logger_runner: &CustomLoggerRunner, + model_call_details: &ModelCallDetails, + timing: CallbackTiming, + ) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> { + match self.run_pre_call(context, request).await { + Ok(result) => Ok(result), + Err(error) => { + let failure_details = model_call_details.clone().with_failure_error(LoggingError { + message: error.message.clone(), + kind: error.kind.clone(), + }); + let response_obj = CallbackValue::new( + "guardrail_error", + serde_json::json!({ + "message": error.message, + "kind": error.kind, + }), + ); + logger_runner + .async_log_failure_event(&failure_details, Some(&response_obj), timing) + .await; + Err(error) + } + } + } + + async fn run_hook( + &self, + event_hook: GuardrailEventHook, + context: &GuardrailContext, + mut request: GuardrailRequest, + ) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> { + if self.guardrails.is_empty() { + return Ok((request, GuardrailDispatchReport::default())); + } + + let mut report = GuardrailDispatchReport::default(); + for guardrail in &self.guardrails { + if !self.should_run(guardrail.as_ref(), event_hook, context) { + continue; + } + + report.invoked += 1; + let decision = match event_hook { + GuardrailEventHook::PreCall => { + guardrail + .async_pre_call_hook(context, request.clone()) + .await? + } + GuardrailEventHook::DuringCall => { + guardrail + .async_moderation_hook(context, request.clone()) + .await? + } + }; + match decision.into_request() { + Ok(next_request) => request = next_request, + Err(error) => return Err(error), + } + } + + Ok((request, report)) + } + + fn should_run( + &self, + guardrail: &dyn CustomGuardrail, + event_hook: GuardrailEventHook, + context: &GuardrailContext, + ) -> bool { + let supports_hook = guardrail.supported_event_hooks().contains(&event_hook); + let selected = context.selected_guardrails.is_empty() + || context + .selected_guardrails + .iter() + .any(|name| name == guardrail.guardrail_name()); + supports_hook && selected + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::integrations::custom_logger::{CallType, CallbackValue, CustomLogger, LogFuture}; + use crate::integrations::types::{StandardLoggingMetadata, StandardLoggingPayload}; + use serde_json::json; + use std::sync::Mutex; + + #[derive(Clone)] + enum TestDecision { + Allow, + Mask, + Block, + } + + struct RecordingCustomGuardrail { + name: String, + hooks: Vec, + decision: TestDecision, + calls: Mutex>, + } + + impl RecordingCustomGuardrail { + fn new(name: &str, hooks: Vec, decision: TestDecision) -> Self { + Self { + name: name.to_string(), + hooks, + decision, + calls: Mutex::new(Vec::new()), + } + } + + fn calls(&self) -> Vec<&'static str> { + self.calls.lock().unwrap().clone() + } + + fn decision(&self, mut request: GuardrailRequest) -> GuardrailDecision { + match self.decision { + TestDecision::Allow => GuardrailDecision::Allow(request), + TestDecision::Mask => { + request.data["masked"] = json!(true); + GuardrailDecision::Mask(request) + } + TestDecision::Block => { + GuardrailDecision::Block(GuardrailError::blocked("blocked by guardrail")) + } + } + } + } + + impl CustomGuardrail for RecordingCustomGuardrail { + fn guardrail_name(&self) -> &str { + &self.name + } + + fn supported_event_hooks(&self) -> &[GuardrailEventHook] { + &self.hooks + } + + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a GuardrailContext, + request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + Box::pin(async move { + self.calls.lock().unwrap().push("async_pre_call_hook"); + Ok(self.decision(request)) + }) + } + + fn async_moderation_hook<'a>( + &'a self, + _context: &'a GuardrailContext, + request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + Box::pin(async move { + self.calls.lock().unwrap().push("async_moderation_hook"); + Ok(self.decision(request)) + }) + } + } + + #[tokio::test] + async fn pre_call_dispatches_to_async_pre_call_hook() { + let guardrail = Arc::new(RecordingCustomGuardrail::new( + "pre", + vec![GuardrailEventHook::PreCall], + TestDecision::Allow, + )); + let runner = CustomGuardrailRunner::new(vec![guardrail.clone()]); + let context = + GuardrailContext::new(CallType::Ocr).with_selected_guardrails(vec!["pre".to_string()]); + let request = GuardrailRequest::new(json!({"messages": ["hello"]})); + + let (result, report) = runner + .run_pre_call(&context, request) + .await + .expect("guardrail allows request"); + + assert_eq!(report.invoked, 1); + assert_eq!(result.data["messages"], json!(["hello"])); + assert_eq!(guardrail.calls(), vec!["async_pre_call_hook"]); + } + + #[tokio::test] + async fn during_call_dispatches_to_async_moderation_hook() { + let guardrail = Arc::new(RecordingCustomGuardrail::new( + "during", + vec![GuardrailEventHook::DuringCall], + TestDecision::Allow, + )); + let runner = CustomGuardrailRunner::new(vec![guardrail.clone()]); + let context = GuardrailContext::new(CallType::Completion) + .with_selected_guardrails(vec!["during".to_string()]); + let request = GuardrailRequest::new(json!({"prompt": "hello"})); + + let (_result, report) = runner + .run_during_call(&context, request) + .await + .expect("guardrail allows request"); + + assert_eq!(report.invoked, 1); + assert_eq!(guardrail.calls(), vec!["async_moderation_hook"]); + } + + #[tokio::test] + async fn mask_decision_continues_with_updated_request() { + let guardrail = Arc::new(RecordingCustomGuardrail::new( + "masker", + vec![GuardrailEventHook::PreCall], + TestDecision::Mask, + )); + let runner = CustomGuardrailRunner::new(vec![guardrail]); + let context = GuardrailContext::new(CallType::Ocr); + let request = GuardrailRequest::new(json!({"document": "secret"})); + + let (result, report) = runner + .run_pre_call(&context, request) + .await + .expect("mask continues"); + + assert_eq!(report.invoked, 1); + assert_eq!(result.data["masked"], json!(true)); + } + + #[tokio::test] + async fn block_decision_short_circuits_and_logs_failure() { + struct RecordingFailureLogger { + errors: Mutex>, + } + + impl CustomLogger for RecordingFailureLogger { + fn async_log_failure_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + _response_obj: Option<&'a CallbackValue>, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + self.errors.lock().unwrap().push( + model_call_details + .failure_error + .as_ref() + .map(|error| error.kind.clone()) + .unwrap_or_default(), + ); + Ok(()) + }) + } + } + + let guardrail = Arc::new(RecordingCustomGuardrail::new( + "blocker", + vec![GuardrailEventHook::PreCall], + TestDecision::Block, + )); + let guardrail_runner = CustomGuardrailRunner::new(vec![guardrail]); + let logger = Arc::new(RecordingFailureLogger { + errors: Mutex::new(Vec::new()), + }); + let logger_runner = CustomLoggerRunner::new(vec![logger.clone()]); + let context = GuardrailContext::new(CallType::Ocr); + let details = ModelCallDetails::from_standard_logging_payload(StandardLoggingPayload { + id: "req_ocr".to_string(), + litellm_call_id: "req_ocr".to_string(), + call_type: "ocr".to_string(), + model: "mistral-ocr-latest".to_string(), + custom_llm_provider: "mistral".to_string(), + response_cost: 0.0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + start_time: 1.0, + end_time: 1.0, + stream: false, + metadata: StandardLoggingMetadata::default(), + messages: None, + }); + + let err = guardrail_runner + .run_pre_call_with_failure_logging( + &context, + GuardrailRequest::new(json!({"document": "bad"})), + &logger_runner, + &details, + CallbackTiming::new(1.0, 2.0), + ) + .await + .expect_err("guardrail blocks request"); + + assert_eq!(err.kind, "GuardrailBlocked"); + assert_eq!( + logger.errors.lock().unwrap().as_slice(), + ["GuardrailBlocked"] + ); + } + + #[tokio::test] + async fn block_decision_short_circuits_later_guardrails_and_provider_work() { + let blocking_guardrail = Arc::new(RecordingCustomGuardrail::new( + "blocker", + vec![GuardrailEventHook::PreCall], + TestDecision::Block, + )); + let later_guardrail = Arc::new(RecordingCustomGuardrail::new( + "later", + vec![GuardrailEventHook::PreCall], + TestDecision::Allow, + )); + let runner = + CustomGuardrailRunner::new(vec![blocking_guardrail.clone(), later_guardrail.clone()]); + let provider_called = Arc::new(Mutex::new(false)); + let provider_called_for_closure = provider_called.clone(); + + let result = runner + .run_before_provider( + GuardrailEventHook::PreCall, + &GuardrailContext::new(CallType::Completion), + GuardrailRequest::new(json!({"prompt": "blocked"})), + move |_request| async move { + *provider_called_for_closure.lock().unwrap() = true; + Ok("provider response") + }, + ) + .await; + + assert!(result.is_err()); + assert_eq!(blocking_guardrail.calls(), vec!["async_pre_call_hook"]); + assert_eq!(later_guardrail.calls(), Vec::<&'static str>::new()); + assert!(!*provider_called.lock().unwrap()); + } + + #[tokio::test] + async fn run_before_provider_returns_provider_guardrail_error_directly() { + let guardrail = Arc::new(RecordingCustomGuardrail::new( + "allow", + vec![GuardrailEventHook::PreCall], + TestDecision::Allow, + )); + let runner = CustomGuardrailRunner::new(vec![guardrail]); + + let result = runner + .run_before_provider( + GuardrailEventHook::PreCall, + &GuardrailContext::new(CallType::Completion), + GuardrailRequest::new(json!({"prompt": "allowed"})), + |_request| async move { + Err::<&'static str, GuardrailError>(GuardrailError::blocked( + "provider-side guardrail error", + )) + }, + ) + .await; + + let err = result.expect_err("provider error is returned directly"); + assert_eq!(err.kind, "GuardrailBlocked"); + assert_eq!(err.message, "provider-side guardrail error"); + } + + #[tokio::test] + async fn no_guardrails_fast_path_dispatches_nothing() { + let runner = CustomGuardrailRunner::new(Vec::new()); + let context = GuardrailContext::new(CallType::Ocr); + let request = GuardrailRequest::new(json!({"document": "ok"})); + + let (result, report) = runner + .run_pre_call(&context, request) + .await + .expect("no guardrails allow request"); + + assert!(runner.is_empty()); + assert_eq!(report, GuardrailDispatchReport::default()); + assert_eq!(result.data["document"], json!("ok")); + } +} diff --git a/litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/types.rs b/litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/types.rs new file mode 100644 index 00000000000..825e56cc0d7 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/types.rs @@ -0,0 +1,110 @@ +use std::collections::HashMap; +use std::future::Future; +use std::pin::Pin; + +use serde_json::Value; + +use crate::integrations::custom_logger::CallType; + +pub type GuardrailFuture<'a> = + Pin> + Send + 'a>>; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum GuardrailEventHook { + PreCall, + DuringCall, +} + +impl GuardrailEventHook { + pub fn as_str(&self) -> &'static str { + match self { + Self::PreCall => "pre_call", + Self::DuringCall => "during_call", + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct GuardrailError { + pub message: String, + pub kind: String, +} + +impl GuardrailError { + pub fn blocked(message: impl Into) -> Self { + Self { + message: message.into(), + kind: "GuardrailBlocked".to_string(), + } + } +} + +impl std::fmt::Display for GuardrailError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}: {}", self.kind, self.message) + } +} + +impl std::error::Error for GuardrailError {} + +#[derive(Clone, Debug)] +pub struct GuardrailContext { + pub call_type: CallType, + pub selected_guardrails: Vec, + pub metadata: HashMap, + pub user_api_key_hash: Option, + pub user_api_key_user_id: Option, + pub user_api_key_team_id: Option, + pub trace_parent: Option, +} + +impl GuardrailContext { + pub fn new(call_type: CallType) -> Self { + Self { + call_type, + selected_guardrails: Vec::new(), + metadata: HashMap::new(), + user_api_key_hash: None, + user_api_key_user_id: None, + user_api_key_team_id: None, + trace_parent: None, + } + } + + pub fn with_selected_guardrails(mut self, selected_guardrails: Vec) -> Self { + self.selected_guardrails = selected_guardrails; + self + } +} + +#[derive(Clone, Debug, PartialEq)] +pub struct GuardrailRequest { + pub data: Value, +} + +impl GuardrailRequest { + pub fn new(data: Value) -> Self { + Self { data } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub enum GuardrailDecision { + Allow(GuardrailRequest), + Mask(GuardrailRequest), + Block(GuardrailError), +} + +impl GuardrailDecision { + pub(super) fn into_request(self) -> Result { + match self { + Self::Allow(request) | Self::Mask(request) => Ok(request), + Self::Block(error) => Err(error), + } + } +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct GuardrailDispatchReport { + pub invoked: usize, +} diff --git a/litellm-rust/crates/ai-gateway/src/integrations/custom_logger.rs b/litellm-rust/crates/ai-gateway/src/integrations/custom_logger.rs deleted file mode 100644 index 53b599d8c98..00000000000 --- a/litellm-rust/crates/ai-gateway/src/integrations/custom_logger.rs +++ /dev/null @@ -1,24 +0,0 @@ -//! The `CustomLogger` trait — the Rust mirror of Python -//! `litellm/integrations/custom_logger.py::CustomLogger`. -//! -//! Synchronous (no `async_trait`): callbacks are O(1) enqueue-and-return so the -//! realtime splice never blocks on a logger. Default bodies are no-ops so a -//! logger can implement only the events it cares about. - -use crate::integrations::types::{LogError, LoggingError, StandardLoggingPayload}; - -pub trait CustomLogger: Send + Sync { - /// Record a successful call. Default: no-op. - fn log_success_event(&self, _payload: &StandardLoggingPayload) -> Result<(), LogError> { - Ok(()) - } - - /// Record a failed call. Default: no-op. - fn log_failure_event( - &self, - _payload: &StandardLoggingPayload, - _error: &LoggingError, - ) -> Result<(), LogError> { - Ok(()) - } -} diff --git a/litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs b/litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs new file mode 100644 index 00000000000..792717dacfc --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs @@ -0,0 +1,317 @@ +//! The `CustomLogger` trait — the Rust mirror of Python +//! `litellm/integrations/custom_logger.py::CustomLogger`. +//! +//! The Python-named async terminal methods are the public Rust callback shape. + +use std::sync::Arc; + +pub mod types; + +pub use types::{ + CallType, CallbackDispatchReport, CallbackTiming, CallbackValue, LogError, LogFuture, + LoggingError, ModelCallDetails, +}; + +pub trait CustomLogger: Send + Sync { + /// Python 1:1 name: `async_log_success_event(model_call_details, response_obj, start_time, end_time)`. + fn async_log_success_event<'a>( + &'a self, + _model_call_details: &'a ModelCallDetails, + _response_obj: &'a CallbackValue, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async { Ok(()) }) + } + + /// Python 1:1 name: `async_log_failure_event(model_call_details, response_obj, start_time, end_time)`. + fn async_log_failure_event<'a>( + &'a self, + _model_call_details: &'a ModelCallDetails, + _response_obj: Option<&'a CallbackValue>, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async { Ok(()) }) + } +} + +pub struct CustomLoggerRunner { + loggers: Vec>, +} + +impl CustomLoggerRunner { + pub fn new(loggers: Vec>) -> Self { + Self { loggers } + } + + pub fn is_empty(&self) -> bool { + self.loggers.is_empty() + } + + pub async fn async_log_success_event( + &self, + model_call_details: &ModelCallDetails, + response_obj: &CallbackValue, + timing: CallbackTiming, + ) -> CallbackDispatchReport { + if self.loggers.is_empty() { + return CallbackDispatchReport::default(); + } + + let mut report = CallbackDispatchReport::default(); + for logger in &self.loggers { + report.invoked += 1; + if let Err(err) = logger + .async_log_success_event(model_call_details, response_obj, timing) + .await + { + report.dropped += 1; + eprintln!("litellm-ai-gateway: async_log_success_event dropped: {err}"); + } + } + report + } + + pub async fn async_log_failure_event( + &self, + model_call_details: &ModelCallDetails, + response_obj: Option<&CallbackValue>, + timing: CallbackTiming, + ) -> CallbackDispatchReport { + if self.loggers.is_empty() { + return CallbackDispatchReport::default(); + } + + let mut report = CallbackDispatchReport::default(); + for logger in &self.loggers { + report.invoked += 1; + if let Err(err) = logger + .async_log_failure_event(model_call_details, response_obj, timing) + .await + { + report.dropped += 1; + eprintln!("litellm-ai-gateway: async_log_failure_event dropped: {err}"); + } + } + report + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::integrations::types::{StandardLoggingMetadata, StandardLoggingPayload}; + use serde_json::json; + use std::sync::Mutex; + + #[derive(Clone, Debug, PartialEq)] + struct RecordedEvent { + hook: &'static str, + model: String, + provider: String, + call_type: String, + request_id: Option, + litellm_call_id: Option, + user_id: Option, + response_object: Option, + error_kind: Option, + start_time: f64, + end_time: f64, + standard_logging_model: Option, + } + + #[derive(Default)] + struct RecordingCustomLogger { + events: Mutex>, + } + + impl RecordingCustomLogger { + fn events(&self) -> Vec { + self.events.lock().unwrap().clone() + } + } + + impl CustomLogger for RecordingCustomLogger { + fn async_log_success_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + response_obj: &'a CallbackValue, + timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push(RecordedEvent { + hook: "async_log_success_event", + model: model_call_details.model.clone(), + provider: model_call_details.custom_llm_provider.clone(), + call_type: model_call_details.call_type.to_string(), + request_id: model_call_details.request_id.clone(), + litellm_call_id: model_call_details.litellm_call_id.clone(), + user_id: model_call_details.metadata.user_api_key_user_id.clone(), + response_object: Some(response_obj.object.clone()), + error_kind: None, + start_time: timing.start_time, + end_time: timing.end_time, + standard_logging_model: model_call_details + .standard_logging_payload + .as_ref() + .map(|payload| payload.model.clone()), + }); + Ok(()) + }) + } + + fn async_log_failure_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + response_obj: Option<&'a CallbackValue>, + timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push(RecordedEvent { + hook: "async_log_failure_event", + model: model_call_details.model.clone(), + provider: model_call_details.custom_llm_provider.clone(), + call_type: model_call_details.call_type.to_string(), + request_id: model_call_details.request_id.clone(), + litellm_call_id: model_call_details.litellm_call_id.clone(), + user_id: model_call_details.metadata.user_api_key_user_id.clone(), + response_object: response_obj.map(|value| value.object.clone()), + error_kind: model_call_details + .failure_error + .as_ref() + .map(|error| error.kind.clone()), + start_time: timing.start_time, + end_time: timing.end_time, + standard_logging_model: model_call_details + .standard_logging_payload + .as_ref() + .map(|payload| payload.model.clone()), + }); + Ok(()) + }) + } + } + + fn payload(call_type: &str, model: &str, provider: &str) -> StandardLoggingPayload { + StandardLoggingPayload { + id: format!("req_{call_type}"), + litellm_call_id: format!("call_{call_type}"), + call_type: call_type.to_string(), + model: model.to_string(), + custom_llm_provider: provider.to_string(), + response_cost: 0.25, + prompt_tokens: 3, + completion_tokens: 4, + total_tokens: 7, + start_time: 10.0, + end_time: 11.5, + stream: false, + metadata: StandardLoggingMetadata { + user_api_key_hash: Some("hash".to_string()), + user_api_key_user_id: Some("user".to_string()), + user_api_key_team_id: Some("team".to_string()), + ..Default::default() + }, + messages: Some(json!([{"role": "user", "content": "read this"}])), + } + } + + #[tokio::test] + async fn rust_custom_logger_reads_success_payload_for_ocr() { + let logger = Arc::new(RecordingCustomLogger::default()); + let runner = CustomLoggerRunner::new(vec![logger.clone()]); + let details = ModelCallDetails::from_standard_logging_payload(payload( + "ocr", + "mistral-ocr-latest", + "mistral", + )); + let response = CallbackValue::new("ocr", json!({"pages": [{"markdown": "ok"}]})); + let report = runner + .async_log_success_event(&details, &response, CallbackTiming::new(10.0, 11.5)) + .await; + + assert_eq!(report.invoked, 1); + assert_eq!(report.dropped, 0); + assert_eq!( + logger.events(), + vec![RecordedEvent { + hook: "async_log_success_event", + model: "mistral-ocr-latest".to_string(), + provider: "mistral".to_string(), + call_type: "ocr".to_string(), + request_id: Some("req_ocr".to_string()), + litellm_call_id: Some("call_ocr".to_string()), + user_id: Some("user".to_string()), + response_object: Some("ocr".to_string()), + error_kind: None, + start_time: 10.0, + end_time: 11.5, + standard_logging_model: Some("mistral-ocr-latest".to_string()), + }] + ); + } + + #[tokio::test] + async fn rust_custom_logger_reads_failure_payload_for_non_ocr_call_type() { + let logger = Arc::new(RecordingCustomLogger::default()); + let runner = CustomLoggerRunner::new(vec![logger.clone()]); + let details = ModelCallDetails::from_standard_logging_payload(payload( + "acompletion", + "gpt-4.1-mini", + "openai", + )) + .with_failure_error(LoggingError { + message: "provider failed".to_string(), + kind: "ProviderError".to_string(), + }); + let response = CallbackValue::new("error", json!({"message": "provider failed"})); + let report = runner + .async_log_failure_event(&details, Some(&response), CallbackTiming::new(2.0, 3.0)) + .await; + + assert_eq!(report.invoked, 1); + assert_eq!(report.dropped, 0); + assert_eq!( + logger.events(), + vec![RecordedEvent { + hook: "async_log_failure_event", + model: "gpt-4.1-mini".to_string(), + provider: "openai".to_string(), + call_type: "acompletion".to_string(), + request_id: Some("req_acompletion".to_string()), + litellm_call_id: Some("call_acompletion".to_string()), + user_id: Some("user".to_string()), + response_object: Some("error".to_string()), + error_kind: Some("ProviderError".to_string()), + start_time: 2.0, + end_time: 3.0, + standard_logging_model: Some("gpt-4.1-mini".to_string()), + }] + ); + } + + #[tokio::test] + async fn no_callback_fast_path_dispatches_nothing() { + let runner = CustomLoggerRunner::new(Vec::new()); + let details = ModelCallDetails::new("mistral-ocr-latest", "mistral", CallType::Ocr); + let response = CallbackValue::new("ocr", json!({})); + + let report = runner + .async_log_success_event(&details, &response, CallbackTiming::new(1.0, 1.5)) + .await; + + assert!(runner.is_empty()); + assert_eq!(report, CallbackDispatchReport::default()); + } + + #[test] + fn with_standard_logging_payload_keeps_top_level_fields_in_sync() { + let details = ModelCallDetails::new("old-model", "old-provider", CallType::Completion) + .with_standard_logging_payload(payload("ocr", "mistral-ocr-latest", "mistral")); + + assert_eq!(details.model, "mistral-ocr-latest"); + assert_eq!(details.custom_llm_provider, "mistral"); + assert_eq!(details.call_type, CallType::Ocr); + assert_eq!(details.request_id, Some("req_ocr".to_string())); + assert_eq!(details.litellm_call_id, Some("call_ocr".to_string())); + } +} diff --git a/litellm-rust/crates/ai-gateway/src/integrations/custom_logger/types.rs b/litellm-rust/crates/ai-gateway/src/integrations/custom_logger/types.rs new file mode 100644 index 00000000000..ba7d67bd46e --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/custom_logger/types.rs @@ -0,0 +1,194 @@ +use std::collections::HashMap; +use std::future::Future; +use std::pin::Pin; + +use serde_json::Value; + +use crate::integrations::types::{StandardLoggingMetadata, StandardLoggingPayload}; + +pub type LogFuture<'a> = Pin> + Send + 'a>>; + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct CallbackDispatchReport { + pub invoked: usize, + pub dropped: usize, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum CallType { + Ocr, + Realtime, + Completion, + Acompletion, + ChatCompletion, + Other(String), +} + +impl CallType { + pub fn as_str(&self) -> &str { + match self { + Self::Ocr => "ocr", + Self::Realtime => "realtime", + Self::Completion => "completion", + Self::Acompletion => "acompletion", + Self::ChatCompletion => "chat_completion", + Self::Other(value) => value.as_str(), + } + } +} + +impl From<&str> for CallType { + fn from(value: &str) -> Self { + match value { + "ocr" => Self::Ocr, + "realtime" => Self::Realtime, + "completion" => Self::Completion, + "acompletion" => Self::Acompletion, + "chat_completion" => Self::ChatCompletion, + other => Self::Other(other.to_string()), + } + } +} + +impl std::fmt::Display for CallType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct CallbackTiming { + pub start_time: f64, + pub end_time: f64, +} + +impl CallbackTiming { + pub fn new(start_time: f64, end_time: f64) -> Self { + Self { + start_time, + end_time, + } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub struct CallbackValue { + pub object: String, + pub value: Value, +} + +impl CallbackValue { + pub fn new(object: impl Into, value: Value) -> Self { + Self { + object: object.into(), + value, + } + } +} + +#[derive(Clone, Debug)] +pub struct ModelCallDetails { + pub model: String, + pub custom_llm_provider: String, + pub call_type: CallType, + pub metadata: StandardLoggingMetadata, + pub extra_metadata: HashMap, + pub request_id: Option, + pub litellm_call_id: Option, + pub response_cost: Option, + pub standard_logging_payload: Option, + pub failure_error: Option, +} + +impl ModelCallDetails { + pub fn new( + model: impl Into, + custom_llm_provider: impl Into, + call_type: CallType, + ) -> Self { + Self { + model: model.into(), + custom_llm_provider: custom_llm_provider.into(), + call_type, + metadata: StandardLoggingMetadata::default(), + extra_metadata: HashMap::new(), + request_id: None, + litellm_call_id: None, + response_cost: None, + standard_logging_payload: None, + failure_error: None, + } + } + + pub fn from_standard_logging_payload(payload: StandardLoggingPayload) -> Self { + let request_id = Some(payload.id.clone()); + let litellm_call_id = Some(payload.litellm_call_id.clone()); + let response_cost = Some(payload.response_cost); + let metadata = payload.metadata.clone(); + Self { + model: payload.model.clone(), + custom_llm_provider: payload.custom_llm_provider.clone(), + call_type: CallType::from(payload.call_type.as_str()), + metadata, + extra_metadata: HashMap::new(), + request_id, + litellm_call_id, + response_cost, + standard_logging_payload: Some(payload), + failure_error: None, + } + } + + pub fn with_standard_logging_payload(mut self, payload: StandardLoggingPayload) -> Self { + self.model = payload.model.clone(); + self.custom_llm_provider = payload.custom_llm_provider.clone(); + self.call_type = CallType::from(payload.call_type.as_str()); + self.request_id = Some(payload.id.clone()); + self.litellm_call_id = Some(payload.litellm_call_id.clone()); + self.response_cost = Some(payload.response_cost); + self.metadata = payload.metadata.clone(); + self.standard_logging_payload = Some(payload); + self + } + + pub fn with_failure_error(mut self, error: LoggingError) -> Self { + self.failure_error = Some(error); + self + } +} + +#[derive(Clone, Debug)] +pub struct LoggingError { + pub message: String, + pub kind: String, +} + +#[derive(Clone, Debug)] +pub struct LogError { + pub message: String, + pub kind: String, +} + +impl LogError { + pub fn channel_full() -> Self { + Self { + message: "logging channel is full; dropping record".to_string(), + kind: "ChannelFull".to_string(), + } + } + + pub fn channel_closed() -> Self { + Self { + message: "logging channel is closed; worker has shut down".to_string(), + kind: "ChannelClosed".to_string(), + } + } +} + +impl std::fmt::Display for LogError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}: {}", self.kind, self.message) + } +} + +impl std::error::Error for LogError {} diff --git a/litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api.rs b/litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api/mod.rs similarity index 68% rename from litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api.rs rename to litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api/mod.rs index 165a90d8dbe..3dad18cb7a3 100644 --- a/litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api.rs +++ b/litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api/mod.rs @@ -1,7 +1,8 @@ //! A `CustomLogger` that ships finished events to the LiteLLM Python proxy's //! `/v1/rust_control_plane/logs` endpoint. //! -//! The callback path is non-blocking: `log_success_event` / `log_failure_event` +//! The callback path is non-blocking: `async_log_success_event` / +//! `async_log_failure_event` //! build a `LogRecord` and `try_send` it onto a bounded channel, returning a //! `LogError` (never panicking, never awaiting) if the channel is full or the //! worker has gone away. A spawned background worker drains the channel, batches @@ -15,54 +16,14 @@ use reqwest::Client; use tokio::sync::mpsc::{self, Receiver, Sender}; use tokio::time::interval; -use crate::constants::{ - DEFAULT_CHANNEL_CAPACITY, DEFAULT_FLUSH_INTERVAL_MS, DEFAULT_MAX_BATCH_SIZE, - DEFAULT_PROXY_BASE_URL, RUST_CONTROL_PLANE_LOGS_PATH, -}; -use crate::integrations::custom_logger::CustomLogger; -use crate::integrations::types::{ - CallbackLogsRequest, LogError, LogRecord, LoggingError, StandardLoggingPayload, +use crate::constants::{DEFAULT_PROXY_BASE_URL, RUST_CONTROL_PLANE_LOGS_PATH}; +use crate::integrations::custom_logger::{ + CallbackTiming, CallbackValue, CustomLogger, LogError, LogFuture, LoggingError, + ModelCallDetails, }; +use types::{CallbackLogsRequest, EgressTunables, LogRecord}; -/// Egress worker tunables. Each field defaults to the matching `DEFAULT_*` const -/// in `crate::constants` and is overridable via an env var (read once at logger -/// construction). -struct EgressTunables { - channel_capacity: usize, - max_batch_size: usize, - flush_interval: Duration, -} - -impl EgressTunables { - fn from_env() -> Self { - Self { - channel_capacity: env_positive( - "LITELLM_LOG_CHANNEL_CAPACITY", - DEFAULT_CHANNEL_CAPACITY, - ), - max_batch_size: env_positive("LITELLM_LOG_BATCH_SIZE", DEFAULT_MAX_BATCH_SIZE), - flush_interval: Duration::from_millis(env_positive( - "LITELLM_LOG_FLUSH_INTERVAL_MS", - DEFAULT_FLUSH_INTERVAL_MS, - )), - } - } -} - -/// Parse a positive integer env var, falling back to `default` on missing, -/// unparseable, or non-positive values. Generic over the integer type so one -/// helper serves both the `usize` capacities and the `u64` interval. -fn env_positive(name: &str, default: T) -> T -where - T: std::str::FromStr + PartialOrd + From, -{ - let zero = T::from(0u8); - std::env::var(name) - .ok() - .and_then(|value| value.trim().parse::().ok()) - .filter(|n| *n > zero) - .unwrap_or(default) -} +pub mod types; /// Ships realtime logging events to the LiteLLM Python proxy. pub struct LiteLLMPythonProxyAPILogger { @@ -118,23 +79,50 @@ impl LiteLLMPythonProxyAPILogger { } impl CustomLogger for LiteLLMPythonProxyAPILogger { - fn log_success_event(&self, payload: &StandardLoggingPayload) -> Result<(), LogError> { - self.enqueue(LogRecord { - status: "success".to_string(), - payload: payload.clone(), - error: None, + fn async_log_success_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + _response_obj: &'a CallbackValue, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + if let Some(payload) = &model_call_details.standard_logging_payload { + self.enqueue(LogRecord { + status: "success".to_string(), + payload: payload.clone(), + error: None, + })?; + } + Ok(()) }) } - fn log_failure_event( - &self, - payload: &StandardLoggingPayload, - error: &LoggingError, - ) -> Result<(), LogError> { - self.enqueue(LogRecord { - status: "failure".to_string(), - payload: payload.clone(), - error: Some(format!("{}: {}", error.kind, error.message)), + fn async_log_failure_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + _response_obj: Option<&'a CallbackValue>, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + if let Some(payload) = &model_call_details.standard_logging_payload { + let fallback_error; + let error = match &model_call_details.failure_error { + Some(error) => error, + None => { + fallback_error = LoggingError { + message: "callback failure event".to_string(), + kind: "CallbackFailure".to_string(), + }; + &fallback_error + } + }; + self.enqueue(LogRecord { + status: "failure".to_string(), + payload: payload.clone(), + error: Some(format!("{}: {}", error.kind, error.message)), + })?; + } + Ok(()) }) } } diff --git a/litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api/types.rs b/litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api/types.rs new file mode 100644 index 00000000000..481a437747f --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api/types.rs @@ -0,0 +1,72 @@ +use std::time::Duration; + +use serde::Serialize; + +use crate::constants::{ + DEFAULT_CHANNEL_CAPACITY, DEFAULT_FLUSH_INTERVAL_MS, DEFAULT_MAX_BATCH_SIZE, +}; +use crate::integrations::types::StandardLoggingPayload; + +#[derive(Serialize)] +pub struct CallbackLogsRequest { + pub records: Vec, +} + +#[derive(Serialize)] +pub struct CallbackLogRecord { + pub status: String, + pub standard_logging_payload: StandardLoggingPayload, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +#[derive(Clone, Debug)] +pub struct LogRecord { + pub status: String, + pub payload: StandardLoggingPayload, + pub error: Option, +} + +impl LogRecord { + pub fn into_callback_record(self) -> CallbackLogRecord { + CallbackLogRecord { + status: self.status, + standard_logging_payload: self.payload, + error: self.error, + } + } +} + +pub(super) struct EgressTunables { + pub channel_capacity: usize, + pub max_batch_size: usize, + pub flush_interval: Duration, +} + +impl EgressTunables { + pub fn from_env() -> Self { + Self { + channel_capacity: env_positive( + "LITELLM_LOG_CHANNEL_CAPACITY", + DEFAULT_CHANNEL_CAPACITY, + ), + max_batch_size: env_positive("LITELLM_LOG_BATCH_SIZE", DEFAULT_MAX_BATCH_SIZE), + flush_interval: Duration::from_millis(env_positive( + "LITELLM_LOG_FLUSH_INTERVAL_MS", + DEFAULT_FLUSH_INTERVAL_MS, + )), + } + } +} + +fn env_positive(name: &str, default: T) -> T +where + T: std::str::FromStr + PartialOrd + From, +{ + let zero = T::from(0u8); + std::env::var(name) + .ok() + .and_then(|value| value.trim().parse::().ok()) + .filter(|n| *n > zero) + .unwrap_or(default) +} diff --git a/litellm-rust/crates/ai-gateway/src/integrations/mod.rs b/litellm-rust/crates/ai-gateway/src/integrations/mod.rs index 8799be0c040..c62f1821ef8 100644 --- a/litellm-rust/crates/ai-gateway/src/integrations/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/integrations/mod.rs @@ -1,10 +1,12 @@ //! Pure-Rust logging integrations. Names map 1:1 to Python //! `litellm/integrations/`: +//! - [`custom_guardrail::CustomGuardrail`] — the guardrail callback trait //! - [`custom_logger::CustomLogger`] — the callback trait //! - [`litellm_python_proxy_api::LiteLLMPythonProxyAPILogger`] — ships events -//! to the Python proxy's `/v1/callbacks/logs` endpoint +//! to the Python proxy's `/v1/rust_control_plane/logs` endpoint //! - [`types`] — the typed `StandardLoggingPayload` wire contract +pub mod custom_guardrail; pub mod custom_logger; pub mod litellm_python_proxy_api; pub mod types; diff --git a/litellm-rust/crates/ai-gateway/src/integrations/types.rs b/litellm-rust/crates/ai-gateway/src/integrations/types.rs index d61a1f816a7..34dce93d8e0 100644 --- a/litellm-rust/crates/ai-gateway/src/integrations/types.rs +++ b/litellm-rust/crates/ai-gateway/src/integrations/types.rs @@ -28,68 +28,6 @@ pub struct RequestMetadata { pub user_api_key_team_id: Option, } -/// A logging-callback failure (e.g. a custom logger raised). Mirrors the Python -/// failure-event shape: a message plus an exception kind/class name. -#[derive(Clone, Debug)] -pub struct LoggingError { - pub message: String, - pub kind: String, -} - -/// A non-fatal error returned by a `CustomLogger` when it cannot enqueue an -/// event (channel full or the background worker has shut down). -#[derive(Clone, Debug)] -pub struct LogError { - pub message: String, - pub kind: String, -} - -impl LogError { - pub fn channel_full() -> Self { - Self { - message: "logging channel is full; dropping record".to_string(), - kind: "ChannelFull".to_string(), - } - } - - pub fn channel_closed() -> Self { - Self { - message: "logging channel is closed; worker has shut down".to_string(), - kind: "ChannelClosed".to_string(), - } - } -} - -impl std::fmt::Display for LogError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{}: {}", self.kind, self.message) - } -} - -impl std::error::Error for LogError {} - -/// Batch wrapper — the top-level request body. -/// Matches Python `CallbackLogsRequest { records: list[CallbackLogRecord] }`. -#[derive(Serialize)] -pub struct CallbackLogsRequest { - pub records: Vec, -} - -/// One finished logging event. -/// Matches `CallbackLogRecord { status, standard_logging_payload, error? }`. -#[derive(Serialize)] -pub struct CallbackLogRecord { - /// "success" | "failure". On "failure", `error` (or payload.error_str) - /// becomes the replayed exception string. - pub status: String, - - pub standard_logging_payload: StandardLoggingPayload, - - /// Only meaningful when status == "failure". Omitted on success. - #[serde(skip_serializing_if = "Option::is_none")] - pub error: Option, -} - /// The self-describing payload. Field names are the EXACT JSON keys the Python /// replay path + spend-logs builder read. #[derive(Clone, Debug, Serialize)] @@ -143,22 +81,3 @@ pub struct StandardLoggingMetadata { #[serde(skip_serializing_if = "Option::is_none")] pub spend_logs_metadata: Option>, } - -/// The unit handed to a `CustomLogger` sink: a finished payload plus its status -/// and (on failure) the replayed error string. -#[derive(Clone, Debug)] -pub struct LogRecord { - pub status: String, - pub payload: StandardLoggingPayload, - pub error: Option, -} - -impl LogRecord { - pub fn into_callback_record(self) -> CallbackLogRecord { - CallbackLogRecord { - status: self.status, - standard_logging_payload: self.payload, - error: self.error, - } - } -} diff --git a/litellm-rust/crates/ai-gateway/src/io/ocr.rs b/litellm-rust/crates/ai-gateway/src/io/ocr.rs index 35e511fa982..55e02839c4e 100644 --- a/litellm-rust/crates/ai-gateway/src/io/ocr.rs +++ b/litellm-rust/crates/ai-gateway/src/io/ocr.rs @@ -1,406 +1 @@ -//! End-to-end OCR orchestration. -//! -//! Owns supported OCR provider calls so the Python side stays a thin bridge: -//! resolve the API key, build the URL + body via the pure transforms, POST it, -//! and normalize the response. The HTTP client is built once and reused. - -use std::sync::OnceLock; -use std::time::Duration; - -use litellm_core::error::CoreError; -use litellm_core::ocr::transformation::{OcrAuthStrategy, OcrResponseHandling}; -use litellm_core::CoreResult; -use serde_json::{Map, Value}; - -mod common_utils; - -use common_utils::{ - convert_document_url_to_data_uri, has_header, ocr_provider_config, poll_document_intelligence, - string_headers, truncate_error_body, -}; - -/// OCR over large documents can take a while; bound it generously rather than -/// hanging forever on an unresponsive upstream. The client-level limit is the -/// outer ceiling; callers can tighten it per request via ``run_ocr``'s ``timeout``. -const OCR_TIMEOUT_SECS: u64 = 600; - -/// Process-wide async HTTP client (connection pool + TLS reused across calls). -/// -/// The Python fallback path uses LiteLLM's standard `BaseLLMHTTPHandler`. This -/// Rust path is opt-in and owns end-to-end OCR I/O, so it cannot call the -/// Python handler directly; keep this route-scoped until litellm-rust has a -/// shared HTTP abstraction. -fn http_client() -> &'static reqwest::Client { - static CLIENT: OnceLock = OnceLock::new(); - CLIENT.get_or_init(|| { - reqwest::Client::builder() - .timeout(Duration::from_secs(OCR_TIMEOUT_SECS)) - .build() - .expect("failed to build reqwest client") - }) -} - -fn upstream_headers( - headers: &[(String, String)], - auth_strategy: OcrAuthStrategy, - api_key: Option<&str>, -) -> Vec<(String, String)> { - let auth_header = api_key.map(|api_key| match auth_strategy { - OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")), - OcrAuthStrategy::Header(header_name) => (header_name.to_string(), api_key.to_string()), - }); - auth_header - .into_iter() - .chain(headers.iter().cloned()) - .collect() -} - -pub struct OcrRequest<'a> { - pub model: &'a str, - pub document: Value, - pub api_key: Option<&'a str>, - pub api_base: Option<&'a str>, - pub custom_llm_provider: &'a str, - pub extra_headers: Option>, - pub optional_params: Map, - pub timeout: Option, -} - -/// Perform an OCR call end to end and return the normalized response as -/// JSON (the shape the Python `OCRResponse` model expects). -/// -/// Async: intended to be awaited directly by the Python bridge's async entrypoint. -pub async fn ocr(request: OcrRequest<'_>) -> CoreResult { - let model = request.model; - let config = ocr_provider_config(request.custom_llm_provider, model) - .ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.to_string()))?; - let env_lookup = |key: &str| std::env::var(key).ok(); - - let headers = string_headers(request.extra_headers)?; - let auth_strategy = config.auth_strategy(); - let api_key = (!has_header(&headers, auth_strategy.header_name())) - .then(|| config.resolve_api_key(request.api_key, &env_lookup)) - .transpose()?; - let url = config.complete_url( - request.api_base, - model, - &request.optional_params, - &env_lookup, - )?; - let filtered_params = config.map_ocr_params(&request.optional_params); - let document = if config.requires_data_uri_document() { - convert_document_url_to_data_uri(request.document).await? - } else { - request.document - }; - let body = config - .transform_ocr_request(model, document, filtered_params)? - .data; - let upstream_headers = upstream_headers(&headers, auth_strategy, api_key.as_deref()); - - let mut request_builder = http_client().post(&url).json(&body); - for (key, value) in &upstream_headers { - request_builder = request_builder.header(key, value); - } - if let Some(duration) = request.timeout { - request_builder = request_builder.timeout(duration); - } - - let response = request_builder - .send() - .await - .map_err(|err| CoreError::Network(err.to_string()))?; - - let status = response.status(); - if config.response_handling() == OcrResponseHandling::AzureDocumentIntelligencePoll - && status.as_u16() == 202 - { - let operation_url = response - .headers() - .get("operation-location") - .and_then(|value| value.to_str().ok()) - .map(str::to_string) - .ok_or_else(|| { - CoreError::InvalidResponse( - "Azure Document Intelligence returned 202 but no Operation-Location header found" - .to_string(), - ) - })?; - let response_json = - poll_document_intelligence(&operation_url, &url, &upstream_headers, request.timeout) - .await?; - return Ok(config - .transform_ocr_response(model, response_json)? - .into_json()); - } - - let text = response - .text() - .await - .map_err(|err| CoreError::Network(err.to_string()))?; - - if !status.is_success() { - return Err(CoreError::Http { - status: status.as_u16(), - body: truncate_error_body(&text), - }); - } - - let response_json: Value = serde_json::from_str(&text) - .map_err(|err| CoreError::InvalidResponse(format!("invalid OCR response JSON: {err}")))?; - - Ok(config - .transform_ocr_response(model, response_json)? - .into_json()) -} - -#[cfg(test)] -mod tests { - use super::*; - use serde_json::json; - use tokio::io::{AsyncReadExt, AsyncWriteExt}; - use tokio::net::{TcpListener, TcpStream}; - - async fn read_http_headers(socket: &mut TcpStream) -> String { - let mut request = Vec::new(); - let mut buffer = [0_u8; 1024]; - loop { - let n = socket.read(&mut buffer).await.expect("reads request"); - if n == 0 { - break; - } - request.extend_from_slice(&buffer[..n]); - if request.windows(4).any(|window| window == b"\r\n\r\n") { - break; - } - } - String::from_utf8(request).expect("request is utf8") - } - - #[test] - fn truncate_error_body_passes_short_strings_through() { - let body = "Unauthorized"; - assert_eq!(truncate_error_body(body), "Unauthorized"); - } - - #[test] - fn truncate_error_body_caps_long_payloads() { - let body = "x".repeat(306); - let truncated = truncate_error_body(&body); - - assert!(truncated.ends_with("... (truncated)")); - let prefix_chars = truncated - .strip_suffix("... (truncated)") - .expect("truncated marker present") - .chars() - .count(); - assert_eq!(prefix_chars, 256); - } - - #[test] - fn truncate_error_body_does_not_split_multibyte_chars() { - let body = "é".repeat(266); - let truncated = truncate_error_body(&body); - assert!(truncated.is_char_boundary(truncated.len())); - } - - #[test] - fn ocr_dispatch_supports_migrated_providers() { - assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some()); - assert!(ocr_provider_config("azure_ai", "pixtral-12b-2409") - .expect("azure ai config resolves") - .requires_data_uri_document()); - assert_eq!( - ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read") - .expect("document intelligence config resolves") - .response_handling(), - OcrResponseHandling::AzureDocumentIntelligencePoll - ); - assert!(ocr_provider_config("vertex_ai", "deepseek-ocr-maas") - .expect("vertex deepseek config resolves") - .supported_ocr_params() - .contains(&"temperature")); - assert!(ocr_provider_config("openai", "gpt-4o").is_none()); - } - - #[test] - fn string_headers_accepts_string_values() { - let headers = json!({ - "x-trace-id": "trace-1" - }) - .as_object() - .unwrap() - .clone(); - - assert_eq!( - string_headers(Some(headers)).expect("string headers accepted"), - vec![("x-trace-id".to_string(), "trace-1".to_string())] - ); - } - - #[test] - fn auth_header_detection_is_case_insensitive() { - let headers = vec![ - ("x-trace-id".to_string(), "trace-1".to_string()), - ("authorization".to_string(), "Bearer sk-test".to_string()), - ]; - - assert!(has_header(&headers, "authorization")); - - let headers = vec![("Authorization".to_string(), "Bearer sk-test".to_string())]; - assert!(has_header(&headers, "authorization")); - - let headers = vec![("x-trace-id".to_string(), "trace-1".to_string())]; - assert!(!has_header(&headers, "authorization")); - } - - #[tokio::test] - async fn ocr_does_not_duplicate_authorization_header_when_header_is_supplied() { - let listener = TcpListener::bind("127.0.0.1:0") - .await - .expect("test listener binds"); - let addr = listener.local_addr().expect("listener has local addr"); - - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts one request"); - let request = read_http_headers(&mut socket).await; - - let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#; - let response = format!( - "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - response_body.len(), - response_body - ); - socket - .write_all(response.as_bytes()) - .await - .expect("writes response"); - request - }); - - let mut headers = Map::new(); - headers.insert( - "Authorization".to_string(), - Value::String("Bearer sk-from-python".to_string()), - ); - headers.insert( - "x-trace-id".to_string(), - Value::String("trace-1".to_string()), - ); - - let response = ocr(OcrRequest { - model: "mistral-ocr-latest", - document: json!({ - "type": "document_url", - "document_url": "https://example.com/doc.pdf" - }), - api_key: Some("sk-for-rust-fallback"), - api_base: Some(&format!("http://{addr}")), - custom_llm_provider: "mistral", - extra_headers: Some(headers), - optional_params: Map::new(), - timeout: Some(Duration::from_secs(5)), - }) - .await - .expect("ocr request succeeds"); - - assert_eq!(response["pages"][0]["markdown"], "ok"); - - let request = server.await.expect("server task completes"); - let authorization_count = request - .lines() - .filter(|line| line.to_ascii_lowercase().starts_with("authorization:")) - .count(); - assert_eq!(authorization_count, 1, "{request}"); - assert!( - request.contains("authorization: Bearer sk-from-python") - || request.contains("Authorization: Bearer sk-from-python"), - "{request}" - ); - } - - #[tokio::test] - async fn document_intelligence_poll_uses_resolved_subscription_key() { - let listener = TcpListener::bind("127.0.0.1:0") - .await - .expect("test listener binds"); - let addr = listener.local_addr().expect("listener has local addr"); - let operation_url = format!("http://{addr}/operations/1"); - - let server = tokio::spawn(async move { - let (mut post_socket, _) = listener.accept().await.expect("accepts post request"); - let post_request = read_http_headers(&mut post_socket).await; - let post_response = format!( - "HTTP/1.1 202 Accepted\r\noperation-location: {operation_url}\r\ncontent-length: 0\r\nconnection: close\r\n\r\n" - ); - post_socket - .write_all(post_response.as_bytes()) - .await - .expect("writes post response"); - - let (mut poll_socket, _) = listener.accept().await.expect("accepts poll request"); - let poll_request = read_http_headers(&mut poll_socket).await; - let response_body = r#"{"status":"succeeded","analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"ok"}]}]}}"#; - let poll_response = format!( - "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - response_body.len(), - response_body - ); - poll_socket - .write_all(poll_response.as_bytes()) - .await - .expect("writes poll response"); - (post_request, poll_request) - }); - - let response = ocr(OcrRequest { - model: "prebuilt-read", - document: json!({ - "type": "document_url", - "document_url": "https://example.com/doc.pdf" - }), - api_key: Some("di-key"), - api_base: Some(&format!("http://{addr}")), - custom_llm_provider: "azure_ai/doc-intelligence", - extra_headers: None, - optional_params: Map::new(), - timeout: Some(Duration::from_secs(5)), - }) - .await - .expect("document intelligence request succeeds"); - - assert_eq!(response["pages"][0]["markdown"], "ok"); - - let (post_request, poll_request) = server.await.expect("server task completes"); - assert!( - post_request - .to_ascii_lowercase() - .contains("ocp-apim-subscription-key: di-key"), - "{post_request}" - ); - assert!( - poll_request - .to_ascii_lowercase() - .contains("ocp-apim-subscription-key: di-key"), - "{poll_request}" - ); - } - - #[test] - fn string_headers_rejects_non_string_values() { - let headers = json!({ - "x-retry-count": 3 - }) - .as_object() - .unwrap() - .clone(); - - let err = string_headers(Some(headers)).expect_err("non-string header rejected"); - assert_eq!( - err, - CoreError::InvalidRequest( - "OCR extra_headers.x-retry-count must be a string, got number".to_string() - ) - ); - } -} +pub use crate::ocr::{ocr, OcrRequest}; diff --git a/litellm-rust/crates/ai-gateway/src/io/realtime.rs b/litellm-rust/crates/ai-gateway/src/io/realtime.rs index 4047de5cb26..40a38c1579a 100644 --- a/litellm-rust/crates/ai-gateway/src/io/realtime.rs +++ b/litellm-rust/crates/ai-gateway/src/io/realtime.rs @@ -1,8 +1,8 @@ //! End-to-end OpenAI realtime invocation. //! -//! The host-facing entry point, mirroring `crate::io::ocr::run_ocr`: open the -//! WebSocket to OpenAI, then splice a client realtime stream to the upstream, -//! driving typed events through the pure `OPENAI_REALTIME_CONFIG` transforms. +//! The host-facing entry point opens the WebSocket to OpenAI, then splices a +//! client realtime stream to the upstream, driving typed events through the pure +//! `OPENAI_REALTIME_CONFIG` transforms. //! Network, auth header, key resolution, and wire (de)serialization live here so //! the `transformation` module stays pure and typed. //! diff --git a/litellm-rust/crates/ai-gateway/src/lib.rs b/litellm-rust/crates/ai-gateway/src/lib.rs index 6c04fbb7626..d8ef7bb5ba1 100644 --- a/litellm-rust/crates/ai-gateway/src/lib.rs +++ b/litellm-rust/crates/ai-gateway/src/lib.rs @@ -3,15 +3,16 @@ //! Two layers, split by feature so the Python `cdylib` can depend on the I/O //! without pulling in the HTTP server: //! -//! - [`io`]: all network I/O (OCR HTTP call, realtime WebSocket splice, the -//! pre-warmed realtime pool). Always available — no feature required. The -//! Python bridge links this for `run_ocr`. +//! - Call-type modules such as [`ocr`]: provider transforms, lifecycle hooks, +//! and provider I/O. Always available — no feature required. +//! - [`io`]: compatibility exports and realtime WebSocket splice helpers. //! - The server modules ([`auth`], [`routes`], [`state`]) and anything pulling //! `axum` are gated behind the `server` feature, which the `litellm-ai-gateway` //! binary turns on. The `python-config` feature additionally pulls in [`python`] //! for the load-time config reader. pub mod io; +pub mod ocr; /// GIL-activity tracking. Pure (atomics only); shared by the `server` routes and /// the `python-config` reader, so it is available without either feature. @@ -27,9 +28,7 @@ pub mod state; // Realtime request logging. Only the server serves realtime, so these are // `server`-gated; `io::realtime` exposes the generic `observe` hook while the // collector and callback fan-out live here. -#[cfg(feature = "server")] mod constants; -#[cfg(feature = "server")] pub mod integrations; #[cfg(feature = "server")] mod realtime; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/client.rs b/litellm-rust/crates/ai-gateway/src/ocr/client.rs new file mode 100644 index 00000000000..79cc7816227 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/ocr/client.rs @@ -0,0 +1,14 @@ +use std::sync::OnceLock; +use std::time::Duration; + +const OCR_TIMEOUT_SECS: u64 = 600; + +pub(super) fn http_client() -> &'static reqwest::Client { + static CLIENT: OnceLock = OnceLock::new(); + CLIENT.get_or_init(|| { + reqwest::Client::builder() + .timeout(Duration::from_secs(OCR_TIMEOUT_SECS)) + .build() + .expect("failed to build reqwest client") + }) +} diff --git a/litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs similarity index 99% rename from litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs rename to litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs index ee0d86c3000..d4b4d9338e7 100644 --- a/litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs @@ -18,7 +18,7 @@ use litellm_core::providers::vertex_ai::ocr::transformation::{ VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG, }; -use super::http_client; +use super::client::http_client; const ERROR_BODY_MAX_CHARS: usize = 256; const AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120; @@ -42,7 +42,6 @@ pub(super) fn ocr_provider_config( "azure_ai" if is_azure_document_intelligence_model(model) => { Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG) } - "azure_ai/doc-intelligence" => Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG), "azure_ai" => Some(&AZURE_AI_OCR_CONFIG), "vertex_ai" if vertex_ai::is_deepseek_model(model) => Some(&VERTEX_AI_DEEPSEEK_OCR_CONFIG), "vertex_ai" => Some(&VERTEX_AI_OCR_CONFIG), diff --git a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs new file mode 100644 index 00000000000..4d93c2a25db --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs @@ -0,0 +1,71 @@ +use litellm_core::error::CoreError; +use litellm_core::ocr::transformation::OcrResponseHandling; +use litellm_core::CoreResult; +use serde_json::Value; + +use super::client::http_client; +use super::common_utils::{poll_document_intelligence, truncate_error_body}; +use super::types::ProviderOcrRequest; + +pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> CoreResult { + let mut request_builder = http_client().post(&request.url).json(&request.body); + for (key, value) in &request.upstream_headers { + request_builder = request_builder.header(key, value); + } + if let Some(duration) = request.timeout { + request_builder = request_builder.timeout(duration); + } + + let response = request_builder + .send() + .await + .map_err(|err| CoreError::Network(err.to_string()))?; + + let status = response.status(); + if request.config.response_handling() == OcrResponseHandling::AzureDocumentIntelligencePoll + && status.as_u16() == 202 + { + let operation_url = response + .headers() + .get("operation-location") + .and_then(|value| value.to_str().ok()) + .map(str::to_string) + .ok_or_else(|| { + CoreError::InvalidResponse( + "Azure Document Intelligence returned 202 but no Operation-Location header found" + .to_string(), + ) + })?; + let response_json = poll_document_intelligence( + &operation_url, + &request.url, + &request.upstream_headers, + request.timeout, + ) + .await?; + return Ok(request + .config + .transform_ocr_response(&request.model, response_json)? + .into_json()); + } + + let text = response + .text() + .await + .map_err(|err| CoreError::Network(err.to_string()))?; + + if !status.is_success() { + return Err(CoreError::Http { + status: status.as_u16(), + body: truncate_error_body(&text), + }); + } + + let response_json: Value = serde_json::from_str(&text) + .map_err(|err| CoreError::InvalidResponse(format!("invalid OCR response JSON: {err}")))?; + + Ok(request + .config + .transform_ocr_response(&request.model, response_json)? + .into_json()) +} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs new file mode 100644 index 00000000000..6be74ed2714 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -0,0 +1,329 @@ +use std::future::Future; +use std::pin::Pin; + +use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; +use litellm_core::error::CoreError; +use litellm_core::ocr::transformation::OcrAuthStrategy; +use litellm_core::CoreResult; +use serde_json::{json, Map, Value}; + +use super::common_utils::{ + convert_document_url_to_data_uri, has_header, ocr_provider_config, string_headers, +}; +use super::types::{PreparedOcrRequest, ProviderOcrRequest}; +use crate::integrations::custom_guardrail::{ + CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest, +}; +use crate::integrations::custom_logger::{ + CallType, CallbackTiming, CallbackValue, CustomLoggerRunner, LoggingError, ModelCallDetails, +}; +use crate::integrations::types::{ + RequestMetadata, StandardLoggingMetadata, StandardLoggingPayload, +}; + +pub(crate) struct OcrLifecycleHooks { + logger_runner: CustomLoggerRunner, + guardrail_runner: CustomGuardrailRunner, + request_metadata: RequestMetadata, +} + +type OcrFuture<'a, T> = Pin> + Send + 'a>>; +type OcrLogFuture<'a> = Pin + Send + 'a>>; + +impl OcrLifecycleHooks { + pub(crate) fn new( + logger_runner: CustomLoggerRunner, + guardrail_runner: CustomGuardrailRunner, + request_metadata: RequestMetadata, + ) -> Self { + Self { + logger_runner, + guardrail_runner, + request_metadata, + } + } + + async fn run_pre_call_guardrails( + &self, + request: PreparedOcrRequest, + ) -> CoreResult { + if self.guardrail_runner.is_empty() { + return Ok(request); + } + + let context = guardrail_context(&self.request_metadata); + let guardrail_request = GuardrailRequest::new(json!({ + "model": request.model, + "custom_llm_provider": request.custom_llm_provider, + "document": request.document, + "optional_params": request.optional_params, + })); + let (guardrail_request, _) = self + .guardrail_runner + .run_pre_call(&context, guardrail_request) + .await + .map_err(guardrail_error_to_core_error)?; + let (document, optional_params) = parse_ocr_pre_call_guardrail_request(guardrail_request)?; + Ok(PreparedOcrRequest { + document, + optional_params, + ..request + }) + } + + async fn prepare_provider_request( + &self, + request: PreparedOcrRequest, + ) -> CoreResult { + let config = ocr_provider_config(&request.custom_llm_provider, &request.model) + .ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.clone()))?; + let env_lookup = |key: &str| std::env::var(key).ok(); + let headers = string_headers(request.extra_headers)?; + let auth_strategy = config.auth_strategy(); + let api_key = (!has_header(&headers, auth_strategy.header_name())) + .then(|| config.resolve_api_key(request.api_key.as_deref(), &env_lookup)) + .transpose()?; + let url = config.complete_url( + request.api_base.as_deref(), + &request.model, + &request.optional_params, + &env_lookup, + )?; + let filtered_params = config.map_ocr_params(&request.optional_params); + let model = request.model.clone(); + let custom_llm_provider = request.custom_llm_provider.clone(); + let document = if config.requires_data_uri_document() { + convert_document_url_to_data_uri(request.document).await? + } else { + request.document + }; + let body = config + .transform_ocr_request(&request.model, document, filtered_params)? + .data; + let upstream_headers = upstream_headers(&headers, auth_strategy, api_key.as_deref()); + let body = self + .run_during_call_guardrails(&model, &custom_llm_provider, &url, body) + .await?; + Ok(ProviderOcrRequest { + model, + config, + url, + body, + upstream_headers, + timeout: request.timeout, + }) + } + + async fn run_during_call_guardrails( + &self, + model: &str, + custom_llm_provider: &str, + url: &str, + body: Value, + ) -> CoreResult { + if self.guardrail_runner.is_empty() { + return Ok(body); + } + + let context = guardrail_context(&self.request_metadata); + let guardrail_request = GuardrailRequest::new(json!({ + "model": model, + "custom_llm_provider": custom_llm_provider, + "url": url, + "body": body, + })); + let (guardrail_request, _) = self + .guardrail_runner + .run_during_call(&context, guardrail_request) + .await + .map_err(guardrail_error_to_core_error)?; + parse_ocr_during_call_guardrail_request(guardrail_request) + } + + fn standard_logging_payload( + &self, + context: &CallLifecycleContext, + timing: &CallLifecycleTiming, + ) -> StandardLoggingPayload { + StandardLoggingPayload { + id: context.litellm_call_id.clone(), + litellm_call_id: context.litellm_call_id.clone(), + call_type: context.call_type.clone(), + model: context.model.clone(), + custom_llm_provider: context.custom_llm_provider.clone(), + response_cost: 0.0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + start_time: timing.start_time, + end_time: timing.end_time, + stream: false, + metadata: StandardLoggingMetadata { + user_api_key_hash: self.request_metadata.user_api_key_hash.clone(), + user_api_key_user_id: self.request_metadata.user_api_key_user_id.clone(), + user_api_key_team_id: self.request_metadata.user_api_key_team_id.clone(), + ..Default::default() + }, + messages: None, + } + } +} + +impl CallLifecycleHooks for OcrLifecycleHooks { + type PreCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>; + type DuringCallFuture<'a> = OcrFuture<'a, ProviderOcrRequest>; + type SuccessFuture<'a> = OcrLogFuture<'a>; + type FailureFuture<'a> = OcrLogFuture<'a>; + + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: PreparedOcrRequest, + ) -> Self::PreCallFuture<'a> { + Box::pin(async move { self.run_pre_call_guardrails(request).await }) + } + + fn async_during_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: PreparedOcrRequest, + ) -> Self::DuringCallFuture<'a> { + Box::pin(async move { self.prepare_provider_request(request).await }) + } + + fn async_log_success_event<'a>( + &'a self, + context: &'a CallLifecycleContext, + response: &'a Value, + timing: &'a CallLifecycleTiming, + ) -> Self::SuccessFuture<'a> { + Box::pin(async move { + if self.logger_runner.is_empty() { + return; + } + let response_obj = CallbackValue::new("ocr", response.clone()); + self.logger_runner + .async_log_success_event( + &ModelCallDetails::from_standard_logging_payload( + self.standard_logging_payload(context, timing), + ), + &response_obj, + CallbackTiming::new(timing.start_time, timing.end_time), + ) + .await; + }) + } + + fn async_log_failure_event<'a>( + &'a self, + context: &'a CallLifecycleContext, + error: &'a CoreError, + timing: &'a CallLifecycleTiming, + ) -> Self::FailureFuture<'a> { + Box::pin(async move { + if self.logger_runner.is_empty() { + return; + } + let logging_error = LoggingError { + message: error.to_string(), + kind: core_error_kind(error).to_string(), + }; + let response_obj = CallbackValue::new( + "error", + json!({ + "message": logging_error.message, + "kind": logging_error.kind, + }), + ); + self.logger_runner + .async_log_failure_event( + &ModelCallDetails::from_standard_logging_payload( + self.standard_logging_payload(context, timing), + ) + .with_failure_error(logging_error), + Some(&response_obj), + CallbackTiming::new(timing.start_time, timing.end_time), + ) + .await; + }) + } +} + +fn upstream_headers( + headers: &[(String, String)], + auth_strategy: OcrAuthStrategy, + api_key: Option<&str>, +) -> Vec<(String, String)> { + api_key + .map(|api_key| match auth_strategy { + OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")), + OcrAuthStrategy::Header(header_name) => (header_name.to_string(), api_key.to_string()), + }) + .into_iter() + .chain(headers.iter().cloned()) + .collect() +} + +fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext { + GuardrailContext { + call_type: CallType::Ocr, + selected_guardrails: Vec::new(), + metadata: std::collections::HashMap::new(), + user_api_key_hash: metadata.user_api_key_hash.clone(), + user_api_key_user_id: metadata.user_api_key_user_id.clone(), + user_api_key_team_id: metadata.user_api_key_team_id.clone(), + trace_parent: None, + } +} + +fn parse_ocr_pre_call_guardrail_request( + request: GuardrailRequest, +) -> CoreResult<(Value, Map)> { + let Value::Object(mut data) = request.data else { + return Err(CoreError::InvalidRequest( + "OCR pre_call guardrail must return an object".to_string(), + )); + }; + let document = data.remove("document").ok_or_else(|| { + CoreError::InvalidRequest("OCR pre_call guardrail removed document".to_string()) + })?; + let optional_params = match data.remove("optional_params") { + Some(Value::Object(params)) => params, + Some(_) => { + return Err(CoreError::InvalidRequest( + "OCR pre_call guardrail optional_params must be an object".to_string(), + )) + } + None => Map::new(), + }; + Ok((document, optional_params)) +} + +fn parse_ocr_during_call_guardrail_request(request: GuardrailRequest) -> CoreResult { + let Value::Object(mut data) = request.data else { + return Err(CoreError::InvalidRequest( + "OCR during_call guardrail must return an object".to_string(), + )); + }; + data.remove("body").ok_or_else(|| { + CoreError::InvalidRequest("OCR during_call guardrail removed body".to_string()) + }) +} + +fn guardrail_error_to_core_error(error: GuardrailError) -> CoreError { + CoreError::InvalidRequest(format!("{}: {}", error.kind, error.message)) +} + +fn core_error_kind(error: &CoreError) -> &'static str { + match error { + CoreError::Auth(_) => "AuthError", + CoreError::InvalidProvider(_) => "InvalidProvider", + CoreError::InvalidRequest(_) => "InvalidRequest", + CoreError::InvalidType { .. } => "InvalidType", + CoreError::MissingField(_) => "MissingField", + CoreError::Http { .. } => "HttpError", + CoreError::InvalidResponse(_) => "InvalidResponse", + CoreError::Network(_) => "NetworkError", + CoreError::Routing(_) => "RoutingError", + } +} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs new file mode 100644 index 00000000000..b54ee39b21d --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs @@ -0,0 +1,25 @@ +use litellm_core::call_lifecycle::CallLifecycle; +use litellm_core::CoreResult; +use serde_json::Value; + +mod client; +mod common_utils; +mod handler; +mod hooks; +mod prepare; +mod types; + +pub use types::OcrRequest; + +use handler::execute_ocr_provider_call; +use prepare::{prepare_ocr_call, PreparedOcrCall}; + +pub async fn ocr(request: OcrRequest<'_>) -> CoreResult { + let PreparedOcrCall { request, hooks } = prepare_ocr_call(request); + CallLifecycle::default() + .run_request(request, &hooks, execute_ocr_provider_call) + .await +} + +#[cfg(test)] +mod tests; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs new file mode 100644 index 00000000000..5a4b350a4c4 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs @@ -0,0 +1,57 @@ +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{SystemTime, UNIX_EPOCH}; + +use litellm_core::routing_utils::provider::{get_custom_llm_provider, CustomLlmProvider}; + +use super::hooks::OcrLifecycleHooks; +use super::types::{OcrRequest, PreparedOcrRequest}; +use crate::integrations::custom_guardrail::CustomGuardrailRunner; +use crate::integrations::custom_logger::CustomLoggerRunner; + +pub(crate) struct PreparedOcrCall { + pub(crate) request: PreparedOcrRequest, + pub(crate) hooks: OcrLifecycleHooks, +} + +pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall { + let call_id = request + .litellm_call_id + .map(str::to_string) + .unwrap_or_else(new_ocr_call_id); + let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider) + .unwrap_or(CustomLlmProvider { + model: request.model, + custom_llm_provider: "mistral", + }); + let model = provider_info.model.to_string(); + let custom_llm_provider = provider_info.custom_llm_provider.to_string(); + + PreparedOcrCall { + request: PreparedOcrRequest { + model, + custom_llm_provider, + litellm_call_id: call_id, + document: request.document, + api_key: request.api_key.map(str::to_string), + api_base: request.api_base.map(str::to_string), + extra_headers: request.extra_headers, + optional_params: request.optional_params, + timeout: request.timeout, + }, + hooks: OcrLifecycleHooks::new( + CustomLoggerRunner::new(request.callbacks), + CustomGuardrailRunner::new(request.guardrails), + request.request_metadata, + ), + } +} + +fn new_ocr_call_id() -> String { + static COUNTER: AtomicU64 = AtomicU64::new(1); + let sequence = COUNTER.fetch_add(1, Ordering::Relaxed); + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_nanos()) + .unwrap_or(0); + format!("ocr-{timestamp}-{sequence}") +} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs new file mode 100644 index 00000000000..35747dc6985 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs @@ -0,0 +1,610 @@ +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use litellm_core::error::CoreError; +use litellm_core::ocr::transformation::OcrResponseHandling; +use serde_json::{json, Map, Value}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; + +use super::common_utils::{has_header, ocr_provider_config, string_headers, truncate_error_body}; +use super::{ocr, OcrRequest}; +use crate::integrations::custom_guardrail::{ + CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook, + GuardrailFuture, GuardrailRequest, +}; +use crate::integrations::custom_logger::{ + CallbackTiming, CallbackValue, CustomLogger, LogFuture, ModelCallDetails, +}; +use crate::integrations::types::RequestMetadata; + +async fn read_http_headers(socket: &mut TcpStream) -> String { + let mut request = Vec::new(); + let mut buffer = [0_u8; 1024]; + loop { + let n = socket.read(&mut buffer).await.expect("reads request"); + if n == 0 { + break; + } + request.extend_from_slice(&buffer[..n]); + if request.windows(4).any(|window| window == b"\r\n\r\n") { + break; + } + } + String::from_utf8(request).expect("request is utf8") +} + +async fn read_http_request(socket: &mut TcpStream) -> String { + let mut request = Vec::new(); + let mut buffer = [0_u8; 1024]; + let header_end = loop { + let n = socket.read(&mut buffer).await.expect("reads request"); + if n == 0 { + break request.len(); + } + request.extend_from_slice(&buffer[..n]); + if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") { + break position + 4; + } + }; + let headers = String::from_utf8_lossy(&request[..header_end]); + let content_length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) + .unwrap_or(0); + while request.len().saturating_sub(header_end) < content_length { + let n = socket.read(&mut buffer).await.expect("reads body"); + if n == 0 { + break; + } + request.extend_from_slice(&buffer[..n]); + } + String::from_utf8(request).expect("request is utf8") +} + +#[derive(Clone, Debug, PartialEq)] +struct RecordedLogEvent { + hook: &'static str, + model: String, + call_type: String, + user_id: Option, + response_object: Option, + error_kind: Option, +} + +#[derive(Default)] +struct RecordingOcrLogger { + events: Mutex>, +} + +impl RecordingOcrLogger { + fn events(&self) -> Vec { + self.events.lock().unwrap().clone() + } +} + +impl CustomLogger for RecordingOcrLogger { + fn async_log_success_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + response_obj: &'a CallbackValue, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push(RecordedLogEvent { + hook: "async_log_success_event", + model: model_call_details.model.clone(), + call_type: model_call_details.call_type.to_string(), + user_id: model_call_details.metadata.user_api_key_user_id.clone(), + response_object: Some(response_obj.object.clone()), + error_kind: None, + }); + Ok(()) + }) + } + + fn async_log_failure_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + response_obj: Option<&'a CallbackValue>, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push(RecordedLogEvent { + hook: "async_log_failure_event", + model: model_call_details.model.clone(), + call_type: model_call_details.call_type.to_string(), + user_id: model_call_details.metadata.user_api_key_user_id.clone(), + response_object: response_obj.map(|value| value.object.clone()), + error_kind: model_call_details + .failure_error + .as_ref() + .map(|error| error.kind.clone()), + }); + Ok(()) + }) + } +} + +struct RecordingOcrGuardrail { + hooks: Vec, + events: Mutex>, + block_pre_call: bool, +} + +impl RecordingOcrGuardrail { + fn new(hooks: Vec) -> Self { + Self { + hooks, + events: Mutex::new(Vec::new()), + block_pre_call: false, + } + } + + fn blocking_pre_call() -> Self { + Self { + hooks: vec![GuardrailEventHook::PreCall], + events: Mutex::new(Vec::new()), + block_pre_call: true, + } + } + + fn events(&self) -> Vec<&'static str> { + self.events.lock().unwrap().clone() + } +} + +impl CustomGuardrail for RecordingOcrGuardrail { + fn guardrail_name(&self) -> &str { + "recording-ocr-guardrail" + } + + fn supported_event_hooks(&self) -> &[GuardrailEventHook] { + &self.hooks + } + + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a GuardrailContext, + mut request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("async_pre_call_hook"); + if self.block_pre_call { + return Ok(GuardrailDecision::Block(GuardrailError::blocked( + "blocked before provider", + ))); + } + request.data["document"]["guarded_pre"] = json!(true); + Ok(GuardrailDecision::Mask(request)) + }) + } + + fn async_moderation_hook<'a>( + &'a self, + _context: &'a GuardrailContext, + mut request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("async_moderation_hook"); + request.data["body"]["guarded_during"] = json!(true); + Ok(GuardrailDecision::Mask(request)) + }) + } +} + +#[test] +fn truncate_error_body_passes_short_strings_through() { + let body = "Unauthorized"; + assert_eq!(truncate_error_body(body), "Unauthorized"); +} + +#[test] +fn truncate_error_body_caps_long_payloads() { + let body = "x".repeat(306); + let truncated = truncate_error_body(&body); + + assert!(truncated.ends_with("... (truncated)")); + let prefix_chars = truncated + .strip_suffix("... (truncated)") + .expect("truncated marker present") + .chars() + .count(); + assert_eq!(prefix_chars, 256); +} + +#[test] +fn truncate_error_body_does_not_split_multibyte_chars() { + let body = "é".repeat(266); + let truncated = truncate_error_body(&body); + assert!(truncated.is_char_boundary(truncated.len())); +} + +#[test] +fn ocr_dispatch_supports_migrated_providers() { + assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some()); + assert!(ocr_provider_config("azure_ai", "pixtral-12b-2409") + .expect("azure ai config resolves") + .requires_data_uri_document()); + assert_eq!( + ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read") + .expect("document intelligence config resolves") + .response_handling(), + OcrResponseHandling::AzureDocumentIntelligencePoll + ); + assert!(ocr_provider_config("vertex_ai", "deepseek-ocr-maas") + .expect("vertex deepseek config resolves") + .supported_ocr_params() + .contains(&"temperature")); + assert!(ocr_provider_config("openai", "gpt-4o").is_none()); +} + +#[test] +fn string_headers_accepts_string_values() { + let headers = json!({ + "x-trace-id": "trace-1" + }) + .as_object() + .unwrap() + .clone(); + + assert_eq!( + string_headers(Some(headers)).expect("string headers accepted"), + vec![("x-trace-id".to_string(), "trace-1".to_string())] + ); +} + +#[test] +fn auth_header_detection_is_case_insensitive() { + let headers = vec![ + ("x-trace-id".to_string(), "trace-1".to_string()), + ("authorization".to_string(), "Bearer sk-test".to_string()), + ]; + + assert!(has_header(&headers, "authorization")); + + let headers = vec![("Authorization".to_string(), "Bearer sk-test".to_string())]; + assert!(has_header(&headers, "authorization")); + + let headers = vec![("x-trace-id".to_string(), "trace-1".to_string())]; + assert!(!has_header(&headers, "authorization")); +} + +#[tokio::test] +async fn ocr_lifecycle_runs_pre_during_and_success_hooks() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let addr = listener.local_addr().expect("listener has local addr"); + + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts one request"); + let request = read_http_request(&mut socket).await; + let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + response_body.len(), + response_body + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes response"); + request + }); + + let logger = Arc::new(RecordingOcrLogger::default()); + let guardrail = Arc::new(RecordingOcrGuardrail::new(vec![ + GuardrailEventHook::PreCall, + GuardrailEventHook::DuringCall, + ])); + let response = ocr(OcrRequest { + model: "mistral-ocr-latest", + document: json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }), + api_key: Some("sk-test"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("mistral"), + extra_headers: None, + optional_params: Map::new(), + timeout: Some(Duration::from_secs(5)), + callbacks: vec![logger.clone()], + guardrails: vec![guardrail.clone()], + request_metadata: RequestMetadata { + user_api_key_user_id: Some("user-1".to_string()), + ..Default::default() + }, + litellm_call_id: Some("ocr-call-1"), + }) + .await + .expect("ocr request succeeds"); + + assert_eq!(response["pages"][0]["markdown"], "ok"); + assert_eq!( + guardrail.events(), + vec!["async_pre_call_hook", "async_moderation_hook"] + ); + assert_eq!( + logger.events(), + vec![RecordedLogEvent { + hook: "async_log_success_event", + model: "mistral-ocr-latest".to_string(), + call_type: "ocr".to_string(), + user_id: Some("user-1".to_string()), + response_object: Some("ocr".to_string()), + error_kind: None, + }] + ); + + let request = server.await.expect("server task completes"); + assert!(request.contains(r#""guarded_pre":true"#), "{request}"); + assert!(request.contains(r#""guarded_during":true"#), "{request}"); +} + +#[tokio::test] +async fn ocr_lifecycle_runs_failure_hook_on_provider_error() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let addr = listener.local_addr().expect("listener has local addr"); + + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts one request"); + let _request = read_http_request(&mut socket).await; + let response_body = "provider failed"; + let response = format!( + "HTTP/1.1 500 Internal Server Error\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + response_body.len(), + response_body + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes response"); + }); + + let logger = Arc::new(RecordingOcrLogger::default()); + let err = ocr(OcrRequest { + model: "mistral-ocr-latest", + document: json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }), + api_key: Some("sk-test"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("mistral"), + extra_headers: None, + optional_params: Map::new(), + timeout: Some(Duration::from_secs(5)), + callbacks: vec![logger.clone()], + guardrails: Vec::new(), + request_metadata: RequestMetadata::default(), + litellm_call_id: Some("ocr-call-2"), + }) + .await + .expect_err("provider error propagates"); + + assert!(matches!(err, CoreError::Http { status: 500, .. })); + server.await.expect("server task completes"); + assert_eq!( + logger.events(), + vec![RecordedLogEvent { + hook: "async_log_failure_event", + model: "mistral-ocr-latest".to_string(), + call_type: "ocr".to_string(), + user_id: None, + response_object: Some("error".to_string()), + error_kind: Some("HttpError".to_string()), + }] + ); +} + +#[tokio::test] +async fn ocr_lifecycle_pre_call_block_skips_provider_socket() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let addr = listener.local_addr().expect("listener has local addr"); + let logger = Arc::new(RecordingOcrLogger::default()); + let guardrail = Arc::new(RecordingOcrGuardrail::blocking_pre_call()); + + let err = ocr(OcrRequest { + model: "mistral-ocr-latest", + document: json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }), + api_key: Some("sk-test"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("mistral"), + extra_headers: None, + optional_params: Map::new(), + timeout: Some(Duration::from_millis(100)), + callbacks: vec![logger.clone()], + guardrails: vec![guardrail.clone()], + request_metadata: RequestMetadata::default(), + litellm_call_id: Some("ocr-call-3"), + }) + .await + .expect_err("guardrail blocks request"); + + assert!(matches!(err, CoreError::InvalidRequest(_))); + assert_eq!(guardrail.events(), vec!["async_pre_call_hook"]); + assert_eq!( + logger.events(), + vec![RecordedLogEvent { + hook: "async_log_failure_event", + model: "mistral-ocr-latest".to_string(), + call_type: "ocr".to_string(), + user_id: None, + response_object: Some("error".to_string()), + error_kind: Some("InvalidRequest".to_string()), + }] + ); + let accepted = tokio::time::timeout(Duration::from_millis(100), listener.accept()).await; + assert!(accepted.is_err(), "provider socket should not be touched"); +} + +#[tokio::test] +async fn ocr_does_not_duplicate_authorization_header_when_header_is_supplied() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let addr = listener.local_addr().expect("listener has local addr"); + + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts one request"); + let request = read_http_headers(&mut socket).await; + let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + response_body.len(), + response_body + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes response"); + request + }); + + let mut headers = Map::new(); + headers.insert( + "Authorization".to_string(), + Value::String("Bearer sk-from-python".to_string()), + ); + headers.insert( + "x-trace-id".to_string(), + Value::String("trace-1".to_string()), + ); + + let response = ocr(OcrRequest { + model: "mistral-ocr-latest", + document: json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }), + api_key: Some("sk-for-rust-fallback"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("mistral"), + extra_headers: Some(headers), + optional_params: Map::new(), + timeout: Some(Duration::from_secs(5)), + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: RequestMetadata::default(), + litellm_call_id: None, + }) + .await + .expect("ocr request succeeds"); + + assert_eq!(response["pages"][0]["markdown"], "ok"); + + let request = server.await.expect("server task completes"); + let authorization_count = request + .lines() + .filter(|line| line.to_ascii_lowercase().starts_with("authorization:")) + .count(); + assert_eq!(authorization_count, 1, "{request}"); + assert!( + request.contains("authorization: Bearer sk-from-python") + || request.contains("Authorization: Bearer sk-from-python"), + "{request}" + ); +} + +#[tokio::test] +async fn document_intelligence_poll_uses_resolved_subscription_key() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let addr = listener.local_addr().expect("listener has local addr"); + let operation_url = format!("http://{addr}/operations/1"); + + let server = tokio::spawn(async move { + let (mut post_socket, _) = listener.accept().await.expect("accepts post request"); + let post_request = read_http_headers(&mut post_socket).await; + let post_response = format!( + "HTTP/1.1 202 Accepted\r\noperation-location: {operation_url}\r\ncontent-length: 0\r\nconnection: close\r\n\r\n" + ); + post_socket + .write_all(post_response.as_bytes()) + .await + .expect("writes post response"); + + let (mut poll_socket, _) = listener.accept().await.expect("accepts poll request"); + let poll_request = read_http_headers(&mut poll_socket).await; + let response_body = r#"{"status":"succeeded","analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"ok"}]}]}}"#; + let poll_response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + response_body.len(), + response_body + ); + poll_socket + .write_all(poll_response.as_bytes()) + .await + .expect("writes poll response"); + (post_request, poll_request) + }); + + let response = ocr(OcrRequest { + model: "doc-intelligence/prebuilt-read", + document: json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }), + api_key: Some("di-key"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("azure_ai"), + extra_headers: None, + optional_params: Map::new(), + timeout: Some(Duration::from_secs(5)), + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: RequestMetadata::default(), + litellm_call_id: None, + }) + .await + .expect("document intelligence request succeeds"); + + assert_eq!(response["pages"][0]["markdown"], "ok"); + + let (post_request, poll_request) = server.await.expect("server task completes"); + assert!( + post_request + .to_ascii_lowercase() + .contains("ocp-apim-subscription-key: di-key"), + "{post_request}" + ); + assert!( + poll_request + .to_ascii_lowercase() + .contains("ocp-apim-subscription-key: di-key"), + "{poll_request}" + ); +} + +#[test] +fn string_headers_rejects_non_string_values() { + let headers = json!({ + "x-retry-count": 3 + }) + .as_object() + .unwrap() + .clone(); + + let err = string_headers(Some(headers)).expect_err("non-string header rejected"); + assert_eq!( + err, + CoreError::InvalidRequest( + "OCR extra_headers.x-retry-count must be a string, got number".to_string() + ) + ); +} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/types.rs b/litellm-rust/crates/ai-gateway/src/ocr/types.rs new file mode 100644 index 00000000000..bde734a4dd1 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/ocr/types.rs @@ -0,0 +1,57 @@ +use std::sync::Arc; +use std::time::Duration; + +use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest}; +use litellm_core::ocr::transformation::OcrProviderConfig; +use serde_json::{Map, Value}; + +use crate::integrations::custom_guardrail::CustomGuardrail; +use crate::integrations::custom_logger::CustomLogger; +use crate::integrations::types::RequestMetadata; + +pub struct OcrRequest<'a> { + pub model: &'a str, + pub document: Value, + pub api_key: Option<&'a str>, + pub api_base: Option<&'a str>, + pub custom_llm_provider: Option<&'a str>, + pub extra_headers: Option>, + pub optional_params: Map, + pub timeout: Option, + pub callbacks: Vec>, + pub guardrails: Vec>, + pub request_metadata: RequestMetadata, + pub litellm_call_id: Option<&'a str>, +} + +pub(crate) struct PreparedOcrRequest { + pub(crate) model: String, + pub(crate) custom_llm_provider: String, + pub(crate) litellm_call_id: String, + pub(crate) document: Value, + pub(crate) api_key: Option, + pub(crate) api_base: Option, + pub(crate) extra_headers: Option>, + pub(crate) optional_params: Map, + pub(crate) timeout: Option, +} + +impl CallLifecycleRequest for PreparedOcrRequest { + fn lifecycle_context(&self) -> CallLifecycleContext { + CallLifecycleContext::new( + "ocr", + self.model.clone(), + self.custom_llm_provider.clone(), + self.litellm_call_id.clone(), + ) + } +} + +pub(crate) struct ProviderOcrRequest { + pub(crate) model: String, + pub(crate) config: &'static dyn OcrProviderConfig, + pub(crate) url: String, + pub(crate) body: Value, + pub(crate) upstream_headers: Vec<(String, String)>, + pub(crate) timeout: Option, +} diff --git a/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs b/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs index 34c82897808..c32e727de54 100644 --- a/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs +++ b/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs @@ -13,7 +13,9 @@ use litellm_core::realtime::types::RealtimeEvent; use serde_json::Value; use crate::constants::DEFAULT_PROVIDER; -use crate::integrations::custom_logger::CustomLogger; +use crate::integrations::custom_logger::{ + CallbackTiming, CallbackValue, CustomLogger, CustomLoggerRunner, LoggingError, ModelCallDetails, +}; use crate::integrations::types::{ RequestMetadata, StandardLoggingMetadata, StandardLoggingPayload, Usage, }; @@ -183,30 +185,45 @@ impl RealTimeStreaming { /// Finish the session: stamp the end time and fan the payload out to every /// callback. On a logger enqueue error we bump a non-fatal counter (the /// realtime session has already ended; a dropped log must never propagate). - pub fn log_messages(&mut self, status: SessionStatus) { + pub async fn log_messages(&mut self, status: SessionStatus) { self.end_time = epoch_seconds(); let payload = self.build_payload(); + let timing = CallbackTiming::new(payload.start_time, payload.end_time); + let runner = CustomLoggerRunner::new(self.callbacks.clone()); match status { SessionStatus::Success => { - for callback in &self.callbacks { - if let Err(err) = callback.log_success_event(&payload) { - self.dropped += 1; - eprintln!("litellm-ai-gateway: log_success_event dropped: {err}"); - } - } + let response = CallbackValue::new("realtime", serde_json::Value::Null); + let report = runner + .async_log_success_event( + &ModelCallDetails::from_standard_logging_payload(payload), + &response, + timing, + ) + .await; + self.dropped += report.dropped as u64; } SessionStatus::Failure => { - let error = crate::integrations::types::LoggingError { + let error = LoggingError { message: "realtime session ended in failure".to_string(), kind: "RealtimeSessionError".to_string(), }; - for callback in &self.callbacks { - if let Err(err) = callback.log_failure_event(&payload, &error) { - self.dropped += 1; - eprintln!("litellm-ai-gateway: log_failure_event dropped: {err}"); - } - } + let response = CallbackValue::new( + "error", + serde_json::json!({ + "message": error.message, + "kind": error.kind, + }), + ); + let report = runner + .async_log_failure_event( + &ModelCallDetails::from_standard_logging_payload(payload) + .with_failure_error(error), + Some(&response), + timing, + ) + .await; + self.dropped += report.dropped as u64; } } } @@ -215,7 +232,8 @@ impl RealTimeStreaming { #[cfg(test)] mod tests { use super::*; - use crate::integrations::types::{LogError, LoggingError}; + use crate::integrations::custom_logger::LogError; + use crate::integrations::custom_logger::LogFuture; use std::sync::atomic::{AtomicU64, Ordering}; fn event(raw: &str) -> RealtimeEvent { @@ -231,17 +249,28 @@ mod tests { } impl CustomLogger for CapturingLogger { - fn log_success_event(&self, payload: &StandardLoggingPayload) -> Result<(), LogError> { - self.calls.fetch_add(1, Ordering::SeqCst); - *self.last_model.lock().unwrap() = Some(payload.model.clone()); - self.last_total_tokens - .store(payload.total_tokens, Ordering::SeqCst); - Ok(()) + fn async_log_success_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + _response_obj: &'a CallbackValue, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + let payload = model_call_details + .standard_logging_payload + .as_ref() + .expect("standard logging payload"); + self.calls.fetch_add(1, Ordering::SeqCst); + *self.last_model.lock().unwrap() = Some(payload.model.clone()); + self.last_total_tokens + .store(payload.total_tokens, Ordering::SeqCst); + Ok(()) + }) } } - #[test] - fn observe_accumulates_model_and_tokens_then_logs() { + #[tokio::test] + async fn observe_accumulates_model_and_tokens_then_logs() { let logger = Arc::new(CapturingLogger::default()); let callbacks: Vec> = vec![logger.clone()]; let mut streaming = RealTimeStreaming::new( @@ -284,7 +313,7 @@ mod tests { Some("hash123") ); - streaming.log_messages(SessionStatus::Success); + streaming.log_messages(SessionStatus::Success).await; assert_eq!(logger.calls.load(Ordering::SeqCst), 1); assert_eq!( logger.last_model.lock().unwrap().as_deref(), @@ -324,19 +353,26 @@ mod tests { /// A logger whose enqueue always fails should bump the dropped counter, not /// panic or propagate. - #[test] - fn failing_logger_bumps_dropped_counter() { + #[tokio::test] + async fn failing_logger_bumps_dropped_counter() { struct FailingLogger; impl CustomLogger for FailingLogger { - fn log_success_event(&self, _p: &StandardLoggingPayload) -> Result<(), LogError> { - Err(LogError::channel_full()) + fn async_log_success_event<'a>( + &'a self, + _model_call_details: &'a ModelCallDetails, + _response_obj: &'a CallbackValue, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async { Err(LogError::channel_full()) }) } - fn log_failure_event( - &self, - _p: &StandardLoggingPayload, - _e: &LoggingError, - ) -> Result<(), LogError> { - Err(LogError::channel_closed()) + + fn async_log_failure_event<'a>( + &'a self, + _model_call_details: &'a ModelCallDetails, + _response_obj: Option<&'a CallbackValue>, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async { Err(LogError::channel_closed()) }) } } let callbacks: Vec> = vec![Arc::new(FailingLogger)]; @@ -346,7 +382,7 @@ mod tests { "gpt-realtime".to_string(), RequestMetadata::default(), ); - streaming.log_messages(SessionStatus::Success); + streaming.log_messages(SessionStatus::Success).await; assert_eq!(streaming.dropped(), 1); } } diff --git a/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs index 899ad73829f..c3f929f5f0b 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs @@ -162,5 +162,5 @@ async fn bridge( } else { SessionStatus::Failure }; - collector.log_messages(status); + collector.log_messages(status).await; } diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 1881bcfa602..9bd4634cc2a 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -10,3 +10,6 @@ rand.workspace = true serde.workspace = true serde_json.workspace = true thiserror.workspace = true + +[dev-dependencies] +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/litellm-rust/crates/core/src/call_lifecycle/README.md b/litellm-rust/crates/core/src/call_lifecycle/README.md new file mode 100644 index 00000000000..692e249ef27 --- /dev/null +++ b/litellm-rust/crates/core/src/call_lifecycle/README.md @@ -0,0 +1,167 @@ +# Call lifecycle + +`litellm_core::call_lifecycle` is the shared execution wrapper for LiteLLM call +types migrated to Rust. It owns lifecycle ordering, phase timing, and trace +observer calls. It must not know about OCR, chat, messages, responses, +completions, provider auth, request transforms, or response normalization. + +Call-type modules own their domain behavior. For example, OCR owns document +payloads, OCR provider transforms, safe document fetch, guardrail payload shape, +callback payload shape, and provider HTTP execution. + +## Runtime order + +Every wrapped call runs in this order: + +1. `async_pre_call_hook` +2. `async_during_call_hook` +3. provider call +4. `async_log_success_event` or `async_log_failure_event` + +`async_pre_call_hook` receives the initial LiteLLM request shape. It is where +pre-call custom guardrails run. + +`async_during_call_hook` converts the initial request into the provider-ready +request. It is where provider config selection, parameter mapping, auth/header +resolution, request transforms, and during-call guardrails belong. + +The provider call receives only the provider-ready request. It should execute +I/O and call the provider response transform. + +Success and failure callbacks receive `CallLifecycleTiming`. Callback failures +must not replace the original provider or guardrail result. + +## Trace contract + +The lifecycle runner records: + +- full call start and end time +- `pre_call` phase timing +- `during_call` phase timing +- `provider_call` phase timing +- `success_callback` phase timing +- `failure_callback` phase timing + +`CallLifecycleObserver` receives phase start and end events. The default +observer is a no-op. Future OTEL support should implement this observer instead +of editing OCR, chat, messages, responses, completions, or provider modules. + +## Required shape + +Each migrated call type should use this folder shape: + +```text +litellm-rust/crates/ai-gateway/src// + mod.rs # thin public entrypoint + types.rs # public request, prepared request, provider request, response types + prepare.rs # model/provider/callback/guardrail setup + hooks.rs # CallLifecycleHooks implementation + handler.rs # provider I/O and response normalization + tests.rs # call-type lifecycle and handler tests +``` + +Provider transforms can live in `litellm-rust/crates/core/src/providers/...`. +Shared call-type helpers can live beside the call type, but generic lifecycle +code stays in this folder. + +## Core API + +The prepared request implements `CallLifecycleRequest`: + +```rust +impl CallLifecycleRequest for PreparedMessagesRequest { + fn lifecycle_context(&self) -> CallLifecycleContext { + CallLifecycleContext::new( + "messages", + self.model.clone(), + self.custom_llm_provider.clone(), + self.litellm_call_id.clone(), + ) + } +} +``` + +The call-type hooks implement `CallLifecycleHooks`: + +```rust +impl CallLifecycleHooks< + PreparedMessagesRequest, + ProviderMessagesRequest, + MessagesResponse, +> for MessagesLifecycleHooks { + fn async_pre_call_hook(...) { + // run pre-call custom guardrails against the LiteLLM request shape + } + + fn async_during_call_hook(...) { + // map params, validate env, transform request, run during-call guardrails + } + + fn async_log_success_event(...) { + // call async_log_success_event on configured custom loggers + } + + fn async_log_failure_event(...) { + // call async_log_failure_event without swallowing the original error + } +} +``` + +The public entrypoint stays thin: + +```rust +pub async fn messages(request: MessagesRequest<'_>) -> CoreResult { + let PreparedMessagesCall { request, hooks } = prepare_messages_call(request)?; + + CallLifecycle::default() + .run_request(request, &hooks, execute_messages_provider_call) + .await +} +``` + +Use `run_request` for new call types. Keep `run` available only for specialized +tests or existing code that already has a `CallLifecycleContext`. + +## Adding a new call type + +1. Add `/types.rs` + +Define the public request accepted by the bridge, the prepared request used by +the lifecycle runner, and the provider request consumed by the handler. + +2. Implement `CallLifecycleRequest` + +Return `call_type`, `model`, `custom_llm_provider`, and `litellm_call_id`. +Do not put provider-specific logic here. + +3. Add `/prepare.rs` + +Resolve model/provider once, generate or preserve `litellm_call_id`, construct +callback and guardrail runners, and return `PreparedCall`. + +4. Add `/hooks.rs` + +Implement `CallLifecycleHooks`. Put pre-call guardrail payload construction, +provider config selection, param mapping, request transform, during-call +guardrail payload construction, and callback payload construction here. + +5. Add `/handler.rs` + +Execute the provider request and normalize the provider response. Do not repeat +provider-specific transforms here; call the provider config. + +6. Add tests + +Cover hook order, success callback payload, failure callback payload, pre-call +guardrail blocking before provider I/O, during-call body mutation, and provider +error mapping. + +## Review checklist + +- Core lifecycle has no call-type or provider-specific branches +- Public call-type entrypoint only prepares and calls `run_request` +- Provider behavior lives behind provider config/transformation code +- Hook method names map to the Python custom logger and guardrail concepts +- Phase timing is recorded once in lifecycle, not separately per call type +- Callback failures never hide the original provider or guardrail error +- Tests prove the provider socket is not touched when pre-call guardrails block diff --git a/litellm-rust/crates/core/src/call_lifecycle/mod.rs b/litellm-rust/crates/core/src/call_lifecycle/mod.rs new file mode 100644 index 00000000000..d9b68a1b726 --- /dev/null +++ b/litellm-rust/crates/core/src/call_lifecycle/mod.rs @@ -0,0 +1,414 @@ +use std::future::Future; +use std::time::{Instant, SystemTime, UNIX_EPOCH}; + +use crate::{CoreError, CoreResult}; + +pub mod types; + +pub use types::{ + CallLifecycleContext, CallLifecyclePhase, CallLifecyclePhaseTiming, CallLifecycleRequest, + CallLifecycleTiming, +}; + +pub trait CallLifecycleHooks: Send + Sync { + type PreCallFuture<'a>: Future> + Send + 'a + where + Self: 'a, + InitialReq: 'a, + ProviderReq: 'a, + Resp: 'a; + + type DuringCallFuture<'a>: Future> + Send + 'a + where + Self: 'a, + InitialReq: 'a, + ProviderReq: 'a, + Resp: 'a; + + type SuccessFuture<'a>: Future + Send + 'a + where + Self: 'a, + Resp: 'a; + + type FailureFuture<'a>: Future + Send + 'a + where + Self: 'a; + + fn async_pre_call_hook<'a>( + &'a self, + context: &'a CallLifecycleContext, + request: InitialReq, + ) -> Self::PreCallFuture<'a>; + + fn async_during_call_hook<'a>( + &'a self, + context: &'a CallLifecycleContext, + request: InitialReq, + ) -> Self::DuringCallFuture<'a>; + + fn async_log_success_event<'a>( + &'a self, + context: &'a CallLifecycleContext, + response: &'a Resp, + timing: &'a CallLifecycleTiming, + ) -> Self::SuccessFuture<'a>; + + fn async_log_failure_event<'a>( + &'a self, + context: &'a CallLifecycleContext, + error: &'a CoreError, + timing: &'a CallLifecycleTiming, + ) -> Self::FailureFuture<'a>; +} + +pub trait CallLifecycleObserver: Send + Sync { + fn on_phase_start(&self, _context: &CallLifecycleContext, _phase: CallLifecyclePhase) {} + + fn on_phase_end(&self, _context: &CallLifecycleContext, _timing: &CallLifecyclePhaseTiming) {} +} + +#[derive(Default)] +pub struct NoopCallLifecycleObserver; + +impl CallLifecycleObserver for NoopCallLifecycleObserver {} + +pub struct CallLifecycle<'a> { + observer: &'a dyn CallLifecycleObserver, +} + +impl<'a> CallLifecycle<'a> { + pub fn new(observer: &'a dyn CallLifecycleObserver) -> Self { + Self { observer } + } + + pub async fn run_request( + &self, + request: InitialReq, + hooks: &Hooks, + provider_call: ProviderCall, + ) -> CoreResult + where + InitialReq: CallLifecycleRequest, + Hooks: CallLifecycleHooks, + ProviderCall: FnOnce(ProviderReq) -> ProviderFuture, + ProviderFuture: Future>, + { + let context = request.lifecycle_context(); + self.run(context, request, hooks, provider_call).await + } + + pub async fn run( + &self, + context: CallLifecycleContext, + request: InitialReq, + hooks: &Hooks, + provider_call: ProviderCall, + ) -> CoreResult + where + Hooks: CallLifecycleHooks, + ProviderCall: FnOnce(ProviderReq) -> ProviderFuture, + ProviderFuture: Future>, + { + let call_start = epoch_seconds(); + let mut phases = Vec::new(); + + let pre_call = self.start_phase(&context, CallLifecyclePhase::PreCall); + let request = match hooks.async_pre_call_hook(&context, request).await { + Ok(request) => { + phases.push(self.finish_phase(&context, pre_call)); + request + } + Err(error) => { + phases.push(self.finish_phase(&context, pre_call)); + self.log_failure(&context, hooks, &error, call_start, &mut phases) + .await; + return Err(error); + } + }; + + let during_call = self.start_phase(&context, CallLifecyclePhase::DuringCall); + let provider_request = match hooks.async_during_call_hook(&context, request).await { + Ok(request) => { + phases.push(self.finish_phase(&context, during_call)); + request + } + Err(error) => { + phases.push(self.finish_phase(&context, during_call)); + self.log_failure(&context, hooks, &error, call_start, &mut phases) + .await; + return Err(error); + } + }; + + let provider_phase = self.start_phase(&context, CallLifecyclePhase::ProviderCall); + let result = provider_call(provider_request).await; + phases.push(self.finish_phase(&context, provider_phase)); + + match &result { + Ok(response) => { + let success_phase = self.start_phase(&context, CallLifecyclePhase::SuccessCallback); + let timing = CallLifecycleTiming::new(call_start, epoch_seconds(), phases.clone()); + hooks + .async_log_success_event(&context, response, &timing) + .await; + phases.push(self.finish_phase(&context, success_phase)); + } + Err(error) => { + self.log_failure(&context, hooks, error, call_start, &mut phases) + .await; + } + } + + result + } + + async fn log_failure( + &self, + context: &CallLifecycleContext, + hooks: &Hooks, + error: &CoreError, + call_start: f64, + phases: &mut Vec, + ) where + Hooks: CallLifecycleHooks, + { + let failure_phase = self.start_phase(context, CallLifecyclePhase::FailureCallback); + let timing = CallLifecycleTiming::new(call_start, epoch_seconds(), phases.clone()); + hooks.async_log_failure_event(context, error, &timing).await; + phases.push(self.finish_phase(context, failure_phase)); + } + + fn start_phase(&self, context: &CallLifecycleContext, phase: CallLifecyclePhase) -> PhaseStart { + self.observer.on_phase_start(context, phase); + PhaseStart { + phase, + start_time: epoch_seconds(), + started_at: Instant::now(), + } + } + + fn finish_phase( + &self, + context: &CallLifecycleContext, + phase_start: PhaseStart, + ) -> CallLifecyclePhaseTiming { + let timing = CallLifecyclePhaseTiming { + phase: phase_start.phase, + start_time: phase_start.start_time, + end_time: epoch_seconds(), + duration: phase_start.started_at.elapsed(), + }; + self.observer.on_phase_end(context, &timing); + timing + } +} + +impl Default for CallLifecycle<'static> { + fn default() -> Self { + static OBSERVER: NoopCallLifecycleObserver = NoopCallLifecycleObserver; + Self::new(&OBSERVER) + } +} + +struct PhaseStart { + phase: CallLifecyclePhase, + start_time: f64, + started_at: Instant, +} + +fn epoch_seconds() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs_f64()) + .unwrap_or(0.0) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::pin::Pin; + use std::sync::Mutex; + + type BoxFuture<'a, T> = Pin + Send + 'a>>; + + #[derive(Default)] + struct RecordingHooks { + events: Mutex>, + } + + struct RecordingRequest(String); + + impl CallLifecycleRequest for RecordingRequest { + fn lifecycle_context(&self) -> CallLifecycleContext { + CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1") + } + } + + impl RecordingHooks { + fn events(&self) -> Vec<&'static str> { + self.events.lock().unwrap().clone() + } + } + + impl CallLifecycleHooks for RecordingHooks { + type PreCallFuture<'a> = BoxFuture<'a, CoreResult>; + type DuringCallFuture<'a> = BoxFuture<'a, CoreResult>; + type SuccessFuture<'a> = BoxFuture<'a, ()>; + type FailureFuture<'a> = BoxFuture<'a, ()>; + + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: String, + ) -> Self::PreCallFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("pre_call"); + Ok(format!("{request}:pre")) + }) + } + + fn async_during_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: String, + ) -> Self::DuringCallFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("during_call"); + Ok(format!("{request}:during")) + }) + } + + fn async_log_success_event<'a>( + &'a self, + _context: &'a CallLifecycleContext, + _response: &'a String, + timing: &'a CallLifecycleTiming, + ) -> Self::SuccessFuture<'a> { + Box::pin(async move { + assert!(timing.end_time >= timing.start_time); + assert_eq!(timing.phases.len(), 3); + self.events.lock().unwrap().push("success"); + }) + } + + fn async_log_failure_event<'a>( + &'a self, + _context: &'a CallLifecycleContext, + _error: &'a CoreError, + _timing: &'a CallLifecycleTiming, + ) -> Self::FailureFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("failure"); + }) + } + } + + impl CallLifecycleHooks for RecordingHooks { + type PreCallFuture<'a> = BoxFuture<'a, CoreResult>; + type DuringCallFuture<'a> = BoxFuture<'a, CoreResult>; + type SuccessFuture<'a> = BoxFuture<'a, ()>; + type FailureFuture<'a> = BoxFuture<'a, ()>; + + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: RecordingRequest, + ) -> Self::PreCallFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("pre_call"); + Ok(RecordingRequest(format!("{}:pre", request.0))) + }) + } + + fn async_during_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: RecordingRequest, + ) -> Self::DuringCallFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("during_call"); + Ok(format!("{}:during", request.0)) + }) + } + + fn async_log_success_event<'a>( + &'a self, + _context: &'a CallLifecycleContext, + _response: &'a String, + _timing: &'a CallLifecycleTiming, + ) -> Self::SuccessFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("success"); + }) + } + + fn async_log_failure_event<'a>( + &'a self, + _context: &'a CallLifecycleContext, + _error: &'a CoreError, + _timing: &'a CallLifecycleTiming, + ) -> Self::FailureFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("failure"); + }) + } + } + + #[tokio::test] + async fn lifecycle_runs_hooks_around_provider_call() { + let hooks = RecordingHooks::default(); + let response = CallLifecycle::default() + .run( + CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1"), + "request".to_string(), + &hooks, + |request| async move { + assert_eq!(request, "request:pre:during"); + Ok("response".to_string()) + }, + ) + .await + .expect("call succeeds"); + + assert_eq!(response, "response"); + assert_eq!(hooks.events(), vec!["pre_call", "during_call", "success"]); + } + + #[tokio::test] + async fn lifecycle_logs_failure_when_provider_fails() { + let hooks = RecordingHooks::default(); + let error = CallLifecycle::default() + .run( + CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1"), + "request".to_string(), + &hooks, + |_request| async move { + Err::(CoreError::Network("provider down".to_string())) + }, + ) + .await + .expect_err("call fails"); + + assert_eq!(error, CoreError::Network("provider down".to_string())); + assert_eq!(hooks.events(), vec!["pre_call", "during_call", "failure"]); + } + + #[tokio::test] + async fn lifecycle_can_run_any_request_with_embedded_context() { + let hooks = RecordingHooks::default(); + let response = CallLifecycle::default() + .run_request( + RecordingRequest("request".to_string()), + &hooks, + |request| async move { + assert_eq!(request, "request:pre:during"); + Ok("response".to_string()) + }, + ) + .await + .expect("call succeeds"); + + assert_eq!(response, "response"); + assert_eq!(hooks.events(), vec!["pre_call", "during_call", "success"]); + } +} diff --git a/litellm-rust/crates/core/src/call_lifecycle/types.rs b/litellm-rust/crates/core/src/call_lifecycle/types.rs new file mode 100644 index 00000000000..8819c8830d2 --- /dev/null +++ b/litellm-rust/crates/core/src/call_lifecycle/types.rs @@ -0,0 +1,75 @@ +use std::time::Duration; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CallLifecycleContext { + pub call_type: String, + pub model: String, + pub custom_llm_provider: String, + pub litellm_call_id: String, +} + +impl CallLifecycleContext { + pub fn new( + call_type: impl Into, + model: impl Into, + custom_llm_provider: impl Into, + litellm_call_id: impl Into, + ) -> Self { + Self { + call_type: call_type.into(), + model: model.into(), + custom_llm_provider: custom_llm_provider.into(), + litellm_call_id: litellm_call_id.into(), + } + } +} + +pub trait CallLifecycleRequest { + fn lifecycle_context(&self) -> CallLifecycleContext; +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum CallLifecyclePhase { + PreCall, + DuringCall, + ProviderCall, + SuccessCallback, + FailureCallback, +} + +impl CallLifecyclePhase { + pub fn as_str(self) -> &'static str { + match self { + Self::PreCall => "pre_call", + Self::DuringCall => "during_call", + Self::ProviderCall => "provider_call", + Self::SuccessCallback => "success_callback", + Self::FailureCallback => "failure_callback", + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct CallLifecyclePhaseTiming { + pub phase: CallLifecyclePhase, + pub start_time: f64, + pub end_time: f64, + pub duration: Duration, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct CallLifecycleTiming { + pub start_time: f64, + pub end_time: f64, + pub phases: Vec, +} + +impl CallLifecycleTiming { + pub fn new(start_time: f64, end_time: f64, phases: Vec) -> Self { + Self { + start_time, + end_time, + phases, + } + } +} diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 2ac479cc725..555a04ce853 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -1,7 +1,9 @@ +pub mod call_lifecycle; pub mod error; pub mod ocr; pub mod providers; pub mod realtime; pub mod router; +pub mod routing_utils; pub use error::{CoreError, CoreResult}; diff --git a/litellm-rust/crates/core/src/routing_utils/README.md b/litellm-rust/crates/core/src/routing_utils/README.md new file mode 100644 index 00000000000..8585c18e421 --- /dev/null +++ b/litellm-rust/crates/core/src/routing_utils/README.md @@ -0,0 +1,7 @@ +# Routing Utils + +Shared helpers for deciding how a LiteLLM model routes to an LLM provider. +Keep provider-name parsing, explicit `custom_llm_provider` handling, and model-prefix normalization here. +Do not put deployment selection or load-balancing logic here; that belongs in `router`. +Do not put provider HTTP transformation logic here; that belongs in `providers`. +Helpers in this folder should be deterministic and easy to unit test without network calls. diff --git a/litellm-rust/crates/core/src/routing_utils/mod.rs b/litellm-rust/crates/core/src/routing_utils/mod.rs new file mode 100644 index 00000000000..8336397f870 --- /dev/null +++ b/litellm-rust/crates/core/src/routing_utils/mod.rs @@ -0,0 +1 @@ +pub mod provider; diff --git a/litellm-rust/crates/core/src/routing_utils/provider.rs b/litellm-rust/crates/core/src/routing_utils/provider.rs new file mode 100644 index 00000000000..6333eedebfc --- /dev/null +++ b/litellm-rust/crates/core/src/routing_utils/provider.rs @@ -0,0 +1,77 @@ +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct CustomLlmProvider<'a> { + pub model: &'a str, + pub custom_llm_provider: &'a str, +} + +pub fn get_custom_llm_provider<'a>( + model: &'a str, + custom_llm_provider: Option<&'a str>, +) -> Option> { + if let Some(custom_llm_provider) = custom_llm_provider.filter(|provider| !provider.is_empty()) { + return Some(CustomLlmProvider { + model: strip_custom_llm_provider_prefix(model, custom_llm_provider), + custom_llm_provider, + }); + } + + let (custom_llm_provider, model) = model.split_once('/')?; + if custom_llm_provider.is_empty() || model.is_empty() { + return None; + } + Some(CustomLlmProvider { + model, + custom_llm_provider, + }) +} + +fn strip_custom_llm_provider_prefix<'a>(model: &'a str, custom_llm_provider: &str) -> &'a str { + model + .strip_prefix(custom_llm_provider) + .and_then(|model| model.strip_prefix('/')) + .unwrap_or(model) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn gets_custom_llm_provider_from_model_prefix() { + assert_eq!( + get_custom_llm_provider("mistral/mistral-ocr-latest", None), + Some(CustomLlmProvider { + model: "mistral-ocr-latest", + custom_llm_provider: "mistral", + }) + ); + assert_eq!( + get_custom_llm_provider("azure_ai/doc-intelligence/prebuilt-layout", None), + Some(CustomLlmProvider { + model: "doc-intelligence/prebuilt-layout", + custom_llm_provider: "azure_ai", + }) + ); + assert_eq!(get_custom_llm_provider("mistral-ocr-latest", None), None); + assert_eq!(get_custom_llm_provider("/model", None), None); + assert_eq!(get_custom_llm_provider("provider/", None), None); + } + + #[test] + fn explicit_custom_llm_provider_strips_matching_model_prefix() { + assert_eq!( + get_custom_llm_provider("mistral/mistral-ocr-latest", Some("mistral")), + Some(CustomLlmProvider { + model: "mistral-ocr-latest", + custom_llm_provider: "mistral", + }) + ); + assert_eq!( + get_custom_llm_provider("mistral/mistral-ocr-latest", Some("vertex_ai")), + Some(CustomLlmProvider { + model: "mistral/mistral-ocr-latest", + custom_llm_provider: "vertex_ai", + }) + ); + } +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 82024e1bf47..946a99f990c 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -96,7 +96,6 @@ fn ocr( optional_params: Option>, timeout_seconds: Option, ) -> PyResult> { - let custom_llm_provider = custom_llm_provider.unwrap_or_else(|| "mistral".to_string()); let (document, extra_headers, optional_params, timeout) = marshal_inputs( py, document, @@ -111,10 +110,14 @@ fn ocr( document, api_key: api_key.as_deref(), api_base: api_base.as_deref(), - custom_llm_provider: &custom_llm_provider, + custom_llm_provider: custom_llm_provider.as_deref(), extra_headers, optional_params, timeout, + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: Default::default(), + litellm_call_id: None, })) }); @@ -138,7 +141,6 @@ fn aocr( optional_params: Option>, timeout_seconds: Option, ) -> PyResult> { - let custom_llm_provider = custom_llm_provider.unwrap_or_else(|| "mistral".to_string()); let (document, extra_headers, optional_params, timeout) = marshal_inputs( py, document, @@ -153,10 +155,14 @@ fn aocr( document, api_key: api_key.as_deref(), api_base: api_base.as_deref(), - custom_llm_provider: &custom_llm_provider, + custom_llm_provider: custom_llm_provider.as_deref(), extra_headers, optional_params, timeout, + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: Default::default(), + litellm_call_id: None, }) .await .map_err(core_error_to_pyerr)?; diff --git a/litellm/__init__.py b/litellm/__init__.py index 9650dc12c97..5ae5942f32f 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1394,7 +1394,7 @@ from .skills.main import ( ) from .containers.main import * from .ocr.main import * -from .ocr.rust_bridge import use_litellm_rust +from .rust_bridge.ocr import use_litellm_rust from .rag.main import * from .sandbox.main import * from .search.main import * diff --git a/litellm/caching/_embedding_router.py b/litellm/caching/_embedding_router.py new file mode 100644 index 00000000000..1ec898012e9 --- /dev/null +++ b/litellm/caching/_embedding_router.py @@ -0,0 +1,45 @@ +"""Shared selection of the embedding path for semantic caches. + +Both the Redis and qdrant semantic caches need the same decision: when the +configured embedding model is a proxy Router deployment, embeddings must run +through the Router so per-deployment auth (e.g. Bedrock aws_role_name) is +applied. Otherwise fall back to a direct litellm embedding call. + +This module is dependency-injected: callers pass the proxy ``llm_router`` and +``llm_model_list`` in, so the decision logic is unit-testable without importing +``litellm.proxy.proxy_server``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from litellm.router import Router + + +def resolve_embedding_router( + embedding_model: str, + llm_router: Router | None, + llm_model_list: list[dict[str, Any]] | None, +) -> Router | None: + """Return ``llm_router`` iff it serves ``embedding_model`` as a deployment.""" + if llm_router is None: + return None + router_model_names: list[str] = ( + [m["model_name"] for m in llm_model_list if "model_name" in m] + if llm_model_list is not None + else [] + ) + if embedding_model in router_model_names: + return llm_router + return None + + +def build_router_embedding_metadata( + request_metadata: dict[str, Any] | None, +) -> dict[str, Any]: + """Forward the caller's full metadata, flagged as a semantic-cache embedding.""" + metadata: dict[str, Any] = dict(request_metadata or {}) + metadata["semantic-cache-embedding"] = True + return metadata diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index cb122e90102..5f2269d1945 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -574,8 +574,9 @@ class Cache: if prompt_kwarg in kwargs: cache_lookup_kwargs[prompt_kwarg] = kwargs[prompt_kwarg] - if isinstance(kwargs.get("metadata"), dict): - cache_lookup_kwargs["metadata"] = {} + metadata = kwargs.get("metadata") + if isinstance(metadata, dict): + cache_lookup_kwargs["metadata"] = dict(metadata) return cache_lookup_kwargs diff --git a/litellm/caching/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py index 68d3b8c20b3..504ef8a54eb 100644 --- a/litellm/caching/qdrant_semantic_cache.py +++ b/litellm/caching/qdrant_semantic_cache.py @@ -22,6 +22,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.types.utils import EmbeddingResponse +from ._embedding_router import build_router_embedding_metadata, resolve_embedding_router from .base_cache import BaseCache @@ -219,37 +220,50 @@ class QdrantSemanticCache(BaseCache): cached_key = payload.get(self.CACHE_KEY_FIELD_NAME) return cached_key is not None and str(cached_key) == str(key) - async def _get_async_embedding(self, prompt: str, **kwargs) -> Any: - llm_model_list = None - llm_router = None - + def _get_embedding( + self, prompt: str, metadata: Dict[str, Any] | None = None + ) -> EmbeddingResponse: + """Embed via the proxy Router when it serves the model, else direct.""" try: - from litellm.proxy.proxy_server import ( - llm_model_list as proxy_llm_model_list, - llm_router as proxy_llm_router, - ) - - llm_model_list = proxy_llm_model_list - llm_router = proxy_llm_router + from litellm.proxy.proxy_server import llm_model_list, llm_router except ImportError: - pass + llm_model_list = None + llm_router = None - router_model_names = ( - [m["model_name"] for m in llm_model_list] - if llm_model_list is not None - else [] + router = resolve_embedding_router( + self.embedding_model, llm_router, llm_model_list ) - if llm_router is not None and self.embedding_model in router_model_names: - user_api_key = kwargs.get("metadata", {}).get("user_api_key", "") - return await llm_router.aembedding( + if router is not None: + return router.embedding( model=self.embedding_model, input=prompt, cache={"no-store": True, "no-cache": True}, - metadata={ - "user_api_key": user_api_key, - "semantic-cache-embedding": True, - "trace_id": kwargs.get("metadata", {}).get("trace_id", None), - }, + metadata=build_router_embedding_metadata(metadata), + ) + return litellm.embedding( + model=self.embedding_model, + input=prompt, + cache={"no-store": True, "no-cache": True}, + ) + + async def _get_async_embedding( + self, prompt: str, metadata: Dict[str, Any] | None = None + ) -> EmbeddingResponse: + try: + from litellm.proxy.proxy_server import llm_model_list, llm_router + except ImportError: + llm_model_list = None + llm_router = None + + router = resolve_embedding_router( + self.embedding_model, llm_router, llm_model_list + ) + if router is not None: + return await router.aembedding( + model=self.embedding_model, + input=prompt, + cache={"no-store": True, "no-cache": True}, + metadata=build_router_embedding_metadata(metadata), ) return await litellm.aembedding( @@ -269,11 +283,7 @@ class QdrantSemanticCache(BaseCache): # create an embedding for prompt embedding_response = cast( EmbeddingResponse, - litellm.embedding( - model=self.embedding_model, - input=prompt, - cache={"no-store": True, "no-cache": True}, - ), + self._get_embedding(prompt, metadata=kwargs.get("metadata")), ) # get the embedding @@ -312,11 +322,7 @@ class QdrantSemanticCache(BaseCache): # convert to embedding embedding_response = cast( EmbeddingResponse, - litellm.embedding( - model=self.embedding_model, - input=prompt, - cache={"no-store": True, "no-cache": True}, - ), + self._get_embedding(prompt, metadata=kwargs.get("metadata")), ) # get the embedding @@ -388,7 +394,9 @@ class QdrantSemanticCache(BaseCache): # get the prompt messages = kwargs["messages"] prompt = get_str_from_messages(messages) - embedding_response = await self._get_async_embedding(prompt, **kwargs) + embedding_response = await self._get_async_embedding( + prompt, metadata=kwargs.get("metadata") + ) # get the embedding embedding = embedding_response["data"][0]["embedding"] @@ -424,7 +432,9 @@ class QdrantSemanticCache(BaseCache): messages = kwargs["messages"] prompt = get_str_from_messages(messages) - embedding_response = await self._get_async_embedding(prompt, **kwargs) + embedding_response = await self._get_async_embedding( + prompt, metadata=kwargs.get("metadata") + ) # get the embedding embedding = embedding_response["data"][0]["embedding"] diff --git a/litellm/caching/redis_semantic_cache.py b/litellm/caching/redis_semantic_cache.py index cce4b75795f..e79392bb7f0 100644 --- a/litellm/caching/redis_semantic_cache.py +++ b/litellm/caching/redis_semantic_cache.py @@ -16,12 +16,13 @@ import os from typing import Any, Dict, List, Optional, Tuple, cast import litellm -from litellm._logging import print_verbose +from litellm._logging import print_verbose, verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) from litellm.types.utils import EmbeddingResponse +from ._embedding_router import build_router_embedding_metadata, resolve_embedding_router from .base_cache import BaseCache @@ -67,9 +68,6 @@ class RedisSemanticCache(BaseCache): Exception: If similarity_threshold is not provided or required Redis connection information is missing """ - from redisvl.extensions.llmcache import SemanticCache # type: ignore[import-not-found, import-untyped] - from redisvl.utils.vectorize import CustomTextVectorizer # type: ignore[import-not-found, import-untyped] - if index_name is None: index_name = self.DEFAULT_REDIS_INDEX_NAME @@ -107,15 +105,42 @@ class RedisSemanticCache(BaseCache): print_verbose(f"Redis semantic-cache redis_url: {redis_url}") - # Initialize the Redis vectorizer and cache - cache_vectorizer = CustomTextVectorizer(self._get_embedding) + # Defer redisvl index construction until first use. redisvl's + # CustomTextVectorizer eagerly embeds a probe string at construction; + # building lazily ensures that probe runs after llm_router is wired so + # per-deployment auth (e.g. Bedrock aws_role_name) is applied. + self._index_name = index_name + self._redis_url = redis_url + self._llmcache = None - self.llmcache = self._init_semantic_cache( - semantic_cache_cls=SemanticCache, - index_name=index_name, - redis_url=redis_url, - cache_vectorizer=cache_vectorizer, - ) + @property + def llmcache(self) -> object: + if getattr(self, "_llmcache", None) is None: + self._llmcache = self._build_llmcache() + return self._llmcache + + @llmcache.setter + def llmcache(self, value: object) -> None: + self._llmcache = value + + def _build_llmcache(self) -> object: + # CustomTextVectorizer probes its embedding dimension at construction by + # embedding "dimension test", so the first cache request issues one extra + # billable embedding on top of the request's own. + from redisvl.extensions.llmcache import SemanticCache # type: ignore[import-not-found, import-untyped] + from redisvl.utils.vectorize import CustomTextVectorizer # type: ignore[import-not-found, import-untyped] + + try: + cache_vectorizer = CustomTextVectorizer(self._get_embedding) + return self._init_semantic_cache( + semantic_cache_cls=SemanticCache, + index_name=self._index_name, + redis_url=self._redis_url, + cache_vectorizer=cache_vectorizer, + ) + except Exception as e: + verbose_logger.error(f"Redis semantic-cache index build failed: {e}") + raise @classmethod def _cache_key_filterable_field(cls) -> Dict[str, str]: @@ -285,27 +310,43 @@ class RedisSemanticCache(BaseCache): return dict_method() return value - def _get_embedding(self, prompt: str) -> List[float]: + def _get_embedding( + self, prompt: str, metadata: Dict[str, Any] | None = None + ) -> List[float]: """ - Generate an embedding vector for the given prompt using the configured embedding model. - - Args: - prompt: The text to generate an embedding for - - Returns: - List[float]: The embedding vector + Routes through the proxy Router when the embedding model is a Router + deployment so per-deployment auth (e.g. Bedrock aws_role_name) applies, + mirroring ``_get_async_embedding``; otherwise embeds directly. """ - # Create an embedding from prompt - embedding_response = cast( - EmbeddingResponse, - litellm.embedding( - model=self.embedding_model, - input=prompt, - cache={"no-store": True, "no-cache": True}, - ), + try: + from litellm.proxy.proxy_server import llm_model_list, llm_router + except ImportError: + llm_model_list = None + llm_router = None + + router = resolve_embedding_router( + self.embedding_model, llm_router, llm_model_list ) - embedding = embedding_response["data"][0]["embedding"] - return embedding + if router is not None: + embedding_response = cast( + EmbeddingResponse, + router.embedding( + model=self.embedding_model, + input=prompt, + cache={"no-store": True, "no-cache": True}, + metadata=build_router_embedding_metadata(metadata), + ), + ) + else: + embedding_response = cast( + EmbeddingResponse, + litellm.embedding( + model=self.embedding_model, + input=prompt, + cache={"no-store": True, "no-cache": True}, + ), + ) + return embedding_response["data"][0]["embedding"] def _get_cache_logic(self, cached_response: Any) -> Any: """ @@ -357,7 +398,12 @@ class RedisSemanticCache(BaseCache): value_str = str(value) - store_kwargs: Dict[str, Any] = { + prompt_embedding = self._get_embedding( + prompt, metadata=kwargs.get("metadata") + ) + + store_kwargs: dict[str, Any] = { + "vector": prompt_embedding, "filters": self._get_cache_filters(key), } @@ -393,8 +439,12 @@ class RedisSemanticCache(BaseCache): # Check the cache for semantically similar prompts in this exact # LiteLLM cache-key scope. - check_kwargs: Dict[str, Any] = { + prompt_embedding = self._get_embedding( + prompt, metadata=kwargs.get("metadata") + ) + check_kwargs: dict[str, Any] = { "prompt": prompt, + "vector": prompt_embedding, "filter_expression": self._get_cache_key_filter_expression(key), } results = self.llmcache.check(**check_kwargs) @@ -435,49 +485,42 @@ class RedisSemanticCache(BaseCache): print_verbose(f"Error retrieving from Redis semantic cache: {str(e)}") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 - async def _get_async_embedding(self, prompt: str, **kwargs) -> List[float]: + async def _get_async_embedding( + self, prompt: str, metadata: Dict[str, Any] | None = None + ) -> List[float]: """ Asynchronously generate an embedding for the given prompt. Args: prompt: The text to generate an embedding for - **kwargs: Additional arguments that may contain metadata + metadata: Request metadata forwarded to the Router embedding call Returns: List[float]: The embedding vector """ - from litellm.proxy.proxy_server import llm_model_list, llm_router - - # Route the embedding request through the proxy if appropriate - router_model_names = ( - [m["model_name"] for m in llm_model_list] - if llm_model_list is not None - else [] - ) - try: - if llm_router is not None and self.embedding_model in router_model_names: - # Use the router for embedding generation - user_api_key = kwargs.get("metadata", {}).get("user_api_key", "") - embedding_response = await llm_router.aembedding( + from litellm.proxy.proxy_server import llm_model_list, llm_router + except ImportError: + llm_model_list = None + llm_router = None + + router = resolve_embedding_router( + self.embedding_model, llm_router, llm_model_list + ) + try: + if router is not None: + embedding_response = await router.aembedding( model=self.embedding_model, input=prompt, cache={"no-store": True, "no-cache": True}, - metadata={ - "user_api_key": user_api_key, - "semantic-cache-embedding": True, - "trace_id": kwargs.get("metadata", {}).get("trace_id", None), - }, + metadata=build_router_embedding_metadata(metadata), ) else: - # Generate embedding directly embedding_response = await litellm.aembedding( model=self.embedding_model, input=prompt, cache={"no-store": True, "no-cache": True}, ) - - # Extract and return the embedding vector return embedding_response["data"][0]["embedding"] except Exception as e: print_verbose(f"Error generating async embedding: {str(e)}") @@ -504,9 +547,11 @@ class RedisSemanticCache(BaseCache): value_str = str(value) # Generate embedding for the value (response) to cache - prompt_embedding = await self._get_async_embedding(prompt, **kwargs) + prompt_embedding = await self._get_async_embedding( + prompt, metadata=kwargs.get("metadata") + ) - store_kwargs: Dict[str, Any] = { + store_kwargs: dict[str, Any] = { "vector": prompt_embedding, "filters": self._get_cache_filters(key), } @@ -544,11 +589,13 @@ class RedisSemanticCache(BaseCache): return None # Generate embedding for the prompt - prompt_embedding = await self._get_async_embedding(prompt, **kwargs) + prompt_embedding = await self._get_async_embedding( + prompt, metadata=kwargs.get("metadata") + ) # Check the cache for semantically similar prompts in this exact # LiteLLM cache-key scope. - check_kwargs: Dict[str, Any] = { + check_kwargs: dict[str, Any] = { "prompt": prompt, "vector": prompt_embedding, "filter_expression": self._get_cache_key_filter_expression(key), diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py index bdbb483dcf6..d35601c7e13 100644 --- a/litellm/google_genai/main.py +++ b/litellm/google_genai/main.py @@ -49,6 +49,7 @@ class GenerateContentSetupResult(BaseModel): custom_llm_provider: str generate_content_provider_config: Optional[BaseGoogleGenAIGenerateContentConfig] generate_content_config_dict: Dict[str, Any] + native_request_fields: dict[str, object] litellm_params: GenericLiteLLMParams litellm_logging_obj: LiteLLMLoggingObj litellm_call_id: Optional[str] @@ -152,6 +153,7 @@ class GenerateContentHelper: request_body={}, # Will be handled by adapter generate_content_provider_config=None, # type: ignore generate_content_config_dict=dict(config or {}), + native_request_fields={}, litellm_params=litellm_params, litellm_logging_obj=litellm_logging_obj, litellm_call_id=litellm_call_id, @@ -171,6 +173,12 @@ class GenerateContentHelper: system_instruction = kwargs.get("systemInstruction") or kwargs.get( "system_instruction" ) + # Native top-level REST fields arrive as loose kwargs and are otherwise dropped. + native_request_fields: dict[str, object] = { + field: kwargs[field] + for field in generate_content_provider_config.get_generate_content_request_top_level_fields() + if field in kwargs + } request_body = ( generate_content_provider_config.transform_generate_content_request( model=model, @@ -201,12 +209,29 @@ class GenerateContentHelper: request_body=request_body, generate_content_provider_config=generate_content_provider_config, generate_content_config_dict=generate_content_config_dict, + native_request_fields=native_request_fields, litellm_params=litellm_params, litellm_logging_obj=litellm_logging_obj, litellm_call_id=litellm_call_id, ) +def _merge_native_request_fields( + native_request_fields: dict[str, object], + extra_body: dict[str, object] | None, +) -> dict[str, object] | None: + """ + Merge native top-level request fields into ``extra_body`` so the HTTP handler + forwards them verbatim onto the outgoing request body. An explicit ``extra_body`` + value wins on conflict. Returns ``None`` only when there is genuinely nothing to + forward (no native fields and no caller-supplied ``extra_body``), preserving the + prior behavior without discarding an explicit ``extra_body={}``. + """ + if not native_request_fields and extra_body is None: + return None + return {**native_request_fields, **(extra_body or {})} + + @client async def agenerate_content( model: str, @@ -350,7 +375,9 @@ def generate_content( litellm_params=setup_result.litellm_params, logging_obj=setup_result.litellm_logging_obj, extra_headers=extra_headers, - extra_body=extra_body, + extra_body=_merge_native_request_fields( + setup_result.native_request_fields, extra_body + ), timeout=timeout or request_timeout, _is_async=_is_async, client=kwargs.get("client"), @@ -447,7 +474,9 @@ async def agenerate_content_stream( litellm_params=setup_result.litellm_params, logging_obj=setup_result.litellm_logging_obj, extra_headers=extra_headers, - extra_body=extra_body, + extra_body=_merge_native_request_fields( + setup_result.native_request_fields, extra_body + ), timeout=timeout or request_timeout, _is_async=True, client=kwargs.get("client"), @@ -503,6 +532,11 @@ def generate_content_stream( **kwargs, ) + # Extract systemInstruction from kwargs to pass to handler + system_instruction = kwargs.get("systemInstruction") or kwargs.get( + "system_instruction" + ) + # Check if we should use the adapter (when provider config is None) if setup_result.generate_content_provider_config is None: if "stream" in kwargs: @@ -531,12 +565,15 @@ def generate_content_stream( litellm_params=setup_result.litellm_params, logging_obj=setup_result.litellm_logging_obj, extra_headers=extra_headers, - extra_body=extra_body, + extra_body=_merge_native_request_fields( + setup_result.native_request_fields, extra_body + ), timeout=timeout or request_timeout, _is_async=_is_async, client=kwargs.get("client"), stream=True, litellm_metadata=kwargs.get("litellm_metadata", {}), + system_instruction=system_instruction, ) except Exception as e: diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 19b711f6748..d2c97a5d6e7 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -485,6 +485,7 @@ class PromptTokensDetailsResult(TypedDict): character_count: int image_count: int video_length_seconds: float + audio_length_seconds: float def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: @@ -535,6 +536,13 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: ) or 0.0 ) + audio_length_seconds = ( + cast( + Optional[float], + getattr(usage.prompt_tokens_details, "audio_length_seconds", 0), + ) + or 0.0 + ) return PromptTokensDetailsResult( cache_hit_tokens=cache_hit_tokens, @@ -546,6 +554,7 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: character_count=character_count, image_count=image_count, video_length_seconds=float(video_length_seconds), + audio_length_seconds=float(audio_length_seconds), ) @@ -667,6 +676,14 @@ def _calculate_input_cost( prompt_tokens_details["video_length_seconds"], ) + ### AUDIO LENGTH COST + if prompt_tokens_details["audio_length_seconds"]: + prompt_cost += calculate_cost_component( + model_info, + "input_cost_per_audio_per_second", + prompt_tokens_details["audio_length_seconds"], + ) + return prompt_cost @@ -743,6 +760,7 @@ def generic_cost_per_token( character_count=0, image_count=0, video_length_seconds=0.0, + audio_length_seconds=0.0, ) if usage.prompt_tokens_details: prompt_tokens_details = _parse_prompt_tokens_details(usage) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index bf9ce3b0acb..4bdda6de2c8 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -466,7 +466,7 @@ def get_format_from_file_id(file_id: Optional[str]) -> Optional[str]: def update_messages_with_model_file_ids( messages: List[AllMessageValues], - model_id: str, + model_id: str | None, model_file_id_mapping: Dict[str, Dict[str, str]], ) -> List[AllMessageValues]: """ @@ -519,7 +519,7 @@ def update_messages_with_model_file_ids( if file_id: provider_file_id = ( model_file_id_mapping.get(file_id, {}).get(model_id) - if model_file_id_mapping + if model_file_id_mapping and model_id is not None else None ) if ( diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 66abb824610..1b1b652eaa0 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -585,6 +585,17 @@ class ChunkProcessor: # # Update usage information if needed prompt_tokens = 0 completion_tokens = 0 + # Anthropic's `message_start` SSE event carries usage.output_tokens=1 as a + # cursor/placeholder; the real value only arrives in `message_delta`. + # If a stream is cancelled before `message_delta` lands, the last-wins + # accumulator below leaves completion_tokens stuck at 1 — which then + # bypasses the `completion_tokens or token_counter(...)` fallback in + # calculate_usage() because 1 is truthy. Count the completion-bearing + # usage events so `_reset_anthropic_cursor_completion_tokens` can tell a + # legitimate single-token reply (Anthropic emits 1 in BOTH message_start + # AND message_delta, so >=2 events is positive evidence message_delta + # arrived) from a stale lone cursor. + completion_usage_updates = 0 ## anthropic prompt caching information ## cache_creation_input_tokens: Optional[int] = None cache_read_input_tokens: Optional[int] = None @@ -617,6 +628,7 @@ class ChunkProcessor: and usage_chunk_dict["completion_tokens"] > 0 ): completion_tokens = usage_chunk_dict["completion_tokens"] + completion_usage_updates += 1 if usage_chunk_dict["cache_creation_input_tokens"] is not None and ( usage_chunk_dict["cache_creation_input_tokens"] > 0 or cache_creation_input_tokens is None @@ -667,6 +679,12 @@ class ChunkProcessor: prompt_tokens_details = usage_chunk_dict["prompt_tokens_details"] + completion_tokens = self._reset_anthropic_cursor_completion_tokens( + chunks=chunks, + completion_tokens=completion_tokens, + completion_usage_updates=completion_usage_updates, + ) + return UsagePerChunk( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, @@ -678,6 +696,47 @@ class ChunkProcessor: prompt_tokens_details=prompt_tokens_details, ) + @staticmethod + def _reset_anthropic_cursor_completion_tokens( + chunks: list[dict[str, Any] | ModelResponse], + completion_tokens: int, + completion_usage_updates: int, + ) -> int: + """Reset a stale Anthropic ``message_start`` cursor placeholder to 0. + + See the ``completion_usage_updates`` comment in + ``_calculate_usage_per_chunk``. The accumulated value is NOT a stale + cursor when either it is > 1 (definitely not a placeholder) or we saw + >= 2 completion-bearing usage events (positive evidence ``message_delta`` + arrived). Otherwise — the only completion update we ever saw was the + Anthropic ``message_start`` cursor (=1) — reset to 0 so + ``calculate_usage()``'s ``or token_counter(text=...)`` fallback estimates + from the actually-received completion text instead of trusting the + placeholder. Gated on ``custom_llm_provider == "anthropic"`` so the + heuristic (which encodes Anthropic's specific message_start SSE shape) + does not silently affect other providers that may legitimately report + ``completion_tokens=1`` from a single usage event. + """ + saw_non_cursor_completion = ( + completion_tokens > 1 or completion_usage_updates >= 2 + ) + if saw_non_cursor_completion: + return completion_tokens + + custom_llm_provider: Optional[str] = None + if chunks: + first_chunk = chunks[0] + if isinstance(first_chunk, dict): + hp = first_chunk.get("_hidden_params") + else: + hp = getattr(first_chunk, "_hidden_params", None) + if isinstance(hp, dict): + custom_llm_provider = hp.get("custom_llm_provider") + + if custom_llm_provider == "anthropic" and completion_tokens == 1: + return 0 + return completion_tokens + def calculate_usage( self, chunks: List[Union[Dict[str, Any], ModelResponse]], diff --git a/litellm/llms/anthropic/batches/transformation.py b/litellm/llms/anthropic/batches/transformation.py index fd67a7fbaf1..7c4e9386d5f 100644 --- a/litellm/llms/anthropic/batches/transformation.py +++ b/litellm/llms/anthropic/batches/transformation.py @@ -43,7 +43,9 @@ class AnthropicBatchesConfig(BaseBatchesConfig): api_base: Optional[str] = None, ) -> dict: """Validate and prepare environment-specific headers and parameters.""" - auth_header = self.anthropic_model_info.get_auth_header(api_key) + if api_base is None and isinstance(litellm_params, dict): + api_base = litellm_params.get("api_base") + auth_header = self.anthropic_model_info.get_auth_header(api_key, api_base) if auth_header is None: raise ValueError( "Missing Anthropic API Key - A call is being made to anthropic but no key is set either in the environment variables or via params" diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index e7e0f12c455..b1212f93059 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -364,6 +364,7 @@ class AnthropicChatCompletion(BaseLLM): messages=messages, optional_params={**optional_params, "is_vertex_request": is_vertex_request}, litellm_params=litellm_params, + api_base=api_base, ) config = ProviderConfigManager.get_provider_chat_config( diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index eaa762204f2..d27187cf2e7 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -548,6 +548,19 @@ class AnthropicModelInfo(BaseLLMModelInfo): return list(set(betas)) + @staticmethod + def _make_api_key_auth_header( + api_key: str, api_base: str | None, use_bearer_for_custom_base: bool = False + ) -> dict: + if use_bearer_for_custom_base and ( + api_base + and "api.anthropic.com" not in api_base + and not api_key.startswith("sk-ant-") + ): + value = api_key if api_key.startswith("Bearer ") else f"Bearer {api_key}" + return {"authorization": value} + return {"x-api-key": api_key} + def get_anthropic_headers( self, api_key: Optional[str] = None, @@ -567,6 +580,8 @@ class AnthropicModelInfo(BaseLLMModelInfo): user_anthropic_beta_headers: Optional[List[str]] = None, code_execution_tool_used: bool = False, container_with_skills_used: bool = False, + api_base: str | None = None, + use_bearer_for_custom_base: bool = False, ) -> dict: betas = set() # Anthropic no longer requires the prompt-caching beta header @@ -615,7 +630,11 @@ class AnthropicModelInfo(BaseLLMModelInfo): elif auth_token and not api_key: headers["authorization"] = f"Bearer {auth_token}" elif api_key: - headers["x-api-key"] = api_key + headers.update( + self._make_api_key_auth_header( + api_key, api_base, use_bearer_for_custom_base + ) + ) if user_anthropic_beta_headers is not None: betas.update(user_anthropic_beta_headers) @@ -644,6 +663,12 @@ class AnthropicModelInfo(BaseLLMModelInfo): api_key: Optional[str] = None, api_base: Optional[str] = None, ) -> Dict: + if api_base is None and isinstance(litellm_params, dict): + api_base = litellm_params.get("api_base") + use_bearer_for_custom_base: bool = bool( + isinstance(litellm_params, dict) + and litellm_params.get("use_bearer_for_custom_base", False) + ) # Check for Anthropic OAuth token in headers headers, api_key = optionally_handle_anthropic_oauth( headers=headers, api_key=api_key @@ -699,6 +724,8 @@ class AnthropicModelInfo(BaseLLMModelInfo): effort_used=effort_used, code_execution_tool_used=code_execution_tool_used, container_with_skills_used=container_with_skills_used, + api_base=api_base, + use_bearer_for_custom_base=use_bearer_for_custom_base, ) headers = {**headers, **anthropic_headers} @@ -734,18 +761,24 @@ class AnthropicModelInfo(BaseLLMModelInfo): return auth_token or get_secret_str("ANTHROPIC_AUTH_TOKEN") @staticmethod - def get_auth_header(api_key: Optional[str] = None) -> Optional[dict]: + def get_auth_header( + api_key: str | None = None, + api_base: str | None = None, + use_bearer_for_custom_base: bool = False, + ) -> dict | None: """Resolve Anthropic credentials and return the appropriate auth header dict. - Checks ANTHROPIC_API_KEY first (-> x-api-key), then - ANTHROPIC_AUTH_TOKEN (-> Authorization: Bearer). + Checks ANTHROPIC_API_KEY first (-> x-api-key or Bearer depending on + use_bearer_for_custom_base), then ANTHROPIC_AUTH_TOKEN (-> Authorization: Bearer). Returns None if neither is available. """ resolved_key = AnthropicModelInfo.get_api_key(api_key) if resolved_key is not None: if is_anthropic_oauth_key(resolved_key): return {"authorization": f"Bearer {resolved_key}"} - return {"x-api-key": resolved_key} + return AnthropicModelInfo._make_api_key_auth_header( + resolved_key, api_base, use_bearer_for_custom_base + ) auth_token = AnthropicModelInfo.get_auth_token() if auth_token is not None: return {"authorization": f"Bearer {auth_token}"} @@ -759,7 +792,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): self, api_key: Optional[str] = None, api_base: Optional[str] = None ) -> List[str]: api_base = AnthropicModelInfo.get_api_base(api_base) - auth_header = AnthropicModelInfo.get_auth_header(api_key) + auth_header = AnthropicModelInfo.get_auth_header(api_key, api_base) if api_base is None or auth_header is None: raise ValueError( "ANTHROPIC_API_BASE/ANTHROPIC_BASE_URL or ANTHROPIC_API_KEY/ANTHROPIC_AUTH_TOKEN is not set. Please set the environment variable, to query Anthropic's `/models` endpoint." diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index 9860997b1cd..a92834b5d71 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -205,32 +205,11 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): if "delta" not in merged_chunk: merged_chunk["delta"] = {} - uncached_input_tokens = chunk.usage.prompt_tokens or 0 - if ( - hasattr(chunk.usage, "prompt_tokens_details") - and chunk.usage.prompt_tokens_details - ): - cached_tokens = ( - getattr(chunk.usage.prompt_tokens_details, "cached_tokens", 0) or 0 - ) - uncached_input_tokens -= cached_tokens + from .transformation import LiteLLMAnthropicMessagesAdapter - usage_dict: UsageDelta = { - "input_tokens": uncached_input_tokens, - "output_tokens": chunk.usage.completion_tokens or 0, - } - if ( - hasattr(chunk.usage, "_cache_creation_input_tokens") - and chunk.usage._cache_creation_input_tokens > 0 - ): - usage_dict["cache_creation_input_tokens"] = ( - chunk.usage._cache_creation_input_tokens - ) - if ( - hasattr(chunk.usage, "_cache_read_input_tokens") - and chunk.usage._cache_read_input_tokens > 0 - ): - usage_dict["cache_read_input_tokens"] = chunk.usage._cache_read_input_tokens + usage_dict: UsageDelta = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta( + chunk.usage + ) merged_chunk["usage"] = usage_dict if self.applied_edits and "context_management" not in merged_chunk: merged_chunk["context_management"] = ContextManagementResponse( diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 46564a565ce..1fd593f63c8 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -1399,6 +1399,95 @@ class LiteLLMAnthropicMessagesAdapter: return "tool_use" return "end_turn" + @staticmethod + def _positive_int(value: object) -> int: + if isinstance(value, bool): + return 0 + if isinstance(value, int) and value > 0: + return value + if isinstance(value, float) and value.is_integer() and value > 0: + return int(value) + return 0 + + @classmethod + def _first_positive_usage_value( + cls, usage: Usage, field_names: tuple[str, ...] + ) -> int: + for field_name in field_names: + value = cls._positive_int(getattr(usage, field_name, None)) + if value > 0: + return value + return 0 + + @classmethod + def _first_positive_prompt_tokens_detail_value( + cls, usage: Usage, field_names: tuple[str, ...] + ) -> int: + prompt_tokens_details = getattr(usage, "prompt_tokens_details", None) + if prompt_tokens_details is None: + return 0 + + for field_name in field_names: + if isinstance(prompt_tokens_details, dict): + value = cls._positive_int(prompt_tokens_details.get(field_name)) + else: + value = cls._positive_int( + getattr(prompt_tokens_details, field_name, None) + ) + if value > 0: + return value + return 0 + + @classmethod + def _get_cache_read_input_tokens(cls, usage: Usage) -> int: + explicit_value = cls._first_positive_usage_value( + usage, ("cache_read_input_tokens", "_cache_read_input_tokens") + ) + if explicit_value > 0: + return explicit_value + return cls._first_positive_prompt_tokens_detail_value(usage, ("cached_tokens",)) + + @classmethod + def _get_cache_creation_input_tokens(cls, usage: Usage) -> int: + explicit_value = cls._first_positive_usage_value( + usage, ("cache_creation_input_tokens", "_cache_creation_input_tokens") + ) + if explicit_value > 0: + return explicit_value + return cls._first_positive_prompt_tokens_detail_value( + usage, ("cache_creation_tokens", "cache_write_tokens") + ) + + @classmethod + def _translate_openai_usage_to_anthropic_usage_delta( + cls, usage: Usage + ) -> UsageDelta: + cache_read_input_tokens = cls._get_cache_read_input_tokens(usage) + cache_creation_input_tokens = cls._get_cache_creation_input_tokens(usage) + input_tokens = max( + (usage.prompt_tokens or 0) + - cache_read_input_tokens + - cache_creation_input_tokens, + 0, + ) + + usage_delta = UsageDelta( + input_tokens=input_tokens, + output_tokens=usage.completion_tokens or 0, + ) + if cache_creation_input_tokens > 0: + usage_delta["cache_creation_input_tokens"] = cache_creation_input_tokens + if cache_read_input_tokens > 0: + usage_delta["cache_read_input_tokens"] = cache_read_input_tokens + return usage_delta + + @classmethod + def _translate_openai_usage_to_anthropic_usage(cls, usage: Usage) -> AnthropicUsage: + return cast( + AnthropicUsage, + cls._translate_openai_usage_to_anthropic_usage_delta(usage), + ) + def translate_openai_response_to_anthropic( self, response: ModelResponse, @@ -1430,32 +1519,12 @@ class LiteLLMAnthropicMessagesAdapter: ) # extract usage usage: Usage = getattr(response, "usage") - uncached_input_tokens = usage.prompt_tokens or 0 - cached_tokens = 0 - if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details: - cached_tokens = ( - getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0 - ) - uncached_input_tokens -= cached_tokens - - anthropic_usage = AnthropicUsage( - input_tokens=uncached_input_tokens, - output_tokens=usage.completion_tokens or 0, - ) - if ( - hasattr(usage, "_cache_creation_input_tokens") - and usage._cache_creation_input_tokens > 0 - ): - anthropic_usage["cache_creation_input_tokens"] = ( - usage._cache_creation_input_tokens - ) - if cached_tokens > 0: - anthropic_usage["cache_read_input_tokens"] = cached_tokens + anthropic_usage = self._translate_openai_usage_to_anthropic_usage(usage) if polyfill_result is not None and polyfill_result.iterations_usage is not None: message_iteration: UsageIteration = { "type": "message", - "input_tokens": uncached_input_tokens, + "input_tokens": anthropic_usage["input_tokens"], "output_tokens": usage.completion_tokens or 0, } anthropic_usage["iterations"] = list(polyfill_result.iterations_usage) + [ @@ -1644,35 +1713,9 @@ class LiteLLMAnthropicMessagesAdapter: else: litellm_usage_chunk = None if litellm_usage_chunk is not None: - uncached_input_tokens = litellm_usage_chunk.prompt_tokens or 0 - cached_tokens = 0 - if ( - hasattr(litellm_usage_chunk, "prompt_tokens_details") - and litellm_usage_chunk.prompt_tokens_details - ): - cached_tokens = ( - getattr( - litellm_usage_chunk.prompt_tokens_details, - "cached_tokens", - 0, - ) - or 0 - ) - uncached_input_tokens -= cached_tokens - - usage_delta = UsageDelta( - input_tokens=uncached_input_tokens, - output_tokens=litellm_usage_chunk.completion_tokens or 0, + usage_delta = self._translate_openai_usage_to_anthropic_usage_delta( + litellm_usage_chunk ) - if ( - hasattr(litellm_usage_chunk, "_cache_creation_input_tokens") - and litellm_usage_chunk._cache_creation_input_tokens > 0 - ): - usage_delta["cache_creation_input_tokens"] = ( - litellm_usage_chunk._cache_creation_input_tokens - ) - if cached_tokens > 0: - usage_delta["cache_read_input_tokens"] = cached_tokens else: usage_delta = UsageDelta(input_tokens=0, output_tokens=0) message_block = MessageBlockDelta( diff --git a/litellm/llms/anthropic/files/handler.py b/litellm/llms/anthropic/files/handler.py index 56296df94a1..170cb086bb0 100644 --- a/litellm/llms/anthropic/files/handler.py +++ b/litellm/llms/anthropic/files/handler.py @@ -84,7 +84,7 @@ class AnthropicFilesHandler: # Get Anthropic API credentials api_base = self.anthropic_model_info.get_api_base(api_base) - auth_header = self.anthropic_model_info.get_auth_header(api_key) + auth_header = self.anthropic_model_info.get_auth_header(api_key, api_base) if auth_header is None: raise ValueError("Missing Anthropic API Key") diff --git a/litellm/llms/anthropic/files/transformation.py b/litellm/llms/anthropic/files/transformation.py index ea9bf00f505..7ffb6beb4c7 100644 --- a/litellm/llms/anthropic/files/transformation.py +++ b/litellm/llms/anthropic/files/transformation.py @@ -95,7 +95,9 @@ class AnthropicFilesConfig(BaseFilesConfig): api_key: Optional[str] = None, api_base: Optional[str] = None, ) -> dict: - auth_header = AnthropicModelInfo.get_auth_header(api_key) + if api_base is None and isinstance(litellm_params, dict): + api_base = litellm_params.get("api_base") + auth_header = AnthropicModelInfo.get_auth_header(api_key, api_base) if auth_header is None: raise ValueError( "Anthropic API key is required. Set ANTHROPIC_API_KEY or ANTHROPIC_AUTH_TOKEN environment variable or pass api_key parameter." diff --git a/litellm/llms/anthropic/skills/transformation.py b/litellm/llms/anthropic/skills/transformation.py index 4ea768b02af..2bfdf7ef4ba 100644 --- a/litellm/llms/anthropic/skills/transformation.py +++ b/litellm/llms/anthropic/skills/transformation.py @@ -38,10 +38,12 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): # Get API key from litellm_params if available api_key = None + api_base = None if litellm_params is not None: api_key = litellm_params.api_key + api_base = litellm_params.api_base - auth_header = AnthropicModelInfo.get_auth_header(api_key) + auth_header = AnthropicModelInfo.get_auth_header(api_key, api_base) if auth_header is None: raise ValueError( "ANTHROPIC_API_KEY or ANTHROPIC_AUTH_TOKEN is required for Skills API" diff --git a/litellm/llms/base_llm/google_genai/transformation.py b/litellm/llms/base_llm/google_genai/transformation.py index e8b3bf1a576..8fb7eb9fde0 100644 --- a/litellm/llms/base_llm/google_genai/transformation.py +++ b/litellm/llms/base_llm/google_genai/transformation.py @@ -62,6 +62,19 @@ class BaseGoogleGenAIGenerateContentConfig(ABC): "get_supported_generate_content_optional_params is not implemented" ) + def get_generate_content_request_top_level_fields(self) -> tuple[str, ...]: + """ + Native Google ``GenerateContentRequest`` fields that sit at the top level + (siblings of ``generationConfig``) rather than inside it. The proxy forwards + these verbatim from a native request so ``generateContent`` is a drop-in for + Google's REST API. + + Excludes ``contents``, ``model`` and ``tools`` (dedicated params), + ``systemInstruction`` (dedicated extraction) and ``generationConfig`` (mapped + to ``config``). + """ + return ("safetySettings", "toolConfig", "cachedContent", "labels") + @abstractmethod def map_generate_content_optional_params( self, diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 5e97394f459..69c0a8529b4 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -904,9 +904,12 @@ class BedrockModelInfo(BaseLLMModelInfo): "mantle/": "mantle", } - # Check explicit routes first + # Check explicit routes first. Match each prefix only as a leading path + # segment so the `bedrock_mantle/` provider prefix is never mistaken for + # the `mantle/` invoke route (which would mangle + # `bedrock_mantle/openai.gpt-5.5` into `bedrock_openai.gpt-5.5`). for prefix, route_type in route_mappings.items(): - if prefix in model: + if BedrockModelInfo._model_has_route_prefix(model, prefix): return route_type # Check for nova spec prefixes (nova/ and nova-2/) @@ -930,14 +933,14 @@ class BedrockModelInfo(BaseLLMModelInfo): """ Check if the model is an explicit converse route. """ - return "converse/" in model + return BedrockModelInfo._model_has_route_prefix(model, "converse/") @staticmethod def _explicit_claude_platform_route(model: str) -> bool: """ Check if the model is an explicit Claude Platform on AWS route. """ - return "claude_platform/" in model + return BedrockModelInfo._model_has_route_prefix(model, "claude_platform/") @staticmethod def get_claude_platform_model(model: str) -> str: @@ -967,42 +970,58 @@ class BedrockModelInfo(BaseLLMModelInfo): """ Check if the model is an explicit invoke route. """ - return "invoke/" in model + return BedrockModelInfo._model_has_route_prefix(model, "invoke/") @staticmethod def _explicit_agent_route(model: str) -> bool: """ Check if the model is an explicit agent route. """ - return "agent/" in model + return BedrockModelInfo._model_has_route_prefix(model, "agent/") @staticmethod def _explicit_agentcore_route(model: str) -> bool: """ Check if the model is an explicit agentcore route. """ - return "agentcore/" in model + return BedrockModelInfo._model_has_route_prefix(model, "agentcore/") + + @staticmethod + def _model_has_route_prefix(model: str, prefix: str) -> bool: + """Whether a route prefix (e.g. ``mantle/``) appears as a leading path segment. + + A route token is only valid at the start of the model id or immediately + after a ``/``. A plain substring check matches the ``bedrock_mantle/`` + provider prefix against the ``mantle/`` route, so the body model gets + mangled to ``bedrock_openai.gpt-5.5``; anchoring to a segment boundary + keeps the bare model id intact. + + ``f"/{prefix}" in model`` matches the token as a segment at any path + depth, not just the second segment; that is intentional and acceptable + for these short, unambiguous route tokens. + """ + return model.startswith(prefix) or f"/{prefix}" in model @staticmethod def _explicit_mantle_route(model: str) -> bool: """ Check if the model is an explicit mantle route (bedrock-mantle endpoint). """ - return "mantle/" in model + return BedrockModelInfo._model_has_route_prefix(model, "mantle/") @staticmethod def _explicit_converse_like_route(model: str) -> bool: """ Check if the model is an explicit converse like route. """ - return "converse_like/" in model + return BedrockModelInfo._model_has_route_prefix(model, "converse_like/") @staticmethod def _explicit_async_invoke_route(model: str) -> bool: """ Check if the model is an explicit async invoke route. """ - return "async_invoke/" in model + return BedrockModelInfo._model_has_route_prefix(model, "async_invoke/") @staticmethod def _explicit_openai_route(model: str) -> bool: @@ -1010,7 +1029,7 @@ class BedrockModelInfo(BaseLLMModelInfo): Check if the model is an explicit openai route. Used for Bedrock imported models that use OpenAI Chat Completions format. """ - return "openai/" in model + return BedrockModelInfo._model_has_route_prefix(model, "openai/") @staticmethod def get_bedrock_provider_config_for_messages_api( diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 79f006a83de..3bdcbd25949 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -1046,9 +1046,9 @@ def _gemini_convert_messages_with_history( if invocation.get("tool_type"): tr_dict["toolType"] = invocation["tool_type"] tr_part: Dict[str, Any] = {"toolResponse": tr_dict} - if "thought_signature" in invocation: + if "response_thought_signature" in invocation: tr_part["thoughtSignature"] = invocation[ - "thought_signature" + "response_thought_signature" ] assistant_content.append(tr_part) # type: ignore diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 0de34ec257c..32bdb17840d 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -1634,13 +1634,16 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): resp = tool_responses_by_id.pop(call_id, None) if resp is not None: merged["response"] = resp.get("response") - # Keep response signature if call didn't have one - if "thought_signature" not in merged and "thought_signature" in resp: - merged["thought_signature"] = resp["thought_signature"] + if "thought_signature" in resp: + merged["response_thought_signature"] = resp["thought_signature"] invocations.append(merged) # Any orphan responses (shouldn't happen, but be safe) for resp_id, resp_entry in tool_responses_by_id.items(): + if "thought_signature" in resp_entry: + resp_entry["response_thought_signature"] = resp_entry[ + "thought_signature" + ] invocations.append(resp_entry) return invocations if invocations else None @@ -3598,16 +3601,22 @@ class ModelResponseIterator: chunk = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or "" message = chunk.replace("\n\n", "") - # Accumulate JSON data self.accumulated_json += message - # Try to parse the accumulated JSON + # json.loads on the whole buffer after every fragment is O(n^2) and + # holds the GIL, freezing the event loop for seconds on large responses + # (https://github.com/BerriAI/litellm/issues/26181). A complete Gemini + # chunk is a JSON object/array, so only attempt the parse once the + # buffer's last non-whitespace byte can close one. + stripped = self.accumulated_json.rstrip() + if not stripped or stripped[-1] not in "}]": + return None + try: _data = json.loads(self.accumulated_json) self.accumulated_json = "" # reset after successful parsing return self.chunk_parser(chunk=_data) except json.JSONDecodeError: - # If it's not valid JSON yet, continue to the next event return None def _common_chunk_parsing_logic( diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index 165dac24903..53f5fb464df 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py @@ -274,6 +274,7 @@ class GoogleBatchEmbeddings(VertexLLM): model_response=model_response, model=model, response_json=_json_response, + resolved_files=resolved_files, ) else: _predictions = VertexAIBatchEmbeddingsResponseObject(**_json_response) # type: ignore @@ -377,6 +378,7 @@ class GoogleBatchEmbeddings(VertexLLM): model_response=model_response, model=model, response_json=_json_response, + resolved_files=resolved_files, ) else: _predictions = VertexAIBatchEmbeddingsResponseObject(**_json_response) # type: ignore diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index ba6e6f0c056..27ca1bd92a4 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -4,7 +4,10 @@ Transformation logic from OpenAI /v1/embeddings format to Google AI Studio /batc Why separate file? Make it easy to see how transformation works """ -from typing import Dict, List, Optional, Tuple +from collections.abc import Mapping +from typing import Dict, List, Optional, Sequence, Tuple + +from pydantic import TypeAdapter, ValidationError from litellm.types.llms.vertex_ai import ( BlobType, @@ -13,10 +16,17 @@ from litellm.types.llms.vertex_ai import ( FileDataType, GeminiEmbeddingInput, PartType, + PromptTokensDetails, + UsageMetadata, VertexAIBatchEmbeddingsRequestBody, VertexAIBatchEmbeddingsResponseObject, ) -from litellm.types.utils import Embedding, EmbeddingResponse, Usage +from litellm.types.utils import ( + Embedding, + EmbeddingResponse, + PromptTokensDetailsWrapper, + Usage, +) from litellm.utils import get_formatted_prompt, token_counter SUPPORTED_EMBEDDING_MIME_TYPES = { @@ -294,11 +304,133 @@ def transform_openai_input_gemini_embed_content( return request_body +_IMAGE_MIME_TYPES = frozenset({"image/png", "image/jpeg"}) +_VIDEO_TOKENS_PER_SECOND = 258.0 +_AUDIO_TOKENS_PER_SECOND = 32.0 +_usage_metadata_adapter = TypeAdapter(UsageMetadata) + + +def _parse_usage_metadata(raw_usage_metadata: object) -> Optional[UsageMetadata]: + if not isinstance(raw_usage_metadata, dict): + return None + try: + return _usage_metadata_adapter.validate_python(raw_usage_metadata) + except ValidationError: + return None + + +def _flatten_input(input: GeminiEmbeddingInput) -> tuple[str, ...]: + if isinstance(input, str): + return (input,) + return tuple( + sub + for element in input + for sub in (element if isinstance(element, list) else [element]) + ) + + +def _is_image_element( + element: str, + resolved_files: Mapping[str, Mapping[str, str]], +) -> bool: + if element.startswith("data:") and ";base64," in element: + try: + mime_type, _ = _parse_data_url(element) + except ValueError: + return False + return mime_type in _IMAGE_MIME_TYPES + if _is_gcs_url(element): + try: + return _infer_mime_type_from_gcs_url(element) in _IMAGE_MIME_TYPES + except ValueError: + return False + if _is_file_reference(element): + file_info = resolved_files.get(element) + return file_info is not None and file_info.get("mime_type") in _IMAGE_MIME_TYPES + return False + + +def _count_input_images( + input: GeminiEmbeddingInput, + resolved_files: Mapping[str, Mapping[str, str]], +) -> int: + return sum( + 1 + for element in _flatten_input(input) + if _is_image_element(element, resolved_files) + ) + + +def _tokens_for_modality(details: Sequence[PromptTokensDetails], modality: str) -> int: + return sum( + detail["tokenCount"] for detail in details if detail["modality"] == modality + ) + + +def _fallback_usage(input: GeminiEmbeddingInput, model: str) -> Usage: + if _is_multimodal_input(input): + return Usage(prompt_tokens=0, total_tokens=0) + input_text = get_formatted_prompt(data={"input": input}, call_type="embedding") + prompt_tokens = token_counter(model=model, text=input_text) + return Usage(prompt_tokens=prompt_tokens, total_tokens=prompt_tokens) + + +def _usage_from_embed_content_response( + input: GeminiEmbeddingInput, + model: str, + raw_usage_metadata: object, + resolved_files: Mapping[str, Mapping[str, str]], +) -> Usage: + usage_metadata = _parse_usage_metadata(raw_usage_metadata) + if usage_metadata is None: + return _fallback_usage(input, model) + + prompt_tokens = usage_metadata.get("promptTokenCount", 0) + total_tokens = usage_metadata.get("totalTokenCount") or prompt_tokens + + details: Sequence[PromptTokensDetails] = ( + usage_metadata.get("promptTokensDetails") or () + ) + text_tokens = _tokens_for_modality(details, "TEXT") + audio_tokens = _tokens_for_modality(details, "AUDIO") + video_tokens = _tokens_for_modality(details, "VIDEO") + image_count = _count_input_images(input, resolved_files) + + video_length_seconds = ( + video_tokens / _VIDEO_TOKENS_PER_SECOND if video_tokens > 0 else 0.0 + ) + audio_length_seconds = ( + audio_tokens / _AUDIO_TOKENS_PER_SECOND if audio_tokens > 0 else 0.0 + ) + + # generic_cost_per_token rewrites text_tokens to the full prompt minus + # other modalities when both text_tokens and image_count are zero. For + # video, that misallocates video tokens to text; a 1-token floor sidesteps + # the rewrite and keeps billing on input_cost_per_video_per_second. + needs_video_text_floor = ( + video_length_seconds > 0 and text_tokens == 0 and image_count == 0 + ) + resolved_text_tokens = 1 if needs_video_text_floor else text_tokens + + return Usage( + prompt_tokens=prompt_tokens, + total_tokens=total_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=resolved_text_tokens, + audio_tokens=audio_tokens, + image_count=image_count, + video_length_seconds=video_length_seconds, + audio_length_seconds=audio_length_seconds, + ), + ) + + def process_embed_content_response( input: GeminiEmbeddingInput, model_response: EmbeddingResponse, model: str, response_json: dict, + resolved_files: Mapping[str, Mapping[str, str]] | None = None, ) -> EmbeddingResponse: """ Process Gemini embedContent response (single embedding for multimodal input). @@ -308,6 +440,8 @@ def process_embed_content_response( model_response: EmbeddingResponse to populate model: Model name response_json: Raw JSON response from embedContent endpoint + resolved_files: Mapping of file references (files/abc) to {mime_type, uri}, + used to bill resolved image references at the per-image rate Returns: EmbeddingResponse with single embedding @@ -327,14 +461,11 @@ def process_embed_content_response( model_response.data = [openai_embedding] model_response.model = model - - if _is_multimodal_input(input): - prompt_tokens = 0 - else: - input_text = get_formatted_prompt(data={"input": input}, call_type="embedding") - prompt_tokens = token_counter(model=model, text=input_text) - model_response.usage = Usage( - prompt_tokens=prompt_tokens, total_tokens=prompt_tokens + model_response.usage = _usage_from_embed_content_response( + input=input, + model=model, + raw_usage_metadata=response_json.get("usageMetadata"), + resolved_files=resolved_files or {}, ) return model_response diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py index 8a92e7ec4a5..a633ef3298a 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py @@ -35,27 +35,21 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert Validate the environment for the request """ + # Work on a local copy — router shallow-copies litellm_params so the caller's + # headers dict may be the shared deployment extra_headers object. + headers = dict(headers) vertex_ai_project = VertexBase.safe_get_vertex_ai_project(litellm_params) vertex_ai_location = VertexBase.safe_get_vertex_ai_location(litellm_params) - project_id: Optional[str] = None - if "Authorization" not in headers: - vertex_credentials = VertexBase.safe_get_vertex_ai_credentials( - litellm_params - ) + vertex_credentials = VertexBase.safe_get_vertex_ai_credentials(litellm_params) + access_token, project_id = self._ensure_access_token( + credentials=vertex_credentials, + project_id=vertex_ai_project, + custom_llm_provider="vertex_ai", + ) + headers["Authorization"] = f"Bearer {access_token}" - access_token, project_id = self._ensure_access_token( - credentials=vertex_credentials, - project_id=vertex_ai_project, - custom_llm_provider="vertex_ai", - ) - - headers["Authorization"] = f"Bearer {access_token}" - else: - # Authorization already in headers, but we still need project_id - project_id = vertex_ai_project - - # Always calculate api_base if not provided, regardless of Authorization header + # Calculate api_base if not provided if api_base is None: api_base = self.get_complete_vertex_url( custom_api_base=api_base, diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index 960d3483848..78669f1e789 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -194,9 +194,11 @@ class VertexAIPartnerModels(VertexBase): encoding=encoding, ) elif "claude" in model: - if headers is None: - headers = {} - headers.update({"Authorization": "Bearer {}".format(access_token)}) + # Build a new dict so we never mutate the shared deployment extra_headers object. + headers = { + **(headers or {}), + "Authorization": "Bearer {}".format(access_token), + } optional_params.update( { diff --git a/litellm/main.py b/litellm/main.py index 1960c456f80..cdea7dbd684 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -633,6 +633,7 @@ async def acompletion( try: # Use a partial function to pass your keyword arguments + kwargs.pop("acompletion", None) func = partial(completion, **completion_kwargs, **kwargs) # Add the context to the function @@ -5063,6 +5064,12 @@ def completion( # type: ignore ######### unpacking kwargs ##################### args = locals() + # Set by the responses->completion fallback so completion() does not bridge + # back to the Responses API: that round-trip mutually recurses forever for a + # model whose model_cost mode is "responses" but whose provider has no + # Responses API config (get_provider_responses_api_config -> None). + skip_responses_api_bridge = kwargs.pop("_skip_responses_api_bridge", False) + skip_mcp_handler = kwargs.pop("_skip_mcp_handler", False) if not skip_mcp_handler and tools: from litellm.responses.mcp.chat_completions_handler import acompletion_with_mcp @@ -5358,7 +5365,7 @@ def completion( # type: ignore messages = update_messages_with_model_file_ids( messages=messages, - model_id=kwargs.get("model_info", {}).get("id", None), + model_id=(kwargs.get("model_info") or {}).get("id", None), model_file_id_mapping=cast( Dict[str, Dict[str, str]], kwargs.get("model_file_id_mapping") or {}, @@ -5560,7 +5567,10 @@ def completion( # type: ignore # detection when the deployment name differs from the model name. _azure_detection_model = base_model or model - if responses_api_model_info.get("mode") == "responses": + if ( + responses_api_model_info.get("mode") == "responses" + and not skip_responses_api_bridge + ): from litellm.completion_extras import responses_api_bridge optional_params, rs_val = ( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1393c683860..d44fc654a56 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -25793,7 +25793,7 @@ "supports_response_schema": true, "supports_tool_choice": true }, - "mistral/mistral-medium-latest": { + "mistral/mistral-medium-2508": { "input_cost_per_token": 4e-07, "litellm_provider": "mistral", "max_input_tokens": 131072, @@ -25801,12 +25801,45 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 2e-06, + "source": "https://mistral.ai/news/mistral-medium-3", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true }, + "mistral/mistral-medium-2604": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/mistral-medium-latest": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "mistral/mistral-medium-3-1-2508": { "input_cost_per_token": 4e-07, "litellm_provider": "mistral", @@ -25833,6 +25866,7 @@ "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 6a196d41768..93b3a892659 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -19,13 +19,7 @@ from litellm.constants import request_timeout from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler -from litellm.ocr.rust_bridge import ( - RustAocr, - RustOcr, - load_rust_aocr, - load_rust_ocr, - rust_ocr_enabled, -) +from litellm.rust_bridge import ocr as rust_ocr_bridge from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -60,27 +54,10 @@ class _PreparedRustOCRCall: _RUST_OCR_PROVIDERS = { "mistral", "azure_ai", - "azure_ai/doc-intelligence", "vertex_ai", } -def _timeout_to_seconds( - timeout: Union[float, httpx.Timeout] | None, -) -> float | None: - """Convert the Python OCR timeout to a single seconds value for the Rust bridge. - - The Rust HTTP client takes one duration; ``httpx.Timeout`` carries separate - connect/read/write/pool values, so pick the read deadline as the closest - analog to a total-request timeout. - """ - if timeout is None: - return None - if isinstance(timeout, httpx.Timeout): - return timeout.read - return float(timeout) - - def _prepare_ocr_request( model: str, document: dict[str, Any], @@ -218,13 +195,9 @@ def _rust_bridge_api_base( ) -> str | None: if prepared_request.api_base is not None: return prepared_request.api_base - if prepared_request.custom_llm_provider == "azure_ai/doc-intelligence": - return resolve_secret("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT") if prepared_request.custom_llm_provider == "azure_ai": - if ( - "doc-intelligence" in prepared_request.model - or "documentintelligence" in prepared_request.model - ): + model = prepared_request.model.lower() + if "doc-intelligence" in model or "documentintelligence" in model: return resolve_secret("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT") return resolve_secret("AZURE_AI_API_BASE") return None @@ -278,57 +251,53 @@ def _prepare_rust_ocr_call( def _run_rust_ocr( - rust_ocr: RustOcr, prepared_request: _PreparedOCRRequest, resolve_api_key: Callable[[str], str | None], -) -> OCRResponse: - """Run the Mistral OCR call through the Rust bridge and wrap the result. - - Resolves the key the same way the Python path does so secret-manager backends - (AWS/Azure/GCP/Vault) work; the Rust bridge's own fallback only reads the - process environment. The request that Rust actually sends (resolved URL and - headers) is mirrored into pre_call so logs match the wire. Dependencies are - injected so this stays unit-testable without patching module globals. - """ +) -> OCRResponse | None: + if rust_ocr_bridge.load_rust_ocr() is None: + return None prepared = _prepare_rust_ocr_call( prepared_request=prepared_request, resolve_api_key=resolve_api_key, ) - return OCRResponse.model_validate( - rust_ocr( - model=prepared_request.model, - document=cast(dict[str, object], prepared_request.document), - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared_request.custom_llm_provider, - extra_headers=prepared.headers, - optional_params=prepared.optional_params, - timeout_seconds=_timeout_to_seconds(prepared_request.effective_timeout), - ) + rust_response = rust_ocr_bridge.ocr( + model=prepared_request.model, + document=prepared_request.document, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared_request.custom_llm_provider, + extra_headers=prepared.headers, + optional_params=prepared.optional_params, + timeout=prepared_request.effective_timeout, ) + if rust_response is None: + return None + return OCRResponse.model_validate(rust_response) async def _run_rust_aocr( - rust_aocr: RustAocr, prepared_request: _PreparedOCRRequest, resolve_api_key: Callable[[str], str | None], -) -> OCRResponse: +) -> OCRResponse | None: + if rust_ocr_bridge.load_rust_aocr() is None: + return None prepared = _prepare_rust_ocr_call( prepared_request=prepared_request, resolve_api_key=resolve_api_key, ) - return OCRResponse.model_validate( - await rust_aocr( - model=prepared_request.model, - document=cast(dict[str, object], prepared_request.document), - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared_request.custom_llm_provider, - extra_headers=prepared.headers, - optional_params=prepared.optional_params, - timeout_seconds=_timeout_to_seconds(prepared_request.effective_timeout), - ) + rust_response = await rust_ocr_bridge.aocr( + model=prepared_request.model, + document=prepared_request.document, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared_request.custom_llm_provider, + extra_headers=prepared.headers, + optional_params=prepared.optional_params, + timeout=prepared_request.effective_timeout, ) + if rust_response is None: + return None + return OCRResponse.model_validate(rust_response) @client @@ -427,21 +396,19 @@ async def aocr( {"model": model, "custom_llm_provider": custom_llm_provider} ) - if _rust_ocr_supported(prepared) and rust_ocr_enabled(): - rust_aocr = load_rust_aocr() - if rust_aocr is None: + if _rust_ocr_supported(prepared) and rust_ocr_bridge.rust_ocr_enabled(): + from litellm.secret_managers.main import get_secret_str + + rust_response = await _run_rust_aocr( + prepared_request=prepared, + resolve_api_key=get_secret_str, + ) + if rust_response is None: verbose_logger.debug( "Async Rust OCR bridge unavailable; falling back to Python path" ) else: - from litellm.secret_managers.main import get_secret_str - - response = await _run_rust_aocr( - rust_aocr=rust_aocr, - prepared_request=prepared, - resolve_api_key=get_secret_str, - ) - return response + return rust_response response = base_llm_http_handler.ocr( model=prepared.model, @@ -704,21 +671,19 @@ def ocr( {"model": model, "custom_llm_provider": custom_llm_provider} ) - # Optional Rust path: hand supported OCR calls to the Rust bridge. - if _rust_ocr_supported(prepared) and rust_ocr_enabled(): - rust_ocr = load_rust_ocr() - if rust_ocr is None: + if _rust_ocr_supported(prepared) and rust_ocr_bridge.rust_ocr_enabled(): + from litellm.secret_managers.main import get_secret_str + + rust_response = _run_rust_ocr( + prepared_request=prepared, + resolve_api_key=get_secret_str, + ) + if rust_response is None: verbose_logger.debug( "Rust OCR bridge unavailable; falling back to Python path" ) else: - from litellm.secret_managers.main import get_secret_str - - return _run_rust_ocr( - rust_ocr=rust_ocr, - prepared_request=prepared, - resolve_api_key=get_secret_str, - ) + return rust_response response = base_llm_http_handler.ocr( model=prepared.model, diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 19f508e6af9..b7ac6a8c325 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -624,7 +624,9 @@ class MCPRequestHandler: Permission hierarchy (all rules are intersections): 1. Get allowed servers from key permissions - 2. Get allowed servers from team permissions (key inherits from team, or intersection) + 2. Get allowed servers from team permissions (key inherits from team, or + intersection; or inherits nothing when require_key_mcp_access_defined + is enabled, making the team a ceiling rather than a default) 3. Get allowed servers from end_user permissions (intersected if set) 4. Get allowed servers from agent permissions (intersected if set) 5. Get allowed servers from org permissions — org acts as a ceiling: if the org @@ -677,7 +679,16 @@ class MCPRequestHandler: if not team_set: base = key_set # no team restriction elif not key_set: - base = team_set # key has no own perms → inherits team + # A key that grants no MCP servers of its own inherits the + # team's by default. With require_key_mcp_access_defined the + # team is a ceiling rather than a default, so the key must + # grant servers explicitly (or via an access group) to reach + # any — it inherits none. + base = ( + set() + if general_settings.get("require_key_mcp_access_defined", False) + else team_set + ) else: base = key_set & team_set # both restrict → intersect diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 2149f079a3d..7b4f1e13a52 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -79,6 +79,7 @@ if MCP_AVAILABLE: ) from litellm.proxy._experimental.mcp_server.server import ( ListMCPToolsRestAPIResponseObject, + MCPInfo, MCPServer, _tool_name_matches, execute_mcp_tool, @@ -238,14 +239,24 @@ if MCP_AVAILABLE: ) return {} - def _create_tool_response_objects(tools, server_mcp_info): - """Helper function to create tool response objects.""" + def _create_tool_response_objects(tools, server: MCPServer): + """Helper function to create tool response objects. + + Enriches the server's ``mcp_info`` with ``server_id`` and ``alias`` so + REST clients can map the internal ``server_name`` to the user-facing + alias without needing access to the ``mcp_routes``-gated server listing. + """ + enriched_mcp_info: MCPInfo = { + **(server.mcp_info or {}), + "server_id": server.server_id, + "alias": server.alias, + } return [ ListMCPToolsRestAPIResponseObject( name=tool.name, description=tool.description, inputSchema=tool.inputSchema, - mcp_info=server_mcp_info, + mcp_info=enriched_mcp_info, ) for tool in tools ] @@ -405,7 +416,7 @@ if MCP_AVAILABLE: ) if not apply_tool_filters: - return _create_tool_response_objects(tools, server.mcp_info) + return _create_tool_response_objects(tools, server) # Always apply allowed_tools/disallowed_tools so the blacklist is # enforced even when no allowlist is set (matches the SSE/HTTP path). @@ -436,7 +447,7 @@ if MCP_AVAILABLE: if _tool_name_matches(tool.name, allowed_tools_for_server) ] - return _create_tool_response_objects(tools, server.mcp_info) + return _create_tool_response_objects(tools, server) async def _resolve_allowed_mcp_servers_for_tool_call( user_api_key_dict: UserAPIKeyAuth, @@ -587,6 +598,8 @@ if MCP_AVAILABLE: "mcp_info": { "server_name": "zapier", "logo_url": "https://www.zapier.com/logo.png", + "server_id": "a1b2c3d4-...", + "alias": "zapier_prod", } } ], diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7db0156bbe0..be5d7a2db78 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3181,6 +3181,7 @@ class SpendLogsMetadata(TypedDict): dict ] # special param to log k,v pairs to spendlogs for a call requester_ip_address: Optional[str] + litellm_call_id: Optional[str] applied_guardrails: Optional[List[str]] mcp_tool_call_metadata: Optional[StandardLoggingMCPToolCall] vector_store_request_metadata: Optional[List[StandardLoggingVectorStoreRequest]] diff --git a/litellm/proxy/client/exceptions.py b/litellm/proxy/client/exceptions.py index fffd1b78b7e..c4089381e30 100644 --- a/litellm/proxy/client/exceptions.py +++ b/litellm/proxy/client/exceptions.py @@ -2,18 +2,30 @@ from typing import Union import requests +from litellm.litellm_core_utils.secret_redaction import redact_string + + +def _redact_orig_exception( + orig_exception: Union[requests.exceptions.HTTPError, str], +) -> Union[requests.exceptions.HTTPError, str]: + if isinstance(orig_exception, requests.exceptions.HTTPError): + return requests.exceptions.HTTPError( + redact_string(str(orig_exception)), response=orig_exception.response + ) + return redact_string(str(orig_exception)) + class UnauthorizedError(Exception): """Exception raised when the API returns a 401 Unauthorized response.""" def __init__(self, orig_exception: Union[requests.exceptions.HTTPError, str]): - self.orig_exception = orig_exception - super().__init__(str(orig_exception)) + self.orig_exception = _redact_orig_exception(orig_exception) + super().__init__(str(self.orig_exception)) class NotFoundError(Exception): """Exception raised when the API returns a 404 Not Found response or indicates a resource was not found.""" def __init__(self, orig_exception: Union[requests.exceptions.HTTPError, str]): - self.orig_exception = orig_exception - super().__init__(str(orig_exception)) + self.orig_exception = _redact_orig_exception(orig_exception) + super().__init__(str(self.orig_exception)) diff --git a/litellm/proxy/client/keys.py b/litellm/proxy/client/keys.py index d8687cbad16..845b49d1581 100644 --- a/litellm/proxy/client/keys.py +++ b/litellm/proxy/client/keys.py @@ -2,6 +2,8 @@ from typing import Any, Dict, List, Optional, Union import requests +from litellm.litellm_core_utils.secret_redaction import redact_string + from .exceptions import UnauthorizedError @@ -314,6 +316,9 @@ class KeysManagementClient: response.raise_for_status() return response.json() except requests.exceptions.HTTPError as e: + redacted_message = redact_string(str(e)) if e.response.status_code == 401: - raise UnauthorizedError(e) - raise + raise UnauthorizedError(e) from None + raise requests.exceptions.HTTPError( + redacted_message, response=e.response + ) from None diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 92f8dd8df98..c2ffd9b077d 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -516,12 +516,13 @@ def _override_openai_response_model( LiteLLM internally prefixes some provider/deployment model identifiers (e.g. `hosted_vllm/...`). That internal identifier should not be returned to clients in the OpenAI `model` field. - Note: This is intentionally verbose. A model mismatch is a useful signal that an internal - model identifier is being stamped/preserved somewhere in the request/response pipeline. - We log mismatches as warnings (and then restamp to the client-requested value) so these - paths stay observable for maintainers/operators without breaking client compatibility. + Note: This is intentionally verbose at debug level. A model mismatch is a useful signal that an + internal model identifier is being stamped/preserved somewhere in the request/response pipeline. + We log mismatches as debug (and then restamp to the client-requested value) so these paths stay + observable for maintainers without breaking client compatibility or alarming operators. - Errors are reserved for cases where the proxy cannot read/override the response model field. + Responses that omit an OpenAI-style `model` field are left unchanged (silent return), + including dict responses with no `model` key. Exceptions: 1. If a fallback occurred (indicated by x-litellm-attempted-fallbacks header), @@ -577,6 +578,8 @@ def _override_openai_response_model( return if isinstance(response_obj, dict): + if "model" not in response_obj: + return downstream_model = response_obj.get("model") if downstream_model != requested_model: verbose_proxy_logger.debug( @@ -589,11 +592,6 @@ def _override_openai_response_model( return if not hasattr(response_obj, "model"): - verbose_proxy_logger.error( - "%s: cannot override response model; missing `model` attribute. response_type=%s", - log_context, - type(response_obj), - ) return downstream_model = getattr(response_obj, "model", None) @@ -608,7 +606,7 @@ def _override_openai_response_model( try: setattr(response_obj, "model", requested_model) except Exception as e: - verbose_proxy_logger.error( + verbose_proxy_logger.debug( "%s: failed to override response.model=%r on response_type=%s. error=%s", log_context, requested_model, diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 8a432eb2f42..928d00109fa 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -1849,7 +1849,10 @@ async def test_model_connection( "responses", "ocr", ] - ] = fastapi.Body("chat", description="The mode to test the model with"), + ] = fastapi.Body( + None, + description="The mode to test the model with. If not provided, auto-detected from model capabilities.", + ), litellm_params: Dict = fastapi.Body( None, description="Parameters for litellm.completion, litellm.embedding for the health check", diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index f28bcc2bcd4..bc2da33672f 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -1,8 +1,27 @@ +import math from typing import TYPE_CHECKING, Any, Dict, Optional, Union from fastapi import HTTPException, status from pydantic import BaseModel + +# Defined above the `litellm.proxy.*` imports so the name is bound even when +# this module is imported first through the proxy import cycle (CodeQL: +# module-level cyclic import). Depends only on `math` + `HTTPException`. +def validate_finite_spend(spend: float | None) -> None: + """Reject NaN/±inf spend before it reaches the DB / spend counter. + + A non-finite spend would otherwise slip past `spend >= max_budget` + enforcement, since any comparison with NaN (and `-inf >= max_budget`) + is False, letting the entity keep spending past its configured budget. + """ + if spend is not None and not math.isfinite(spend): + raise HTTPException( + status_code=400, + detail={"error": f"spend must be a finite number. Received: {spend}"}, + ) + + from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache from litellm.proxy._types import ( diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 8548d1c672c..99e276cfbb4 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -36,6 +36,7 @@ from litellm.proxy.management_endpoints.common_utils import ( _is_user_team_admin, _user_has_admin_view, require_caller_user_id_for_non_admin, + validate_finite_spend, ) from litellm.proxy.management_endpoints.key_management_endpoints import ( generate_key_helper_fn, @@ -1256,6 +1257,25 @@ def _check_user_update_authz( ) +async def _invalidate_user_spend_counter_if_changed( + non_default_values: dict[str, Any], +) -> None: + """Invalidate the cross-pod spend counter after a direct ``spend`` change. + + A direct ``spend`` change must also invalidate the cross-pod spend counter + enforcement reads; the DB write alone leaves a warm counter at the stale + value. ``non_default_values["user_id"]`` is populated in every branch of the + caller (incl. the email-new-user insert path, whose response is a bare model + and not safely subscriptable). + """ + if non_default_values.get("spend") is not None: + from litellm.proxy.proxy_server import _invalidate_spend_counter + + await _invalidate_spend_counter( + counter_key=f"spend:user:{non_default_values['user_id']}" + ) + + async def _update_single_user_helper( user_request: UpdateUserRequest, user_api_key_dict: UserAPIKeyAuth, @@ -1336,6 +1356,9 @@ async def _update_single_user_helper( existing_metadata=existing_metadata or {}, ) + # Reject NaN/±inf spend before it can reach the DB / spend counter. + validate_finite_spend(non_default_values.get("spend")) + # Perform the update response: Optional[Dict[str, Any]] = None @@ -1384,6 +1407,8 @@ async def _update_single_user_helper( litellm_proxy_admin_name=litellm_proxy_admin_name, ) + await _invalidate_user_spend_counter_if_changed(non_default_values) + if response is None: raise HTTPException( status_code=400, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index db205803a47..9f316256ac9 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -69,6 +69,7 @@ from litellm.proxy.management_endpoints.common_utils import ( _is_user_team_admin, _set_object_metadata_field, _team_member_has_permission, + validate_finite_spend, ) from litellm.proxy.management_endpoints.model_management_endpoints import ( _add_model_to_db, @@ -2210,6 +2211,9 @@ async def _validate_update_key_data( user_api_key_cache: Any, ) -> None: """Validate permissions and constraints for key update.""" + # Reject NaN/±inf spend before it can reach the DB / spend counter. + validate_finite_spend(data.spend) + _is_proxy_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value _check_allowed_routes_caller_permission( @@ -2269,12 +2273,14 @@ async def _validate_update_key_data( # existing admin-only budget semantics). budget_limits uses # model_fields_set because an explicit null/[] clears the field # and must gate the same as setting or changing it. + # - spend gates on presence alone (not a value diff): the DB spend + # lags the live cross-pod counter, so letting an "unchanged" spend + # through the non-admin path would let a key owner / team member + # overwrite the live counter below real usage and silently weaken + # enforcement. _is_budget_change = ( (data.max_budget is not None and data.max_budget != existing_key_row.max_budget) - or ( - data.spend is not None - and data.spend != getattr(existing_key_row, "spend", None) - ) + or data.spend is not None or "budget_limits" in data.model_fields_set ) @@ -2609,15 +2615,24 @@ async def update_key_fn( ) if data.spend is not None: - try: - from litellm.proxy.proxy_server import _invalidate_spend_counter + from litellm.proxy.proxy_server import spend_counter_cache - token_to_invalidate = _hash_token_if_needed(key) - await _invalidate_spend_counter( - counter_key=f"spend:key:{token_to_invalidate}" - ) - except Exception: - pass + counter_key = f"spend:key:{_hash_token_if_needed(key)}" + spend_counter_cache.in_memory_cache.set_cache( + key=counter_key, value=data.spend, ttl=60 + ) + if spend_counter_cache.redis_cache is not None: + try: + await spend_counter_cache.redis_cache.async_set_cache( + key=counter_key, value=data.spend, ttl=60 + ) + except Exception as redis_err: + verbose_proxy_logger.warning( + "Failed to update spend counter %s in Redis after key spend update: %s. " + "Budget checks may use stale value until counter expires.", + counter_key, + redis_err, + ) asyncio.create_task( KeyManagementEventHooks.async_key_updated_hook( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b0beaae02cf..3c91f1bc1d7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1114,6 +1114,33 @@ _OPENAPI_HTTP_METHODS = { # `_SSO_SENSITIVE_FIELDS` / `_CACHE_SENSITIVE_FIELDS` constants in the SSO # and cache endpoint files. _ALERTING_SENSITIVE_VARS: Set[str] = {"SLACK_WEBHOOK_URL", "SMTP_PASSWORD"} +_DB_LITELLM_PARAM_ENV_REF_KEYS = frozenset( + { + "api_key", + "client_secret", + "vertex_credentials", + "vertex_ai_credentials", + "aws_access_key_id", + "aws_secret_access_key", + } +) + + +def _db_model_is_team_scoped(model: object) -> bool: + model_info = getattr(model, "model_info", None) + if isinstance(model_info, BaseModel): + return getattr(model_info, "team_id", None) is not None + if isinstance(model_info, str): + try: + model_info = json.loads(model_info) + except (TypeError, ValueError): + model_info = None + if isinstance(model_info, dict) and model_info.get("team_id") is not None: + return True + if getattr(model_info, "team_id", None) is not None: + return True + model_name = getattr(model, "model_name", None) + return isinstance(model_name, str) and model_name.startswith("model_name_") def _strip_operation_id_method_suffix(operation_id: str) -> str: @@ -5255,6 +5282,24 @@ class ProxyConfig: deleted_deployments += 1 return deleted_deployments + def _resolve_db_litellm_param( + self, key: str, value: object, resolve_env_refs: bool = True + ) -> object: + if not isinstance(value, str): + return value + + decrypted_value = decrypt_value_helper( + value=value, key=key, return_original_value=True + ) + if ( + resolve_env_refs + and key in _DB_LITELLM_PARAM_ENV_REF_KEYS + and isinstance(decrypted_value, str) + and decrypted_value.startswith("os.environ/") + ): + return get_secret(decrypted_value) + return decrypted_value + def _add_deployment(self, db_models: list) -> int: """ Iterate through db models @@ -5272,15 +5317,13 @@ class ProxyConfig: ## ADD MODEL LOGIC for m in db_models: _litellm_params = m.litellm_params + resolve_env_refs = not _db_model_is_team_scoped(m) if isinstance(_litellm_params, dict): # decrypt values for k, v in _litellm_params.items(): - if isinstance(v, str): - # decrypt value - returns original value if decryption fails or no key is set - _value = decrypt_value_helper( - value=v, key=k, return_original_value=True - ) - _litellm_params[k] = _value + _litellm_params[k] = self._resolve_db_litellm_param( + key=k, value=v, resolve_env_refs=resolve_env_refs + ) _litellm_params = LiteLLM_Params(**_litellm_params) else: @@ -5308,15 +5351,15 @@ class ProxyConfig: _model_list: list = [] for m in new_models: _litellm_params = m.litellm_params + resolve_env_refs = not _db_model_is_team_scoped(m) if isinstance(_litellm_params, BaseModel): _litellm_params = _litellm_params.model_dump() if isinstance(_litellm_params, dict): # decrypt values for k, v in _litellm_params.items(): - decrypted_value = decrypt_value_helper( - value=v, key=k, return_original_value=True + _litellm_params[k] = self._resolve_db_litellm_param( + key=k, value=v, resolve_env_refs=resolve_env_refs ) - _litellm_params[k] = decrypted_value _litellm_params = LiteLLM_Params(**_litellm_params) else: verbose_proxy_logger.error( diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 8d89ff4a1ff..f6e0303e349 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -79,6 +79,7 @@ def _get_spend_logs_metadata( cold_storage_object_key: Optional[str] = None, litellm_overhead_time_ms: Optional[float] = None, cost_breakdown: Optional[CostBreakdown] = None, + litellm_call_id: Optional[str] = None, ) -> SpendLogsMetadata: if metadata is None: return SpendLogsMetadata( @@ -109,6 +110,7 @@ def _get_spend_logs_metadata( attempted_retries=None, max_retries=None, cost_breakdown=None, + litellm_call_id=litellm_call_id, ) verbose_proxy_logger.debug( "getting payload for SpendLogs, available keys in metadata: " @@ -133,6 +135,7 @@ def _get_spend_logs_metadata( clean_metadata["cold_storage_object_key"] = cold_storage_object_key clean_metadata["litellm_overhead_time_ms"] = litellm_overhead_time_ms clean_metadata["cost_breakdown"] = cost_breakdown + clean_metadata["litellm_call_id"] = litellm_call_id return clean_metadata @@ -383,6 +386,10 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs if standard_logging_payload is not None else None ), + litellm_call_id=cast( + Optional[str], + kwargs.get("litellm_call_id") or litellm_params.get("litellm_call_id"), + ), ) special_usage_fields = ["completion_tokens", "prompt_tokens", "total_tokens"] diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index a29f0a01b7a..1187640bf0c 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -59,6 +59,7 @@ class LiteLLMCompletionTransformationHandler: completion_args = {} completion_args.update(kwargs) completion_args.update(litellm_completion_request) + completion_args["_skip_responses_api_bridge"] = True litellm_completion_response: Union[ ModelResponse, litellm.CustomStreamWrapper @@ -107,6 +108,7 @@ class LiteLLMCompletionTransformationHandler: acompletion_args = {} acompletion_args.update(kwargs) acompletion_args.update(litellm_completion_request) + acompletion_args["_skip_responses_api_bridge"] = True litellm_completion_response: Union[ ModelResponse, litellm.CustomStreamWrapper diff --git a/litellm/router.py b/litellm/router.py index 0008fcaa30b..1aba259a328 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -2323,6 +2323,14 @@ class Router: verbose_router_logger.error( f"Fallback also failed: {fallback_error}" ) + # No fallback handled the mid-stream error, so surface the + # real provider exception (e.g. RateLimitError) instead of + # leaking the internal MidStreamFallbackError to the client + if ( + isinstance(fallback_error, MidStreamFallbackError) + and fallback_error.original_exception is not None + ): + raise fallback_error.original_exception from fallback_error raise fallback_error finally: # Close the underlying streams to release HTTP connections @@ -2754,6 +2762,11 @@ class Router: verbose_router_logger.error( f"Responses streaming fallback also failed: {fallback_error}" ) + if ( + isinstance(fallback_error, MidStreamFallbackError) + and fallback_error.original_exception is not None + ): + raise fallback_error.original_exception from fallback_error raise fallback_error finally: with anyio.CancelScope(shield=True): @@ -2890,6 +2903,11 @@ class Router: verbose_router_logger.error( f"Fallback also failed: {fallback_error}" ) + if ( + isinstance(fallback_error, MidStreamFallbackError) + and fallback_error.original_exception is not None + ): + raise fallback_error.original_exception from fallback_error raise fallback_error finally: if hasattr(model_response, "close"): diff --git a/litellm/rust_bridge/__init__.py b/litellm/rust_bridge/__init__.py index ec89e3b65b4..3da5b98449b 100644 --- a/litellm/rust_bridge/__init__.py +++ b/litellm/rust_bridge/__init__.py @@ -4,5 +4,6 @@ from litellm.rust_bridge.loader import ( get_native_bridge, native_bridge_available, ) +from litellm.rust_bridge.ocr import use_litellm_rust -__all__ = ["get_native_bridge", "native_bridge_available"] +__all__ = ["get_native_bridge", "native_bridge_available", "use_litellm_rust"] diff --git a/litellm/ocr/rust_bridge.py b/litellm/rust_bridge/ocr.py similarity index 52% rename from litellm/ocr/rust_bridge.py rename to litellm/rust_bridge/ocr.py index 631fd4c63c5..36a088b6b1a 100644 --- a/litellm/ocr/rust_bridge.py +++ b/litellm/rust_bridge/ocr.py @@ -1,30 +1,21 @@ -""" -Optional Rust-backed OCR path. - -Enable with ``litellm.use_litellm_rust()``; the sync ``litellm.ocr()`` entrypoint -then routes supported Mistral calls through the compiled ``litellm.rust_bridge._native`` -extension, which performs the whole OCR call (URL, headers, HTTP, parse) in Rust. - -No module-level ``litellm`` imports keep this a leaf so ``litellm/ocr/main.py`` -can import it statically without forming an import cycle. -""" +"""Thin Python wrapper for the native Rust OCR bridge.""" from __future__ import annotations import os -from typing import Awaitable, Final, Protocol, cast +from typing import Any, Awaitable, Final, Protocol, Union, cast + +import httpx class RustOcr(Protocol): - """Signature of the compiled Rust OCR entrypoint.""" - def __call__( self, model: str, document: dict[str, object], api_key: str | None, api_base: str | None, - custom_llm_provider: str, + custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, @@ -33,15 +24,13 @@ class RustOcr(Protocol): class RustAocr(Protocol): - """Signature of the compiled ``litellm_python_bridge.aocr`` entrypoint.""" - def __call__( self, model: str, document: dict[str, object], api_key: str | None, api_base: str | None, - custom_llm_provider: str, + custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, @@ -50,7 +39,7 @@ class RustAocr(Protocol): class _Unset: - """Sentinel type so ``ocr=None`` can clear a prior injection while omission preserves it.""" + pass _UNSET: Final[_Unset] = _Unset() @@ -76,12 +65,6 @@ def use_litellm_rust( ocr: RustOcr | None | _Unset = _UNSET, aocr: RustAocr | None | _Unset = _UNSET, ) -> None: - """Route supported OCR calls through the packaged Rust extension. - - ``ocr`` and ``aocr`` inject bridge callables; when omitted the compiled - extension is loaded on demand and any previously injected bridge is - preserved. Pass ``None`` explicitly to clear a prior injection. - """ global _rust_ocr_enabled, _rust_ocr_impl, _rust_aocr_impl _rust_ocr_enabled = enabled if not isinstance(ocr, _Unset): @@ -91,17 +74,10 @@ def use_litellm_rust( def rust_ocr_enabled() -> bool: - """Whether the Rust OCR path has been turned on via ``use_litellm_rust()``.""" return _rust_ocr_enabled def load_rust_ocr() -> RustOcr | None: - """Return the Rust OCR callable, or ``None`` when no bridge is available. - - Prefers an injected implementation, otherwise loads the compiled - ``litellm.rust_bridge._native`` extension; a missing extension yields ``None`` so - the caller can fall back to the Python path instead of hard-failing. - """ if _rust_ocr_impl is not None: return _rust_ocr_impl from litellm.rust_bridge import get_native_bridge @@ -113,7 +89,6 @@ def load_rust_ocr() -> RustOcr | None: def load_rust_aocr() -> RustAocr | None: - """Return the async Rust OCR callable, or ``None`` when unavailable.""" if _rust_aocr_impl is not None: return _rust_aocr_impl from litellm.rust_bridge import get_native_bridge @@ -122,3 +97,63 @@ def load_rust_aocr() -> RustAocr | None: if native_bridge is None: return None return cast(RustAocr, getattr(native_bridge, "aocr", None)) + + +def _timeout_to_seconds(timeout: Union[float, httpx.Timeout] | None) -> float | None: + if timeout is None: + return None + if isinstance(timeout, httpx.Timeout): + return timeout.read + return float(timeout) + + +def ocr( + *, + model: str, + document: dict[str, Any], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, Any] | None, + optional_params: dict[str, object], + timeout: Union[float, httpx.Timeout] | None, +) -> dict[str, object] | None: + rust_ocr = load_rust_ocr() + if rust_ocr is None: + return None + return rust_ocr( + model=model, + document=cast(dict[str, object], document), + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=cast(dict[str, object] | None, extra_headers), + optional_params=optional_params, + timeout_seconds=_timeout_to_seconds(timeout), + ) + + +async def aocr( + *, + model: str, + document: dict[str, Any], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, Any] | None, + optional_params: dict[str, object], + timeout: Union[float, httpx.Timeout] | None, +) -> dict[str, object] | None: + rust_aocr = load_rust_aocr() + if rust_aocr is None: + return None + return await rust_aocr( + model=model, + document=cast(dict[str, object], document), + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=cast(dict[str, object] | None, extra_headers), + optional_params=optional_params, + timeout_seconds=_timeout_to_seconds(timeout), + ) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 24d6e84fba7..c3f99b1d18b 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1539,6 +1539,9 @@ class PromptTokensDetailsWrapper( video_length_seconds: Optional[float] = None """Length of videos sent to the model. Used for Vertex AI multimodal embeddings.""" + audio_length_seconds: Optional[float] = None + """Length of audio sent to the model. Used for multimodal embeddings priced per audio-second.""" + cache_creation_tokens: Optional[int] = None """Number of cache creation tokens sent to the model. Used for Anthropic prompt caching.""" @@ -1553,6 +1556,8 @@ class PromptTokensDetailsWrapper( del self.image_count if self.video_length_seconds is None: del self.video_length_seconds + if self.audio_length_seconds is None: + del self.audio_length_seconds if self.web_search_requests is None: del self.web_search_requests if self.cache_creation_tokens is None: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9867874cb65..a174c1b5efd 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -25965,7 +25965,7 @@ "supports_response_schema": true, "supports_tool_choice": true }, - "mistral/mistral-medium-latest": { + "mistral/mistral-medium-2508": { "input_cost_per_token": 4e-07, "litellm_provider": "mistral", "max_input_tokens": 131072, @@ -25973,12 +25973,45 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 2e-06, + "source": "https://mistral.ai/news/mistral-medium-3", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true }, + "mistral/mistral-medium-2604": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/mistral-medium-latest": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "mistral/mistral-medium-3-1-2508": { "input_cost_per_token": 4e-07, "litellm_provider": "mistral", @@ -26005,6 +26038,7 @@ "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true diff --git a/tests/logging_callback_tests/test_gcs_pub_sub.py b/tests/logging_callback_tests/test_gcs_pub_sub.py index f4cc9735177..c37a2e3f65d 100644 --- a/tests/logging_callback_tests/test_gcs_pub_sub.py +++ b/tests/logging_callback_tests/test_gcs_pub_sub.py @@ -31,6 +31,7 @@ verbose_logger.setLevel(logging.DEBUG) ignored_keys = [ "request_id", + "metadata.litellm_call_id", "session_id", "startTime", "endTime", diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 5f8fcbf835e..bb13a7ce8cc 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -1883,10 +1883,11 @@ def test_get_server_auth_header_no_auth_headers(): def test_create_tool_response_objects(): - """Test _create_tool_response_objects function.""" + """Test _create_tool_response_objects enriches mcp_info with server_id and alias.""" from litellm.proxy._experimental.mcp_server.rest_endpoints import ( _create_tool_response_objects, ) + from litellm.types.mcp_server.mcp_server_manager import MCPServer from mcp.types import Tool as MCPTool # Create mock tools @@ -1903,20 +1904,32 @@ def test_create_tool_response_objects(): ), ] - server_mcp_info = { - "server_name": "zapier", + server = MCPServer( + server_id="a1b2c3d4", + name="zapier_internal", + alias="zapier", + transport="http", + mcp_info={ + "server_name": "zapier_internal", + "logo_url": "https://zapier.com/logo.png", + }, + ) + + result = _create_tool_response_objects(mock_tools, server) + + expected_mcp_info = { + "server_name": "zapier_internal", "logo_url": "https://zapier.com/logo.png", + "server_id": "a1b2c3d4", + "alias": "zapier", } - - result = _create_tool_response_objects(mock_tools, server_mcp_info) - assert len(result) == 2 assert result[0].name == "send_email" assert result[0].description == "Send an email" - assert result[0].mcp_info == server_mcp_info + assert result[0].mcp_info == expected_mcp_info assert result[1].name == "create_event" assert result[1].description == "Create a calendar event" - assert result[1].mcp_info == server_mcp_info + assert result[1].mcp_info == expected_mcp_info @pytest.mark.asyncio @@ -1930,6 +1943,8 @@ async def test_get_tools_for_single_server(): # Create a mock server (pin allowlist fields; MagicMock auto-attrs are truthy) mock_server = MagicMock() mock_server.mcp_info = {"server_name": "zapier"} + mock_server.server_id = "zapier_id" + mock_server.alias = "zapier_alias" mock_server.allowed_tools = None mock_server.disallowed_tools = None @@ -1963,7 +1978,11 @@ async def test_get_tools_for_single_server(): # Verify the result assert len(result) == 1 assert result[0].name == "send_email" - assert result[0].mcp_info == {"server_name": "zapier"} + assert result[0].mcp_info == { + "server_name": "zapier", + "server_id": "zapier_id", + "alias": "zapier_alias", + } @pytest.mark.asyncio diff --git a/tests/test_litellm/caching/test_embedding_router.py b/tests/test_litellm/caching/test_embedding_router.py new file mode 100644 index 00000000000..550095a112a --- /dev/null +++ b/tests/test_litellm/caching/test_embedding_router.py @@ -0,0 +1,67 @@ +import os +import sys +from unittest.mock import MagicMock + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.caching._embedding_router import ( + build_router_embedding_metadata, + resolve_embedding_router, +) + + +def test_resolve_returns_router_when_model_is_a_deployment(): + router = MagicMock() + assert ( + resolve_embedding_router("sem-embed", router, [{"model_name": "sem-embed"}]) + is router + ) + + +def test_resolve_returns_none_when_model_not_in_router(): + router = MagicMock() + assert ( + resolve_embedding_router("sem-embed", router, [{"model_name": "other"}]) is None + ) + + +def test_resolve_returns_none_when_router_is_none(): + assert ( + resolve_embedding_router("sem-embed", None, [{"model_name": "sem-embed"}]) + is None + ) + + +def test_resolve_returns_none_when_model_list_is_none(): + router = MagicMock() + assert resolve_embedding_router("sem-embed", router, None) is None + + +def test_resolve_skips_entries_missing_model_name(): + router = MagicMock() + model_list = [ + {"litellm_params": {"model": "bedrock/x"}}, + {"model_name": "sem-embed"}, + ] + assert resolve_embedding_router("sem-embed", router, model_list) is router + assert resolve_embedding_router("other", router, [{"litellm_params": {}}]) is None + + +def test_build_metadata_preserves_request_fields_and_adds_flag(): + md = build_router_embedding_metadata( + {"user_api_key": "sk-x", "user_api_key_team_id": "team-1", "trace_id": "t-1"} + ) + assert md == { + "user_api_key": "sk-x", + "user_api_key_team_id": "team-1", + "trace_id": "t-1", + "semantic-cache-embedding": True, + } + + +def test_build_metadata_handles_none_and_does_not_mutate_input(): + original = {"user_api_key": "sk-x"} + md = build_router_embedding_metadata(original) + assert md == {"user_api_key": "sk-x", "semantic-cache-embedding": True} + assert original == {"user_api_key": "sk-x"} + assert build_router_embedding_metadata(None) == {"semantic-cache-embedding": True} diff --git a/tests/test_litellm/caching/test_qdrant_semantic_cache.py b/tests/test_litellm/caching/test_qdrant_semantic_cache.py index 949e6ccc292..67d4e2d9892 100644 --- a/tests/test_litellm/caching/test_qdrant_semantic_cache.py +++ b/tests/test_litellm/caching/test_qdrant_semantic_cache.py @@ -1,5 +1,6 @@ import os import sys +import types from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -806,3 +807,104 @@ def test_qdrant_semantic_cache_large_vector_size(): ) create_payload = put_call.kwargs["json"] assert create_payload["vectors"]["size"] == 4096 + + +def _router_proxy_module(router, model_name): + mod = types.ModuleType("litellm.proxy.proxy_server") + mod.llm_router = router + mod.llm_model_list = [{"model_name": model_name}] + return mod + + +def test_qdrant_sync_get_cache_routes_through_router(monkeypatch): + from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache + + cache = QdrantSemanticCache.__new__(QdrantSemanticCache) + cache.embedding_model = "sem-embed" + cache.qdrant_api_base = "http://test.qdrant.local" + cache.collection_name = "test_collection" + cache.headers = {"Content-Type": "application/json", "api-key": "test_key"} + cache.similarity_threshold = 0.8 + cache.sync_client = MagicMock() + search_response = MagicMock() + search_response.status_code = 200 + search_response.json.return_value = {"result": []} + cache.sync_client.post.return_value = search_response + + router = MagicMock() + router.embedding = MagicMock( + return_value={"data": [{"embedding": [0.3, 0.3, 0.3]}]} + ) + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + _router_proxy_module(router, "sem-embed"), + ) + + with patch("litellm.embedding") as direct_embed: + result = cache.get_cache( + key="test_key", + messages=[{"content": "What is the capital of France?"}], + metadata={}, + ) + + assert result is None + router.embedding.assert_called_once() + assert router.embedding.call_args.kwargs["model"] == "sem-embed" + direct_embed.assert_not_called() + + +def test_qdrant_sync_set_cache_falls_back_to_direct(monkeypatch): + from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache + + cache = QdrantSemanticCache.__new__(QdrantSemanticCache) + cache.embedding_model = "text-embedding-ada-002" + cache.qdrant_api_base = "http://test.qdrant.local" + cache.collection_name = "test_collection" + cache.headers = {"Content-Type": "application/json", "api-key": "test_key"} + cache.sync_client = MagicMock() + put_response = MagicMock() + put_response.status_code = 200 + cache.sync_client.put.return_value = put_response + + fake_proxy = types.ModuleType("litellm.proxy.proxy_server") + fake_proxy.llm_router = None + fake_proxy.llm_model_list = None + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy) + + with patch( + "litellm.embedding", return_value={"data": [{"embedding": [0.1, 0.1, 0.1]}]} + ) as direct_embed: + cache.set_cache( + key="test_key", + value={"content": "Paris"}, + messages=[{"content": "What is the capital of France?"}], + ) + + direct_embed.assert_called_once() + + +@pytest.mark.asyncio +async def test_qdrant_async_embedding_forwards_full_metadata(monkeypatch): + from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache + + cache = QdrantSemanticCache.__new__(QdrantSemanticCache) + cache.embedding_model = "sem-embed" + + router = MagicMock() + router.aembedding = AsyncMock(return_value={"data": [{"embedding": [0.1, 0.2]}]}) + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + _router_proxy_module(router, "sem-embed"), + ) + + await cache._get_async_embedding( + "hello", + metadata={"user_api_key": "sk-x", "user_api_key_team_id": "team-1"}, + ) + + md = router.aembedding.call_args.kwargs["metadata"] + assert md["user_api_key"] == "sk-x" + assert md["user_api_key_team_id"] == "team-1" + assert md["semantic-cache-embedding"] is True diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 6aaadce93ce..37c938bb38d 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -3,7 +3,6 @@ import sys from unittest.mock import MagicMock, patch import pytest -from fastapi.testclient import TestClient sys.path.insert( 0, os.path.abspath("../../..") @@ -94,6 +93,42 @@ async def test_redis_cache_async_increment_default_does_not_bump_existing_ttl( mock_redis_instance.expire.assert_not_awaited() +@pytest.mark.parametrize("namespace", [None, "litellm"]) +@pytest.mark.asyncio +async def test_async_delete_cache_applies_namespace( + namespace, monkeypatch, redis_no_ping +): + """async_delete_cache must prefix keys with the namespace, matching every + other cache operation. Without this, Redis NOPERM errors occur when an + ACL restricts DEL to the litellm:* pattern.""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + mock_redis_instance = AsyncMock() + + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + await redis_cache.async_delete_cache(key="3997c4abcdef") + + expected_key = "litellm:3997c4abcdef" if namespace else "3997c4abcdef" + mock_redis_instance.delete.assert_awaited_once_with(expected_key) + + +@pytest.mark.parametrize("namespace", [None, "litellm"]) +def test_delete_cache_applies_namespace(namespace, monkeypatch, redis_no_ping): + """delete_cache must prefix keys with the namespace, matching every other + cache operation.""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + mock_redis_client = MagicMock() + redis_cache.redis_client = mock_redis_client + + redis_cache.delete_cache(key="3997c4abcdef") + + expected_key = "litellm:3997c4abcdef" if namespace else "3997c4abcdef" + mock_redis_client.delete.assert_called_once_with(expected_key) + + @pytest.mark.asyncio async def test_redis_client_init_with_socket_timeout(monkeypatch, redis_no_ping): monkeypatch.setenv("REDIS_HOST", "my-fake-host") diff --git a/tests/test_litellm/caching/test_redis_semantic_cache.py b/tests/test_litellm/caching/test_redis_semantic_cache.py index 13f9d00136d..1d3129d6467 100644 --- a/tests/test_litellm/caching/test_redis_semantic_cache.py +++ b/tests/test_litellm/caching/test_redis_semantic_cache.py @@ -104,6 +104,7 @@ def test_redis_semantic_cache_get_cache(monkeypatch): # Verify llmcache.check was called redis_semantic_cache.llmcache.check.assert_called_once_with( prompt="What is the capital of France?", + vector=[0.1, 0.2, 0.3], filter_expression="cache-key-filter", ) @@ -138,10 +139,16 @@ def test_redis_semantic_cache_rejects_unscoped_cache_hit(monkeypatch): ] ) - with patch.object( - redis_semantic_cache, - "_get_cache_key_filter_expression", - return_value="cache-key-filter", + with ( + patch( + "litellm.embedding", + return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]}, + ), + patch.object( + redis_semantic_cache, + "_get_cache_key_filter_expression", + return_value="cache-key-filter", + ), ): metadata = {} result = redis_semantic_cache.get_cache( @@ -176,16 +183,21 @@ def test_redis_semantic_cache_set_cache_stores_cache_key_filter(monkeypatch): redis_semantic_cache = RedisSemanticCache(similarity_threshold=0.8) redis_semantic_cache.llmcache.store = MagicMock() - redis_semantic_cache.set_cache( - key="test_key", - value={"content": "Paris"}, - messages=[{"content": "What is the capital of France?"}], - ttl=60, - ) + with patch( + "litellm.embedding", + return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]}, + ): + redis_semantic_cache.set_cache( + key="test_key", + value={"content": "Paris"}, + messages=[{"content": "What is the capital of France?"}], + ttl=60, + ) redis_semantic_cache.llmcache.store.assert_called_once_with( "What is the capital of France?", "{'content': 'Paris'}", + vector=[0.1, 0.2, 0.3], filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}, ttl=60, ) @@ -299,10 +311,11 @@ def test_redis_semantic_cache_reraises_unexpected_isolated_index_error(monkeypat monkeypatch.setenv("REDIS_PASSWORD", "test_password") with pytest.raises(ValueError, match="connection failed"): - RedisSemanticCache( + cache = RedisSemanticCache( similarity_threshold=0.8, index_name="existing_index", ) + _ = cache.llmcache def test_redis_semantic_cache_reraises_unexpected_index_error(): @@ -534,6 +547,7 @@ def test_redis_semantic_cache_set_cache_uses_responses_string_input(): return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"} ) redis_semantic_cache._get_ttl = MagicMock(return_value=None) + redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) redis_semantic_cache.set_cache( key="test_key", @@ -544,6 +558,7 @@ def test_redis_semantic_cache_set_cache_uses_responses_string_input(): redis_semantic_cache.llmcache.store.assert_called_once_with( "What is the capital of France?", "{'content': 'Paris'}", + vector=[0.1, 0.2, 0.3], filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}, ) @@ -564,6 +579,7 @@ def test_redis_semantic_cache_get_cache_uses_responses_string_input(): } ] ) + redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) with patch.object( redis_semantic_cache, @@ -581,6 +597,7 @@ def test_redis_semantic_cache_get_cache_uses_responses_string_input(): assert metadata["semantic-similarity"] == pytest.approx(0.9) redis_semantic_cache.llmcache.check.assert_called_once_with( prompt="What is the capital of France?", + vector=[0.1, 0.2, 0.3], filter_expression="cache-key-filter", ) @@ -594,6 +611,7 @@ def test_redis_semantic_cache_set_cache_flattens_structured_responses_input(): return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"} ) redis_semantic_cache._get_ttl = MagicMock(return_value=None) + redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) redis_semantic_cache.set_cache( key="test_key", @@ -616,6 +634,7 @@ def test_redis_semantic_cache_set_cache_flattens_structured_responses_input(): redis_semantic_cache.llmcache.store.assert_called_once_with( "What is the capital of France?\nAnswer briefly.", "{'content': 'Paris'}", + vector=[0.1, 0.2, 0.3], filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}, ) @@ -740,6 +759,7 @@ def test_redis_semantic_cache_get_cache_sets_similarity_when_no_results(): redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) redis_semantic_cache.llmcache = MagicMock() redis_semantic_cache.llmcache.check = MagicMock(return_value=[]) + redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) with patch.object( redis_semantic_cache, @@ -757,6 +777,7 @@ def test_redis_semantic_cache_get_cache_sets_similarity_when_no_results(): assert metadata["semantic-similarity"] == 0.0 redis_semantic_cache.llmcache.check.assert_called_once_with( prompt="What is the capital of France?", + vector=[0.1, 0.2, 0.3], filter_expression="cache-key-filter", ) @@ -870,6 +891,63 @@ async def test_redis_semantic_cache_async_paths_set_similarity_on_misses(): ) +def test_redis_get_embedding_routes_through_router(monkeypatch): + import sys + import types + + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + cache = RedisSemanticCache.__new__(RedisSemanticCache) + cache.embedding_model = "sem-embed" + + router = MagicMock() + router.embedding = MagicMock(return_value={"data": [{"embedding": [0.5, 0.6]}]}) + fake_proxy = types.ModuleType("litellm.proxy.proxy_server") + fake_proxy.llm_router = router + fake_proxy.llm_model_list = [{"model_name": "sem-embed"}] + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy) + + with patch("litellm.embedding") as direct_embed: + vec = cache._get_embedding("hello", metadata={"user_api_key": "sk-x"}) + + assert vec == [0.5, 0.6] + router.embedding.assert_called_once() + assert router.embedding.call_args.kwargs["model"] == "sem-embed" + assert router.embedding.call_args.kwargs["input"] == "hello" + assert router.embedding.call_args.kwargs["cache"] == { + "no-store": True, + "no-cache": True, + } + assert router.embedding.call_args.kwargs["metadata"] == { + "user_api_key": "sk-x", + "semantic-cache-embedding": True, + } + direct_embed.assert_not_called() + + +def test_redis_get_embedding_falls_back_to_direct(monkeypatch): + import sys + import types + + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + cache = RedisSemanticCache.__new__(RedisSemanticCache) + cache.embedding_model = "text-embedding-ada-002" + + fake_proxy = types.ModuleType("litellm.proxy.proxy_server") + fake_proxy.llm_router = None + fake_proxy.llm_model_list = None + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy) + + with patch( + "litellm.embedding", return_value={"data": [{"embedding": [0.1, 0.2]}]} + ) as direct_embed: + vec = cache._get_embedding("hello") + + assert vec == [0.1, 0.2] + direct_embed.assert_called_once() + + def test_cache_get_cache_passes_responses_input_to_backend_cache(): from litellm.caching.caching import Cache @@ -893,7 +971,7 @@ def test_cache_get_cache_passes_responses_input_to_backend_cache(): ) -def test_cache_get_cache_filters_sensitive_kwargs_from_backend_cache(): +def test_cache_get_cache_filters_non_lookup_kwargs_from_backend_cache(): from litellm.caching.caching import Cache cache = Cache.__new__(Cache) @@ -927,7 +1005,11 @@ def test_cache_get_cache_filters_sensitive_kwargs_from_backend_cache(): forwarded_kwargs = cache.cache.get_cache.call_args.kwargs assert forwarded_kwargs == { "input": "What is the capital of France?", - "metadata": {"semantic-similarity": 0.7}, + "metadata": { + "user_api_key": "sk-secret", + "trace_id": "trace-id", + "semantic-similarity": 0.7, + }, } assert forwarded_kwargs["metadata"] is not metadata cache._get_cache_logic.assert_called_once_with( @@ -988,3 +1070,166 @@ def test_cache_get_cache_passes_responses_input_to_dynamic_cache(): cached_result={"content": "Paris"}, max_age=float("inf"), ) + + +def test_redis_sync_set_cache_passes_precomputed_vector(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + cache = RedisSemanticCache.__new__(RedisSemanticCache) + cache.llmcache = MagicMock() + cache._get_cache_filters = MagicMock( + return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"} + ) + cache._get_ttl = MagicMock(return_value=None) + cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) + + cache.set_cache( + key="test_key", + value={"content": "Paris"}, + messages=[{"content": "What is the capital of France?"}], + ) + + cache._get_embedding.assert_called_once() + cache.llmcache.store.assert_called_once_with( + "What is the capital of France?", + "{'content': 'Paris'}", + vector=[0.1, 0.2, 0.3], + filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}, + ) + + +def test_redis_sync_get_cache_passes_precomputed_vector(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + cache = RedisSemanticCache.__new__(RedisSemanticCache) + cache.similarity_threshold = 0.8 + cache.llmcache = MagicMock() + cache.llmcache.check = MagicMock( + return_value=[ + { + "prompt": "What is the capital of France?", + "response": '{"content": "Paris"}', + "vector_distance": 0.1, + RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key", + } + ] + ) + cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) + + with patch.object( + cache, "_get_cache_key_filter_expression", return_value="cache-key-filter" + ): + result = cache.get_cache( + key="test_key", + messages=[{"content": "What is the capital of France?"}], + metadata={}, + ) + + assert result == {"content": "Paris"} + cache._get_embedding.assert_called_once() + cache.llmcache.check.assert_called_once_with( + prompt="What is the capital of France?", + vector=[0.1, 0.2, 0.3], + filter_expression="cache-key-filter", + ) + + +@pytest.mark.asyncio +async def test_redis_async_embedding_forwards_full_metadata(monkeypatch): + import sys + import types + + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + cache = RedisSemanticCache.__new__(RedisSemanticCache) + cache.embedding_model = "sem-embed" + + router = MagicMock() + router.aembedding = AsyncMock(return_value={"data": [{"embedding": [0.1, 0.2]}]}) + fake_proxy = types.ModuleType("litellm.proxy.proxy_server") + fake_proxy.llm_router = router + fake_proxy.llm_model_list = [{"model_name": "sem-embed"}] + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy) + + await cache._get_async_embedding( + "hello", + metadata={"user_api_key": "sk-x", "user_api_key_team_id": "team-1"}, + ) + + md = router.aembedding.call_args.kwargs["metadata"] + assert md["user_api_key"] == "sk-x" + assert md["user_api_key_team_id"] == "team-1" # FAILS today: team_id is dropped + assert md["semantic-cache-embedding"] is True + + +def test_redis_init_defers_redisvl_construction(monkeypatch): + semantic_cache_mock = MagicMock() + custom_vectorizer_mock = MagicMock() + + with patch.dict( + "sys.modules", + { + "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), + "redisvl.utils.vectorize": MagicMock( + CustomTextVectorizer=custom_vectorizer_mock + ), + }, + ): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + monkeypatch.setenv("REDIS_HOST", "localhost") + monkeypatch.setenv("REDIS_PORT", "6379") + monkeypatch.setenv("REDIS_PASSWORD", "test_password") + + cache = RedisSemanticCache(similarity_threshold=0.8) + + semantic_cache_mock.assert_not_called() + custom_vectorizer_mock.assert_not_called() + + first = cache.llmcache + semantic_cache_mock.assert_called_once() + custom_vectorizer_mock.assert_called_once() + + second = cache.llmcache + assert first is second + semantic_cache_mock.assert_called_once() + + +def test_redis_failed_llmcache_build_is_not_memoized(monkeypatch): + built_cache = MagicMock() + semantic_cache_mock = MagicMock( + side_effect=[ConnectionError("redis down"), built_cache] + ) + custom_vectorizer_mock = MagicMock() + + with patch.dict( + "sys.modules", + { + "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), + "redisvl.utils.vectorize": MagicMock( + CustomTextVectorizer=custom_vectorizer_mock + ), + }, + ): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + monkeypatch.setenv("REDIS_HOST", "localhost") + monkeypatch.setenv("REDIS_PORT", "6379") + monkeypatch.setenv("REDIS_PASSWORD", "test_password") + + cache = RedisSemanticCache(similarity_threshold=0.8) + + with pytest.raises(ConnectionError, match="redis down"): + _ = cache.llmcache + + assert cache.llmcache is built_cache + assert semantic_cache_mock.call_count == 2 + + +def test_redis_llmcache_setter_supported(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + cache = RedisSemanticCache.__new__(RedisSemanticCache) + sentinel = MagicMock() + cache.llmcache = sentinel + assert cache.llmcache is sentinel diff --git a/tests/test_litellm/google_genai/test_google_genai_main.py b/tests/test_litellm/google_genai/test_google_genai_main.py index 5854e4b55af..8f56b4e4bc0 100644 --- a/tests/test_litellm/google_genai/test_google_genai_main.py +++ b/tests/test_litellm/google_genai/test_google_genai_main.py @@ -2,6 +2,7 @@ """ Test to verify the Google GenAI generate_content adapter functionality """ + import json import os import sys @@ -42,4 +43,228 @@ async def test_agenerate_content_stream(): stream=True, ) mock_post.assert_called_once() - mock_post.call_args.kwargs["stream"] == True + assert mock_post.call_args.kwargs["stream"] is True + + +def _mock_gemini_post_response(): + """A minimal stand-in for a successful Gemini generateContent HTTP response.""" + from unittest.mock import MagicMock + + resp = MagicMock() + resp.status_code = 200 + resp.headers = {} + resp.json.return_value = { + "candidates": [ + { + "content": {"parts": [{"text": "hi"}], "role": "model"}, + "finishReason": "STOP", + } + ] + } + return resp + + +NATIVE_TOP_LEVEL_FIELD_CASES = [ + ( + "safetySettings", + [{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"}], + ), + ("toolConfig", {"functionCallingConfig": {"mode": "AUTO"}}), + ("cachedContent", "cachedContents/abc123"), + ("labels", {"team": "search"}), +] + + +@pytest.mark.parametrize("field_name, field_value", NATIVE_TOP_LEVEL_FIELD_CASES) +def test_native_top_level_field_forwarded_to_request_body(field_name, field_value): + """ + Regression for https://github.com/BerriAI/litellm/issues/12671 + + Google's native generateContent body carries fields like safetySettings at the + top level (siblings of generationConfig). The proxy spreads them as loose kwargs + into generate_content. They must reach Google's request body at the top level and + must NOT be silently dropped nor nested under generationConfig. + """ + from unittest.mock import patch + + from litellm.google_genai.main import generate_content + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + with patch.object( + HTTPHandler, "post", return_value=_mock_gemini_post_response() + ) as mock_post: + generate_content( + model="gemini/gemini-2.0-flash", + contents=[{"role": "user", "parts": [{"text": "Say hi"}]}], + custom_llm_provider="gemini", + api_key="test-key", + **{field_name: field_value}, + ) + + assert mock_post.called, "expected the request to reach the HTTP client" + body = mock_post.call_args.kwargs["json"] + assert ( + body[field_name] == field_value + ), f"{field_name} should be forwarded to Google at the top level" + assert field_name not in body.get("generationConfig", {}), ( + f"{field_name} must be a top-level sibling of generationConfig, " + "not nested inside it" + ) + + +@pytest.mark.asyncio +async def test_native_safety_settings_forwarded_async(): + """The async path (used by the proxy's :generateContent route) must also forward + native top-level fields.""" + from unittest.mock import AsyncMock, patch + + from litellm.google_genai.main import agenerate_content + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + safety_settings = [ + {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"} + ] + + with patch.object( + AsyncHTTPHandler, + "post", + new_callable=AsyncMock, + return_value=_mock_gemini_post_response(), + ) as mock_post: + await agenerate_content( + model="gemini/gemini-2.0-flash", + contents=[{"role": "user", "parts": [{"text": "Say hi"}]}], + custom_llm_provider="gemini", + api_key="test-key", + safetySettings=safety_settings, + ) + + assert mock_post.called + body = mock_post.call_args.kwargs["json"] + assert body["safetySettings"] == safety_settings + assert "safetySettings" not in body.get("generationConfig", {}) + + +def test_native_fields_coexist_with_generation_config(): + """Forwarding native top-level fields must not regress the already-working + generationConfig path; both must land in their correct positions.""" + from unittest.mock import patch + + from litellm.google_genai.main import generate_content + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + safety_settings = [ + {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"} + ] + + with patch.object( + HTTPHandler, "post", return_value=_mock_gemini_post_response() + ) as mock_post: + generate_content( + model="gemini/gemini-2.0-flash", + contents=[{"role": "user", "parts": [{"text": "Say hi"}]}], + custom_llm_provider="gemini", + api_key="test-key", + safetySettings=safety_settings, + generationConfig={"temperature": 0, "responseMimeType": "application/json"}, + ) + + body = mock_post.call_args.kwargs["json"] + assert body["safetySettings"] == safety_settings + generation_config = body["generationConfig"] + assert generation_config["temperature"] == 0 + assert generation_config["responseMimeType"] == "application/json" + + +def test_explicit_extra_body_overrides_native_top_level_field(): + """An explicit extra_body value takes precedence over the same top-level field.""" + from unittest.mock import patch + + from litellm.google_genai.main import generate_content + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + native = [{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"}] + override = [ + {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_ONLY_HIGH"} + ] + + with patch.object( + HTTPHandler, "post", return_value=_mock_gemini_post_response() + ) as mock_post: + generate_content( + model="gemini/gemini-2.0-flash", + contents=[{"role": "user", "parts": [{"text": "Say hi"}]}], + custom_llm_provider="gemini", + api_key="test-key", + safetySettings=native, + extra_body={"safetySettings": override}, + ) + + body = mock_post.call_args.kwargs["json"] + assert body["safetySettings"] == override + + +def test_native_fields_and_system_instruction_forwarded_on_sync_stream(): + """The sync streaming path (generate_content_stream) must forward native top-level + fields AND systemInstruction. The PR changed the merge here and newly added the + systemInstruction kwarg; without coverage a regression on either ships green.""" + from unittest.mock import patch + + from litellm.google_genai.main import generate_content_stream + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + safety_settings = [ + {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"} + ] + system_instruction = {"parts": [{"text": "Be terse"}]} + + with patch.object( + HTTPHandler, "post", return_value=_mock_gemini_post_response() + ) as mock_post: + generate_content_stream( + model="gemini/gemini-2.0-flash", + contents=[{"role": "user", "parts": [{"text": "Say hi"}]}], + custom_llm_provider="gemini", + api_key="test-key", + safetySettings=safety_settings, + systemInstruction=system_instruction, + ) + + assert mock_post.called + body = mock_post.call_args.kwargs["json"] + assert body["safetySettings"] == safety_settings + assert body["systemInstruction"] == system_instruction + assert "safetySettings" not in body.get("generationConfig", {}) + + +@pytest.mark.asyncio +async def test_native_fields_forwarded_on_async_stream(): + """The async streaming path (agenerate_content_stream) backs the proxy's + :streamGenerateContent route and must forward native top-level fields too.""" + from unittest.mock import AsyncMock, patch + + from litellm.google_genai.main import agenerate_content_stream + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + safety_settings = [ + {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"} + ] + + with patch.object( + AsyncHTTPHandler, + "post", + new_callable=AsyncMock, + return_value=_mock_gemini_post_response(), + ) as mock_post: + await agenerate_content_stream( + model="gemini/gemini-2.0-flash", + contents=[{"role": "user", "parts": [{"text": "Say hi"}]}], + custom_llm_provider="gemini", + api_key="test-key", + safetySettings=safety_settings, + ) + + assert mock_post.called + body = mock_post.call_args.kwargs["json"] + assert body["safetySettings"] == safety_settings + assert "safetySettings" not in body.get("generationConfig", {}) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 7f3d5a959a1..a5d1934237f 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -951,6 +951,7 @@ def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details(): "character_count": 0, "image_count": 0, "video_length_seconds": 0.0, + "audio_length_seconds": 0.0, } model_info: ModelInfo = {} diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index f0db0409bd7..5a8ac0313a7 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -11,7 +11,9 @@ sys.path.insert( import time +import litellm from litellm.constants import SENTRY_DENYLIST, SENTRY_PII_DENYLIST +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.litellm_core_utils.litellm_logging import set_callbacks from litellm.types.utils import ModelResponse, TextCompletionResponse @@ -3408,6 +3410,140 @@ def test_handle_anthropic_messages_response_logging_degrades_on_unparseable_resp assert result.usage.prompt_tokens == 4 # type: ignore[attr-defined] +class _SuccessCapturingLogger(CustomLogger): + """Records the success payload. success_payload is populated only in + async_log_success_event, so it stays None when the buggy no-op + async_log_stream_event path runs for streaming.""" + + def __init__(self): + super().__init__() + self.success_payload = None + self.success_calls = 0 + self.stream_event_calls = 0 + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.success_calls += 1 + self.success_payload = kwargs.get("standard_logging_object") + + async def async_log_stream_event(self, kwargs, response_obj, start_time, end_time): + self.stream_event_calls += 1 + + +def _responses_stream_sse_bytes(): + """A full Responses stream: an opened message item, two text deltas, then the + terminal response.completed carrying usage. Exercises mid-stream delta handling + in addition to end-of-stream success logging.""" + import json + + events = [ + { + "type": "response.output_item.added", + "output_index": 0, + "item": { + "id": "msg-1", + "type": "message", + "status": "in_progress", + "role": "assistant", + "content": [], + }, + }, + { + "type": "response.output_text.delta", + "item_id": "msg-1", + "output_index": 0, + "content_index": 0, + "delta": "hello ", + "sequence_number": 2, + }, + { + "type": "response.output_text.delta", + "item_id": "msg-1", + "output_index": 0, + "content_index": 0, + "delta": "world", + "sequence_number": 3, + }, + { + "type": "response.completed", + "sequence_number": 4, + "response": _responses_api_response_with_text("hello world").model_dump(), + }, + ] + return [f"data: {json.dumps(e)}\n\n".encode("utf-8") for e in events] + + +def _fake_streaming_responses_http_response(): + sse_chunks = _responses_stream_sse_bytes() + + async def aiter_bytes(*args, **kwargs): + for chunk in sse_chunks: + yield chunk + + resp = MagicMock() + resp.status_code = 200 + resp.headers = {} + resp.aiter_bytes = aiter_bytes + return resp + + +def _chunk_text(chunk): + if isinstance(chunk, (bytes, bytearray)): + return chunk.decode("utf-8", "ignore") + return str(chunk) + + +async def _drain_until_logged(logger, max_iter=30): + for _ in range(max_iter): + if logger.success_payload is not None: + break + await asyncio.sleep(0.1) + + +@pytest.mark.asyncio +async def test_streaming_anthropic_messages_openai_bridge_fires_success_logging( + monkeypatch, +): + """Regression for #28595 / #28943. The existing tests above call + _handle_anthropic_messages_response_logging directly; they do not cover the + streaming wiring that originally broke. Drive a real streaming + anthropic_messages call routed to the OpenAI Responses backend (upstream SSE + mocked) and assert the bridge surfaces delta chunks and fires success logging + exactly once with real cost. On the broken version the stream ran but only the + no-op async_log_stream_event was called, so success_payload stayed None and the + SpendLogs row never landed.""" + logger = _SuccessCapturingLogger() + monkeypatch.setattr(litellm, "callbacks", [logger]) + + chunks = [] + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new=AsyncMock(return_value=_fake_streaming_responses_http_response()), + ): + stream = await litellm.anthropic_messages( + model="openai/gpt-4o", + api_key="sk-test-28595", + messages=[{"role": "user", "content": "ping"}], + max_tokens=16, + stream=True, + ) + async for chunk in stream: # logging fires on stream end; must drain fully + chunks.append(chunk) + + await _drain_until_logged(logger) + + assert chunks, "stream yielded no chunks" + assert any("content_block_delta" in _chunk_text(c) for c in chunks), ( + "no delta chunks surfaced; the streaming text deltas were not forwarded" + ) + assert logger.success_payload is not None, ( + "async_log_success_event never fired for streaming /v1/messages -> openai " + "Responses bridge; the no-op stream path dropped the spend row" + ) + assert logger.success_calls == 1, "bridge call must log success exactly once" + assert logger.success_payload["response_cost"] > 0 + assert logger.success_payload["call_type"] == "anthropic_messages" + + def test_failure_handler_records_recovered_partial_spend(logging_obj): """A stream interrupted mid-flight still billed the provider for the chunks already delivered. When the router stashes that recovered usage as diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py new file mode 100644 index 00000000000..453c7490d98 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py @@ -0,0 +1,321 @@ +""" +Regression tests for the Anthropic message_start cursor=1 bug in +ChunkProcessor._calculate_usage_per_chunk. + +Background +---------- +Anthropic streams a `message_start` event that carries +`usage.output_tokens=1` as a placeholder ("cursor"). The real cumulative +output count only arrives in the final `message_delta` event. When a +stream is cancelled before `message_delta` lands (very common for +thinking models on long-tail prompts), the last-wins accumulator in +ChunkProcessor leaves completion_tokens stuck at 1. Because 1 is +truthy, the `completion_tokens or token_counter(text=...)` fallback in +calculate_usage() never fires, and the request is billed for 1 output +token even when several thousand tokens of text were actually streamed. + +These tests pin the post-fix behavior: completion_tokens should reset +to 0 when the only update we saw was the cursor, allowing the +text-based fallback to estimate from the real completion text. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor +from litellm.types.utils import ( + Delta, + ModelResponseStream, + StreamingChoices, + Usage, +) + + +def _make_chunk( + *, + content: str = "", + usage: Usage = None, + finish_reason: str = None, + custom_llm_provider: str = "anthropic", +) -> ModelResponseStream: + chunk = ModelResponseStream( + id="msg_test", + created=1738900000, + model="claude-sonnet-4-6", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=finish_reason, + index=0, + delta=Delta(content=content, role="assistant"), + ) + ], + usage=usage, + ) + # The cursor reset is now gated on provider; populate the same field the + # real streaming_handler sets (see litellm/litellm_core_utils/streaming_handler.py). + chunk._hidden_params = {"custom_llm_provider": custom_llm_provider} + return chunk + + +class TestAnthropicCursorBug: + """The core regression: completion_tokens=1 cursor must not leak through.""" + + def test_only_message_start_cursor_resets_completion_to_zero(self): + """ + Stream cancelled before message_delta — only the message_start cursor + (output_tokens=1) was seen. Per-chunk accumulator must reset to 0 so + token_counter fallback can estimate from completion text. + """ + # Anthropic message_start: input_tokens accurate, output_tokens=1 cursor + message_start = _make_chunk( + usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025) + ) + # Several content_block_delta chunks (no usage attached) + text_chunks = [ + _make_chunk(content="Hello"), + _make_chunk(content=" world"), + _make_chunk(content=" this is partial."), + ] + chunks = [message_start, *text_chunks] + + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + + assert result["prompt_tokens"] == 1024 + # The cursor value of 1 must NOT leak through — should be reset to 0 + # so the text-based fallback estimates the real completion length. + assert result["completion_tokens"] == 0, ( + "completion_tokens=1 from message_start cursor leaked through. " + "Should reset to 0 when only cursor was seen, so token_counter " + "fallback in calculate_usage() can estimate from completion text." + ) + + def test_message_start_plus_message_delta_uses_delta_value(self): + """ + Normal complete stream: message_start cursor=1, then message_delta=3847. + Last-wins must give 3847 (the real value). + """ + message_start = _make_chunk( + usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025) + ) + text_chunks = [_make_chunk(content=t) for t in ["Hello", " world", "!"]] + # message_delta with the real cumulative output_tokens + message_delta = _make_chunk( + usage=Usage(prompt_tokens=1024, completion_tokens=3847, total_tokens=4871), + finish_reason="stop", + ) + chunks = [message_start, *text_chunks, message_delta] + + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + + assert result["prompt_tokens"] == 1024 + assert result["completion_tokens"] == 3847 + + def test_calculate_usage_falls_back_to_token_counter_for_cursor_only(self): + """ + End-to-end via calculate_usage(): cursor-only stream + real completion + text should produce a token-counter estimate, NOT 1. + """ + message_start = _make_chunk( + usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025) + ) + # ~50 visible chars ≈ ~12 tokens (anthropic-style tokenizer ballpark) + text_chunks = [ + _make_chunk(content="Based on your question, I think the answer is "), + _make_chunk(content="forty-two. Here is my reasoning: "), + ] + chunks = [message_start, *text_chunks] + completion_output = ( + "Based on your question, I think the answer is forty-two. " + "Here is my reasoning: " + ) + + processor = ChunkProcessor(chunks=chunks, messages=[]) + usage = processor.calculate_usage( + chunks=chunks, + model="claude-sonnet-4-6", + completion_output=completion_output, + messages=[], + ) + + # Should be a token_counter estimate of the text, not the cursor 1 + assert usage.completion_tokens > 1, ( + f"Expected token_counter estimate of completion text, got " + f"completion_tokens={usage.completion_tokens} (likely stuck at cursor)" + ) + + def test_cache_fields_preserved_from_message_start(self): + """cache_read / cache_creation come from message_start and must survive.""" + message_start_usage = Usage( + prompt_tokens=1024, completion_tokens=1, total_tokens=1025 + ) + # Anthropic puts these in message_start + message_start_usage.cache_read_input_tokens = 512 + message_start_usage.cache_creation_input_tokens = 128 + message_start = _make_chunk(usage=message_start_usage) + + chunks = [message_start, _make_chunk(content="hi")] + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + + assert result["cache_read_input_tokens"] == 512 + assert result["cache_creation_input_tokens"] == 128 + + def test_openai_streaming_unaffected(self): + """ + OpenAI's only usage chunk is the penultimate one (with + stream_options.include_usage=true), and it carries the real value + directly. Our cursor fix must not break this path — output > 1 + means saw_non_cursor_completion=True so no reset happens. + """ + # Simulate OpenAI: content chunks first, then ONE usage chunk at the end + text_chunks = [_make_chunk(content=t) for t in ["The", " answer", " is 42"]] + usage_chunk = _make_chunk( + usage=Usage(prompt_tokens=42, completion_tokens=15, total_tokens=57), + finish_reason="stop", + ) + chunks = [*text_chunks, usage_chunk] + + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + + assert result["prompt_tokens"] == 42 + assert result["completion_tokens"] == 15 + + def test_single_token_completion_legitimate_case(self): + """ + Edge case: a stream that legitimately completes with output_tokens=1 + (e.g., model returns just "Yes."). Without saw_non_cursor_completion + we'd reset to 0 and fall through to token_counter — but token_counter + on a 1-token string also gives ~1, so billing is still approximately + correct. This test pins that the result is sane (1 or 0). + """ + message_start = _make_chunk( + usage=Usage(prompt_tokens=20, completion_tokens=1, total_tokens=21) + ) + text_chunk = _make_chunk(content="Yes.") + # Anthropic's message_delta also gives output_tokens=1 in this case + message_delta = _make_chunk( + usage=Usage(prompt_tokens=20, completion_tokens=1, total_tokens=21), + finish_reason="stop", + ) + chunks = [message_start, text_chunk, message_delta] + + processor = ChunkProcessor(chunks=chunks, messages=[]) + usage = processor.calculate_usage( + chunks=chunks, + model="claude-sonnet-4-6", + completion_output="Yes.", + messages=[], + ) + + # Two completion-bearing usage events (message_start AND message_delta + # both with output_tokens=1) is positive evidence that message_delta + # arrived — saw_non_cursor_completion goes True via the count >= 2 + # branch and the reset is suppressed. Result: completion_tokens stays + # at the legitimate value of 1. + assert usage.completion_tokens == 1, ( + f"Legitimate single-token completion should bill exactly 1 token " + f"(message_start + message_delta both saw output_tokens=1, " + f"confirming message_delta arrived), got {usage.completion_tokens}" + ) + + def test_anthropic_cache_only_chunks_after_message_start_still_resets(self): + """ + Cache-only chunks (cache_read_input_tokens > 0 but completion_tokens=0) + following message_start should not be mistaken for completion progress. + The cursor=1 from message_start stays the only completion update; reset + must fire so token_counter estimates from completion text instead of + billing the placeholder. + """ + message_start_usage = Usage( + prompt_tokens=1024, completion_tokens=1, total_tokens=1025 + ) + message_start_usage.cache_read_input_tokens = 4096 + message_start = _make_chunk(usage=message_start_usage) + # Subsequent chunks with cache fields but no completion_tokens + cache_chunk_usage = Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0) + cache_chunk_usage.cache_read_input_tokens = 4096 + cache_chunk = _make_chunk(content="partial", usage=cache_chunk_usage) + # No message_delta — stream was cancelled + chunks = [message_start, cache_chunk] + + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + + assert result["cache_read_input_tokens"] == 4096 + assert result["completion_tokens"] == 0, ( + "cache chunks alone don't count as completion progress — only " + "completion_tokens > 0 in a usage event proves real output happened. " + "Reset to 0 forces token_counter fallback." + ) + + +class TestProviderGuard: + """Class A: the cursor-reset heuristic must NOT silently affect non-Anthropic + providers, even if they happen to report completion_tokens=1.""" + + def test_non_anthropic_provider_completion_tokens_one_not_reset(self): + """ + Some non-Anthropic provider legitimately reports completion_tokens=1 + in its single usage chunk. Without the provider guard the cursor + heuristic would silently reset it to 0 and bill via token_counter, + producing a different (often inflated) number than what the provider + actually charged. + """ + chunks = [ + _make_chunk( + usage=Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11), + finish_reason="stop", + custom_llm_provider="openai", + ), + ] + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + assert result["completion_tokens"] == 1, ( + "Non-Anthropic providers must not be subject to the message_start " + "cursor reset — their completion_tokens=1 is the real value." + ) + + def test_unknown_provider_completion_tokens_one_not_reset(self): + """No custom_llm_provider on hidden_params (older path or custom + plugin) — heuristic must not fire.""" + chunk = _make_chunk( + usage=Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11), + ) + # Explicitly clear hidden_params to simulate the unknown-provider case + chunk._hidden_params = {} + processor = ChunkProcessor(chunks=[chunk], messages=[]) + result = processor._calculate_usage_per_chunk(chunks=[chunk]) + assert result["completion_tokens"] == 1 + + +class TestNonAnthropicStreamingIntact: + """Make sure providers without cursor pattern still work.""" + + def test_completion_tokens_above_one_never_resets(self): + """Any chunk reporting completion_tokens > 1 sets saw_non_cursor + and prevents the reset.""" + chunks = [ + _make_chunk( + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + ), + ] + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + assert result["completion_tokens"] == 5 + + def test_no_usage_chunks_leaves_zero(self): + """Stream with zero usage info → completion_tokens stays 0 + (token_counter fallback will handle it).""" + chunks = [_make_chunk(content="hi"), _make_chunk(content=" there")] + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + assert result["prompt_tokens"] == 0 + assert result["completion_tokens"] == 0 diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index f93f2404fa4..45820b9833f 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -2175,6 +2175,283 @@ def test_translate_openai_response_to_anthropic_cache_tokens_from_prompt_tokens_ assert anthropic_response["usage"]["cache_read_input_tokens"] == 30 +def test_translate_openai_usage_to_anthropic_cache_tokens_from_dict_details_with_integral_floats(): + usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + ) + usage.prompt_tokens_details = { + "cached_tokens": 30.0, + "cache_write_tokens": 20.0, + } + + anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta( + usage + ) + + assert anthropic_usage["input_tokens"] == 70 + assert anthropic_usage["output_tokens"] == 50 + assert anthropic_usage["cache_read_input_tokens"] == 30 + assert anthropic_usage["cache_creation_input_tokens"] == 20 + + +def test_translate_openai_usage_to_anthropic_ignores_fractional_cache_tokens(): + usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + ) + usage.prompt_tokens_details = { + "cached_tokens": 30.5, + "cache_creation_tokens": 20.25, + } + + anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta( + usage + ) + + assert anthropic_usage["input_tokens"] == 120 + assert anthropic_usage["output_tokens"] == 50 + assert "cache_read_input_tokens" not in anthropic_usage + assert "cache_creation_input_tokens" not in anthropic_usage + + +def test_translate_openai_usage_to_anthropic_ignores_bool_cache_tokens(): + usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + ) + usage.cache_read_input_tokens = True + usage.cache_creation_input_tokens = True + + anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta( + usage + ) + + assert anthropic_usage["input_tokens"] == 120 + assert anthropic_usage["output_tokens"] == 50 + assert "cache_read_input_tokens" not in anthropic_usage + assert "cache_creation_input_tokens" not in anthropic_usage + + +def test_translate_openai_response_to_anthropic_cache_creation_from_prompt_tokens_details(): + from litellm.types.utils import PromptTokensDetailsWrapper + + usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=30, + cache_creation_tokens=20, + ), + ) + + response = ModelResponse( + id="test-id", + choices=[ + Choices( + index=0, + finish_reason="stop", + message=Message( + role="assistant", + content="Test response", + ), + ) + ], + model="gpt-4o-2024-08-06", + usage=usage, + ) + + adapter = LiteLLMAnthropicMessagesAdapter() + anthropic_response = adapter.translate_openai_response_to_anthropic( + response=response, + tool_name_mapping=None, + ) + + assert anthropic_response["usage"]["input_tokens"] == 70 + assert anthropic_response["usage"]["output_tokens"] == 50 + assert anthropic_response["usage"]["cache_read_input_tokens"] == 30 + assert anthropic_response["usage"]["cache_creation_input_tokens"] == 20 + + +def test_translate_openai_response_to_anthropic_cache_tokens_from_usage_fields(): + usage = Usage(prompt_tokens=120, completion_tokens=50, total_tokens=170) + usage.cache_read_input_tokens = 30 + usage.cache_creation_input_tokens = 20 + + response = ModelResponse( + id="test-id", + choices=[ + Choices( + index=0, + finish_reason="stop", + message=Message( + role="assistant", + content="Test response", + ), + ) + ], + model="claude-3-sonnet-20240229", + usage=usage, + ) + + adapter = LiteLLMAnthropicMessagesAdapter() + anthropic_response = adapter.translate_openai_response_to_anthropic( + response=response, + tool_name_mapping=None, + ) + + assert anthropic_response["usage"]["input_tokens"] == 70 + assert anthropic_response["usage"]["output_tokens"] == 50 + assert anthropic_response["usage"]["cache_read_input_tokens"] == 30 + assert anthropic_response["usage"]["cache_creation_input_tokens"] == 20 + + +def test_translate_openai_response_to_anthropic_cache_tokens_from_private_usage_fields(): + usage = Usage(prompt_tokens=120, completion_tokens=50, total_tokens=170) + + response = ModelResponse( + id="test-id", + choices=[ + Choices( + index=0, + finish_reason="stop", + message=Message( + role="assistant", + content="Test response", + ), + ) + ], + model="claude-3-sonnet-20240229", + usage=usage, + ) + response.usage._cache_read_input_tokens = 30 + response.usage._cache_creation_input_tokens = 20 + + adapter = LiteLLMAnthropicMessagesAdapter() + anthropic_response = adapter.translate_openai_response_to_anthropic( + response=response, + tool_name_mapping=None, + ) + + assert anthropic_response["usage"]["input_tokens"] == 70 + assert anthropic_response["usage"]["output_tokens"] == 50 + assert anthropic_response["usage"]["cache_read_input_tokens"] == 30 + assert anthropic_response["usage"]["cache_creation_input_tokens"] == 20 + + +def test_translate_streaming_openai_response_to_anthropic_cache_tokens_from_prompt_tokens_details(): + from litellm.types.utils import PromptTokensDetailsWrapper + + usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=30, + cache_creation_tokens=20, + ), + ) + response = ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta(), + finish_reason="stop", + ) + ], + usage=usage, + ) + + adapter = LiteLLMAnthropicMessagesAdapter() + message_delta = adapter.translate_streaming_openai_response_to_anthropic( + response=response, + current_content_block_index=0, + ) + + assert message_delta["usage"]["input_tokens"] == 70 + assert message_delta["usage"]["output_tokens"] == 50 + assert message_delta["usage"]["cache_read_input_tokens"] == 30 + assert message_delta["usage"]["cache_creation_input_tokens"] == 20 + + +def test_translate_streaming_openai_response_to_anthropic_cache_tokens_from_hidden_params_usage(): + from litellm.types.utils import PromptTokensDetailsWrapper + + usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=30, + cache_creation_tokens=20, + ), + ) + response = ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta(), + finish_reason="stop", + ) + ], + ) + response._hidden_params = {"usage": usage} + + adapter = LiteLLMAnthropicMessagesAdapter() + message_delta = adapter.translate_streaming_openai_response_to_anthropic( + response=response, + current_content_block_index=0, + ) + + assert message_delta["usage"]["input_tokens"] == 70 + assert message_delta["usage"]["output_tokens"] == 50 + assert message_delta["usage"]["cache_read_input_tokens"] == 30 + assert message_delta["usage"]["cache_creation_input_tokens"] == 20 + + +def test_translate_streaming_openai_response_to_anthropic_cache_tokens_with_applied_edits(): + from litellm.types.utils import PromptTokensDetailsWrapper + + usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=30, + cache_creation_tokens=20, + ), + ) + response = ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta(), + finish_reason="stop", + ) + ], + usage=usage, + ) + + adapter = LiteLLMAnthropicMessagesAdapter() + message_delta = adapter.translate_streaming_openai_response_to_anthropic( + response=response, + current_content_block_index=0, + applied_edits=[{"type": "compact_20260112"}], + ) + + assert message_delta["usage"]["input_tokens"] == 70 + assert message_delta["usage"]["output_tokens"] == 50 + assert message_delta["usage"]["cache_read_input_tokens"] == 30 + assert message_delta["usage"]["cache_creation_input_tokens"] == 20 + assert message_delta["context_management"]["applied_edits"][0]["type"] == ( + "compact_20260112" + ) + + # ===================================================================== # Web Search Tool Transformation Tests # ===================================================================== diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py index f74c5b61300..d67de0dcaf8 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py @@ -24,6 +24,7 @@ from litellm.types.utils import ( Message, ModelResponse, ModelResponseStream, + PromptTokensDetailsWrapper, StreamingChoices, Usage, ) @@ -88,6 +89,59 @@ def test_fake_stream_usage_preserved(): assert message_delta["usage"]["input_tokens"] == 10 +def test_delayed_usage_chunk_preserves_cache_tokens(): + usage = Usage( + prompt_tokens=120, + completion_tokens=5, + total_tokens=125, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=30, + cache_creation_tokens=20, + ), + ) + chunks = [ + ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta(content="Two."), + finish_reason=None, + ) + ], + ), + ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta(), + finish_reason="stop", + ) + ], + ), + ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta(), + finish_reason=None, + ) + ], + usage=usage, + ), + ] + wrapper = AnthropicStreamWrapper(completion_stream=iter(chunks), model="gpt-4o") + events = list(wrapper) + + message_delta = next( + event for event in events if event.get("type") == "message_delta" + ) + + assert message_delta["usage"]["input_tokens"] == 70 + assert message_delta["usage"]["output_tokens"] == 5 + assert message_delta["usage"]["cache_read_input_tokens"] == 30 + assert message_delta["usage"]["cache_creation_input_tokens"] == 20 + + def test_splitter_passes_through_non_combined_chunks(): """A chunk with content but no finish_reason is not split.""" chunk = ModelResponseStream( diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py index 8bc39a6d85e..d9d9474d6f3 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py @@ -30,7 +30,9 @@ from litellm.types.utils import ( ChatCompletionDeltaToolCall, Delta, Function, + PromptTokensDetailsWrapper, StreamingChoices, + Usage, ) @@ -107,6 +109,34 @@ def _input_json_deltas(events: List[dict]) -> List[str]: ] +def test_held_stop_reason_usage_merge_preserves_openai_cache_token_details(): + """OpenAI-compatible usage chunks carry cache reads in prompt_tokens_details.""" + wrapper = AnthropicStreamWrapper(completion_stream=iter([]), model="claude-x") + wrapper.holding_stop_reason_chunk = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"input_tokens": 0, "output_tokens": 0}, + } + + usage_chunk = MagicMock() + usage_chunk.usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=30, + cache_creation_tokens=20, + ), + ) + + merged_chunk = wrapper._merge_usage_into_held_stop_reason_chunk(usage_chunk) + + assert merged_chunk["usage"]["input_tokens"] == 70 + assert merged_chunk["usage"]["output_tokens"] == 50 + assert merged_chunk["usage"]["cache_read_input_tokens"] == 30 + assert merged_chunk["usage"]["cache_creation_input_tokens"] == 20 + + def test_first_text_delta_after_tool_use_is_not_dropped_sync(): """A tool_use -> text transition (text resuming after a tool call) carries the resumed text's first token in the trigger chunk. Without the fix it was diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index 03b72c504b8..b6c055f4390 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -159,6 +159,59 @@ class TestGetAnthropicHeaders: assert "authorization" not in headers assert "anthropic-dangerous-direct-browser-access" not in headers + def test_custom_api_base_uses_bearer_header(self): + """Custom api_base and non-standard API key should produce Authorization: Bearer header when opted in.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + config = AnthropicModelInfo() + headers = config.get_anthropic_headers( + api_key="my-custom-ollama-token", + computer_tool_used=False, + prompt_caching_set=False, + pdf_used=False, + is_vertex_request=False, + api_base="https://ollama.com/", + use_bearer_for_custom_base=True, + ) + + assert headers["authorization"] == "Bearer my-custom-ollama-token" + assert "x-api-key" not in headers + + def test_custom_api_base_uses_bearer_header_already_starts_with_bearer(self): + """If the key already starts with Bearer and Bearer opt-in is enabled, use it directly.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + config = AnthropicModelInfo() + headers = config.get_anthropic_headers( + api_key="Bearer my-custom-ollama-token", + computer_tool_used=False, + prompt_caching_set=False, + pdf_used=False, + is_vertex_request=False, + api_base="https://ollama.com/", + use_bearer_for_custom_base=True, + ) + + assert headers["authorization"] == "Bearer my-custom-ollama-token" + assert "x-api-key" not in headers + + def test_custom_api_base_uses_x_api_key_when_standard_key(self): + """If the key is standard sk-ant- key, use x-api-key even with custom api_base.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + config = AnthropicModelInfo() + headers = config.get_anthropic_headers( + api_key=FAKE_REGULAR_KEY, + computer_tool_used=False, + prompt_caching_set=False, + pdf_used=False, + is_vertex_request=False, + api_base="https://ollama.com/", + ) + + assert headers["x-api-key"] == FAKE_REGULAR_KEY + assert "authorization" not in headers + def test_oauth_includes_standard_headers(self): """OAuth path should still include standard Anthropic headers.""" from litellm.llms.anthropic.common_utils import AnthropicModelInfo @@ -242,6 +295,46 @@ class TestValidateEnvironmentOAuth: assert updated_headers["x-api-key"] == FAKE_REGULAR_KEY assert "authorization" not in updated_headers + + def test_custom_api_base_via_param(self): + """validate_environment uses Bearer when use_bearer_for_custom_base is set in litellm_params.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + config = AnthropicModelInfo() + headers = {} + + updated_headers = config.validate_environment( + headers=headers, + model="claude-sonnet-4-5-20250929", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={"use_bearer_for_custom_base": True}, + api_key="custom-api-key", + api_base="https://custom-gateway.com", + ) + + assert updated_headers["authorization"] == "Bearer custom-api-key" + assert "x-api-key" not in updated_headers + + def test_custom_api_base_via_litellm_params(self): + """validate_environment uses Bearer when api_base and use_bearer_for_custom_base are in litellm_params.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + config = AnthropicModelInfo() + headers = {} + + updated_headers = config.validate_environment( + headers=headers, + model="claude-sonnet-4-5-20250929", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={"api_base": "https://custom-gateway.com", "use_bearer_for_custom_base": True}, + api_key="custom-api-key", + api_base=None, + ) + + assert updated_headers["authorization"] == "Bearer custom-api-key" + assert "x-api-key" not in updated_headers assert "anthropic-dangerous-direct-browser-access" not in updated_headers @@ -1004,6 +1097,20 @@ class TestGetAuthHeader: result = AnthropicModelInfo.get_auth_header() assert result == {"authorization": f"Bearer {FAKE_OAUTH_TOKEN}"} + def test_custom_api_base_get_auth_header_uses_bearer(self): + """Non-standard API key and custom api_base returns Bearer when use_bearer_for_custom_base=True.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + result = AnthropicModelInfo.get_auth_header(api_key="my-custom-key", api_base="https://custom-gateway.com", use_bearer_for_custom_base=True) + assert result == {"authorization": "Bearer my-custom-key"} + + def test_custom_api_base_get_auth_header_uses_x_api_key_when_standard(self): + """Standard sk-ant- key with custom api_base should still return x-api-key.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + result = AnthropicModelInfo.get_auth_header(api_key=FAKE_REGULAR_KEY, api_base="https://custom-gateway.com") + assert result == {"x-api-key": FAKE_REGULAR_KEY} + class TestGetApiBaseFallbackChain: """Tests for AnthropicModelInfo.get_api_base() fallback to ANTHROPIC_BASE_URL.""" diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py index 6298eeb25e9..6b869076044 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -263,3 +263,123 @@ def test_bundled_bedrock_opus_model_info_declares_output_config_effort_ceiling( model_info = GetModelCostMap.load_local_model_cost_map()[model] assert model_info["bedrock_output_config_effort_ceiling"] == expected_ceiling + + +def test_route_prefix_matched_as_path_segment_not_substring(): + """Route tokens like ``mantle/`` must match only at a path-segment boundary. + + The ``bedrock_mantle/`` provider prefix contains the substring ``mantle/``; + a substring match misroutes ``bedrock_mantle/openai.gpt-5.5`` to the Claude + Mythos mantle config, whose request transform strips ``mantle/`` and mangles + the body model into ``bedrock_openai.gpt-5.5``. These assertions fail under + the old substring matching and pass once matching is anchored to ``startswith`` + or a ``/`` boundary. + """ + # The bedrock_mantle/ provider prefix must NOT be read as the mantle/ route. + assert ( + BedrockModelInfo.get_bedrock_route("bedrock_mantle/openai.gpt-5.5") != "mantle" + ) + assert ( + BedrockModelInfo.get_bedrock_route("bedrock_mantle/openai.gpt-5.4") == "invoke" + ) + assert ( + BedrockModelInfo._explicit_mantle_route("bedrock_mantle/openai.gpt-5.5") + is False + ) + + # A genuine mantle route still resolves, via the startswith branch... + assert ( + BedrockModelInfo.get_bedrock_route("mantle/anthropic.claude-mythos-preview") + == "mantle" + ) + # ...and via the mid-path "/mantle/" branch (after the bedrock/ provider prefix). + assert ( + BedrockModelInfo.get_bedrock_route( + "bedrock/mantle/anthropic.claude-mythos-preview" + ) + == "mantle" + ) + + +def test_model_has_route_prefix_exercises_both_branches(): + """``_model_has_route_prefix`` matches on ``startswith`` or a ``/`` boundary only.""" + # startswith branch + assert ( + BedrockModelInfo._model_has_route_prefix( + "mantle/anthropic.claude-mythos-preview", "mantle/" + ) + is True + ) + # f"/{prefix}" boundary branch + assert ( + BedrockModelInfo._model_has_route_prefix( + "bedrock/mantle/anthropic.claude-mythos-preview", "mantle/" + ) + is True + ) + # neither branch: the token only appears glued to another segment + assert ( + BedrockModelInfo._model_has_route_prefix( + "bedrock_mantle/openai.gpt-5.5", "mantle/" + ) + is False + ) + + +@pytest.mark.parametrize( + "route_method, token", + [ + (BedrockModelInfo._explicit_converse_route, "converse"), + (BedrockModelInfo._explicit_converse_like_route, "converse_like"), + (BedrockModelInfo._explicit_invoke_route, "invoke"), + (BedrockModelInfo._explicit_async_invoke_route, "async_invoke"), + (BedrockModelInfo._explicit_agent_route, "agent"), + (BedrockModelInfo._explicit_agentcore_route, "agentcore"), + (BedrockModelInfo._explicit_claude_platform_route, "claude_platform"), + (BedrockModelInfo._explicit_openai_route, "openai"), + ], + ids=[ + "converse", + "converse_like", + "invoke", + "async_invoke", + "agent", + "agentcore", + "claude_platform", + "openai", + ], +) +def test_explicit_route_helpers_match_token_only_as_path_segment(route_method, token): + """Each migrated ``_explicit_*_route`` matches its token only as a path segment. + + A leading segment (start of the id or right after a ``/``) matches; the token + glued onto a preceding segment does not. Reverting any method to the old + ``"/" in model`` substring check makes the non-segment case return True + and fails this test. + """ + # leading-segment forms match + assert route_method(f"{token}/some-model") is True + assert route_method(f"bedrock/{token}/some-model") is True + # the token only as a non-segment substring must not match + assert route_method(f"x{token}/y") is False + + +def test_explicit_invoke_route_does_not_match_async_invoke(): + """``invoke/`` must not substring-match ``async_invoke/`` models. + + This is the concrete improvement of the segment-boundary migration: the old + ``"invoke/" in model`` check wrongly classified async-invoke models as the + invoke route. + """ + async_invoke_model = "async_invoke/twelvelabs.marengo-embed-2-7-v1:0" + assert BedrockModelInfo._explicit_invoke_route(async_invoke_model) is False + assert ( + BedrockModelInfo._explicit_invoke_route(f"bedrock/{async_invoke_model}") + is False + ) + # ...while async_invoke/ is still detected as its own route. + assert BedrockModelInfo._explicit_async_invoke_route(async_invoke_model) is True + assert ( + BedrockModelInfo._explicit_async_invoke_route(f"bedrock/{async_invoke_model}") + is True + ) diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_context_circulation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_context_circulation.py index 6d913ad5d1d..1422531edbf 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_context_circulation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_context_circulation.py @@ -22,7 +22,6 @@ from litellm.llms.vertex_ai.gemini.transformation import ( ) from litellm.types.llms.vertex_ai import HttpxPartType - # --- Response extraction tests --- @@ -63,6 +62,7 @@ class TestExtractServerSideToolInvocations: assert result[0]["args"] == {"queries": ["weather Buenos Aires"]} assert result[0]["response"] == {"weather": "Sunny, 20°C"} assert result[0]["thought_signature"] == "sig_call_1" + assert result[0]["response_thought_signature"] == "sig_resp_1" def test_returns_none_when_no_server_side_tools(self): """No toolCall/toolResponse parts → returns None.""" @@ -145,6 +145,59 @@ class TestExtractServerSideToolInvocations: assert result[0]["id"] == "exec1" assert "response" not in result[0] + def test_extracts_tool_call_and_response_with_different_signatures(self): + """Case where toolCall and toolResponse have different signatures.""" + parts: List[HttpxPartType] = [ + { + "thoughtSignature": "sig_call_1", + "toolCall": { + "toolType": "GOOGLE_SEARCH_WEB", + "id": "abc123", + "args": {"queries": ["weather Buenos Aires"]}, + }, + }, + { + "thoughtSignature": "sig_resp_1", + "toolResponse": { + "toolType": "GOOGLE_SEARCH_WEB", + "id": "abc123", + "response": {"weather": "Sunny, 20°C"}, + }, + }, + ] + + result = VertexGeminiConfig._extract_server_side_tool_invocations(parts) + + assert result is not None + assert len(result) == 1 + assert result[0]["tool_type"] == "GOOGLE_SEARCH_WEB" + assert result[0]["id"] == "abc123" + assert result[0]["args"] == {"queries": ["weather Buenos Aires"]} + assert result[0]["response"] == {"weather": "Sunny, 20°C"} + assert result[0]["thought_signature"] == "sig_call_1" + assert result[0]["response_thought_signature"] == "sig_resp_1" + + def test_orphan_response_signature_extraction(self): + """Orphan toolResponse is captured and has response_thought_signature set.""" + parts: List[HttpxPartType] = [ + { + "thoughtSignature": "sig_resp_1", + "toolResponse": { + "toolType": "GOOGLE_SEARCH_WEB", + "id": "orphan123", + "response": {"result": "some response"}, + }, + }, + ] + + result = VertexGeminiConfig._extract_server_side_tool_invocations(parts) + + assert result is not None + assert len(result) == 1 + assert result[0]["id"] == "orphan123" + assert result[0]["response"] == {"result": "some response"} + assert result[0]["response_thought_signature"] == "sig_resp_1" + # --- Input re-injection tests --- @@ -200,6 +253,76 @@ class TestReInjectServerSideToolInvocations: "weather": "Sunny, 20°C" } + def test_roundtrip_single_invocation_with_different_signatures(self): + """Server-side invocations with different signatures for call and response.""" + messages = [ + {"role": "user", "content": "What's the weather?"}, + { + "role": "assistant", + "content": "It's sunny in Buenos Aires.", + "provider_specific_fields": { + "server_side_tool_invocations": [ + { + "tool_type": "GOOGLE_SEARCH_WEB", + "id": "abc123", + "args": {"queries": ["weather Buenos Aires"]}, + "response": {"weather": "Sunny, 20°C"}, + "thought_signature": "sig_call", + "response_thought_signature": "sig_resp", + } + ] + }, + }, + {"role": "user", "content": "Thanks!"}, + ] + + contents = _gemini_convert_messages_with_history(messages) + + model_turn = [c for c in contents if c["role"] == "model"] + assert len(model_turn) == 1 + + parts = model_turn[0]["parts"] + tool_call_parts = [p for p in parts if "toolCall" in p] + tool_response_parts = [p for p in parts if "toolResponse" in p] + + assert len(tool_call_parts) == 1 + assert tool_call_parts[0]["thoughtSignature"] == "sig_call" + + assert len(tool_response_parts) == 1 + assert tool_response_parts[0]["thoughtSignature"] == "sig_resp" + + def test_roundtrip_orphan_response_signature(self): + """Orphan response signature is preserved and re-injected into toolResponse part.""" + messages = [ + {"role": "user", "content": "What's the weather?"}, + { + "role": "assistant", + "content": "It's sunny in Buenos Aires.", + "provider_specific_fields": { + "server_side_tool_invocations": [ + { + "tool_type": "GOOGLE_SEARCH_WEB", + "id": "orphan123", + "response": {"result": "some response"}, + "response_thought_signature": "sig_orphan_resp", + } + ] + }, + }, + {"role": "user", "content": "Thanks!"}, + ] + + contents = _gemini_convert_messages_with_history(messages) + + model_turn = [c for c in contents if c["role"] == "model"] + assert len(model_turn) == 1 + + parts = model_turn[0]["parts"] + tool_response_parts = [p for p in parts if "toolResponse" in p] + + assert len(tool_response_parts) == 1 + assert tool_response_parts[0]["thoughtSignature"] == "sig_orphan_resp" + def test_no_invocations_no_extra_parts(self): """Without server_side_tool_invocations, no extra parts are added.""" messages = [ diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 1a2d0d86810..20e7aec377b 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2799,6 +2799,77 @@ def test_partial_json_chunk_on_first_chunk(): ), "Should switch to accumulated_json mode" +def test_accumulated_json_does_not_reparse_every_fragment(): + """ + Regression test for https://github.com/BerriAI/litellm/issues/26181 + + handle_accumulated_json_chunk used to call json.loads on the entire + accumulated buffer after EVERY fragment. For a large response fragmented + across many chunks that is O(n^2) work in a single GIL-holding C call, + which freezes the asyncio event loop for seconds and kills liveness probes. + + The buffer only becomes a complete JSON object on the final fragment, so a + correct implementation parses it ~once, not once per fragment. We assert the + full chunk still parses correctly AND that json.loads is not called on every + fragment (which is what made it quadratic). + """ + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + iterator = ModelResponseIterator( + streaming_response=MagicMock(), + sync_stream=True, + logging_obj=MagicMock(), + ) + iterator.chunk_type = "accumulated_json" + + text = "x" * 200_000 # no braces/brackets so only the final fragment closes + blob = json.dumps( + {"candidates": [{"content": {"role": "model", "parts": [{"text": text}]}}]} + ) + fragments = [blob[i : i + 4096] for i in range(0, len(blob), 4096)] + assert len(fragments) > 10, "need a multi-fragment payload to exercise the bug" + + parsed = None + with patch("json.loads", wraps=json.loads) as spy: + for fragment in fragments: + out = iterator.handle_accumulated_json_chunk(chunk=fragment) + if out is not None: + parsed = out + parse_calls = spy.call_count + + assert parsed is not None, "the reassembled chunk must still parse" + assert parsed.choices[0].delta.content == text, "content must be preserved intact" + + assert parse_calls <= 2, ( + f"json.loads was called {parse_calls} times for {len(fragments)} " + "fragments; the O(n^2) per-fragment re-parse has regressed" + ) + + +def test_accumulated_json_partial_fragment_returns_none_without_parsing(): + """A fragment that cannot complete the JSON must not trigger a json.loads + parse of the whole growing buffer (issue #26181).""" + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + iterator = ModelResponseIterator( + streaming_response=MagicMock(), + sync_stream=True, + logging_obj=MagicMock(), + ) + iterator.chunk_type = "accumulated_json" + + with patch("json.loads", wraps=json.loads) as spy: + result = iterator.handle_accumulated_json_chunk( + chunk='{"candidates": [{"content": {"parts": [{"text": "partial' + ) + assert result is None + assert spy.call_count == 0, "incomplete buffer should not be parsed" + + def test_google_ai_studio_presence_penalty_supported(): """ Test that presence_penalty is supported for Google AI Studio Gemini. diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index bb4e6c67e9e..86b3f0976ab 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -10,9 +10,11 @@ Covers: import pytest +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import ( _build_part_for_input, _is_multimodal_input, + process_embed_content_response, process_response, transform_openai_input_gemini_content, transform_openai_input_gemini_embed_content, @@ -72,7 +74,9 @@ class TestBuildPartForInput: assert part["file_data"]["file_uri"] == GCS_URL def test_file_reference_resolved(self): - resolved = {"files/abc": {"mime_type": "image/jpeg", "uri": "https://example.com/abc"}} + resolved = { + "files/abc": {"mime_type": "image/jpeg", "uri": "https://example.com/abc"} + } part = _build_part_for_input("files/abc", resolved_files=resolved) assert part["file_data"] is not None assert part["file_data"]["mime_type"] == "image/jpeg" @@ -94,7 +98,9 @@ class TestTransformOpenaiInputGeminiContent: def test_multiple_texts(self): result = transform_openai_input_gemini_content( - input=["hello", "world"], model="gemini-embedding-2-preview", optional_params={} + input=["hello", "world"], + model="gemini-embedding-2-preview", + optional_params={}, ) assert len(result["requests"]) == 2 assert result["requests"][0]["content"]["parts"][0]["text"] == "hello" @@ -109,7 +115,10 @@ class TestTransformOpenaiInputGeminiContent: ) assert len(result["requests"]) == 2 # First request is text - assert result["requests"][0]["content"]["parts"][0]["text"] == "The food was delicious" + assert ( + result["requests"][0]["content"]["parts"][0]["text"] + == "The food was delicious" + ) # Second request is image assert result["requests"][1]["content"]["parts"][0]["inline_data"] is not None @@ -288,3 +297,207 @@ class TestProcessResponse: model="gemini-embedding-2-preview", optional_params={}, ) + + +class TestProcessEmbedContentResponseUsage: + """Gemini Embedding 2 embedContent usageMetadata must drive spend. + + Regression for multimodal calls recording prompt_tokens=0 / spend=$0. + """ + + MODEL = "gemini-embedding-2" + + def test_multimodal_image_preserves_usage_metadata(self): + response_json = { + "embedding": {"values": [0.1, 0.2, 0.3]}, + "usageMetadata": { + "promptTokenCount": 258, + "totalTokenCount": 258, + "promptTokensDetails": [{"modality": "IMAGE", "tokenCount": 258}], + }, + } + result = process_embed_content_response( + input=[IMAGE_DATA_URI], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens == 258 + assert result.usage.total_tokens == 258 + assert result.usage.prompt_tokens_details.image_count == 1 + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + assert prompt_cost > 0 + + def test_text_modality_detail_populated(self): + response_json = { + "embedding": {"values": [0.1, 0.2]}, + "usageMetadata": { + "promptTokenCount": 12, + "totalTokenCount": 12, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12}], + }, + } + result = process_embed_content_response( + input="a short caption", + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens == 12 + assert result.usage.prompt_tokens_details.text_tokens == 12 + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + assert prompt_cost > 0 + + def test_video_modality_derives_seconds_and_text_floor(self): + response_json = { + "embedding": {"values": [0.1]}, + "usageMetadata": { + "promptTokenCount": 516, + "totalTokenCount": 516, + "promptTokensDetails": [{"modality": "VIDEO", "tokenCount": 516}], + }, + } + result = process_embed_content_response( + input=["gs://bucket/clip.mp4"], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens == 516 + assert result.usage.prompt_tokens_details.video_length_seconds == pytest.approx( + 2.0 + ) + assert result.usage.prompt_tokens_details.text_tokens == 1 + + def test_missing_usage_metadata_does_not_estimate_from_base64(self): + response_json = {"embedding": {"values": [0.1, 0.2]}} + result = process_embed_content_response( + input=[IMAGE_DATA_URI], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens == 0 + assert result.usage.total_tokens == 0 + + def test_missing_usage_metadata_text_falls_back_to_token_counter(self): + response_json = {"embedding": {"values": [0.1, 0.2]}} + result = process_embed_content_response( + input="hello world this is plain text", + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens > 0 + + def test_file_reference_image_billed_per_image_not_text(self): + """files/... image refs must bill per-image, not at the text token rate.""" + response_json = { + "embedding": {"values": [0.1, 0.2, 0.3]}, + "usageMetadata": { + "promptTokenCount": 258, + "totalTokenCount": 258, + "promptTokensDetails": [{"modality": "IMAGE", "tokenCount": 258}], + }, + } + result = process_embed_content_response( + input=["files/img123"], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + resolved_files={ + "files/img123": { + "mime_type": "image/png", + "uri": "https://example.com/img123", + } + }, + ) + assert result.usage.prompt_tokens_details.image_count == 1 + assert result.usage.prompt_tokens_details.text_tokens == 0 + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + assert prompt_cost == pytest.approx(0.00012) + + def test_file_reference_non_image_not_counted_as_image(self): + """A files/... ref resolving to a non-image mime must not be image-counted.""" + response_json = { + "embedding": {"values": [0.1, 0.2]}, + "usageMetadata": { + "promptTokenCount": 64, + "totalTokenCount": 64, + "promptTokensDetails": [{"modality": "AUDIO", "tokenCount": 64}], + }, + } + result = process_embed_content_response( + input=["files/clip1"], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + resolved_files={ + "files/clip1": { + "mime_type": "audio/mpeg", + "uri": "https://example.com/clip1", + } + }, + ) + assert result.usage.prompt_tokens_details.image_count == 0 + assert result.usage.prompt_tokens_details.audio_tokens == 64 + assert result.usage.prompt_tokens_details.audio_length_seconds == pytest.approx( + 2.0 + ) + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + assert prompt_cost == pytest.approx(2.0 * 0.00016) + + def test_video_plus_audio_does_not_double_bill_text(self): + """Video+audio responses must not get video tokens reassigned to text.""" + response_json = { + "embedding": {"values": [0.1]}, + "usageMetadata": { + "promptTokenCount": 580, + "totalTokenCount": 580, + "promptTokensDetails": [ + {"modality": "VIDEO", "tokenCount": 516}, + {"modality": "AUDIO", "tokenCount": 64}, + ], + }, + } + result = process_embed_content_response( + input=["gs://bucket/clip.mp4"], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens_details.text_tokens == 1 + assert result.usage.prompt_tokens_details.video_length_seconds == pytest.approx( + 2.0 + ) + assert result.usage.prompt_tokens_details.audio_length_seconds == pytest.approx( + 2.0 + ) + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + # 1 floor text token at 2e-7 + 2s of video at 7.9e-4 + 2s of audio at 1.6e-4 + assert prompt_cost == pytest.approx(1 * 2e-7 + 2 * 0.00079 + 2 * 0.00016) diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index 6f4bb4e59c2..d64b5c8d742 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -1,10 +1,11 @@ -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation import ( VertexAIPartnerModelsAnthropicMessagesConfig, ) +from litellm.llms.vertex_ai.vertex_ai_partner_models.main import VertexAIPartnerModels from litellm.types.router import GenericLiteLLMParams @@ -233,40 +234,36 @@ def test_both_compact_and_context_management_headers_added(): ), f"anthropic-beta should contain 'context-management-2025-06-27', got: {updated_headers['anthropic-beta']}" -def test_validate_environment_with_authorization_header_calculates_api_base(): - """Test that api_base is calculated even when Authorization header is already present""" +def test_validate_environment_always_refreshes_token_ignoring_stale_bearer(): + """Regression: stale Authorization in shared deployment extra_headers must not + skip token refresh on /v1/messages — _ensure_access_token is always called.""" config = VertexAIPartnerModelsAnthropicMessagesConfig() - # Simulate scenario where Authorization is already in headers (e.g., from cached extra_headers) - headers = {"Authorization": "Bearer existing-token"} + headers = {"Authorization": "Bearer EXPIRED"} litellm_params = { "vertex_project": "test-project", "vertex_location": "us-central1", - "extra_headers": {"anthropic-beta": "context-1m-2025-08-07"}, } - optional_params = {} - with patch.object( - config, "get_complete_vertex_url", return_value="https://mock-vertex-url" - ) as mock_get_url: + with ( + patch.object( + config, "_ensure_access_token", return_value=("fresh-token", "test-project") + ) as mock_ensure, + patch.object( + config, "get_complete_vertex_url", return_value="https://mock-vertex-url" + ), + ): updated_headers, api_base = config.validate_anthropic_messages_environment( headers=headers, model="claude-sonnet-4", messages=[], - optional_params=optional_params, + optional_params={}, litellm_params=litellm_params, api_base=None, ) - # Verify that api_base was calculated even though Authorization was already present - assert ( - api_base == "https://mock-vertex-url" - ), f"api_base should be calculated even with Authorization header. Got: {api_base}" - assert mock_get_url.called, "get_complete_vertex_url should be called" - - # Verify Authorization header is still present - assert ( - "Authorization" in updated_headers - ), "Authorization header should be preserved" + mock_ensure.assert_called_once() + assert updated_headers["Authorization"] == "Bearer fresh-token" + assert api_base == "https://mock-vertex-url" def test_transform_anthropic_messages_request_removes_scope_from_cache_control(): @@ -371,3 +368,77 @@ def test_provider_config_manager_reuses_vertex_anthropic_messages_config_instanc assert first_config is second_config finally: ProviderConfigManager._get_provider_anthropic_messages_config_cached.cache_clear() + + +def test_validate_environment_does_not_mutate_caller_headers(): + """Regression: beta headers (e.g. web-search) must not leak into the caller's + headers dict — which may be the shared deployment extra_headers object.""" + config = VertexAIPartnerModelsAnthropicMessagesConfig() + caller_headers: dict = {} + + with ( + patch.object( + config, "_ensure_access_token", return_value=("token", "test-project") + ), + patch.object( + config, "get_complete_vertex_url", return_value="https://mock-url" + ), + ): + config.validate_anthropic_messages_environment( + headers=caller_headers, + model="claude-sonnet-4", + messages=[], + optional_params={ + "tools": [{"type": "web_search_20250305", "name": "web_search"}] + }, + litellm_params={ + "vertex_ai_project": "p", + "vertex_ai_location": "us-central1", + }, + api_base=None, + ) + + assert ( + caller_headers == {} + ), "validate_anthropic_messages_environment must not mutate the caller's headers dict" + + +def test_vertex_claude_completion_does_not_mutate_shared_extra_headers(): + """Regression: router shallow-copies litellm_params so extra_headers is a shared + reference. Verify that the chat/completions path builds a new headers dict instead + of calling .update() on the shared object.""" + handler = VertexAIPartnerModels() + shared_extra_headers = {} # simulates deployment["litellm_params"]["extra_headers"] + + mock_response = MagicMock() + + with ( + patch.object( + handler, "_ensure_access_token", return_value=("ya29.fresh", "proj") + ), + patch.object( + handler, "get_complete_vertex_url", return_value="https://mock-url" + ), + patch( + "litellm.llms.anthropic.chat.AnthropicChatCompletion.completion", + return_value=mock_response, + ), + ): + handler.completion( + model="claude-haiku-4-5@20251001", + messages=[{"role": "user", "content": "hi"}], + model_response=MagicMock(), + print_verbose=lambda *a, **k: None, + encoding=None, + logging_obj=MagicMock(), + api_base=None, + optional_params={}, + custom_prompt_dict={}, + headers=shared_extra_headers, + timeout=30, + litellm_params={}, + ) + + assert ( + shared_extra_headers == {} + ), "extra_headers must not be mutated by completion()" diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 7e23e441f50..acad249a2bb 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -1,4 +1,4 @@ -"""Tests for the optional Rust-backed OCR path (``litellm/ocr/rust_bridge.py``).""" +"""Tests for the optional Rust-backed OCR path.""" import importlib import builtins @@ -15,7 +15,7 @@ from litellm.llms.base_llm.ocr.transformation import OCRResponse # function onto `litellm.ocr` and shadows the submodule, so import the modules # explicitly via importlib rather than attribute traversal. ocr_main = importlib.import_module("litellm.ocr.main") -rust_bridge = importlib.import_module("litellm.ocr.rust_bridge") +rust_bridge = importlib.import_module("litellm.rust_bridge.ocr") rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader") MODEL = "mistral/mistral-ocr-latest" @@ -49,7 +49,7 @@ class RecordingBridge: document: dict[str, object], api_key: str | None, api_base: str | None, - custom_llm_provider: str, + custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, @@ -81,7 +81,7 @@ class RecordingAsyncBridge: document: dict[str, object], api_key: str | None, api_base: str | None, - custom_llm_provider: str, + custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, @@ -108,7 +108,7 @@ class RaisingBridge: document: dict[str, object], api_key: str | None, api_base: str | None, - custom_llm_provider: str, + custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, @@ -123,7 +123,7 @@ class RaisingAsyncBridge: document: dict[str, object], api_key: str | None, api_base: str | None, - custom_llm_provider: str, + custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, @@ -366,17 +366,78 @@ def test_load_rust_ocr_uses_compiled_extension(monkeypatch): def test_timeout_to_seconds_handles_float_timeout_and_none(): - assert ocr_main._timeout_to_seconds(12.5) == 12.5 - assert ocr_main._timeout_to_seconds(None) is None - assert ocr_main._timeout_to_seconds(httpx.Timeout(30.0, read=42.0)) == 42.0 + assert rust_bridge._timeout_to_seconds(12.5) == 12.5 + assert rust_bridge._timeout_to_seconds(None) is None + assert rust_bridge._timeout_to_seconds(httpx.Timeout(30.0, read=42.0)) == 42.0 -def test_run_rust_ocr_forwards_args_and_wraps_response(): +def test_bridge_wrapper_forwards_prepared_args_and_wraps_response(): + bridge = RecordingBridge() + + litellm.use_litellm_rust(True, ocr=bridge) + response = rust_bridge.ocr( + model="mistral-ocr-latest", + document=DOCUMENT, + api_key="sk-test", + api_base="https://proxy.internal", + custom_llm_provider="mistral", + extra_headers={"Authorization": "Bearer sk-test", "x-trace-id": "trace-1"}, + optional_params={"include_image_base64": True, "pages": [0]}, + timeout=12.5, + ) + + assert response == FAKE_OCR_RESPONSE + call = bridge.calls[0] + assert call == { + "model": "mistral-ocr-latest", + "document": DOCUMENT, + "api_key": "sk-test", + "api_base": "https://proxy.internal", + "custom_llm_provider": "mistral", + "extra_headers": { + "Authorization": "Bearer sk-test", + "x-trace-id": "trace-1", + }, + "optional_params": {"include_image_base64": True, "pages": [0]}, + "timeout_seconds": 12.5, + } + + +@pytest.mark.asyncio +async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response(): + bridge = RecordingAsyncBridge() + + litellm.use_litellm_rust(True, aocr=bridge) + response = await rust_bridge.aocr( + model="mistral-ocr-maas", + document=DOCUMENT, + api_key=None, + api_base=None, + custom_llm_provider="vertex_ai", + extra_headers=None, + optional_params={"vertex_project": "project-1"}, + timeout=httpx.Timeout(30.0, read=42.0), + ) + + assert response == FAKE_OCR_RESPONSE + assert bridge.calls[0] == { + "model": "mistral-ocr-maas", + "document": DOCUMENT, + "api_key": None, + "api_base": None, + "custom_llm_provider": "vertex_ai", + "extra_headers": None, + "optional_params": {"vertex_project": "project-1"}, + "timeout_seconds": 42.0, + } + + +def test_run_rust_ocr_prepares_request_and_wraps_response(): bridge = RecordingBridge() logging_obj = RecordingLogging() + litellm.use_litellm_rust(True, ocr=bridge) response = ocr_main._run_rust_ocr( - rust_ocr=bridge, prepared_request=build_prepared_request( logging_obj=logging_obj, api_base="https://proxy.internal", @@ -389,8 +450,7 @@ def test_run_rust_ocr_forwards_args_and_wraps_response(): assert isinstance(response, OCRResponse) assert response.pages[0].markdown == "hello world" - call = bridge.calls[0] - assert call == { + assert bridge.calls[0] == { "model": "mistral-ocr-latest", "document": DOCUMENT, "api_key": "sk-test", @@ -406,12 +466,10 @@ def test_run_rust_ocr_forwards_args_and_wraps_response(): def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): - """No explicit api_key: the resolver (get_secret_str in production) supplies it, - so secret-manager backends (AWS/Azure/GCP/Vault) work like the Python path.""" bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) ocr_main._run_rust_ocr( - rust_ocr=bridge, prepared_request=build_prepared_request(api_key=None, timeout=None), resolve_api_key=lambda name: ( "sk-from-vault" if name == "MISTRAL_API_KEY" else None @@ -421,16 +479,34 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): assert bridge.calls[0]["api_key"] == "sk-from-vault" +def test_run_rust_ocr_prefers_explicit_key_over_resolver(): + bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) + + def _resolver(name: str) -> str | None: + raise AssertionError(f"resolver should not be called for {name}") + + ocr_main._run_rust_ocr( + prepared_request=build_prepared_request( + api_key="sk-explicit", + timeout=None, + ), + resolve_api_key=_resolver, + ) + + assert bridge.calls[0]["api_key"] == "sk-explicit" + + def test_run_rust_ocr_uses_provider_api_key_env_var(): bridge = RecordingBridge() resolver_calls = [] + litellm.use_litellm_rust(True, ocr=bridge) def _resolver(name): resolver_calls.append(name) return "sk-provider-env" ocr_main._run_rust_ocr( - rust_ocr=bridge, prepared_request=build_prepared_request( provider_config=FakeOCRConfig(api_key_env_var="PROVIDER_OCR_API_KEY"), model="provider-ocr-model", @@ -446,9 +522,9 @@ def test_run_rust_ocr_uses_provider_api_key_env_var(): def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) ocr_main._run_rust_ocr( - rust_ocr=bridge, prepared_request=build_prepared_request( custom_llm_provider="vertex_ai", model="mistral-ocr-maas", @@ -472,6 +548,7 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager(): bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) def _resolver(name: str) -> str | None: return { @@ -480,7 +557,6 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana }.get(name) ocr_main._run_rust_ocr( - rust_ocr=bridge, prepared_request=build_prepared_request( custom_llm_provider="vertex_ai", model="mistral-ocr-maas", @@ -495,9 +571,9 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) ocr_main._run_rust_ocr( - rust_ocr=bridge, prepared_request=build_prepared_request( custom_llm_provider="azure_ai", model="pixtral-12b-2409", @@ -514,12 +590,12 @@ def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) ocr_main._run_rust_ocr( - rust_ocr=bridge, prepared_request=build_prepared_request( - custom_llm_provider="azure_ai/doc-intelligence", - model="prebuilt-layout", + custom_llm_provider="azure_ai", + model="doc-intelligence/prebuilt-layout", api_base=None, timeout=None, ), @@ -533,30 +609,12 @@ def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): assert bridge.calls[0]["api_base"] == "https://document-intelligence.example.com" -def test_run_rust_ocr_prefers_explicit_key_over_resolver(): - bridge = RecordingBridge() - resolver_calls = [] - - def _resolver(name): - resolver_calls.append(name) - return "sk-from-vault" - - ocr_main._run_rust_ocr( - rust_ocr=bridge, - prepared_request=build_prepared_request(api_key="sk-explicit", timeout=None), - resolve_api_key=_resolver, - ) - - assert bridge.calls[0]["api_key"] == "sk-explicit" - assert resolver_calls == [] # resolver never consulted when a key is supplied - - def test_run_rust_ocr_runs_pre_call_logging(): - """The Rust shortcut must run pre_call so callbacks and spend tracking fire.""" logging_obj = RecordingLogging() + bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) ocr_main._run_rust_ocr( - rust_ocr=RecordingBridge(), prepared_request=build_prepared_request( logging_obj=logging_obj, api_base="https://api.mistral.ai/v1", @@ -573,7 +631,6 @@ def test_run_rust_ocr_runs_pre_call_logging(): complete_input = additional_args["complete_input_dict"] assert complete_input["document"] == DOCUMENT assert complete_input["include_image_base64"] is True - # The logged request mirrors what Rust sends: resolved URL + headers. assert additional_args["api_base"] == "https://api.mistral.ai/v1/ocr" assert additional_args["headers"] == { "Authorization": "Bearer sk-test", @@ -594,7 +651,6 @@ def test_ocr_routes_to_rust_when_enabled(fake_bridge): assert response.pages[0].markdown == "hello world" assert len(fake_bridge.calls) == 1 call = fake_bridge.calls[0] - # Provider prefix is stripped before reaching the bridge. assert call["model"] == "mistral-ocr-latest" assert call["document"] == DOCUMENT assert call["api_key"] == "sk-test" @@ -603,7 +659,6 @@ def test_ocr_routes_to_rust_when_enabled(fake_bridge): "Authorization": "Bearer sk-test", "x-trace-id": "trace-1", } - # Raw OCR params ride along in optional_params; Rust filters to supported keys. assert call["optional_params"].get("include_image_base64") is True @@ -621,6 +676,19 @@ def test_ocr_routes_azure_ai_to_rust_when_enabled(fake_bridge): assert fake_bridge.calls[0]["custom_llm_provider"] == "azure_ai" +def test_ocr_rust_path_converts_file_document_before_bridge(fake_bridge): + response = litellm.ocr( + model=MODEL, + document={"type": "file", "file": b"%PDF-1.4", "mime_type": "application/pdf"}, + api_key="sk-test", + ) + + assert isinstance(response, OCRResponse) + document = fake_bridge.calls[0]["document"] + assert document["type"] == "document_url" + assert document["document_url"].startswith("data:application/pdf;base64,") + + def test_ocr_exception_type_uses_resolved_provider_context( monkeypatch: pytest.MonkeyPatch, ): @@ -694,12 +762,10 @@ def test_ocr_forwards_timeout_to_rust(fake_bridge): def test_ocr_passes_default_request_timeout_to_rust(fake_bridge): - """When no explicit timeout is given, the library default (request_timeout) - must still be forwarded so the Rust path matches the Python path's deadline.""" - from litellm.constants import request_timeout - litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") + from litellm.constants import request_timeout + assert fake_bridge.calls[0]["timeout_seconds"] == float(request_timeout) @@ -717,7 +783,7 @@ def test_ocr_does_not_route_to_rust_when_disabled(): def test_ocr_falls_back_to_python_when_bridge_unavailable(monkeypatch): """Rust enabled but no bridge available (no injected impl, no compiled wheel): ocr() must degrade to the Python HTTP handler instead of raising.""" - monkeypatch.setattr(ocr_main, "load_rust_ocr", lambda: None) + monkeypatch.setattr(rust_bridge, "load_rust_ocr", lambda: None) litellm.use_litellm_rust(True) # enabled, but load_rust_ocr() returns None in CI captured = {} diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 20fd5c1d86a..b3b0e8adcf6 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -170,6 +170,104 @@ class TestMCPRequestHandler: mock_key_servers.assert_called_once_with(user_api_key_auth) mock_team_servers.assert_called_once_with(user_api_key_auth) + @pytest.mark.parametrize( + "require_key_mcp_access_defined,expected", + [ + # Default (flag off): a key with no MCP scope of its own inherits + # the team's servers. + (False, ["team_server1", "team_server2"]), + # Flag on: the team is a ceiling, not a default — the key inherits + # nothing and must grant servers explicitly. + (True, []), + ], + ) + async def test_require_key_mcp_access_defined_gates_team_inheritance( + self, require_key_mcp_access_defined, expected + ): + """The require_key_mcp_access_defined general setting flips an empty key + from inheriting its team's MCP servers (default) to inheriting none.""" + auth = UserAPIKeyAuth( + api_key="test-key", user_id="test-user", team_id="test-team" + ) + with ( + patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_key", + new_callable=AsyncMock, + return_value=[], + ), + patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_team", + new_callable=AsyncMock, + return_value=["team_server1", "team_server2"], + ), + patch.object( + MCPRequestHandler, + "_get_key_access_group_mcp_server_extras", + new_callable=AsyncMock, + return_value=[], + ), + patch( + "litellm.proxy.proxy_server.general_settings", + {"require_key_mcp_access_defined": require_key_mcp_access_defined}, + ), + ): + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + + assert sorted(result) == sorted(expected) + + @pytest.mark.parametrize( + "key_servers,grants,expected,scenario", + [ + # Explicit key subset is still honored under the flag (intersected + # with the team ceiling) — the flag only removes empty-key inheritance. + (["team_server1"], [], ["team_server1"], "explicit_subset_survives"), + # An access-group grant is the escape hatch: it surfaces even though + # the key inherits nothing from the team. + ([], ["granted_server"], ["granted_server"], "access_group_grant_survives"), + ], + ) + async def test_require_key_mcp_access_defined_preserves_explicit_grants( + self, key_servers, grants, expected, scenario + ): + """With require_key_mcp_access_defined on, a key still reaches servers it + grants explicitly or via an access group — only blanket team inheritance + is removed.""" + auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="test-team", + access_group_ids=["grp"] if grants else [], + ) + with ( + patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_key", + new_callable=AsyncMock, + return_value=key_servers, + ), + patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_team", + new_callable=AsyncMock, + return_value=["team_server1", "team_server2"], + ), + patch.object( + MCPRequestHandler, + "_get_key_access_group_mcp_server_extras", + new_callable=AsyncMock, + return_value=grants, + ), + patch( + "litellm.proxy.proxy_server.general_settings", + {"require_key_mcp_access_defined": True}, + ), + ): + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + + assert sorted(result) == sorted(expected) + @pytest.mark.parametrize("team_servers", [[], ["team_server1", "team_server2"]]) async def test_no_mcp_servers_sentinel_returns_empty(self, team_servers): """A key scoped to the no-mcp-servers sentinel resolves to zero servers, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 47c9396f121..27f9a311250 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -1800,3 +1800,71 @@ class TestConnectionErrorMessage: message = rest_endpoints._connection_error_message(RuntimeError("weird")) assert "weird" not in message assert "proxy logs" in message.lower() + + +class TestToolResponseMcpInfoEnrichment: + """The REST tools/list response must expose the user-facing alias and the + server_id alongside the internal server_name so clients (agent builder UIs) + can map the internal config key to a friendly name without needing the + mcp_routes-gated server listing. + """ + + def test_enriches_mcp_info_with_alias_and_server_id(self): + from mcp.types import Tool as MCPTool + + from litellm.proxy._experimental.mcp_server.server import MCPServer + from litellm.types.mcp import MCPTransport + + server = MCPServer( + server_id="a1b2c3d4", + name="mcpAtlassian", + alias="atlassian", + server_name="mcpAtlassian", + transport=MCPTransport.http, + mcp_info={"server_name": "mcpAtlassian"}, + ) + tools = [ + MCPTool( + name="get_issue", + description="Fetch a Jira issue", + inputSchema={"type": "object"}, + ) + ] + + result = rest_endpoints._create_tool_response_objects(tools, server) + + assert result[0].mcp_info == { + "server_name": "mcpAtlassian", + "server_id": "a1b2c3d4", + "alias": "atlassian", + } + + def test_alias_none_is_explicit_in_mcp_info(self): + from mcp.types import Tool as MCPTool + + from litellm.proxy._experimental.mcp_server.server import MCPServer + from litellm.types.mcp import MCPTransport + + server = MCPServer( + server_id="server-uuid", + name="no_alias_server", + alias=None, + server_name="no_alias_server", + transport=MCPTransport.http, + mcp_info={"server_name": "no_alias_server"}, + ) + tools = [ + MCPTool( + name="ping", + description="Ping", + inputSchema={"type": "object"}, + ) + ] + + result = rest_endpoints._create_tool_response_objects(tools, server) + + assert result[0].mcp_info == { + "server_name": "no_alias_server", + "server_id": "server-uuid", + "alias": None, + } diff --git a/tests/test_litellm/proxy/client/test_keys.py b/tests/test_litellm/proxy/client/test_keys.py index 136408e01ac..620daefb39e 100644 --- a/tests/test_litellm/proxy/client/test_keys.py +++ b/tests/test_litellm/proxy/client/test_keys.py @@ -1,5 +1,6 @@ import os import sys +import traceback import pytest import requests @@ -11,7 +12,7 @@ sys.path.insert( import responses -from litellm.proxy.client.exceptions import UnauthorizedError +from litellm.proxy.client.exceptions import NotFoundError, UnauthorizedError from litellm.proxy.client.keys import KeysManagementClient @@ -420,3 +421,96 @@ def test_info_server_error(client): ) with pytest.raises(requests.exceptions.HTTPError): client.info(key="test-key") + + +LEAKY_KEY = "sk-1234567890abcdefghijklmnop" + + +def _render_full_traceback(exc: BaseException) -> str: + return "".join(traceback.format_exception(type(exc), exc, exc.__traceback__)) + + +@responses.activate +def test_info_not_found_redacts_key_everywhere(client): + """A 404 must not echo the raw key embedded in the request URL. + + Covers str(exc) and the rendered traceback, since the chain through + __cause__ / __context__ is what logging.exception() and the default + excepthook print. + """ + responses.add( + responses.GET, + f"{client._base_url}/key/info?key={LEAKY_KEY}", + status=404, + json={"error": {"message": "Key not found", "code": "404"}}, + ) + with pytest.raises(requests.exceptions.HTTPError) as excinfo: + client.info(key=LEAKY_KEY) + + exc = excinfo.value + assert LEAKY_KEY not in str(exc) + assert "REDACTED" in str(exc) + assert LEAKY_KEY not in _render_full_traceback(exc) + assert exc.__cause__ is None and exc.__suppress_context__ + assert exc.response is not None + assert exc.response.status_code == 404 + assert exc.request is not None + # Known residual: the live request URL still carries the key, since the + # response is preserved so callers keep status_code / text. str(exc) and the + # traceback are scrubbed; the URL-borne key is the root issue tracked in + # LIT-4013 (move the lookup key out of the query string server-side). + assert LEAKY_KEY in exc.response.request.url + + +@responses.activate +def test_info_unauthorized_redacts_key_everywhere(client): + """A 401 surfaced as UnauthorizedError must not echo the raw key in the + message, the retained original, or the rendered traceback chain.""" + responses.add( + responses.GET, + f"{client._base_url}/key/info?key={LEAKY_KEY}", + status=401, + json={"error": "Unauthorized"}, + ) + with pytest.raises(UnauthorizedError) as excinfo: + client.info(key=LEAKY_KEY) + + exc = excinfo.value + assert LEAKY_KEY not in str(exc) + assert "REDACTED" in str(exc) + assert LEAKY_KEY not in str(exc.orig_exception) + assert LEAKY_KEY not in _render_full_traceback(exc) + assert exc.__cause__ is None and exc.__suppress_context__ + assert isinstance(exc.orig_exception, requests.exceptions.HTTPError) + assert exc.orig_exception.response is not None + assert exc.orig_exception.response.status_code == 401 + + +def _http_error_with_key(prefix: str, status: int) -> requests.exceptions.HTTPError: + resp = requests.Response() + resp.status_code = status + return requests.exceptions.HTTPError( + f"{prefix} for url: http://x/key/info?key={LEAKY_KEY}", response=resp + ) + + +def test_unauthorized_error_redacts_wrapped_key(): + """UnauthorizedError scrubs the key in str(exc) and in the retained + orig_exception, while preserving the response for structured access.""" + wrapped = UnauthorizedError( + _http_error_with_key("401 Client Error: Unauthorized", 401) + ) + assert LEAKY_KEY not in str(wrapped) + assert "REDACTED" in str(wrapped) + assert LEAKY_KEY not in str(wrapped.orig_exception) + assert wrapped.orig_exception.response.status_code == 401 + + +def test_not_found_error_redacts_wrapped_key(): + """NotFoundError scrubs the key in str(exc) and in the retained + orig_exception, while preserving the response for structured access.""" + wrapped = NotFoundError(_http_error_with_key("404 Client Error: Not Found", 404)) + assert LEAKY_KEY not in str(wrapped) + assert "REDACTED" in str(wrapped) + assert LEAKY_KEY not in str(wrapped.orig_exception) + assert wrapped.orig_exception.response.status_code == 404 diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py index 7a8d04507dc..dd9cbcf5232 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py @@ -570,3 +570,41 @@ class TestRequireCallerUserIdForNonAdmin: assert exc_info.value.status_code == 403 assert "Service-account keys" in str(exc_info.value.detail) + + +class TestValidateFiniteSpend: + """`validate_finite_spend` rejects NaN/±inf so a non-finite spend cannot + bypass `spend >= max_budget` enforcement (NaN/-inf compare false).""" + + def test_none_is_allowed(self): + from litellm.proxy.management_endpoints.common_utils import ( + validate_finite_spend, + ) + + assert validate_finite_spend(None) is None + + def test_finite_value_is_allowed(self): + from litellm.proxy.management_endpoints.common_utils import ( + validate_finite_spend, + ) + + assert validate_finite_spend(0.0) is None + assert validate_finite_spend(12.5) is None + # Negative spend is intentionally allowed. Admins may set a negative + # spend counter to grant an entity extra allowance for the current + # budget period only (e.g. a large one-time spend grant), effectively + # raising their headroom without raising the recurring budget ceiling. + # Future changes should continue to allow negative spend counters. + assert validate_finite_spend(-50.0) is None + + @pytest.mark.parametrize("bad", [float("nan"), float("inf"), float("-inf")]) + def test_non_finite_is_rejected(self, bad): + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.common_utils import ( + validate_finite_spend, + ) + + with pytest.raises(HTTPException) as exc_info: + validate_finite_spend(bad) + assert exc_info.value.status_code == 400 diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index b4602e0ad8b..27e82df90c1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -3102,6 +3102,86 @@ async def test_ghsa_wvg4_proxy_admin_can_update_user_budget(mocker): assert result is not None +@pytest.mark.asyncio +async def test_admin_user_update_spend_invalidates_counter(mocker): + """A direct /user/update spend change must invalidate the cross-pod + spend counter so enforcement re-reads the new DB value.""" + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_single_user_helper, + ) + + mock_prisma_client = mocker.MagicMock() + existing_user = mocker.MagicMock() + existing_user.model_dump.return_value = {"user_id": "target-user", "spend": 50.0} + existing_user.user_id = "target-user" + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( + return_value=existing_user + ) + mock_prisma_client.update_data = mocker.AsyncMock( + return_value={"user_id": "target-user", "spend": -25.0} + ) + mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x) + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + mock_invalidate = mocker.patch( + "litellm.proxy.proxy_server._invalidate_spend_counter", + new=mocker.AsyncMock(), + ) + + # Use a negative spend: this also implicitly validates that negative spend + # is allowed, which is desirable. Admins may grant an entity extra + # allowance for the current budget period only (a one-time spend grant) + # without raising the recurring budget ceiling. Future changes should + # continue allowing negative spend counters. + user_request = UpdateUserRequest(user_id="target-user", spend=-25) + admin_caller = UserAPIKeyAuth( + user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + await _update_single_user_helper( + user_request=user_request, user_api_key_dict=admin_caller + ) + mock_invalidate.assert_awaited_once_with(counter_key="spend:user:target-user") + + +@pytest.mark.asyncio +async def test_user_update_rejects_non_finite_spend(mocker): + """NaN/inf spend is rejected before any DB write or counter invalidation.""" + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_single_user_helper, + ) + + mock_prisma_client = mocker.MagicMock() + existing_user = mocker.MagicMock() + existing_user.model_dump.return_value = {"user_id": "target-user", "spend": 50.0} + existing_user.user_id = "target-user" + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( + return_value=existing_user + ) + mock_prisma_client.update_data = mocker.AsyncMock() + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + mock_invalidate = mocker.patch( + "litellm.proxy.proxy_server._invalidate_spend_counter", + new=mocker.AsyncMock(), + ) + + user_request = UpdateUserRequest(user_id="target-user", spend=float("nan")) + admin_caller = UserAPIKeyAuth( + user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + with pytest.raises(HTTPException) as exc: + await _update_single_user_helper( + user_request=user_request, user_api_key_dict=admin_caller + ) + assert exc.value.status_code == 400 + mock_prisma_client.update_data.assert_not_called() + mock_invalidate.assert_not_awaited() + + @pytest.mark.asyncio async def test_resolve_user_email_metadata_maps_page_user_ids_to_email(mocker): """Regression for LIT-3889. diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index b8ec8a8a388..97397fb06be 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -6639,9 +6639,9 @@ async def test_reset_key_spend_success(monkeypatch): @pytest.mark.asyncio -async def test_update_key_spend_invalidates_counter(monkeypatch): +async def test_update_key_spend_updates_counter(monkeypatch): """ - Test that updating a key's spend via update_key_fn immediately invalidates the spend counter. + Test that updating a key's spend via update_key_fn immediately updates the spend counter. """ from litellm.proxy.management_endpoints.key_management_endpoints import ( update_key_fn, @@ -6676,15 +6676,17 @@ async def test_update_key_spend_invalidates_counter(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) monkeypatch.setattr("litellm.store_audit_logs", False) + mock_spend_counter_cache = MagicMock() + mock_spend_counter_cache.redis_cache = MagicMock() + mock_spend_counter_cache.redis_cache.async_set_cache = AsyncMock() + monkeypatch.setattr( + "litellm.proxy.proxy_server.spend_counter_cache", + mock_spend_counter_cache, + ) - with ( - patch( - "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object" - ) as mock_delete_cache, - patch( - "litellm.proxy.proxy_server._invalidate_spend_counter" - ) as mock_invalidate, - ): + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object" + ) as mock_delete_cache: mock_delete_cache.return_value = None user_api_key_dict = UserAPIKeyAuth( @@ -6704,7 +6706,12 @@ async def test_update_key_spend_invalidates_counter(monkeypatch): ) mock_delete_cache.assert_awaited_once() - mock_invalidate.assert_awaited_once_with(counter_key=f"spend:key:{hashed_key}") + mock_spend_counter_cache.in_memory_cache.set_cache.assert_called_once_with( + key=f"spend:key:{hashed_key}", value=0.0, ttl=60 + ) + mock_spend_counter_cache.redis_cache.async_set_cache.assert_awaited_once_with( + key=f"spend:key:{hashed_key}", value=0.0, ttl=60 + ) @pytest.mark.asyncio @@ -9718,6 +9725,39 @@ class TestKeyOwnerPrivilegeEscalation: assert exc_info.value.status_code == 403 mock_check.assert_called_once() + @pytest.mark.asyncio + async def test_creator_cannot_reset_own_spend_to_stale_value(self): + """Submitting `spend` equal to the stale DB value must still require + admin. The DB spend lags the live cross-pod counter, so an + "unchanged" spend on the non-admin path would let the creator + overwrite the live counter below real usage. Any explicit `spend` + is a budget change, regardless of value match.""" + existing = self._make_existing_key(created_by="creator-123") + existing.spend = 0.0 + # spend equals the stale DB value (0.0) — the old `!=` gate skipped + # the admin check here. + data = UpdateKeyRequest(key="sk-test", spend=0.0) + auth = self._make_auth(user_id="creator-123") + + mock_check = AsyncMock( + side_effect=HTTPException(status_code=403, detail="Not authorized") + ) + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints._check_key_admin_access", + mock_check, + ): + with pytest.raises(HTTPException): + await _validate_update_key_data( + data=data, + existing_key_row=existing, + user_api_key_dict=auth, + llm_router=None, + premium_user=False, + prisma_client=AsyncMock(), + user_api_key_cache=MagicMock(), + ) + mock_check.assert_called_once() + @pytest.mark.asyncio async def test_assigned_user_blocked_from_model_escalation(self): data = UpdateKeyRequest(key="sk-test", models=["gpt-4", "claude-opus"]) diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 592232f45f5..196ff045208 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -10,8 +10,8 @@ from __future__ import annotations import os from types import SimpleNamespace -from typing import Any, Dict, List, Optional -from unittest.mock import AsyncMock, MagicMock, patch +from typing import Any, Dict +from unittest.mock import AsyncMock, MagicMock import pytest @@ -887,6 +887,84 @@ def test_ProxyConfig__add_deployment_invalid_litellm_params_skips(monkeypatch): assert pc._add_deployment(db_models=[bad]) == 0 +def test_ProxyConfig__add_deployment_resolves_env_refs_after_db_decrypt(monkeypatch): + monkeypatch.setenv("LITELLM_DB_MODEL_API_KEY", "resolved-secret") + monkeypatch.setenv("LITELLM_MASTER_KEY", "master-secret") + monkeypatch.setattr( + "litellm.proxy.proxy_server.decrypt_value_helper", + lambda value, key, return_original_value: value, + ) + fake_router = MagicMock() + fake_router.upsert_deployment = MagicMock(return_value=True) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router) + pc = ProxyConfig() + db_model = SimpleNamespace( + model_id="model-1", + model_name="env-model", + model_info={"id": "model-1"}, + litellm_params={ + "model": "openai/gpt-4o-mini", + "api_key": "os.environ/LITELLM_DB_MODEL_API_KEY", + "api_base": "os.environ/LITELLM_MASTER_KEY", + }, + blocked=False, + ) + + added = pc._add_deployment(db_models=[db_model]) + deployment = fake_router.upsert_deployment.call_args.kwargs["deployment"] + + assert added == 1 + assert deployment.litellm_params.api_key == "resolved-secret" + assert deployment.litellm_params.api_base == "os.environ/LITELLM_MASTER_KEY" + + +def test_ProxyConfig__add_deployment_keeps_team_env_refs_literal(monkeypatch): + def fail_on_call(secret_name, *args, **kwargs): + raise AssertionError("team DB models should not resolve env refs") + + monkeypatch.setenv("LITELLM_MASTER_KEY", "master-secret") + monkeypatch.setattr( + "litellm.proxy.proxy_server.decrypt_value_helper", + lambda value, key, return_original_value: value, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.get_secret", fail_on_call) + fake_router = MagicMock() + fake_router.upsert_deployment = MagicMock(return_value=True) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router) + pc = ProxyConfig() + db_model = SimpleNamespace( + model_id="model-1", + model_name="model_name_team-1_abc", + model_info={"id": "model-1", "team_id": "team-1"}, + litellm_params={ + "model": "openai/gpt-4o-mini", + "api_key": "os.environ/LITELLM_MASTER_KEY", + "api_base": "https://attacker.example", + }, + blocked=False, + ) + + added = pc._add_deployment(db_models=[db_model]) + deployment = fake_router.upsert_deployment.call_args.kwargs["deployment"] + + assert added == 1 + assert deployment.litellm_params.api_key == "os.environ/LITELLM_MASTER_KEY" + assert deployment.litellm_params.api_base == "https://attacker.example" + + +def test_ProxyConfig__resolve_db_litellm_param_skips_non_string_values(monkeypatch): + def fail_on_call(value, key, return_original_value): + raise AssertionError("decrypt_value_helper should only receive strings") + + monkeypatch.setattr( + "litellm.proxy.proxy_server.decrypt_value_helper", + fail_on_call, + ) + pc = ProxyConfig() + + assert pc._resolve_db_litellm_param(key="tpm", value=100) == 100 + + # --------------------------------------------------------------------------- # ProxyConfig.decrypt_model_list_from_db # --------------------------------------------------------------------------- @@ -919,6 +997,71 @@ def test_ProxyConfig_decrypt_model_list_from_db_returns_decrypted(monkeypatch): } +def test_ProxyConfig_decrypt_model_list_from_db_resolves_env_refs_after_db_decrypt( + monkeypatch, +): + monkeypatch.setenv("LITELLM_DB_MODEL_API_KEY", "resolved-secret") + monkeypatch.setenv("LITELLM_MASTER_KEY", "master-secret") + monkeypatch.setattr( + "litellm.proxy.proxy_server.decrypt_value_helper", + lambda value, key, return_original_value: ( + "os.environ/LITELLM_DB_MODEL_API_KEY" + if key == "api_key" + else "os.environ/LITELLM_MASTER_KEY" if key == "api_base" else value + ), + ) + pc = ProxyConfig() + m = SimpleNamespace( + model_id="model-1", + model_name="env-model", + model_info={"id": "model-1"}, + litellm_params={ + "api_key": "encrypted-env-ref", + "api_base": "encrypted-api-base-env-ref", + "model": "openai/gpt-4o-mini", + }, + blocked=False, + ) + + out = pc.decrypt_model_list_from_db(new_models=[m]) + + assert out[0]["litellm_params"]["api_key"] == "resolved-secret" + assert out[0]["litellm_params"]["api_base"] == "os.environ/LITELLM_MASTER_KEY" + + +def test_ProxyConfig_decrypt_model_list_from_db_keeps_team_env_refs_literal_after_db_decrypt( + monkeypatch, +): + def fail_on_call(secret_name, *args, **kwargs): + raise AssertionError("team DB models should not resolve env refs") + + monkeypatch.setenv("LITELLM_MASTER_KEY", "master-secret") + monkeypatch.setattr( + "litellm.proxy.proxy_server.decrypt_value_helper", + lambda value, key, return_original_value: ( + "os.environ/LITELLM_MASTER_KEY" if key == "api_key" else value + ), + ) + monkeypatch.setattr("litellm.proxy.proxy_server.get_secret", fail_on_call) + pc = ProxyConfig() + m = SimpleNamespace( + model_id="model-1", + model_name="model_name_team-1_abc", + model_info={"id": "model-1", "team_id": "team-1"}, + litellm_params={ + "api_key": "encrypted-env-ref", + "api_base": "https://attacker.example", + "model": "openai/gpt-4o-mini", + }, + blocked=False, + ) + + out = pc.decrypt_model_list_from_db(new_models=[m]) + + assert out[0]["litellm_params"]["api_key"] == "os.environ/LITELLM_MASTER_KEY" + assert out[0]["litellm_params"]["api_base"] == "https://attacker.example" + + def test_ProxyConfig_decrypt_model_list_from_db_invalid_params_skips(): pc = ProxyConfig() bad = SimpleNamespace( diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 0b583129591..6d0b4509b45 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -342,6 +342,7 @@ def test_ui_view_request_response_forbids_non_admin_without_db(client, monkeypat ignored_keys = [ "request_id", + "metadata.litellm_call_id", "session_id", "startTime", "endTime", diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index e305054d075..874e0654a1f 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -2120,3 +2120,112 @@ def test_get_logging_payload_failure_without_recovered_usage_is_zero(): ) assert payload["total_tokens"] == 0 + + +def test_get_logging_payload_sets_litellm_call_id_for_correlation(): + """LIT-3868: a successful spend log must carry the x-litellm-call-id (the + trace id) in its metadata, distinct from request_id, which stays the + provider response id. Without this there is no way to correlate a DB row + with its trace for a successful call. + """ + provider_response_id = "chatcmpl-e6e6f3e9-c392-404e-9a71-5361c79d8470" + trace_call_id = "c6a77556-19ce-4406-b287-53f5fb4b2b55" + + kwargs = { + "model": "openai/gpt-4o-mini", + "call_type": "acompletion", + "litellm_call_id": trace_call_id, + "litellm_params": {"metadata": {"user_api_key": "sk-test"}}, + } + response_obj = { + "id": provider_response_id, + "usage": {"prompt_tokens": 5, "completion_tokens": 7, "total_tokens": 12}, + } + now = datetime.datetime.now(timezone.utc) + + payload = get_logging_payload( + kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now + ) + metadata = json.loads(payload["metadata"]) + + assert payload["request_id"] == provider_response_id + assert metadata["litellm_call_id"] == trace_call_id + assert metadata["litellm_call_id"] != payload["request_id"] + + +def test_get_logging_payload_litellm_call_id_falls_back_to_litellm_params(): + """litellm_call_id may only be present in litellm_params; it must still land + in the spend log metadata so correlation works on that path too. + """ + trace_call_id = "fallback-7a1c-42d9-9f0e-2b6c5d4e3f21" + kwargs = { + "model": "openai/gpt-4o-mini", + "call_type": "acompletion", + "litellm_params": { + "litellm_call_id": trace_call_id, + "metadata": {"user_api_key": "sk-test"}, + }, + } + response_obj = { + "id": "chatcmpl-abc123", + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + now = datetime.datetime.now(timezone.utc) + + payload = get_logging_payload( + kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now + ) + + assert json.loads(payload["metadata"])["litellm_call_id"] == trace_call_id + + +def test_get_logging_payload_litellm_call_id_when_response_has_no_id(): + """When the provider returns no id, request_id falls back to the call id, so + request_id and the metadata call id hold the same value and correlation + still resolves. + """ + trace_call_id = "noid-5b2e-4c7a-9d10-3f8a1c2b4e6d" + kwargs = { + "model": "openai/gpt-4o-mini", + "call_type": "acompletion", + "litellm_call_id": trace_call_id, + "litellm_params": {"metadata": {"user_api_key": "sk-test"}}, + } + response_obj = { + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} + } + now = datetime.datetime.now(timezone.utc) + + payload = get_logging_payload( + kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now + ) + + assert json.loads(payload["metadata"])["litellm_call_id"] == trace_call_id + assert payload["request_id"] == trace_call_id + + +def test_get_logging_payload_cache_hit_keeps_raw_litellm_call_id(): + """On a cache hit request_id is suffixed to stay unique, but the metadata + litellm_call_id stays the raw trace id so the row still points at its trace. + """ + trace_call_id = "cache-9a1c-42d9-9f0e-2b6c5d4e3f21" + kwargs = { + "model": "openai/gpt-4o-mini", + "call_type": "acompletion", + "litellm_call_id": trace_call_id, + "cache_hit": True, + "litellm_params": {"metadata": {"user_api_key": "sk-test"}}, + } + response_obj = { + "id": "chatcmpl-cache-src", + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + now = datetime.datetime.now(timezone.utc) + + payload = get_logging_payload( + kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now + ) + + assert json.loads(payload["metadata"])["litellm_call_id"] == trace_call_id + assert "_cache_hit" in payload["request_id"] + assert json.loads(payload["metadata"])["litellm_call_id"] != payload["request_id"] diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 364861d6e31..3d3a2cc1013 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -2022,6 +2022,52 @@ class TestOverrideOpenAIResponseModel: assert response_obj.model == requested_model + def test_skips_model_override_when_response_has_no_model_attribute(self): + from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult + + response_obj = SearchResponse( + results=[SearchResult(title="t", url="http://x.com", snippet="s")], + object="search", + ) + + _override_openai_response_model( + response_obj=response_obj, + requested_model="my-search-tool", + log_context="test_context", + ) + + assert not hasattr(response_obj, "model") + + def test_skips_model_override_for_dict_without_model_key(self): + response_obj = { + "object": "search", + "results": [{"title": "t", "url": "http://x.com", "snippet": "s"}], + } + + _override_openai_response_model( + response_obj=response_obj, + requested_model="my-search-tool", + log_context="test_context", + ) + + assert "model" not in response_obj + + def test_override_model_swallows_setattr_failure(self): + class ReadOnlyModelResponse: + @property + def model(self) -> str: + return "downstream-model" + + response_obj = ReadOnlyModelResponse() + + _override_openai_response_model( + response_obj=response_obj, + requested_model="my-model", + log_context="test_context", + ) + + assert response_obj.model == "downstream-model" + class TestIsAzureModelRouterRequest: """Tests for _is_azure_model_router_request helper""" diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_handler.py b/tests/test_litellm/responses/litellm_completion_transformation/test_handler.py new file mode 100644 index 00000000000..04db7192364 --- /dev/null +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_handler.py @@ -0,0 +1,73 @@ +"""Regression tests for the responses -> completion fallback bridge guard. + +When the Responses API falls back to chat completions (no native responses +config), it must tag the forwarded ``litellm.completion`` / ``litellm.acompletion`` +call with ``_skip_responses_api_bridge=True`` so ``completion()`` does not bridge +the request straight back to the Responses API and mutually recurse forever. + +Both fallback paths are covered: the sync ``response_api_handler`` (``_is_async`` +False) and the async ``async_response_api_handler`` (``_is_async`` True). The +module-level ``litellm.completion`` / ``litellm.acompletion`` are patched to +capture the forwarded kwargs; if the flag-setting line is removed the captured +kwargs lack the flag and these tests fail. +""" + +import os +import sys +from unittest.mock import patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.responses.litellm_completion_transformation.handler import ( + LiteLLMCompletionTransformationHandler, +) + + +class _StopForwarding(Exception): + """Raised by the mocked (a)completion once the forwarded kwargs are captured.""" + + +def test_sync_fallback_tags_skip_responses_api_bridge(): + handler = LiteLLMCompletionTransformationHandler() + captured: dict = {} + + def fake_completion(**kwargs): + captured.update(kwargs) + raise _StopForwarding() + + with patch("litellm.completion", fake_completion): + with pytest.raises(_StopForwarding): + handler.response_api_handler( + model="gpt-4o", + input="hello", + responses_api_request={}, + custom_llm_provider="openai", + _is_async=False, + ) + + assert captured.get("_skip_responses_api_bridge") is True + + +@pytest.mark.asyncio +async def test_async_fallback_tags_skip_responses_api_bridge(): + handler = LiteLLMCompletionTransformationHandler() + captured: dict = {} + + async def fake_acompletion(**kwargs): + captured.update(kwargs) + raise _StopForwarding() + + with patch("litellm.acompletion", fake_acompletion): + coro = handler.response_api_handler( + model="gpt-4o", + input="hello", + responses_api_request={}, + custom_llm_provider="openai", + _is_async=True, + ) + with pytest.raises(_StopForwarding): + await coro + + assert captured.get("_skip_responses_api_bridge") is True diff --git a/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py b/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py index 496b0276b87..7cc05d6e30a 100644 --- a/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py +++ b/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py @@ -3,19 +3,45 @@ from pathlib import Path import pytest +import litellm from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +REPO_ROOT = Path(__file__).parents[2] +MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" -@pytest.mark.parametrize("model", ["mistral/mistral-medium-3-5"]) -def test_mistral_medium_3_5_model_info(model): - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - with open(json_path) as f: - model_cost = json.load(f) +MEDIUM_3_5_MODELS = ( + "mistral/mistral-medium-3-5", + "mistral/mistral-medium-2604", + "mistral/mistral-medium-latest", +) - info = model_cost.get(model) - assert ( - info is not None - ), f"{model} not found in model_prices_and_context_window.json" +SYNCED_MODELS = MEDIUM_3_5_MODELS + ( + "mistral/mistral-medium-2508", + "mistral/mistral-medium-3-1-2508", +) + + +def _load(path): + with open(path) as f: + return json.load(f) + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Force get_model_info to resolve against the in-repo cost map instead of the + remote one fetched at import time, which still carries the pre-merge pricing.""" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + +@pytest.mark.parametrize("model", MEDIUM_3_5_MODELS) +def test_medium_3_5_specs(model): + info = _load(MAIN_PATH).get(model) + assert info is not None, f"{model} missing from model_prices_and_context_window.json" assert info["litellm_provider"] == "mistral" assert info["mode"] == "chat" @@ -27,10 +53,11 @@ def test_mistral_medium_3_5_model_info(model): assert info["max_output_tokens"] == 262144 assert info["max_tokens"] == 262144 + assert info["supports_reasoning"] is True + assert info["supports_vision"] is True assert info["supports_function_calling"] is True assert info["supports_response_schema"] is True assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True assert info["supports_assistant_prefill"] is True routed_model, provider, _, _ = get_llm_provider(model=model) @@ -38,18 +65,32 @@ def test_mistral_medium_3_5_model_info(model): assert provider == "mistral" -def test_mistral_medium_3_5_backup_matches_main(): - """Ensure the bundled model cost map stays in sync with the canonical file.""" - repo_root = Path(__file__).parents[2] - main_path = repo_root / "model_prices_and_context_window.json" - backup_path = repo_root / "litellm" / "model_prices_and_context_window_backup.json" +def test_mistral_medium_latest_resolves_to_medium_3_5(local_model_cost_map): + """LIT-3883: the -latest alias was retargeted to Medium 3.5; get_model_info must + return the 3.5 pricing/context/reasoning, not the stale Medium 3.1 values.""" + info = litellm.get_model_info(model="mistral/mistral-medium-latest") - with open(main_path) as f: - main_cost = json.load(f) - with open(backup_path) as f: - backup_cost = json.load(f) + assert info["input_cost_per_token"] == 1.5e-06 + assert info["output_cost_per_token"] == 7.5e-06 + assert info["max_input_tokens"] == 262144 + assert info["supports_reasoning"] is True - for model in ("mistral/mistral-medium-3-5",): - assert backup_cost.get(model) == main_cost.get( - model - ), f"{model} differs between main and backup model cost maps" + +def test_mistral_medium_2508_keeps_medium_3_1_specs(): + """The date-pinned 2508 alias is Medium 3.1 and must not inherit 3.5 pricing.""" + info = _load(MAIN_PATH).get("mistral/mistral-medium-2508") + assert info is not None, "mistral/mistral-medium-2508 missing from cost map" + + assert info["input_cost_per_token"] == 4e-07 + assert info["output_cost_per_token"] == 2e-06 + assert info["max_input_tokens"] == 131072 + assert info.get("supports_reasoning") is not True + + +@pytest.mark.parametrize("model", SYNCED_MODELS) +def test_backup_matches_main(model): + """Ensure the bundled (backup) cost map stays in sync with the canonical file.""" + main_cost = _load(MAIN_PATH) + backup_cost = _load(BACKUP_PATH) + + assert backup_cost.get(model) == main_cost.get(model), f"{model} differs between main and backup model cost maps" diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 470d38caf10..7be176fffc7 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1,3 +1,4 @@ +import asyncio import copy import json import os @@ -12,6 +13,7 @@ sys.path.insert( import litellm +from litellm.exceptions import MidStreamFallbackError def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): @@ -2409,6 +2411,177 @@ async def test_aresponses_streaming_iterator_combines_partial_usage(): assert merged.total_tokens == 49 +def _midstream_rate_limit_error(): + rate_limit_error = litellm.RateLimitError( + message="vertex_ai_betaException - Resource exhausted.", + model="gemini", + llm_provider="vertex_ai_beta", + ) + midstream_error = MidStreamFallbackError( + message=str(rate_limit_error), + model="gemini", + llm_provider="vertex_ai_beta", + original_exception=rate_limit_error, + is_pre_first_chunk=True, + ) + return rate_limit_error, midstream_error + + +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_surfaces_rate_limit_without_fallbacks(): + """Regression for #26015: a mid-stream 429 with no fallbacks configured must + surface a clean RateLimitError, not leak the internal MidStreamFallbackError + wrapper to the client, and must terminate instead of hanging.""" + rate_limit_error, midstream_error = _midstream_rate_limit_error() + + router = litellm.Router( + model_list=[ + { + "model_name": "gemini", + "litellm_params": { + "model": "vertex_ai/gemini-2.0-flash", + "api_key": "fake-key", + }, + }, + ], + num_retries=0, + ) + + class _RaisingStream: + def __init__(self): + self.chunks = [] + + def __aiter__(self): + return self + + async def __anext__(self): + raise midstream_error + + stream = _RaisingStream() + setattr(stream, "model", "gemini") + setattr(stream, "custom_llm_provider", "vertex_ai_beta") + setattr(stream, "logging_obj", MagicMock()) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(side_effect=midstream_error), + ): + result = await router._acompletion_streaming_iterator( + model_response=stream, + messages=[{"role": "user", "content": "Hello"}], + initial_kwargs={"model": "gemini", "stream": True}, + ) + + async def _consume(): + async for _ in result: + pass + + with pytest.raises(litellm.RateLimitError) as exc_info: + await asyncio.wait_for(_consume(), timeout=10) + + assert not isinstance(exc_info.value, MidStreamFallbackError) + assert exc_info.value.status_code == 429 + assert exc_info.value is rate_limit_error + + +def test_completion_streaming_iterator_surfaces_rate_limit_without_fallbacks(): + """Sync counterpart of + test_acompletion_streaming_iterator_surfaces_rate_limit_without_fallbacks.""" + rate_limit_error, midstream_error = _midstream_rate_limit_error() + + router = litellm.Router( + model_list=[ + { + "model_name": "gemini", + "litellm_params": { + "model": "vertex_ai/gemini-2.0-flash", + "api_key": "fake-key", + }, + }, + ], + num_retries=0, + ) + + class _RaisingSyncStream: + def __init__(self): + self.model = "gemini" + self.custom_llm_provider = "vertex_ai_beta" + self.logging_obj = MagicMock() + self.chunks = [] + + def __iter__(self): + return self + + def __next__(self): + raise midstream_error + + with patch.object( + router, + "function_with_fallbacks", + side_effect=midstream_error, + ): + result = router._completion_streaming_iterator( + model_response=_RaisingSyncStream(), + messages=[{"role": "user", "content": "Hello"}], + initial_kwargs={"model": "gemini", "stream": True}, + ) + + with pytest.raises(litellm.RateLimitError) as exc_info: + list(result) + + assert not isinstance(exc_info.value, MidStreamFallbackError) + assert exc_info.value.status_code == 429 + assert exc_info.value is rate_limit_error + + +@pytest.mark.asyncio +async def test_aresponses_streaming_iterator_surfaces_rate_limit_without_fallbacks(): + """Responses-API counterpart of + test_acompletion_streaming_iterator_surfaces_rate_limit_without_fallbacks.""" + rate_limit_error, midstream_error = _midstream_rate_limit_error() + + router = litellm.Router( + model_list=[ + { + "model_name": "gemini", + "litellm_params": { + "model": "vertex_ai/gemini-2.0-flash", + "api_key": "fake-key", + }, + }, + ], + num_retries=0, + ) + src = _make_responses_iterator(error=midstream_error, model="gemini") + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(side_effect=midstream_error), + ): + wrapped = await router._aresponses_streaming_iterator( + response=src, + initial_kwargs={ + "model": "gemini", + "stream": True, + "input": "Hello", + "original_generic_function": litellm.aresponses, + }, + ) + + async def _consume(): + async for _ in wrapped: + pass + + with pytest.raises(litellm.RateLimitError) as exc_info: + await asyncio.wait_for(_consume(), timeout=10) + + assert not isinstance(exc_info.value, MidStreamFallbackError) + assert exc_info.value.status_code == 429 + assert exc_info.value is rate_limit_error + + @pytest.mark.asyncio async def test_async_function_with_fallbacks_common_utils(): """Test the async_function_with_fallbacks_common_utils method""" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index ee3d4e835ba..b53acf930f2 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -21125,10 +21125,9 @@ export interface components { }; /** * Mode - * @description The mode to test the model with - * @default chat + * @description The mode to test the model with. If not provided, auto-detected from model capabilities. */ - mode: ("chat" | "completion" | "embedding" | "audio_speech" | "audio_transcription" | "image_generation" | "video_generation" | "batch" | "rerank" | "realtime" | "responses" | "ocr") | null; + mode?: ("chat" | "completion" | "embedding" | "audio_speech" | "audio_transcription" | "image_generation" | "video_generation" | "batch" | "rerank" | "realtime" | "responses" | "ocr") | null; /** * Model Info * @description Model info for the health check