diff --git a/.circleci/scripts/unit_selection.sh b/.circleci/scripts/unit_selection.sh index 5ce8b6c84ba..d56e29fb627 100755 --- a/.circleci/scripts/unit_selection.sh +++ b/.circleci/scripts/unit_selection.sh @@ -7,6 +7,9 @@ legacy_flags=( caching-local enterprise-package enterprise-routing + integrations + llm-other-providers + llm-vertex-ai mcp-integration misc proxy-db-auth-checks @@ -50,6 +53,9 @@ legacy_paths() { echo tests/unit/enterprise/proxy/test_file_deletion_blocking.py echo tests/unit/enterprise/proxy/test_managed_files_access_check.py echo tests/unit/enterprise/proxy/test_managed_files_hook.py ;; + integrations) echo tests/unit/integrations ;; + llm-other-providers) find tests/unit/llms -name 'test_*.py' -not -path 'tests/unit/llms/vertex_ai/*' ;; + llm-vertex-ai) echo tests/unit/llms/vertex_ai ;; mcp-integration) echo tests/unit/experimental_mcp_client echo tests/unit/proxy/_experimental/mcp_server @@ -71,6 +77,7 @@ legacy_paths() { echo tests/unit/messages echo tests/unit/rag echo tests/unit/rerank_api + echo tests/unit/secret_managers echo tests/unit/vector_stores echo tests/unit/videos ;; proxy-db-auth-checks) diff --git a/.circleci/tests.yml b/.circleci/tests.yml index 10ee19f146a..41e9f11cefa 100644 --- a/.circleci/tests.yml +++ b/.circleci/tests.yml @@ -354,6 +354,28 @@ workflows: - proxy-db-endpoints-and-responses base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> + - unit: + name: unit-llm-vertex-ai + flag: llm-vertex-ai + shards: 2 + workers: 1 + reruns: 2 + base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> + pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> + - unit: + name: unit-llm-other-providers + flag: llm-other-providers + shards: 3 + reruns: 2 + base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> + pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> + - unit: + name: unit-integrations + flag: integrations + shards: 2 + reruns: 3 + base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> + pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> - unit: name: unit-misc flag: misc diff --git a/.github/merge-smoke-tests.json b/.github/merge-smoke-tests.json index 8ed7b917460..a563424c230 100644 --- a/.github/merge-smoke-tests.json +++ b/.github/merge-smoke-tests.json @@ -1,8 +1,8 @@ { "cases": { - "CHAT-JSON": "tests/test_litellm/llms/openai/test_openai.py::test_acompletion_returns_json_reply_over_injected_transport", - "CHAT-TEXT-STREAM": "tests/test_litellm/llms/openai/test_openai.py::test_acompletion_streams_text_deltas_over_injected_transport", - "CHAT-TOOL-STREAM": "tests/test_litellm/llms/openai/test_openai.py::test_acompletion_streams_tool_call_arguments_over_injected_transport", + "CHAT-JSON": "tests/unit/llms/openai/test_openai.py::test_acompletion_returns_json_reply_over_injected_transport", + "CHAT-TEXT-STREAM": "tests/unit/llms/openai/test_openai.py::test_acompletion_streams_text_deltas_over_injected_transport", + "CHAT-TOOL-STREAM": "tests/unit/llms/openai/test_openai.py::test_acompletion_streams_tool_call_arguments_over_injected_transport", "MODEL-ALLOW": "tests/test_litellm/proxy/auth/test_auth_checks.py::test_can_object_call_model_allows_listed_model_for_key", "MODEL-DENY": "tests/test_litellm/proxy/auth/test_auth_checks.py::test_can_object_call_model_denials_return_forbidden[key-key_model_access_denied]", "COST-EXPLICIT": "tests/unit/test_cost_calculator.py::test_completion_cost_charges_explicit_per_token_rates_over_registered_ones", diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 91b54f4ee70..4dca8075440 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -80,7 +80,8 @@ jobs: - shard: integrations artifact-name: integrations - test-path: "tests/test_litellm/integrations" + test-path: "" + unit-flag: integrations workers: 2 reruns: 3 timeout-minutes: 20 @@ -89,6 +90,7 @@ jobs: - shard: Vertex AI artifact-name: llm-vertex-ai test-path: "tests/test_litellm/llms/vertex_ai" + unit-flag: llm-vertex-ai workers: 1 reruns: 2 timeout-minutes: 20 @@ -97,6 +99,7 @@ jobs: - shard: All Other Providers artifact-name: llm-other-providers test-path: "tests/test_litellm/llms --ignore=tests/test_litellm/llms/vertex_ai" + unit-flag: llm-other-providers workers: 2 reruns: 2 timeout-minutes: 20 @@ -105,7 +108,6 @@ jobs: - shard: misc artifact-name: misc test-path: >- - tests/test_litellm/secret_managers tests/test_litellm/interactions tests/test_litellm/ocr tests/test_litellm/passthrough diff --git a/Makefile b/Makefile index 62e6ae53275..f27525b58ff 100644 --- a/Makefile +++ b/Makefile @@ -314,7 +314,7 @@ test-unit: install-test-deps # Matrix test targets (matching CI workflow groups) test-unit-llms: install-test-deps - $(UV_RUN) pytest tests/test_litellm/llms --tb=short -vv -n 4 --durations=20 + $(UV_RUN) pytest tests/unit/llms --tb=short -vv -n 4 --durations=20 test-unit-proxy-guardrails: install-test-deps $(UV_RUN) pytest tests/test_litellm/proxy/guardrails tests/test_litellm/proxy/management_endpoints tests/test_litellm/proxy/management_helpers --tb=short -vv -n 4 --durations=20 @@ -326,13 +326,13 @@ test-unit-proxy-misc: install-test-deps $(UV_RUN) pytest tests/test_litellm/proxy/_experimental tests/test_litellm/proxy/agent_endpoints tests/test_litellm/proxy/anthropic_endpoints tests/test_litellm/proxy/common_utils tests/test_litellm/proxy/discovery_endpoints tests/test_litellm/proxy/experimental tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/health_endpoints tests/test_litellm/proxy/image_endpoints tests/test_litellm/proxy/middleware tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/pass_through_endpoints tests/test_litellm/proxy/prompts tests/test_litellm/proxy/public_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/shutdown tests/test_litellm/proxy/spend_tracking tests/test_litellm/proxy/ui_crud_endpoints tests/test_litellm/proxy/vector_store_endpoints tests/test_litellm/proxy/test_*.py --tb=short -vv -n 4 --durations=20 test-unit-integrations: install-test-deps - $(UV_RUN) pytest tests/test_litellm/integrations --tb=short -vv -n 4 --durations=20 + $(UV_RUN) pytest tests/unit/integrations --tb=short -vv -n 4 --durations=20 test-unit-core-utils: install-test-deps $(UV_RUN) pytest tests/test_litellm/litellm_core_utils --tb=short -vv -n 2 --durations=20 test-unit-other: install-test-deps - $(UV_RUN) pytest tests/test_litellm/caching tests/test_litellm/responses tests/test_litellm/secret_managers tests/unit/vector_stores tests/unit/a2a_protocol tests/test_litellm/anthropic_interface tests/unit/completion_extras tests/unit/containers tests/unit/enterprise tests/unit/experimental_mcp_client tests/unit/google_genai tests/unit/images tests/unit/interactions tests/test_litellm/interactions tests/test_litellm/passthrough tests/test_litellm/router_strategy tests/test_litellm/router_utils tests/unit/types --tb=short -vv -n 4 --durations=20 + $(UV_RUN) pytest tests/test_litellm/caching tests/test_litellm/responses tests/unit/secret_managers tests/unit/vector_stores tests/unit/a2a_protocol tests/test_litellm/anthropic_interface tests/unit/completion_extras tests/unit/containers tests/unit/enterprise tests/unit/experimental_mcp_client tests/unit/google_genai tests/unit/images tests/unit/interactions tests/test_litellm/interactions tests/test_litellm/passthrough tests/test_litellm/router_strategy tests/test_litellm/router_utils tests/unit/types --tb=short -vv -n 4 --durations=20 test-unit-root: install-test-deps $(UV_RUN) pytest tests/unit/test_*.py tests/test_litellm/test_*.py --tb=short -vv -n 4 --durations=20 diff --git a/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/readme.md b/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/readme.md index 70562b18aa6..a3869213be2 100644 --- a/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/readme.md +++ b/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/readme.md @@ -2,7 +2,7 @@ Every `litellm_*` metric family the proxy can expose on `/metrics` (136 families across 97 panels), grouped into rows: proxy traffic, latency, spend and tokens, cache, LLM API deployments, key and team rate limits, budgets, guardrails, MCP, managed files and batches, users and teams, the Redis circuit breaker, the spend log cleanup job, and the `prometheus_system` service callback metrics (per-service latency, request and failure rates, spend update queue sizes). Panel titles are the metric names so you can grep the JSON for the metric you care about -Import `grafana_dashboard.json` from **Dashboards > New > Import** and pick your Prometheus data source when prompted (the `DS_PROMETHEUS` variable). Counters are plotted as `rate()` over `$__rate_interval`, histograms as p50 / p95 / p99, gauges as the raw value grouped by the most useful label. Every query names the metric exactly as the proxy emits it (counters carry the `_total` suffix the Prometheus client adds), and `tests/test_litellm/integrations/test_prometheus_metric_name_consistency.py` fails if a metric is renamed without updating this dashboard +Import `grafana_dashboard.json` from **Dashboards > New > Import** and pick your Prometheus data source when prompted (the `DS_PROMETHEUS` variable). Counters are plotted as `rate()` over `$__rate_interval`, histograms as p50 / p95 / p99, gauges as the raw value grouped by the most useful label. Every query names the metric exactly as the proxy emits it (counters carry the `_total` suffix the Prometheus client adds), and `tests/unit/integrations/test_prometheus_metric_name_consistency.py` fails if a metric is renamed without updating this dashboard The first eleven rows need only `callbacks: ["prometheus"]`. The last three rows and the `litellm_admission_*` panels are emitted by other subsystems and stay empty until those are on: the service callback row needs `service_callback: ["prometheus_system"]` in `litellm_settings`, the circuit breaker row needs a Redis cache, the cleanup row needs spend log retention, and admission control needs its middleware enabled. Within the base rows, many panels only fill in once the matching feature is in use: budgets need keys, teams, users or orgs with `max_budget` set, cache panels need caching on, guardrail and MCP panels need those features configured, deployment health needs the router with more than one deployment or a failure to record, and `litellm_in_flight_requests` needs traffic at scrape time. An empty panel for a feature you do not use is expected diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index e7d911f5fd9..3677d1d654f 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3166,10 +3166,11 @@ dependencies = [ "aws-smithy-types", "bytes", "futures-util", + "proptest", "rstest", - "sse-stream", "thiserror 2.0.19", "tokio", + "tokio-util", ] [[package]] @@ -5468,19 +5469,6 @@ dependencies = [ "wasm-bindgen", ] -[[package]] -name = "sse-stream" -version = "0.2.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c25ac7aff0abd1dbc474536e40416e1102c7dd9bfba0b9861c6d357f835dcfb4" -dependencies = [ - "bytes", - "futures-util", - "http-body 1.1.0", - "http-body-util", - "pin-project-lite", -] - [[package]] name = "stable_deref_trait" version = "1.2.1" diff --git a/litellm-rust/crates/core/tests/messages/request.rs b/litellm-rust/crates/core/tests/messages/request.rs index 2927356b773..d37910d4ac4 100644 --- a/litellm-rust/crates/core/tests/messages/request.rs +++ b/litellm-rust/crates/core/tests/messages/request.rs @@ -398,7 +398,7 @@ async fn unsupported_params_are_dropped_under_drop_params_and_rejected_without_i api_key: Some("sk".into()), api_base: Some(upstream.uri()), shaping: MessagesShaping { - capabilities: capabilities.clone(), + capabilities, drop_params, ..MessagesShaping::default() }, diff --git a/litellm-rust/crates/framer/Cargo.toml b/litellm-rust/crates/framer/Cargo.toml index 62bfcc7da3d..e11f2c02a97 100644 --- a/litellm-rust/crates/framer/Cargo.toml +++ b/litellm-rust/crates/framer/Cargo.toml @@ -8,16 +8,17 @@ repository.workspace = true [features] default = ["aws", "sse"] aws = ["dep:aws-smithy-eventstream", "dep:aws-smithy-types"] -sse = ["dep:sse-stream"] +sse = [] [dependencies] aws-smithy-eventstream = { version = "=0.61.4", optional = true } aws-smithy-types = { version = "1.6.1", optional = true } bytes = "1" futures-util.workspace = true -sse-stream = { version = "=0.2.6", optional = true } thiserror.workspace = true +tokio-util = { version = "0.7", features = ["codec", "io"] } [dev-dependencies] +proptest.workspace = true rstest.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/framer/src/aws_event_stream.rs b/litellm-rust/crates/framer/src/aws_event_stream.rs index efd7adeb64b..405ec2d5ad2 100644 --- a/litellm-rust/crates/framer/src/aws_event_stream.rs +++ b/litellm-rust/crates/framer/src/aws_event_stream.rs @@ -1,66 +1,47 @@ -use bytes::{Buf, Bytes, BytesMut}; -use futures_util::{Stream, StreamExt}; +use aws_smithy_eventstream::frame::{read_message_from, write_message_to}; +pub use aws_smithy_types::event_stream::{Header, HeaderValue, Message}; +use bytes::BytesMut; +use tokio_util::codec::{Decoder, Encoder}; -use aws_smithy_eventstream::frame::read_message_from; -use aws_smithy_types::event_stream::Header; - -use crate::{Error, Framer}; +use crate::EventStreamError; +const MIN_FRAME_BYTES: usize = 16; const MAX_FRAME_BYTES: usize = 16 * 1024 * 1024; -#[derive(Clone, Debug, PartialEq)] -pub struct AwsEventStreamFrame { - pub headers: Vec
, - pub payload: Bytes, -} - #[derive(Clone, Copy, Debug, Default)] -pub struct AwsEventStreamFramer; +pub struct AwsEventStreamCodec; -impl Framer for AwsEventStreamFramer { - type Frame = AwsEventStreamFrame; +impl Decoder for AwsEventStreamCodec { + type Item = Message; + type Error = EventStreamError; - fn frame(self, input: S) -> impl Stream> + Send - where - S: Stream> + Send, - B: Buf + Send, - E: std::error::Error + Send + Sync + 'static, - { - futures_util::stream::try_unfold( - (Box::pin(input), BytesMut::new()), - |(mut input, mut buffer)| async move { - loop { - if buffer.len() >= 4 { - let length = (&buffer[..4]).get_u32() as usize; - if !(16..=MAX_FRAME_BYTES).contains(&length) { - return Err(Error::InvalidLength(length)); - } - if buffer.len() >= length { - let raw = buffer.split_to(length).freeze(); - let message = read_message_from(raw)?; - let frame = AwsEventStreamFrame { - headers: message.headers().to_vec(), - payload: message.payload().clone(), - }; - return Ok(Some((frame, (input, buffer)))); - } - } - match input.next().await { - Some(Ok(mut chunk)) => { - while chunk.has_remaining() { - let bytes = chunk.chunk(); - buffer.extend_from_slice(bytes); - let length = bytes.len(); - chunk.advance(length); - } - } - Some(Err(error)) => return Err(Error::Body(Box::new(error))), - None if buffer.is_empty() => return Ok(None), - None => return Err(Error::Truncated), - } - } - }, - ) - .fuse() + fn decode(&mut self, src: &mut BytesMut) -> Result, EventStreamError> { + let Some(prefix) = src.first_chunk::<4>() else { + return Ok(None); + }; + let length = u32::from_be_bytes(*prefix) as usize; + if !(MIN_FRAME_BYTES..=MAX_FRAME_BYTES).contains(&length) { + return Err(EventStreamError::InvalidLength(length)); + } + if src.len() < length { + return Ok(None); + } + Ok(Some(read_message_from(src.split_to(length).freeze())?)) + } + + fn decode_eof(&mut self, src: &mut BytesMut) -> Result, EventStreamError> { + match self.decode(src)? { + Some(message) => Ok(Some(message)), + None if src.is_empty() => Ok(None), + None => Err(EventStreamError::Truncated), + } + } +} + +impl Encoder for AwsEventStreamCodec { + type Error = EventStreamError; + + fn encode(&mut self, message: Message, dst: &mut BytesMut) -> Result<(), EventStreamError> { + Ok(write_message_to(&message, dst)?) } } diff --git a/litellm-rust/crates/framer/src/error.rs b/litellm-rust/crates/framer/src/error.rs index b1f7ed96c5a..879d7557671 100644 --- a/litellm-rust/crates/framer/src/error.rs +++ b/litellm-rust/crates/framer/src/error.rs @@ -1,17 +1,21 @@ +#[cfg(feature = "sse")] #[derive(Debug, thiserror::Error)] -pub enum Error { - #[cfg(feature = "sse")] - #[error("SSE framing failed: {0}")] - Sse(#[from] sse_stream::Error), - #[cfg(feature = "aws")] - #[error("AWS EventStream framing failed: {0}")] - Aws(#[from] aws_smithy_eventstream::error::Error), +pub enum SseError { #[error("body stream failed: {0}")] - Body(#[source] Box), - #[cfg(feature = "aws")] + Body(#[from] std::io::Error), + #[error("SSE field is not UTF-8: {0}")] + InvalidUtf8(#[from] std::str::Utf8Error), +} + +#[cfg(feature = "aws")] +#[derive(Debug, thiserror::Error)] +pub enum EventStreamError { + #[error("body stream failed: {0}")] + Body(#[from] std::io::Error), #[error("invalid AWS EventStream frame length: {0}")] InvalidLength(usize), - #[cfg(feature = "aws")] #[error("truncated AWS EventStream frame")] Truncated, + #[error("malformed AWS EventStream frame: {0}")] + Malformed(#[from] aws_smithy_eventstream::error::Error), } diff --git a/litellm-rust/crates/framer/src/framed.rs b/litellm-rust/crates/framer/src/framed.rs new file mode 100644 index 00000000000..7a19dd40e13 --- /dev/null +++ b/litellm-rust/crates/framer/src/framed.rs @@ -0,0 +1,21 @@ +use std::io; + +use bytes::Buf; +use futures_util::{Stream, StreamExt, TryStreamExt}; +use tokio_util::{ + codec::{Decoder, FramedRead}, + io::StreamReader, +}; + +pub fn frames( + input: S, + codec: D, +) -> impl Stream> + Send +where + S: Stream> + Send, + B: Buf + Send, + E: std::error::Error + Send + Sync + 'static, + D: Decoder + Send, +{ + FramedRead::new(StreamReader::new(input.map_err(io::Error::other)), codec).fuse() +} diff --git a/litellm-rust/crates/framer/src/lib.rs b/litellm-rust/crates/framer/src/lib.rs index 552de419984..223f2f64120 100644 --- a/litellm-rust/crates/framer/src/lib.rs +++ b/litellm-rust/crates/framer/src/lib.rs @@ -1,8 +1,8 @@ mod error; -mod framer; +mod framed; pub use error::*; -pub use framer::*; +pub use framed::frames; #[cfg(feature = "aws")] pub mod aws_event_stream; diff --git a/litellm-rust/crates/framer/src/sse.rs b/litellm-rust/crates/framer/src/sse.rs index 79659f6ce13..6fee1cfab7f 100644 --- a/litellm-rust/crates/framer/src/sse.rs +++ b/litellm-rust/crates/framer/src/sse.rs @@ -1,43 +1,170 @@ -use futures_util::{Stream, StreamExt}; +use std::str; -use crate::{Error, Framer}; +use bytes::{Buf, BufMut, BytesMut}; +use tokio_util::codec::{Decoder, Encoder}; -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct SseFrame { +use crate::SseError; + +const BOM: &[u8] = b"\xEF\xBB\xBF"; + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct SseEvent { pub event: Option, - pub data: Option, + pub data: String, pub id: Option, pub retry: Option, } #[derive(Clone, Copy, Debug, Default)] -pub struct SseFramer; +pub struct SseCodec { + past_bom: bool, +} -impl Framer for SseFramer { - type Frame = SseFrame; +impl Decoder for SseCodec { + type Item = SseEvent; + type Error = SseError; - fn frame(self, input: S) -> impl Stream> + Send - where - S: Stream> + Send, - B: bytes::Buf + Send, - E: std::error::Error + Send + Sync + 'static, - { - let frames = Box::pin(sse_stream::SseStream::from_bytes_stream(input)); - futures_util::stream::try_unfold(frames, |mut frames| async move { - let Some(frame) = frames.next().await else { - return Ok(None); - }; - let frame = frame?; - Ok(Some(( - SseFrame { - event: frame.event, - data: frame.data, - id: frame.id, - retry: frame.retry, - }, - frames, - ))) - }) - .fuse() + fn decode(&mut self, src: &mut BytesMut) -> Result, SseError> { + if !self.skip_bom(src) { + return Ok(None); + } + while let Some(end) = block_end(src) { + let block = src.split_to(end); + let pending = lines(&block) + .map(|(line, _)| line) + .take_while(|line| !line.is_empty()) + .try_fold(Pending::default(), Pending::apply)?; + if let Some(event) = pending.dispatch() { + return Ok(Some(event)); + } + } + Ok(None) + } + + fn decode_eof(&mut self, _pending: &mut BytesMut) -> Result, SseError> { + Ok(None) + } +} + +impl SseCodec { + fn skip_bom(&mut self, src: &mut BytesMut) -> bool { + if self.past_bom { + return true; + } + if src.starts_with(BOM) { + src.advance(BOM.len()); + } else if BOM.starts_with(src) { + return false; + } + self.past_bom = true; + true + } +} + +fn block_end(bytes: &[u8]) -> Option { + lines(bytes) + .find(|(line, _)| line.is_empty()) + .map(|(_, end)| end) +} + +fn lines(bytes: &[u8]) -> impl Iterator { + let mut cursor: usize = 0; + std::iter::from_fn(move || { + let rest = &bytes[cursor..]; + let end = rest.iter().position(|byte| matches!(byte, b'\n' | b'\r'))?; + cursor += end + terminator_len(&rest[end..]); + Some((&rest[..end], cursor)) + }) +} + +fn terminator_len(terminated: &[u8]) -> usize { + match terminated { + [b'\r', b'\n', ..] => 2, + _ => 1, + } +} + +#[derive(Default)] +struct Pending { + event: Option, + data: Option, + id: Option, + retry: Option, +} + +impl Pending { + fn apply(self, line: &[u8]) -> Result { + let (name, value) = split_field(line); + Ok(match name { + b"event" => Self { + event: Some(str::from_utf8(value)?.to_owned()), + ..self + }, + b"data" => Self { + data: Some(append_data(self.data, str::from_utf8(value)?)), + ..self + }, + b"id" if !value.contains(&0) => Self { + id: Some(str::from_utf8(value)?.to_owned()), + ..self + }, + b"retry" => Self { + retry: parse_retry(value).or(self.retry), + ..self + }, + _ => self, + }) + } + + fn dispatch(self) -> Option { + Some(SseEvent { + event: self.event, + data: self.data?, + id: self.id, + retry: self.retry, + }) + } +} + +fn split_field(line: &[u8]) -> (&[u8], &[u8]) { + let Some(colon) = line.iter().position(|byte| *byte == b':') else { + return (line, &[]); + }; + let value = &line[colon + 1..]; + (&line[..colon], value.strip_prefix(b" ").unwrap_or(value)) +} + +fn append_data(buffer: Option, line: &str) -> String { + match buffer { + Some(existing) => format!("{existing}\n{line}"), + None => line.to_owned(), + } +} + +fn parse_retry(value: &[u8]) -> Option { + if !value.iter().all(u8::is_ascii_digit) { + return None; + } + str::from_utf8(value).ok()?.parse().ok() +} + +impl Encoder for SseCodec { + type Error = SseError; + + fn encode(&mut self, event: SseEvent, dst: &mut BytesMut) -> Result<(), SseError> { + if let Some(name) = event.event { + dst.put_slice(format!("event: {name}\n").as_bytes()); + } + for line in event.data.split('\n') { + dst.put_slice(format!("data: {line}\n").as_bytes()); + } + if let Some(id) = event.id { + dst.put_slice(format!("id: {id}\n").as_bytes()); + } + if let Some(retry) = event.retry { + dst.put_slice(format!("retry: {retry}\n").as_bytes()); + } + dst.put_u8(b'\n'); + Ok(()) } } diff --git a/litellm-rust/crates/framer/tests/aws_event_stream.rs b/litellm-rust/crates/framer/tests/aws_event_stream.rs index c90a15a2b0e..d16caa39948 100644 --- a/litellm-rust/crates/framer/tests/aws_event_stream.rs +++ b/litellm-rust/crates/framer/tests/aws_event_stream.rs @@ -4,89 +4,174 @@ mod support; use std::io; -use futures_util::TryStreamExt; -use litellm_framing::aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer}; -use litellm_framing::{Error, Framer}; +use bytes::Bytes; +use futures_util::{StreamExt, TryStreamExt, stream}; +use litellm_framing::{ + EventStreamError, + aws_event_stream::{AwsEventStreamCodec, Header, HeaderValue, Message}, + frames, +}; +use proptest::prelude::*; use rstest::{fixture, rstest}; +use support::{body_cause, cut_at, encode_all, every, input, runtime}; -use support::encode; - -async fn collect_aws(bytes: &[u8], chunk_size: usize) -> Result, Error> { - AwsEventStreamFramer - .frame(futures_util::stream::iter( - bytes.chunks(chunk_size).map(Ok::<_, io::Error>), - )) +async fn collect(pieces: Vec) -> Result, EventStreamError> { + frames(input(pieces), AwsEventStreamCodec) .try_collect() .await } -#[fixture] -fn two_frames() -> Vec { - [encode(b"\xff\x00"), encode(b"second")].concat() +fn message(payload: &[u8]) -> Message { + Message::new(Bytes::copy_from_slice(payload)) + .add_header(Header::new( + ":event-type", + HeaderValue::String("payload".into()), + )) + .add_header(Header::new("sequence", HeaderValue::Int32(7))) } #[fixture] fn payload_frame() -> Vec { - encode(b"payload") + encode_all(AwsEventStreamCodec, [message(b"payload")]) +} + +fn header_value() -> impl Strategy { + prop_oneof![ + "[a-z]{0,8}".prop_map(|text| HeaderValue::String(text.into())), + any::().prop_map(HeaderValue::Int32), + any::().prop_map(HeaderValue::Bool), + proptest::collection::vec(any::(), 0..8) + .prop_map(|bytes| HeaderValue::ByteArray(bytes.into())), + ] +} + +fn arbitrary_message() -> impl Strategy { + ( + proptest::collection::vec(("[a-z:-]{1,12}", header_value()), 0..3), + proptest::collection::vec(any::(), 0..32), + ) + .prop_map(|(headers, payload)| { + headers.into_iter().fold( + Message::new(Bytes::from(payload)), + |message, (name, value)| message.add_header(Header::new(name, value)), + ) + }) +} + +proptest! { + #[test] + fn any_messages_survive_a_round_trip_through_any_cuts( + messages in proptest::collection::vec(arbitrary_message(), 1..4), + cuts in proptest::collection::vec(0_usize..512, 0..4), + ) { + let wire = encode_all(AwsEventStreamCodec, messages.clone()); + let decoded = runtime().block_on(collect(cut_at(&wire, cuts))).unwrap(); + prop_assert_eq!(decoded, messages); + } } #[rstest] -#[case(1)] -#[case(3)] -#[case(12)] -#[case(usize::MAX)] +#[case::prelude_crc(8)] +#[case::message_crc(usize::MAX)] #[tokio::test] -async fn fragmented_and_coalesced_frames_preserve_typed_headers_and_binary_payloads( - two_frames: Vec, - #[case] chunk_size: usize, -) { - let chunk_size = chunk_size.min(two_frames.len()); - let frames = collect_aws(&two_frames, chunk_size).await.unwrap(); - assert_eq!(frames.len(), 2); - assert_eq!(frames[0].payload, &b"\xff\x00"[..]); - assert_eq!(frames[1].payload, "second"); - assert_eq!( - frames[0].headers[0].value().as_string().unwrap().as_str(), - "payload" - ); - assert_eq!(frames[0].headers[1].value().as_int32(), Ok(7)); -} - -#[rstest] -#[case(8)] -#[case(0)] -#[tokio::test] -async fn rejects_corrupt_crcs(payload_frame: Vec, #[case] index: usize) { - let corrupt_index = if index == 0 { - payload_frame.len() - 1 - } else { - index - }; +async fn a_corrupt_crc_is_malformed(payload_frame: Vec, #[case] index: usize) { let mut corrupt = payload_frame; - corrupt[corrupt_index] ^= 1; - assert!(matches!(collect_aws(&corrupt, 3).await, Err(Error::Aws(_)))); -} - -#[rstest] -#[case(0_u32)] -#[case(15)] -#[case(u32::MAX)] -#[tokio::test] -async fn rejects_invalid_lengths(#[case] length: u32) { + let flipped = index.min(corrupt.len() - 1); + corrupt[flipped] ^= 1; assert!(matches!( - collect_aws(&length.to_be_bytes(), 1).await, - Err(Error::InvalidLength(_)) + collect(every(&corrupt, 3)).await, + Err(EventStreamError::Malformed(_)) )); } #[rstest] -#[case(1)] -#[case(3)] -#[case(5)] +#[case::zero(0)] +#[case::below_minimum(15)] +#[case::above_maximum(16 * 1024 * 1024 + 1)] +#[case::u32_max(u32::MAX)] #[tokio::test] -async fn rejects_truncation(payload_frame: Vec, #[case] end: usize) { +async fn a_length_outside_the_frame_bounds_fails_before_buffering(#[case] length: u32) { assert!(matches!( - collect_aws(&payload_frame[..end], 1).await, - Err(Error::Truncated) + collect(every(&length.to_be_bytes(), 1)).await, + Err(EventStreamError::InvalidLength(seen)) if seen == length as usize )); } + +#[rstest] +#[case::before_the_length(1)] +#[case::inside_the_prelude(5)] +#[case::one_byte_short(usize::MAX)] +#[tokio::test] +async fn eof_inside_a_frame_is_truncation(payload_frame: Vec, #[case] end: usize) { + let end = end.min(payload_frame.len() - 1); + assert!(matches!( + collect(every(&payload_frame[..end], 1)).await, + Err(EventStreamError::Truncated) + )); +} + +const FRAME_OVERHEAD_BYTES: usize = 16; +const MAX_FRAME_BYTES: usize = 16 * 1024 * 1024; + +#[tokio::test] +async fn a_frame_at_exactly_the_maximum_length_decodes() { + let largest = Message::new(vec![0xAB; MAX_FRAME_BYTES - FRAME_OVERHEAD_BYTES]); + let wire = encode_all(AwsEventStreamCodec, [largest.clone()]); + assert_eq!(wire.len(), MAX_FRAME_BYTES); + assert_eq!(collect(every(&wire, 1 << 20)).await.unwrap(), vec![largest]); +} + +#[tokio::test] +async fn a_frame_one_byte_over_the_maximum_length_is_rejected_by_its_prelude() { + let oversized = Message::new(vec![0xAB; MAX_FRAME_BYTES - FRAME_OVERHEAD_BYTES + 1]); + let wire = encode_all(AwsEventStreamCodec, [oversized]); + assert!(matches!( + collect(every(&wire[..4], 1)).await, + Err(EventStreamError::InvalidLength(length)) if length == MAX_FRAME_BYTES + 1 + )); +} + +#[tokio::test] +async fn an_empty_body_yields_nothing() { + assert_eq!(collect(vec![]).await.unwrap(), vec![]); +} + +#[tokio::test] +async fn a_complete_frame_precedes_a_truncated_following_frame() { + let wire = encode_all(AwsEventStreamCodec, [message(b"first"), message(b"second")]); + let mut messages = Box::pin(frames( + input(every(&wire[..wire.len() - 1], 3)), + AwsEventStreamCodec, + )); + + assert_eq!(messages.next().await.unwrap().unwrap(), message(b"first")); + assert!(matches!( + messages.next().await, + Some(Err(EventStreamError::Truncated)) + )); + assert!(messages.next().await.is_none()); +} + +#[tokio::test] +async fn a_body_error_after_a_complete_frame_preserves_its_cause() { + let first = encode_all(AwsEventStreamCodec, [message(b"first")]); + let mut messages = Box::pin(frames( + stream::iter([ + Ok(cut_at(&first, [5])[0].clone()), + Ok(cut_at(&first, [5])[1].clone()), + Ok(Bytes::from_static(b"\0\0\0")), + Err(io::Error::new(io::ErrorKind::ConnectionReset, "reset")), + ]), + AwsEventStreamCodec, + )); + + assert_eq!(messages.next().await.unwrap().unwrap(), message(b"first")); + let Some(Err(EventStreamError::Body(body))) = messages.next().await else { + panic!("the body error surfaces"); + }; + assert_eq!( + body_cause::(&body).unwrap().kind(), + io::ErrorKind::ConnectionReset + ); + assert!(messages.next().await.is_none()); +} diff --git a/litellm-rust/crates/framer/tests/chaining.rs b/litellm-rust/crates/framer/tests/chaining.rs index afd24a90704..81884d58ba1 100644 --- a/litellm-rust/crates/framer/tests/chaining.rs +++ b/litellm-rust/crates/framer/tests/chaining.rs @@ -2,28 +2,64 @@ mod support; -use std::io; +use bytes::Bytes; +use futures_util::{StreamExt, TryStreamExt}; +use litellm_framing::{ + EventStreamError, SseError, + aws_event_stream::{AwsEventStreamCodec, Message}, + frames, + sse::{SseCodec, SseEvent}, +}; +use proptest::prelude::*; +use support::{body_cause, cut_at, encode_all, every, input, runtime}; -use futures_util::TryStreamExt; -use litellm_framing::Framer; -use litellm_framing::aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer}; -use litellm_framing::sse::SseFramer; +fn delta(data: &str) -> SseEvent { + SseEvent { + event: Some("delta".into()), + data: data.into(), + id: Some("7".into()), + retry: None, + } +} -use support::encode; +fn envelopes(payloads: Vec) -> Vec { + encode_all(AwsEventStreamCodec, payloads.into_iter().map(Message::new)) +} + +proptest! { + #[test] + fn an_sse_event_cut_anywhere_across_envelopes_is_reassembled(cut in 0_usize..64, chunk in 1_usize..8) { + let sse = encode_all(SseCodec::default(), [delta("hello")]); + let wire = envelopes(cut_at(&sse, [cut.min(sse.len())])); + let events = runtime().block_on(async { + let payloads = frames(input(every(&wire, chunk)), AwsEventStreamCodec) + .map_ok(|message| message.payload().clone()); + frames(payloads, SseCodec::default()).try_collect::>().await + }) + .unwrap(); + prop_assert_eq!(events, vec![delta("hello")]); + } +} #[tokio::test] -async fn hosting_payloads_feed_the_same_sse_framer_across_envelope_boundaries() { - let bytes = [encode(b"event: delta\ndata: hel"), encode(b"lo\nid: 7\n\n")].concat(); - let envelopes = AwsEventStreamFramer.frame(futures_util::stream::iter( - bytes.chunks(3).map(Ok::<_, io::Error>), +async fn a_truncated_envelope_after_an_sse_event_keeps_the_event_and_its_cause() { + let complete = encode_all(SseCodec::default(), [delta("complete")]); + let incomplete = encode_all(SseCodec::default(), [delta("incomplete")]); + let wire = envelopes(vec![complete.into(), incomplete.into()]); + let payloads = frames( + input(every(&wire[..wire.len() - 1], 3)), + AwsEventStreamCodec, + ) + .map_ok(|message| message.payload().clone()); + let mut events = Box::pin(frames(payloads, SseCodec::default())); + + assert_eq!(events.next().await.unwrap().unwrap(), delta("complete")); + let Some(Err(SseError::Body(body))) = events.next().await else { + panic!("the envelope error surfaces through the SSE layer"); + }; + assert!(matches!( + body_cause::(&body), + Some(EventStreamError::Truncated) )); - let frames = SseFramer - .frame(envelopes.map_ok(|frame: AwsEventStreamFrame| frame.payload)) - .try_collect::>() - .await - .unwrap(); - assert_eq!(frames.len(), 1); - assert_eq!(frames[0].event.as_deref(), Some("delta")); - assert_eq!(frames[0].data.as_deref(), Some("hello")); - assert_eq!(frames[0].id.as_deref(), Some("7")); + assert!(events.next().await.is_none()); } diff --git a/litellm-rust/crates/framer/tests/sse.rs b/litellm-rust/crates/framer/tests/sse.rs index 66339dfbfd2..2fa064653a6 100644 --- a/litellm-rust/crates/framer/tests/sse.rs +++ b/litellm-rust/crates/framer/tests/sse.rs @@ -1,67 +1,169 @@ #![cfg(feature = "sse")] +mod support; + use std::io; -use futures_util::{StreamExt, TryStreamExt}; -use litellm_framing::sse::{SseFrame, SseFramer}; -use litellm_framing::{Error, Framer}; +use bytes::Bytes; +use futures_util::{StreamExt, TryStreamExt, stream}; +use litellm_framing::{ + SseError, frames, + sse::{SseCodec, SseEvent}, +}; +use proptest::prelude::*; use rstest::rstest; +use support::{body_cause, cut_at, encode_all, every, input, runtime}; -async fn collect_sse(chunks: &[&[u8]]) -> Result, Error> { - SseFramer - .frame(futures_util::stream::iter( - chunks.iter().copied().map(Ok::<_, io::Error>), - )) +async fn collect(pieces: Vec) -> Result, SseError> { + frames(input(pieces), SseCodec::default()) .try_collect() .await } +fn event(name: Option<&str>, data: &str) -> SseEvent { + SseEvent { + event: name.map(str::to_owned), + data: data.to_owned(), + id: None, + retry: None, + } +} + +fn sse_event() -> impl Strategy { + ( + proptest::option::of("[^\r\n\0]{0,8}"), + "[^\r\0]{0,16}", + proptest::option::of("[^\r\n\0]{0,8}"), + proptest::option::of(any::()), + ) + .prop_map(|(event, data, id, retry)| SseEvent { + event, + data, + id, + retry, + }) +} + +fn terminators() -> impl Strategy { + prop_oneof![Just(&b"\n"[..]), Just(&b"\r\n"[..]), Just(&b"\r"[..])] +} + +proptest! { + #[test] + fn any_events_survive_a_round_trip_through_any_terminator_and_any_cuts( + events in proptest::collection::vec(sse_event(), 1..4), + terminator in terminators(), + cuts in proptest::collection::vec(0_usize..256, 0..4), + bom in any::(), + ) { + let lf_wire = encode_all(SseCodec::default(), events.clone()); + let body: Vec = lf_wire + .iter() + .flat_map(|byte| if *byte == b'\n' { terminator.to_vec() } else { vec![*byte] }) + .collect(); + let wire = if bom { [&b"\xEF\xBB\xBF"[..], &body].concat() } else { body }; + let decoded = runtime().block_on(collect(cut_at(&wire, cuts))).unwrap(); + prop_assert_eq!(decoded, events); + } +} + #[rstest] -#[case( - &[&b":ping\r\nevent: delta\r\nid: 7\r\nretry: 10\r\ndata: \xe2"[..], &b"\x82"[..], &b"\xac\r"[..], &b"\ndata: next\r\n\r"[..], &b"\ndata: [DONE]\n\n"[..]], - vec![ - SseFrame { - event: Some("delta".into()), - data: Some("€\nnext".into()), - id: Some("7".into()), - retry: Some(10), - }, - SseFrame { - event: None, - data: Some("[DONE]".into()), - id: None, - retry: None, - }, - ] -)] +#[case::comment(b":ping\ndata: x\n\n")] +#[case::unknown_field(b"vendor: 1\ndata: x\n\n")] +#[case::field_without_colon(b"garbage\ndata: x\n\n")] +#[case::retry_with_non_digits(b"retry: soon\ndata: x\n\n")] +#[case::retry_with_a_sign(b"retry: +5\ndata: x\n\n")] +#[case::retry_without_a_value(b"retry:\ndata: x\n\n")] +#[case::id_with_nul(b"id: a\0b\ndata: x\n\n")] #[tokio::test] -async fn fragmented_utf8_crlf_and_multiline_data_retain_metadata_and_sentinel( - #[case] chunks: &[&[u8]], - #[case] expected: Vec, +async fn lines_the_spec_ignores_do_not_change_the_event(#[case] wire: &[u8]) { + assert_eq!( + collect(every(wire, 1)).await.unwrap(), + vec![event(None, "x")] + ); +} + +#[rstest] +#[case::no_data_at_all(b"event: ping\nid: 1\n\ndata: x\n\n", vec![event(None, "x")])] +#[case::empty_data_field(b"data:\n\n", vec![event(None, "")])] +#[case::one_leading_space_stripped(b"data: x\n\n", vec![event(None, " x")])] +#[case::multiline_data(b"data: a\ndata: b\ndata:\n\n", vec![event(None, "a\nb\n")])] +#[case::last_event_name_wins(b"event: a\nevent: b\ndata: x\n\n", vec![event(Some("b"), "x")])] +#[case::last_retry_wins(b"retry: 1\nretry: 2\ndata: x\n\n", vec![SseEvent { retry: Some(2), ..event(None, "x") }])] +#[case::split_utf8_across_lines_is_not_joined(b"data: \xe2\x82\xac\ndata: \xe2\x82\xac\n\n", vec![event(None, "€\n€")])] +#[tokio::test] +async fn dispatch_follows_the_data_buffer(#[case] wire: &[u8], #[case] expected: Vec) { + assert_eq!(collect(every(wire, 1)).await.unwrap(), expected); +} + +#[rstest] +#[case::unterminated_single(b"data: partial\n", vec![])] +#[case::unterminated_tail_after_complete(b"data: complete\n\ndata: unfinished\n", vec![event(None, "complete")])] +#[case::lone_cr_terminates_at_eof(b"data: x\r\r", vec![event(None, "x")])] +#[case::lone_cr_line_then_eof(b"data: x\r", vec![])] +#[tokio::test] +async fn eof_dispatches_only_terminated_events( + #[case] wire: &[u8], + #[case] expected: Vec, ) { - assert_eq!(collect_sse(chunks).await.unwrap(), expected); + assert_eq!( + collect(vec![Bytes::copy_from_slice(wire)]).await.unwrap(), + expected + ); +} + +#[rstest] +#[case::inside_the_first_line(vec![&b"data: a\r"[..], &b"\ndata: b\r\n\r\n"[..]])] +#[case::inside_the_blank_line(vec![&b"data: a\r\ndata: b\r\n\r"[..], &b"\n"[..]])] +#[tokio::test] +async fn a_crlf_split_across_chunks_is_one_terminator(#[case] pieces: Vec<&[u8]>) { + let pieces = pieces.into_iter().map(Bytes::copy_from_slice).collect(); + assert_eq!(collect(pieces).await.unwrap(), vec![event(None, "a\nb")]); } #[tokio::test] -async fn eof_does_not_dispatch_an_unterminated_frame() { - assert!(collect_sse(&[b"data: partial\n"]).await.unwrap().is_empty()); +async fn a_bom_is_stripped_only_at_the_start_of_the_stream() { + let wire = b"\xEF\xBB\xBFdata: a\n\n\xEF\xBB\xBFdata: b\ndata: c\n\n"; + let decoded = collect(every(wire, 2)).await.unwrap(); + assert_eq!(decoded, vec![event(None, "a"), event(None, "c")]); +} + +#[tokio::test] +async fn invalid_utf8_in_a_field_fails_after_earlier_events_and_terminates() { + let mut events = Box::pin(frames( + input(every(b"data: ok\n\ndata: \xff\n\n", 3)), + SseCodec::default(), + )); + + assert_eq!(events.next().await.unwrap().unwrap(), event(None, "ok")); + assert!(matches!( + events.next().await, + Some(Err(SseError::InvalidUtf8(_))) + )); + assert!(events.next().await.is_none()); } #[rstest] #[case(io::ErrorKind::ConnectionReset)] #[case(io::ErrorKind::UnexpectedEof)] #[tokio::test] -async fn framing_errors_terminate_and_preserve_input_error_causes(#[case] kind: io::ErrorKind) { - let mut frames = Box::pin(SseFramer.frame(futures_util::stream::iter([ - Err(io::Error::new(kind, "reset")), - Ok(&b"data: later\n\n"[..]), - ]))); - let error = frames.next().await.unwrap().unwrap_err(); - assert!(matches!( - error, - Error::Sse(sse_stream::Error::Body(ref cause)) - if cause.downcast_ref::().unwrap().kind() == kind +async fn a_body_error_keeps_earlier_events_and_its_cause_then_terminates( + #[case] kind: io::ErrorKind, +) { + let mut events = Box::pin(frames( + stream::iter([ + Ok(&b"data: first\n\ndata: partial"[..]), + Err(io::Error::new(kind, "reset")), + Ok(&b"\n\n"[..]), + ]), + SseCodec::default(), )); - assert!(frames.next().await.is_none()); - assert!(frames.next().await.is_none()); + + assert_eq!(events.next().await.unwrap().unwrap(), event(None, "first")); + let Some(Err(SseError::Body(body))) = events.next().await else { + panic!("the body error surfaces"); + }; + assert_eq!(body_cause::(&body).unwrap().kind(), kind); + assert!(events.next().await.is_none()); + assert!(events.next().await.is_none()); } diff --git a/litellm-rust/crates/framer/tests/support/mod.rs b/litellm-rust/crates/framer/tests/support/mod.rs index 9db305af073..9ff67aef149 100644 --- a/litellm-rust/crates/framer/tests/support/mod.rs +++ b/litellm-rust/crates/framer/tests/support/mod.rs @@ -1,15 +1,57 @@ -use aws_smithy_eventstream::frame::write_message_to; -use aws_smithy_types::event_stream::{Header, HeaderValue, Message}; -use bytes::Bytes; +#![allow(dead_code)] -pub fn encode(payload: &'static [u8]) -> Vec { - let message = Message::new(Bytes::from_static(payload)) - .add_header(Header::new( - ":event-type", - HeaderValue::String("payload".into()), - )) - .add_header(Header::new("sequence", HeaderValue::Int32(7))); - let mut bytes = Vec::new(); - write_message_to(&message, &mut bytes).unwrap(); - bytes +use std::{error::Error, io}; + +use bytes::{Bytes, BytesMut}; +use futures_util::{Stream, stream}; +use tokio_util::codec::Encoder; + +pub fn encode_all(mut codec: C, items: impl IntoIterator) -> Vec +where + C: Encoder, + C::Error: std::fmt::Debug, +{ + let mut wire = BytesMut::new(); + for item in items { + codec.encode(item, &mut wire).unwrap(); + } + wire.to_vec() +} + +pub fn cut_at(bytes: &[u8], offsets: impl IntoIterator) -> Vec { + let mut sorted: Vec = offsets + .into_iter() + .filter(|offset| *offset <= bytes.len()) + .collect(); + sorted.sort_unstable(); + sorted.dedup(); + let bounds = std::iter::once(0) + .chain(sorted) + .chain(std::iter::once(bytes.len())) + .collect::>(); + bounds + .windows(2) + .map(|pair| Bytes::copy_from_slice(&bytes[pair[0]..pair[1]])) + .collect() +} + +pub fn every(bytes: &[u8], size: usize) -> Vec { + bytes + .chunks(size.max(1)) + .map(Bytes::copy_from_slice) + .collect() +} + +pub fn input(pieces: Vec) -> impl Stream> + Send { + stream::iter(pieces.into_iter().map(Ok)) +} + +pub fn body_cause(body: &io::Error) -> Option<&T> { + body.get_ref()?.downcast_ref::() +} + +pub fn runtime() -> tokio::runtime::Runtime { + tokio::runtime::Builder::new_current_thread() + .build() + .unwrap() } diff --git a/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/streaming_iterator.rs b/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/streaming_iterator.rs index 35e7d5820b0..3f1b7ed9bcc 100644 --- a/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/streaming_iterator.rs +++ b/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/streaming_iterator.rs @@ -2,9 +2,9 @@ use base64::Engine; use bytes::Buf; use futures_util::{Stream, StreamExt}; use litellm_framing::{ - Framer, - aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer}, - sse::{SseFrame, SseFramer}, + aws_event_stream::{AwsEventStreamCodec, Message}, + frames, + sse::{SseCodec, SseEvent}, }; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; @@ -13,8 +13,6 @@ use serde_json::{Map, Value}; pub enum Error { #[error("stream framing failed: {0}")] StreamFraming(String), - #[error("Anthropic SSE frame has no data")] - MissingStreamData, #[error("Anthropic stream event is invalid: {0}")] InvalidStreamEvent(String), #[error("Bedrock event payload is invalid: {0}")] @@ -165,15 +163,14 @@ struct BedrockChunkPayload { bytes: String, } -pub fn decode_anthropic_sse_frame(frame: SseFrame) -> Result { - let data = frame.data.ok_or(Error::MissingStreamData)?; - serde_json::from_str(&data).map_err(|error| Error::InvalidStreamEvent(error.to_string())) +pub fn decode_anthropic_sse_frame(event: SseEvent) -> Result { + serde_json::from_str(&event.data).map_err(|error| Error::InvalidStreamEvent(error.to_string())) } pub fn decode_bedrock_anthropic_frame( - frame: AwsEventStreamFrame, + message: Message, ) -> Result { - let payload: BedrockChunkPayload = serde_json::from_slice(&frame.payload) + let payload: BedrockChunkPayload = serde_json::from_slice(message.payload()) .map_err(|error| Error::InvalidBedrockPayload(error.to_string()))?; let event = base64::engine::general_purpose::STANDARD .decode(payload.bytes) @@ -189,9 +186,8 @@ where B: Buf + Send, E: std::error::Error + Send + Sync + 'static, { - SseFramer.frame(input).map(|frame| { - let frame = frame.map_err(|error| Error::StreamFraming(error.to_string()))?; - decode_anthropic_sse_frame(frame) + frames(input, SseCodec::default()).map(|event| { + decode_anthropic_sse_frame(event.map_err(|error| Error::StreamFraming(error.to_string()))?) }) } @@ -203,9 +199,10 @@ where B: Buf + Send, E: std::error::Error + Send + Sync + 'static, { - AwsEventStreamFramer.frame(input).map(|frame| { - let frame = frame.map_err(|error| Error::StreamFraming(error.to_string()))?; - decode_bedrock_anthropic_frame(frame) + frames(input, AwsEventStreamCodec).map(|message| { + decode_bedrock_anthropic_frame( + message.map_err(|error| Error::StreamFraming(error.to_string()))?, + ) }) } @@ -247,12 +244,10 @@ mod tests { #[test] fn decodes_citations_delta_events() { - let event = decode_anthropic_sse_frame(SseFrame { + let event = decode_anthropic_sse_frame(SseEvent { event: Some("content_block_delta".into()), - data: Some( - r#"{"type":"content_block_delta","index":0,"delta":{"type":"citations_delta","citation":{"type":"char_location"}}}"# - .into(), - ), + data: r#"{"type":"content_block_delta","index":0,"delta":{"type":"citations_delta","citation":{"type":"char_location"}}}"# + .into(), id: None, retry: None, }) diff --git a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/support.rs b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/support.rs index 30e46f93248..e1797784b84 100644 --- a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/support.rs +++ b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager/support.rs @@ -118,7 +118,7 @@ pub(super) struct ParityCase { pub(super) fn parity_cases() -> Vec { serde_json::from_str(include_str!(concat!( env!("CARGO_MANIFEST_DIR"), - "/../../../tests/test_litellm/secret_managers/hashicorp_vault_parity.json" + "/../../../tests/unit/secret_managers/hashicorp_vault_parity.json" ))) .unwrap() } diff --git a/litellm-rust/crates/secrets/PARITY.md b/litellm-rust/crates/secrets/PARITY.md index aeed4ba4b83..9439da93d65 100644 --- a/litellm-rust/crates/secrets/PARITY.md +++ b/litellm-rust/crates/secrets/PARITY.md @@ -32,7 +32,7 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo `_SecretManagerRuntime` is a private implementation detail, not a replacement SDK class. Its async methods return Futures; public `async def` methods retain lazy coroutine creation and `asyncio.create_task` support. Passing the same names and arguments is insufficient to claim parity until the remaining return-value, error, cache and configuration differences above are closed -## [tests/test_litellm/secret_managers/test_aws_secret_manager_replication.py](../../../tests/test_litellm/secret_managers/test_aws_secret_manager_replication.py) +## [tests/unit/secret_managers/test_aws_secret_manager_replication.py](../../../tests/unit/secret_managers/test_aws_secret_manager_replication.py) | Python test | Rust coverage or boundary | | --- | --- | @@ -48,7 +48,7 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo | `test_replicate_secret_http_error_raises` | [direct_replication_returns_response_or_service_error](../secrets-aws/tests/secret_manager/writes.rs) | | `test_replicate_secret_timeout_raises` | [write_and_replication_timeouts_remain_errors](../secrets-aws/tests/secret_manager/writes.rs) | -## [tests/test_litellm/secret_managers/test_aws_secret_manager_rotation.py](../../../tests/test_litellm/secret_managers/test_aws_secret_manager_rotation.py) +## [tests/unit/secret_managers/test_aws_secret_manager_rotation.py](../../../tests/unit/secret_managers/test_aws_secret_manager_rotation.py) | Python test | Rust coverage or boundary | | --- | --- | @@ -59,7 +59,7 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo | `test_write_secret_to_name_inside_recovery_window_restores_and_stores_new_value` | [recovery_window_alias_is_restored_updated_and_tagged](../secrets-aws/tests/secret_manager/writes.rs) | | `test_write_secret_to_live_existing_name_still_fails_without_overwriting` | [create_failure_does_not_overwrite_an_alias_without_a_deletion_date](../secrets-aws/tests/secret_manager/writes.rs) | -## [tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py](../../../tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py) +## [tests/unit/secret_managers/test_aws_secret_manager_v2.py](../../../tests/unit/secret_managers/test_aws_secret_manager_v2.py) | Python test | Rust coverage or boundary | | --- | --- | @@ -70,14 +70,14 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo | `test_prepare_request_explicit_bedrock_runtime_endpoint_param_still_wins` | [endpoint_overrides_replace_the_service_and_override_the_region](../secrets-aws/tests/secret_manager/configuration.rs) | | `test_prepare_request_env_bedrock_runtime_endpoint_still_wins` | [endpoint_overrides_replace_the_service_and_override_the_region](../secrets-aws/tests/secret_manager/configuration.rs) | -## [tests/test_litellm/secret_managers/test_base_secret_manager.py](../../../tests/test_litellm/secret_managers/test_base_secret_manager.py) +## [tests/unit/secret_managers/test_base_secret_manager.py](../../../tests/unit/secret_managers/test_base_secret_manager.py) | Python test | Rust coverage or boundary | | --- | --- | | `test_raise_if_unsafe_secret_name_rejects_traversal_and_line_breaks` | [names_reject_path_traversal_and_control_characters](../secrets-types/tests/rotation.rs) | | `test_raise_if_unsafe_secret_name_allows_legitimate_aliases` | [names_allow_safe_values](../secrets-types/tests/rotation.rs) | -## [tests/test_litellm/secret_managers/test_custom_secret_manager.py](../../../tests/test_litellm/secret_managers/test_custom_secret_manager.py) +## [tests/unit/secret_managers/test_custom_secret_manager.py](../../../tests/unit/secret_managers/test_custom_secret_manager.py) | Python test | Rust coverage or boundary | | --- | --- | @@ -89,7 +89,7 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo | `test_custom_secret_manager_integration_with_litellm` | [manager_strings_are_coerced_like_literal_eval](../secrets/tests/resolution.rs) | | `test_minimal_custom_secret_manager` | Exercises the Python example subclass itself or Python default methods. Caller-authored Python implementations remain Python callbacks; resolver integration is covered by `manager_strings_are_coerced_like_literal_eval` | -## [tests/test_litellm/secret_managers/test_cyberark_secret_manager.py](../../../tests/test_litellm/secret_managers/test_cyberark_secret_manager.py) +## [tests/unit/secret_managers/test_cyberark_secret_manager.py](../../../tests/unit/secret_managers/test_cyberark_secret_manager.py) | Python test | Rust coverage or boundary | | --- | --- | @@ -97,7 +97,7 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo | `test_async_write_matches_parity_fixture` | [writes_match_python_parity_fixture](../secrets-cyberark/tests/secret_manager/writes.rs) | | `test_missing_credentials_raise_value_error` | [new_validates_credentials_before_license_and_configuration](../secrets-cyberark/tests/secret_manager/configuration.rs) | -## [tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py](../../../tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py) +## [tests/unit/secret_managers/test_get_azure_ad_token_provider.py](../../../tests/unit/secret_managers/test_get_azure_ad_token_provider.py) | Python test | Rust coverage or boundary | | --- | --- | @@ -115,7 +115,7 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo | `test_get_azure_ad_token_provider_prefers_workload_identity_over_managed_identity` | Credential-selection contract belongs to `litellm-auth-azure`, not a secrets crate | | `test_get_azure_ad_token_provider_defaults_to_default_azure_credential` | Credential-selection contract belongs to `litellm-auth-azure`, not a secrets crate | -## [tests/test_litellm/secret_managers/test_hashicorp_secret_manager.py](../../../tests/test_litellm/secret_managers/test_hashicorp_secret_manager.py) +## [tests/unit/secret_managers/test_hashicorp_secret_manager.py](../../../tests/unit/secret_managers/test_hashicorp_secret_manager.py) | Python test | Rust coverage or boundary | | --- | --- | @@ -130,13 +130,13 @@ The rollout decision comes from [catalog.py](../../../litellm/rust_bridge/catalo | `test_tls_login_uses_login_namespace` | [tls_login_posts_the_role_and_uses_the_client_identity](../secrets-hashicorp/tests/secret_manager/configuration.rs) | | `test_configuration_matches_native_parity_fixture` | [configuration_matches_python_parity_fixture](../secrets-hashicorp/tests/secret_manager/configuration.rs) | -## [tests/test_litellm/secret_managers/test_secret_manager_handler.py](../../../tests/test_litellm/secret_managers/test_secret_manager_handler.py) +## [tests/unit/secret_managers/test_secret_manager_handler.py](../../../tests/unit/secret_managers/test_secret_manager_handler.py) | Python test | Rust coverage or boundary | | --- | --- | | `test_azure_key_vault_matches_rust_parity_fixture` | [parity_fixture_matches_python_backend_contract](../secrets-azure/tests/key_vault.rs) | -## [tests/test_litellm/secret_managers/test_secret_managers_main.py](../../../tests/test_litellm/secret_managers/test_secret_managers_main.py) +## [tests/unit/secret_managers/test_secret_managers_main.py](../../../tests/unit/secret_managers/test_secret_managers_main.py) | Python test | Rust coverage or boundary | | --- | --- | diff --git a/litellm/__init__.py b/litellm/__init__.py index e334fbe8ca8..676c735b9e8 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -50,6 +50,7 @@ from litellm.types.integrations.datadog import DatadogInitParams from litellm.types.integrations.newrelic import NewRelicInitParams from litellm.litellm_core_utils.core_helpers import drop_params_env_flag from litellm.types.integrations.pointfive import PointFiveInitParams +from litellm.types.integrations.zerobus import ZerobusInitParams from litellm._logging import ( set_verbose, _turn_on_debug, @@ -157,6 +158,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "deepeval", "s3_v2", "pointfive", + "zerobus", "aws_sqs", "vector_store_pre_call_hook", "dotprompt", @@ -442,6 +444,7 @@ datadog_llm_observability_params: Optional[Union[DatadogLLMObsInitParams, Dict]] datadog_params: Optional[Union[DatadogInitParams, Dict]] = None newrelic_params: Optional[Union[NewRelicInitParams, Dict]] = None pointfive_params: Optional[Union[PointFiveInitParams, Mapping[str, object]]] = None +zerobus_params: Optional[Union[ZerobusInitParams, Mapping[str, object]]] = None aws_sqs_callback_params: Optional[Dict] = None generic_logger_headers: Optional[Dict] = None default_key_generate_params: Optional[Dict] = None diff --git a/litellm/anthropic_interface/exceptions/__init__.py b/litellm/anthropic_interface/exceptions/__init__.py index 7f2de0e60dc..7c3cea0a28a 100644 --- a/litellm/anthropic_interface/exceptions/__init__.py +++ b/litellm/anthropic_interface/exceptions/__init__.py @@ -2,7 +2,9 @@ from .exception_mapping_utils import ( ANTHROPIC_ERROR_TYPE_MAP, + AnthropicErrorSseFrame, AnthropicExceptionMapping, + anthropic_error_sse_frame, ) from .exceptions import ( AnthropicErrorDetail, @@ -14,6 +16,8 @@ __all__ = [ "ANTHROPIC_ERROR_TYPE_MAP", "AnthropicErrorDetail", "AnthropicErrorResponse", + "AnthropicErrorSseFrame", "AnthropicErrorType", "AnthropicExceptionMapping", + "anthropic_error_sse_frame", ] diff --git a/litellm/anthropic_interface/exceptions/exception_mapping_utils.py b/litellm/anthropic_interface/exceptions/exception_mapping_utils.py index d9c9925275b..eb3ec8aaee2 100644 --- a/litellm/anthropic_interface/exceptions/exception_mapping_utils.py +++ b/litellm/anthropic_interface/exceptions/exception_mapping_utils.py @@ -4,11 +4,12 @@ Utilities for mapping exceptions to Anthropic error format. Similar to litellm/litellm_core_utils/exception_mapping_utils.py but for Anthropic response format. """ +import json from typing import Final from litellm.litellm_core_utils.safe_json_loads import safe_json_loads -from .exceptions import AnthropicErrorResponse, AnthropicErrorType +from .exceptions import AnthropicErrorDetail, AnthropicErrorResponse, AnthropicErrorType # HTTP status code -> Anthropic error type # Source: https://docs.anthropic.com/en/api/errors @@ -166,3 +167,36 @@ class AnthropicExceptionMapping: message=message, request_id=request_id, ) + + +class AnthropicErrorSseFrame(str): + """One `event: error` frame, for a stream that fails once the response headers are out. + + Anthropic clients pick stream events by the `event:` name, so a frame carrying only a `data:` + line is skipped and the failure never reaches the caller. The frame remembers the status and + body it was built from, so a stream that fails before its first byte can still answer as a + JSON error with that exact status instead of a 200 that only says `api_error` + """ + + status_code: int + error_response: AnthropicErrorResponse + + def __new__(cls, status_code: int, error_response: AnthropicErrorResponse) -> "AnthropicErrorSseFrame": + frame: Final = super().__new__(cls, f"event: error\ndata: {json.dumps(error_response)}\n\n") + frame.status_code = status_code + frame.error_response = error_response + return frame + + def json_body(self, call_id: str | None) -> AnthropicErrorResponse: + if call_id is None: + return self.error_response + detail: Final[AnthropicErrorDetail] = {**self.error_response["error"], "litellm_call_id": call_id} + body: Final[AnthropicErrorResponse] = {**self.error_response, "error": detail} + return body + + +def anthropic_error_sse_frame(status_code: int, raw_message: str) -> AnthropicErrorSseFrame: + return AnthropicErrorSseFrame( + status_code, + AnthropicExceptionMapping.transform_to_anthropic_error(status_code=status_code, raw_message=raw_message), + ) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 819a279a43c..246ac4fd369 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -706,6 +706,10 @@ def _get_batch_job_usage_from_response_body( if ResponseAPILoggingUtils._is_response_api_usage(_usage_dict): return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(_usage_dict) usage: Final[Usage] = Usage(**_usage_dict) + if custom_llm_provider == "xai": + from litellm.llms.xai.chat.transformation import XAIChatConfig + + XAIChatConfig.fold_reasoning_tokens_into_completion(usage) return usage diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 76b6c73b375..f977fc03891 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -31,6 +31,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.openai.openai import OpenAIBatchesAPI from litellm.llms.vertex_ai.batches.handler import VertexAIBatchPrediction +from litellm.llms.xai.batches.handler import XAIBatchesHandler from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( CancelBatchRequest, @@ -59,6 +60,7 @@ openai_batches_instance: Final = OpenAIBatchesAPI() azure_batches_instance: Final = AzureBatchesAPI() vertex_ai_batches_instance: Final = VertexAIBatchPrediction(gcs_bucket_name="") anthropic_batches_instance: Final = AnthropicBatchesHandler() +xai_batches_instance: Final = XAIBatchesHandler() base_llm_http_handler = BaseLLMHTTPHandler() ################################################# @@ -105,10 +107,22 @@ def _resolve_timeout( @client async def acreate_batch( completion_window: Literal["24h"], - endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses", "/v1/ocr"], + endpoint: Literal[ + "/v1/chat/completions", + "/v1/embeddings", + "/v1/completions", + "/v1/responses", + "/v1/ocr", + "/v1/images/generations", + "/v1/images/edits", + "/v1/videos/generations", + "/v1/videos", + "/v1/videos/edits", + "/v1/videos/extensions", + ], input_file_id: str, custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "mistral" + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "mistral", "xai" ] = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, @@ -157,10 +171,22 @@ async def acreate_batch( @client def create_batch( completion_window: Literal["24h"], - endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses", "/v1/ocr"], + endpoint: Literal[ + "/v1/chat/completions", + "/v1/embeddings", + "/v1/completions", + "/v1/responses", + "/v1/ocr", + "/v1/images/generations", + "/v1/images/edits", + "/v1/videos/generations", + "/v1/videos", + "/v1/videos/edits", + "/v1/videos/extensions", + ], input_file_id: str, custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "mistral" + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "mistral", "xai" ] = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, @@ -243,6 +269,14 @@ def create_batch( model=model, ) return response + if custom_llm_provider == LlmProviders.XAI.value: + return xai_batches_instance.create_batch( + _is_async=_is_async, + create_batch_data=_create_batch_request, + api_base=optional_params.api_base, + api_key=optional_params.api_key, + timeout=timeout, + ) api_base: str | None = None if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there @@ -345,7 +379,7 @@ def create_batch( async def aretrieve_batch( batch_id: str, custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral" + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral", "xai" ] = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, @@ -393,10 +427,18 @@ def _handle_retrieve_batch_providers_without_provider_config( _retrieve_batch_request: RetrieveBatchRequest, _is_async: bool, custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral" + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral", "xai" ] = "openai", logging_obj: LiteLLMLoggingObj | None = None, ): + if custom_llm_provider == LlmProviders.XAI.value: + return xai_batches_instance.retrieve_batch( + _is_async=_is_async, + batch_id=batch_id, + api_base=optional_params.api_base, + api_key=optional_params.api_key, + timeout=timeout, + ) api_base: str | None = None if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there @@ -518,7 +560,7 @@ def _handle_retrieve_batch_providers_without_provider_config( def retrieve_batch( batch_id: str, custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral" + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral", "xai" ] = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, @@ -741,6 +783,15 @@ def list_batches( timeout = 600.0 _is_async: Final = kwargs.pop("alist_batches", False) is True + if custom_llm_provider == LlmProviders.XAI.value: + return xai_batches_instance.list_batches( + _is_async=_is_async, + api_base=optional_params.api_base, + api_key=optional_params.api_key, + timeout=timeout, + after=after, + limit=limit, + ) if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there api_base = ( @@ -837,7 +888,7 @@ def list_batches( async def acancel_batch( batch_id: str, model: str | None = None, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "litellm_proxy"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "litellm_proxy", "xai"] = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, extra_body: dict[str, str] | None = None, @@ -883,7 +934,7 @@ async def acancel_batch( def cancel_batch( batch_id: str, model: str | None = None, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "litellm_proxy"] | str = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "litellm_proxy", "xai"] | str = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, extra_body: dict[str, str] | None = None, @@ -933,6 +984,14 @@ def cancel_batch( ) _is_async: Final = kwargs.pop("acancel_batch", False) is True + if custom_llm_provider == LlmProviders.XAI.value: + return xai_batches_instance.cancel_batch( + _is_async=_is_async, + batch_id=batch_id, + api_base=optional_params.api_base, + api_key=optional_params.api_key, + timeout=timeout, + ) api_base: str | None = None if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: api_base = ( diff --git a/litellm/constants.py b/litellm/constants.py index e7ba1f6b07f..8316761c95b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1952,6 +1952,15 @@ SENTRY_DENYLIST: Final = [ "auth_token", "jwt_token", "private_key", + "authorization", + "api-key", + "x-api-key", + "x-goog-api-key", + "ocp-apim-subscription-key", + "x-litellm-api-key", + "x-mcp-auth", + "cookie", + "set-cookie", "SLACK_WEBHOOK_URL", "ALERTING_WEBHOOK_URL", "webhook_url", @@ -1974,6 +1983,12 @@ SENTRY_DENYLIST: Final = [ ] SENTRY_PII_DENYLIST: Final = [ "user_id", + "user_email", + "end_user_id", + "user_api_key_hash", + "user_api_key_user_id", + "user_api_key_user_email", + "user_api_key_end_user_id", "email", "phone", "address", diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 01670be74c8..4e3b92edc89 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -10,6 +10,7 @@ import os from collections.abc import AsyncIterator, Awaitable, Callable, Generator, Sequence from contextlib import AbstractAsyncContextManager from functools import partial +from importlib.metadata import version from types import MappingProxyType from typing import Final, TypeAlias, TypeVar, cast @@ -34,8 +35,16 @@ _TransportContext: TypeAlias = AbstractAsyncContextManager[_TransportStreams] from mcp.types import ( METHOD_NOT_FOUND, REQUEST_TIMEOUT, + ClientCapabilities, + ElicitationCapability, + FormElicitationCapability, GetPromptRequestParams, GetPromptResult, + Implementation, + InitializedNotification, + InitializeRequest, + InitializeRequestParams, + InitializeResult, InputRequiredResult, ListPromptsResult, ListResourcesResult, @@ -44,12 +53,14 @@ from mcp.types import ( PaginatedResult, Prompt, ResourceTemplate, + SamplingCapability, ServerNotification, + UrlElicitationCapability, ) from mcp.types import CallToolRequestParams as MCPCallToolRequestParams from mcp.types import CallToolResult as MCPCallToolResult from mcp.types import Tool as MCPTool -from pydantic import AnyUrl +from pydantic import AnyUrl, TypeAdapter from litellm._logging import verbose_logger from litellm.constants import ( @@ -64,11 +75,13 @@ from litellm.proxy._experimental.mcp_server.mcp_debug import capture_upstream_er from litellm.proxy._experimental.mcp_server.result_conversion import error_text_result from litellm.types.llms.custom_http import VerifyTypes from litellm.types.mcp import ( + MCP_LEGACY_VERSIONS, MCPAuth, MCPAuthType, MCPStdioConfig, MCPTransport, MCPTransportType, + MCPUpstreamProtocol, credential_redirect_hook, has_header, without_header, @@ -386,7 +399,9 @@ class MCPClient: sampling_callback: Callable | None = None, elicitation_callback: Callable | None = None, logging_callback: Callable | None = None, + protocol_version: MCPUpstreamProtocol = "auto", ): + self.protocol_version: MCPUpstreamProtocol = TypeAdapter(MCPUpstreamProtocol).validate_python(protocol_version) self.server_url: str = server_url self.transport_type: MCPTransport = transport_type self.auth_type: MCPAuthType = auth_type @@ -525,6 +540,35 @@ class MCPClient: return safe_env + async def _initialize_session(self, session: ClientSession) -> InitializeResult: + if self.protocol_version == "auto": + automatic: Final = await session.initialize() + if automatic.protocol_version not in MCP_LEGACY_VERSIONS: + raise MCPError(code=-32022, message="Upstream selected an unsupported MCP protocol version") + return automatic + result: Final = await session.send_request( + InitializeRequest( + params=InitializeRequestParams( + protocol_version=self.protocol_version, + client_info=Implementation(name="litellm", version=version("litellm")), + capabilities=ClientCapabilities( + sampling=SamplingCapability() if self._sampling_callback is not None else None, + elicitation=ElicitationCapability( + form=FormElicitationCapability(), url=UrlElicitationCapability() + ) + if self._elicitation_callback is not None + else None, + ), + ) + ), + InitializeResult, + ) + if result.protocol_version != self.protocol_version: + raise MCPError(code=-32022, message="Upstream did not accept the configured MCP protocol version") + session.adopt(result) + await session.send_notification(InitializedNotification()) + return result + async def _execute_session_operation( self, transport_ctx: _TransportContext, @@ -579,7 +623,7 @@ class MCPClient: ) session: Final = await session_ctx.__aenter__() try: - init_result: Final = await session.initialize() + init_result: Final = await self._initialize_session(session) instructions: Final = getattr(init_result, "instructions", None) self._last_initialize_instructions = ( instructions.strip() or None if isinstance(instructions, str) else None diff --git a/litellm/files/main.py b/litellm/files/main.py index 72832aeccc9..723784795b0 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -28,12 +28,15 @@ FileCreateProvider = Literal[ "manus", "anthropic", "mistral", + "xai", ] FileRetrieveProvider = Literal[ - "openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic", "mistral" + "openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic", "mistral", "xai" ] -FileDeleteProvider = Literal["openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic", "mistral"] -FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic", "mistral"] +FileDeleteProvider = Literal[ + "openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic", "mistral", "xai" +] +FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic", "mistral", "xai"] import litellm from litellm import get_secret_str from litellm.files.streaming import FileContentStreamingResponse @@ -49,6 +52,8 @@ from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.openai.common_utils import get_openai_credentials from litellm.llms.openai.openai import FileDeleted, FileObject, OpenAIFilesAPI from litellm.llms.vertex_ai.files.handler import VertexAIFilesHandler +from litellm.llms.xai.batches.handler import XAIBatchesHandler +from litellm.llms.xai.batches.transformation import is_xai_batch_results_id from litellm.types.llms.openai import ( CreateFileRequest, FileContentRequest, @@ -103,6 +108,7 @@ openai_files_instance: Final = OpenAIFilesAPI() azure_files_instance: Final = AzureOpenAIFilesAPI() vertex_ai_files_instance: Final = VertexAIFilesHandler() bedrock_files_instance: Final = BedrockFilesHandler() +xai_batch_results_instance: Final = XAIBatchesHandler() ################################################# @@ -920,6 +926,15 @@ def file_content( client=client, ) + if custom_llm_provider == LlmProviders.XAI.value and is_xai_batch_results_id(file_id): + return xai_batch_results_instance.batch_results_content( + _is_async=_is_async, + batch_id=file_id, + api_base=optional_params.api_base, + api_key=optional_params.api_key, + timeout=timeout, + ) + # Check if provider has a custom files config (e.g., Anthropic, Manus) provider_config: Final = ProviderConfigManager.get_provider_files_config( model="", diff --git a/litellm/integrations/callback_configs.json b/litellm/integrations/callback_configs.json index 5bd8aca55fa..4e72075dc5c 100644 --- a/litellm/integrations/callback_configs.json +++ b/litellm/integrations/callback_configs.json @@ -406,6 +406,45 @@ }, "description": "PointFive Logging Integration" }, + { + "id": "zerobus", + "displayName": "Databricks Zerobus", + "logo": "databricks.svg", + "supports_key_team_logging": false, + "dynamic_params": { + "ZEROBUS_WORKSPACE_URL": { + "type": "text", + "ui_name": "Workspace URL", + "description": "Databricks workspace URL, e.g. https://dbc-a1b2c3d4-e5f6.cloud.databricks.com", + "required": true + }, + "ZEROBUS_SERVER_ENDPOINT": { + "type": "text", + "ui_name": "Zerobus Endpoint", + "description": "Zerobus ingest endpoint, e.g. https://.zerobus..cloud.databricks.com", + "required": true + }, + "ZEROBUS_CLIENT_ID": { + "type": "text", + "ui_name": "Service Principal Client ID", + "description": "OAuth client id of a service principal with USE CATALOG, USE SCHEMA, SELECT and MODIFY on the table", + "required": true + }, + "ZEROBUS_CLIENT_SECRET": { + "type": "password", + "ui_name": "Service Principal Client Secret", + "description": "OAuth client secret of the service principal", + "required": true + }, + "ZEROBUS_TABLE_NAME": { + "type": "text", + "ui_name": "Table", + "description": "Fully qualified Unity Catalog table, catalog.schema.table, created with the LiteLLM trace schema", + "required": true + } + }, + "description": "Databricks Zerobus Ingest Logging Integration" + }, { "id": "s3", "displayName": "S3", diff --git a/litellm/integrations/levo/README.md b/litellm/integrations/levo/README.md index 5296acb7ff4..1fbd202d9a5 100644 --- a/litellm/integrations/levo/README.md +++ b/litellm/integrations/levo/README.md @@ -92,7 +92,7 @@ litellm/integrations/levo/ ## Testing -See the test files in `tests/test_litellm/integrations/levo/`: +See the test files in `tests/unit/integrations/levo/`: - `test_levo.py`: Unit tests for configuration - `test_levo_integration.py`: Integration tests for callback registration diff --git a/litellm/integrations/zerobus/__init__.py b/litellm/integrations/zerobus/__init__.py new file mode 100644 index 00000000000..b1f5bc2ca40 --- /dev/null +++ b/litellm/integrations/zerobus/__init__.py @@ -0,0 +1,5 @@ +"""Databricks Zerobus logging integration for LiteLLM.""" + +from litellm.integrations.zerobus.logger import ZerobusLogger + +__all__ = ("ZerobusLogger",) diff --git a/litellm/integrations/zerobus/client.py b/litellm/integrations/zerobus/client.py new file mode 100644 index 00000000000..bf3a9e3e269 --- /dev/null +++ b/litellm/integrations/zerobus/client.py @@ -0,0 +1,161 @@ +""" +Writes rows to a Unity Catalog table through the Zerobus Ingest REST API. + +Zerobus only accepts a Databricks OAuth token minted for its own resource and scoped to +the target table's privileges, so the client mints that token itself with the service +principal's client credentials and reuses it until shortly before it expires. +""" + +import asyncio +import base64 +import json +import time +from collections.abc import Callable, Mapping, Sequence +from typing import Final + +import httpx +from pydantic import BaseModel, ValidationError + +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.types.integrations.zerobus import ( + RETRYABLE_INGEST_STATUS_CODES, + TOKEN_REFRESH_LEEWAY_SECONDS, + ZerobusAccessToken, + ZerobusConnection, + ZerobusIngestFailure, +) + +TOKEN_PATH: Final = "/oidc/v1/token" +OAUTH_SCOPE: Final = "all-apis" + + +class _TokenResponse(BaseModel): + access_token: str + expires_in: float = 3600 + + +class ZerobusIngestError(Exception): + """A batch could not be written and the failure is worth retrying.""" + + +def zerobus_resource(workspace_id: str) -> str: + return f"api://databricks/workspaces/{workspace_id}/zerobusDirectWriteApi" + + +def authorization_details(table_name: str) -> str: + """The Unity Catalog privileges Zerobus requires the token to carry, as the token endpoint expects them.""" + catalog, schema, _table = table_name.split(".", 2) + return json.dumps( + ( + { + "type": "unity_catalog_privileges", + "privileges": ("USE CATALOG",), + "object_type": "CATALOG", + "object_full_path": catalog, + }, + { + "type": "unity_catalog_privileges", + "privileges": ("USE SCHEMA",), + "object_type": "SCHEMA", + "object_full_path": f"{catalog}.{schema}", + }, + { + "type": "unity_catalog_privileges", + "privileges": ("SELECT", "MODIFY"), + "object_type": "TABLE", + "object_full_path": table_name, + }, + ) + ) + + +def insert_url(connection: ZerobusConnection) -> str: + return f"{connection.server_endpoint.rstrip('/')}/zerobus/v1/tables/{connection.table_name}/insert" + + +def token_url(connection: ZerobusConnection) -> str: + return f"{connection.workspace_url.rstrip('/')}{TOKEN_PATH}" + + +def _basic_auth(client_id: str, client_secret: str) -> str: + return "Basic " + base64.b64encode(f"{client_id}:{client_secret}".encode()).decode() + + +def _status_failure(what: str, error: httpx.HTTPStatusError) -> ZerobusIngestFailure: + status: Final = error.response.status_code + return ZerobusIngestFailure( + detail=f"{what} returned {status}: {error.response.text}"[:500], + retryable=status in RETRYABLE_INGEST_STATUS_CODES, + ) + + +class ZerobusIngestClient: + def __init__( + self, + connection: ZerobusConnection, + http_client: AsyncHTTPHandler, + clock: Callable[[], float] = time.time, + ) -> None: + self.connection: Final = connection + self.http_client: Final = http_client + self.clock: Final = clock + self._token: ZerobusAccessToken | None = None + self._token_lock: Final = asyncio.Lock() + + async def insert(self, rows: Sequence[Mapping[str, object]]) -> ZerobusIngestFailure | None: + """Write ``rows`` as one request. ``None`` means Zerobus accepted every row.""" + token: Final = await self.access_token() + if isinstance(token, ZerobusIngestFailure): + return token + try: + await self.http_client.post( + insert_url(self.connection), + content=json.dumps([dict(row) for row in rows]).encode(), + headers={"Content-Type": "application/json", "Authorization": f"Bearer {token.value}"}, + ) + except httpx.HTTPStatusError as error: + if error.response.status_code == 401: + self._token = None + return ZerobusIngestFailure(detail="insert returned 401, token discarded", retryable=True) + return _status_failure("insert", error) + except (httpx.HTTPError, litellm.Timeout) as error: + return ZerobusIngestFailure(detail=f"insert failed: {error}", retryable=True) + return None + + async def access_token(self) -> ZerobusAccessToken | ZerobusIngestFailure: + """The cached token while it has more than the leeway left, otherwise a fresh one.""" + async with self._token_lock: + cached: Final = self._token + if cached is not None and cached.expires_at - self.clock() > TOKEN_REFRESH_LEEWAY_SECONDS: + return cached + minted: Final = await self._mint_token() + if isinstance(minted, ZerobusAccessToken): + self._token = minted + return minted + + async def _mint_token(self) -> ZerobusAccessToken | ZerobusIngestFailure: + connection: Final = self.connection + try: + response: Final = await self.http_client.post( + token_url(connection), + data={ + "grant_type": "client_credentials", + "scope": OAUTH_SCOPE, + "resource": zerobus_resource(connection.workspace_id), + "authorization_details": authorization_details(connection.table_name), + }, + headers={ + "Content-Type": "application/x-www-form-urlencoded", + "Authorization": _basic_auth(connection.client_id, connection.client_secret), + }, + ) + except httpx.HTTPStatusError as error: + return _status_failure("token request", error) + except (httpx.HTTPError, litellm.Timeout) as error: + return ZerobusIngestFailure(detail=f"token request failed: {error}", retryable=True) + try: + parsed: Final = _TokenResponse.model_validate_json(response.text) + except ValidationError as error: + return ZerobusIngestFailure(detail=f"token response was not understood: {error}", retryable=False) + return ZerobusAccessToken(value=parsed.access_token, expires_at=self.clock() + parsed.expires_in) diff --git a/litellm/integrations/zerobus/logger.py b/litellm/integrations/zerobus/logger.py new file mode 100644 index 00000000000..e2007218c8e --- /dev/null +++ b/litellm/integrations/zerobus/logger.py @@ -0,0 +1,230 @@ +"""Databricks Zerobus logging integration.""" + +import asyncio +from collections.abc import Mapping +from datetime import datetime +from typing import Final +from urllib.parse import urlsplit + +import litellm +from litellm._logging import verbose_logger +from litellm.integrations.custom_batch_logger import CustomBatchLogger +from litellm.integrations.zerobus.client import ZerobusIngestClient, ZerobusIngestError +from litellm.integrations.zerobus.row import trace_row +from litellm.litellm_core_utils.redact_messages import ( + redacted_standard_logging_payload, + should_redact_message_logging, +) +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client, httpxSpecialProvider +from litellm.secret_managers.main import get_secret_str +from litellm.types.integrations.zerobus import ZerobusConnection, ZerobusInitParams + +_ENV_REFERENCE_PREFIX: Final = "os.environ/" + + +def _resolved_secret(value: str | None) -> str | None: + """Resolve a config value that may name a secret; an unset ``os.environ/NAME`` stays unresolved.""" + if value is None: + return None + resolved: Final = get_secret_str(value) + if resolved: + return resolved + return None if value.startswith(_ENV_REFERENCE_PREFIX) else value + + +def _configured_params() -> ZerobusInitParams: + configured: Final = litellm.zerobus_params + if isinstance(configured, ZerobusInitParams): + return configured + if isinstance(configured, Mapping): + return ZerobusInitParams.model_validate(configured) + return ZerobusInitParams() + + +def _setting(configured: str | None, env_var: str) -> str: + """Prefer the configured value, falling back to the environment the proxy UI writes.""" + value: Final = _resolved_secret(configured) or get_secret_str(env_var) + if not value: + raise ValueError( + f"zerobus logging requires {env_var}. Set it in the environment, or " + f"litellm_settings.zerobus_params.{env_var.removeprefix('ZEROBUS_').lower()} in config.yaml" + ) + return value + + +def _workspace_id(server_endpoint: str) -> str: + """The Zerobus endpoint is ``https://.zerobus..``, so the id is its first label.""" + host: Final = urlsplit(server_endpoint).hostname or "" + workspace_id: Final = host.split(".", 1)[0] + if not workspace_id.isdigit(): + raise ValueError( + f"ZEROBUS_SERVER_ENDPOINT {server_endpoint!r} does not look like " + "https://.zerobus..cloud.databricks.com" + ) + return workspace_id + + +def _table_name(configured: str | None) -> str: + table_name: Final = _setting(configured, "ZEROBUS_TABLE_NAME") + if table_name.count(".") != 2: + raise ValueError(f"ZEROBUS_TABLE_NAME {table_name!r} must be fully qualified as catalog.schema.table") + return table_name + + +def connection_for(params: ZerobusInitParams) -> ZerobusConnection: + """The connection configured right now, so a UI edit takes effect without a restart.""" + server_endpoint: Final = _setting(params.server_endpoint, "ZEROBUS_SERVER_ENDPOINT") + return ZerobusConnection( + workspace_url=_setting(params.workspace_url, "ZEROBUS_WORKSPACE_URL"), + workspace_id=_workspace_id(server_endpoint), + server_endpoint=server_endpoint, + client_id=_setting(params.client_id, "ZEROBUS_CLIENT_ID"), + client_secret=_setting(params.client_secret, "ZEROBUS_CLIENT_SECRET"), + table_name=_table_name(params.table_name), + ) + + +class ZerobusLogger(CustomBatchLogger): + preserve_events_added_during_flush = True + + def __init__( + self, + params: ZerobusInitParams | None = None, + client: ZerobusIngestClient | None = None, + start_periodic_flush: bool = True, + ) -> None: + resolved: Final = params if params is not None else _configured_params() + self.params: Final = resolved + self.given_client: Final = client + self._cached_client: ZerobusIngestClient | None = None + if client is None: + connection_for(resolved) + super().__init__( + flush_lock=asyncio.Lock(), + batch_size=resolved.batch_size, + flush_interval=resolved.flush_interval, + turn_off_message_logging=bool(resolved.turn_off_message_logging), + ) + self._flushing: bool = False + self._batch_flush_task: asyncio.Task[None] | None = None + self._periodic_flush_task: asyncio.Task[None] | None = ( + self._start_periodic_flush_task() if start_periodic_flush else None + ) + + @property + def client(self) -> ZerobusIngestClient: + """A client for the current connection, kept while the connection is unchanged so its token is reused.""" + if self.given_client is not None: + return self.given_client + connection: Final = connection_for(self.params) + cached: Final = self._cached_client + if cached is not None and cached.connection == connection: + return cached + fresh: Final = ZerobusIngestClient( + connection=connection, + http_client=get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback), + ) + self._cached_client = fresh + return fresh + + def _start_periodic_flush_task(self) -> asyncio.Task[None] | None: + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + return None + return loop.create_task(self.periodic_flush()) + + def _start_batch_flush_task(self) -> None: + if self._batch_flush_task is not None and not self._batch_flush_task.done(): + return + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + return + self._batch_flush_task = loop.create_task(self.flush_queue(skip_if_flushing=True)) + + def _flush_task_is_alive(self) -> bool: + task: Final = self._periodic_flush_task + return task is not None and not task.done() and not task.get_loop().is_closed() + + async def async_log_success_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: datetime, + end_time: datetime, + ) -> None: + await self._enqueue(kwargs) + + async def async_log_failure_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: datetime, + end_time: datetime, + ) -> None: + await self._enqueue(kwargs) + + async def _enqueue(self, kwargs: Mapping[str, object]) -> None: + try: + if not self._flush_task_is_alive(): + self._periodic_flush_task = self._start_periodic_flush_task() + + payload: Final = self._payload_for(kwargs) + if payload is None: + verbose_logger.debug("zerobus: event carried no standard_logging_object, skipping") + return + + if self._flushing and len(self.log_queue) >= self.max_queue_size: + verbose_logger.warning("zerobus: queue at %s rows during a flush, dropped a row", self.max_queue_size) + return + + self.log_queue.append(trace_row(payload)) + self._drop_overflow() + if len(self.log_queue) >= self.batch_size: + self._start_batch_flush_task() + except Exception: # noqa: BLE001 # logging must never break the request path + verbose_logger.exception("zerobus: failed to queue an event") + + def _payload_for(self, kwargs: Mapping[str, object]) -> Mapping[str, object] | None: + """The payload to buffer, redacted the way the framework redacts the success path.""" + details: Final = self.redact_standard_logging_payload_from_model_call_details( + dict(kwargs) # mutable-ok: both framework helpers take the call details as a dict + ) + payload: Final = details.get("standard_logging_object") + if not isinstance(payload, dict): + return None + if should_redact_message_logging(details): + return redacted_standard_logging_payload(payload) + return payload + + def _drop_overflow(self) -> None: + """Trim the oldest rows, except mid flush when the in-flight batch is the head of the queue.""" + if self._flushing: + return + overflow: Final = len(self.log_queue) - self.max_queue_size + if overflow <= 0: + return + del self.log_queue[:overflow] + verbose_logger.warning("zerobus: queue over %s rows, dropped %s oldest", self.max_queue_size, overflow) + + async def flush_queue(self, skip_if_flushing: bool = False) -> None: + if skip_if_flushing and self._flushing: + return + self._flushing = True + try: + await super().flush_queue() + finally: + self._flushing = False + + async def async_send_batch(self) -> None: + """A retryable failure propagates so the rows are kept; a permanent one drops them so the queue moves on.""" + rows: Final = tuple(self.log_queue) + if not rows: + return + failure: Final = await self.client.insert(rows) + if failure is None: + return + if failure.retryable: + raise ZerobusIngestError(failure.detail) + verbose_logger.error("zerobus: dropping %s rows, %s", len(rows), failure.detail) diff --git a/litellm/integrations/zerobus/row.py b/litellm/integrations/zerobus/row.py new file mode 100644 index 00000000000..c4da7975c44 --- /dev/null +++ b/litellm/integrations/zerobus/row.py @@ -0,0 +1,156 @@ +""" +Shape of one Delta table row per LiteLLM request. + +Zerobus validates every record against the target table and rejects unknown columns, so +the row is a fixed set of scalar columns for filtering plus JSON-encoded ``VARIANT`` +columns for anything nested. ``create_table_sql`` renders the matching DDL. +""" + +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + +TRACE_TABLE_COLUMNS: Final[Mapping[str, str]] = MappingProxyType( + { + "id": "STRING", + "trace_id": "STRING", + "session_id": "STRING", + "litellm_call_id": "STRING", + "call_type": "STRING", + "status": "STRING", + "model": "STRING", + "model_group": "STRING", + "model_id": "STRING", + "custom_llm_provider": "STRING", + "api_base": "STRING", + "stream": "BOOLEAN", + "cache_hit": "BOOLEAN", + "start_time": "TIMESTAMP", + "end_time": "TIMESTAMP", + "completion_start_time": "TIMESTAMP", + "response_time": "DOUBLE", + "prompt_tokens": "LONG", + "completion_tokens": "LONG", + "total_tokens": "LONG", + "response_cost": "DOUBLE", + "saved_cache_cost": "DOUBLE", + "api_key_hash": "STRING", + "api_key_alias": "STRING", + "team_id": "STRING", + "team_alias": "STRING", + "user_id": "STRING", + "org_id": "STRING", + "end_user": "STRING", + "requester_ip_address": "STRING", + "user_agent": "STRING", + "request_tags": "VARIANT", + "messages": "VARIANT", + "response": "VARIANT", + "error_str": "STRING", + "error_information": "VARIANT", + "metadata": "VARIANT", + "model_parameters": "VARIANT", + "hidden_params": "VARIANT", + "guardrail_information": "VARIANT", + "cost_breakdown": "VARIANT", + } +) + +_MICROSECONDS: Final = 1_000_000 + + +def create_table_sql(table_name: str) -> str: + columns: Final = ",\n".join(f" {name} {delta_type}" for name, delta_type in TRACE_TABLE_COLUMNS.items()) + return f"CREATE TABLE {table_name} (\n{columns}\n);" + + +def _text(payload: Mapping[str, object], key: str) -> str | None: + value: Final = payload.get(key) + return value if isinstance(value, str) else None + + +def _flag(payload: Mapping[str, object], key: str) -> bool | None: + value: Final = payload.get(key) + return value if isinstance(value, bool) else None + + +def _number(payload: Mapping[str, object], key: str) -> float | None: + value: Final = payload.get(key) + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + return float(value) + + +def _count(payload: Mapping[str, object], key: str) -> int | None: + value: Final = _number(payload, key) + return None if value is None else int(value) + + +def _timestamp_micros(payload: Mapping[str, object], key: str) -> int | None: + """Delta ``TIMESTAMP`` over Zerobus is epoch microseconds; LiteLLM keeps epoch seconds.""" + seconds: Final = _number(payload, key) + if seconds is None or seconds <= 0: + return None + return int(seconds * _MICROSECONDS) + + +def _json(payload: Mapping[str, object], key: str) -> str | None: + value: Final = payload.get(key) + return None if value is None else safe_dumps(value) + + +def _metadata(payload: Mapping[str, object]) -> Mapping[str, object]: + value: Final = payload.get("metadata") + return value if isinstance(value, Mapping) else MappingProxyType({}) + + +def trace_row(payload: Mapping[str, object]) -> Mapping[str, object]: + """One ``TRACE_TABLE_COLUMNS`` row for a ``StandardLoggingPayload``.""" + metadata: Final = _metadata(payload) + return MappingProxyType( + { + "id": _text(payload, "id"), + "trace_id": _text(payload, "trace_id"), + "session_id": _text(payload, "session_id"), + "litellm_call_id": _text(payload, "litellm_call_id"), + "call_type": _text(payload, "call_type"), + "status": _text(payload, "status"), + "model": _text(payload, "model"), + "model_group": _text(payload, "model_group"), + "model_id": _text(payload, "model_id"), + "custom_llm_provider": _text(payload, "custom_llm_provider"), + "api_base": _text(payload, "api_base"), + "stream": _flag(payload, "stream"), + "cache_hit": _flag(payload, "cache_hit"), + "start_time": _timestamp_micros(payload, "startTime"), + "end_time": _timestamp_micros(payload, "endTime"), + "completion_start_time": _timestamp_micros(payload, "completionStartTime"), + "response_time": _number(payload, "response_time"), + "prompt_tokens": _count(payload, "prompt_tokens"), + "completion_tokens": _count(payload, "completion_tokens"), + "total_tokens": _count(payload, "total_tokens"), + "response_cost": _number(payload, "response_cost"), + "saved_cache_cost": _number(payload, "saved_cache_cost"), + "api_key_hash": _text(metadata, "user_api_key_hash"), + "api_key_alias": _text(metadata, "user_api_key_alias"), + "team_id": _text(metadata, "user_api_key_team_id"), + "team_alias": _text(metadata, "user_api_key_team_alias"), + "user_id": _text(metadata, "user_api_key_user_id"), + "org_id": _text(metadata, "user_api_key_org_id"), + "end_user": _text(payload, "end_user"), + "requester_ip_address": _text(payload, "requester_ip_address"), + "user_agent": _text(payload, "user_agent"), + "request_tags": _json(payload, "request_tags"), + "messages": _json(payload, "messages"), + "response": _json(payload, "response"), + "error_str": _text(payload, "error_str"), + "error_information": _json(payload, "error_information"), + "metadata": _json(payload, "metadata"), + "model_parameters": _json(payload, "model_parameters"), + "hidden_params": _json(payload, "hidden_params"), + "guardrail_information": _json(payload, "guardrail_information"), + "cost_breakdown": _json(payload, "cost_breakdown"), + } + ) diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index 6294f3bc577..7049fdd1f39 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -52,6 +52,7 @@ from litellm.integrations.vantage.vantage_logger import VantageLogger from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( VectorStorePreCallHook, ) +from litellm.integrations.zerobus import ZerobusLogger from litellm.proxy.hooks.dynamic_rate_limiter import _PROXY_DynamicRateLimitHandler from litellm.proxy.hooks.dynamic_rate_limiter_v3 import _PROXY_DynamicRateLimitHandlerV3 @@ -97,6 +98,7 @@ class CustomLoggerRegistry: "deepeval": DeepEvalLogger, "s3_v2": S3Logger, "pointfive": PointFiveLogger, + "zerobus": ZerobusLogger, "aws_sqs": SQSLogger, "dynamic_rate_limiter": _PROXY_DynamicRateLimitHandler, "dynamic_rate_limiter_v3": _PROXY_DynamicRateLimitHandlerV3, diff --git a/litellm/litellm_core_utils/health_check_helpers.py b/litellm/litellm_core_utils/health_check_helpers.py index 22b8d850c83..a0f027cd58f 100644 --- a/litellm/litellm_core_utils/health_check_helpers.py +++ b/litellm/litellm_core_utils/health_check_helpers.py @@ -114,10 +114,9 @@ class HealthCheckHelpers: """ Health check for batch mode. - Calls list_batches for providers that support it (openai, hosted_vllm, azure, - vertex_ai). For all other providers (e.g. bedrock) the batch API surface doesn't - include list_batches, so we fall back to acompletion to verify connectivity and - credential validity instead. + Calls list_batches for providers that support it. For all other providers (e.g. bedrock) + the batch API surface doesn't include list_batches, so we fall back to acompletion to + verify connectivity and credential validity instead. """ import litellm @@ -132,10 +131,9 @@ class HealthCheckHelpers: litellm_params={"api_base": api_base} if api_base else None, ) - if custom_llm_provider in LIST_BATCHES_SUPPORTED_PROVIDERS: - return await litellm.alist_batches(**filtered_model_params) - else: + if custom_llm_provider not in LIST_BATCHES_SUPPORTED_PROVIDERS: return await litellm.acompletion(**model_params) + return await litellm.alist_batches(**{**filtered_model_params, "custom_llm_provider": custom_llm_provider}) @staticmethod async def _image_edit_health_check(edit_request: Callable[[], Awaitable["ImageResponse"]]) -> "ImageResponse": diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 83ab2bc11a2..0ef5bbf807a 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -43,8 +43,6 @@ from litellm.constants import ( DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT, EMPTY_MAPPING, PROVIDER_REQUEST_ID_HEADERS, - SENTRY_DENYLIST, - SENTRY_PII_DENYLIST, ) from litellm.cost_calculator import ( RealtimeAPITokenUsageProcessor, @@ -213,6 +211,7 @@ from ..integrations.s3 import S3Logger from ..integrations.s3_v2 import S3Logger as S3V2Logger from ..integrations.supabase import Supabase from ..integrations.traceloop import TraceloopLogger +from ..integrations.zerobus import ZerobusLogger from .exception_mapping_utils import _get_response_headers from .initialize_dynamic_callback_params import ( get_trusted_callback_params, @@ -380,9 +379,12 @@ _DEPLOYMENT_PRICING_KEYS: Final = ( "output_cost_per_token", "input_cost_per_token_batches", "output_cost_per_token_batches", + "input_cost_per_token_above_200k_tokens_batches", "input_cost_per_token_above_272k_tokens_batches", + "output_cost_per_token_above_200k_tokens_batches", "output_cost_per_token_above_272k_tokens_batches", "cache_read_input_token_cost_batches", + "cache_read_input_token_cost_above_200k_tokens_batches", "cache_read_input_token_cost_above_272k_tokens_batches", "cache_creation_input_token_cost_batches", "cache_creation_input_token_cost_above_272k_tokens_batches", @@ -4215,6 +4217,9 @@ class Logging(LiteLLMLoggingBaseClass): json_mode=False, litellm_params={}, ) + elif result is None: + verbose_logger.warning("LiteLLM: the anthropic_messages stream assembled no response, logging an empty one") + return litellm.ModelResponse(model=self.model) else: from litellm.types.llms.anthropic import AnthropicResponse @@ -4423,21 +4428,10 @@ def set_callbacks(callback_list, function_id=None): print_verbose("Package 'sentry_sdk' is missing. Installing it...") subprocess.check_call([sys.executable, "-m", "pip", "install", "sentry_sdk"]) import sentry_sdk - from sentry_sdk.scrubber import EventScrubber + from litellm.litellm_core_utils.sentry_scrubbing import build_sentry_init_options sentry_sdk_instance = sentry_sdk - sentry_trace_rate = os.environ.get("SENTRY_API_TRACE_RATE", "1.0") - sentry_sample_rate = ( - os.environ.get("SENTRY_API_SAMPLE_RATE") if "SENTRY_API_SAMPLE_RATE" in os.environ else "1.0" - ) - sentry_sdk_instance.init( - dsn=os.environ.get("SENTRY_DSN"), - traces_sample_rate=float(sentry_trace_rate), - sample_rate=float(sentry_sample_rate if sentry_sample_rate else 1.0), - send_default_pii=False, # Prevent sending Personal Identifiable Information - event_scrubber=EventScrubber(denylist=SENTRY_DENYLIST, pii_denylist=SENTRY_PII_DENYLIST), - environment=os.environ.get("SENTRY_ENVIRONMENT", "production"), - ) + sentry_sdk_instance.init(**build_sentry_init_options(os.environ)) capture_exception = sentry_sdk_instance.capture_exception add_breadcrumb = sentry_sdk_instance.add_breadcrumb elif callback == "slack": @@ -4660,6 +4654,14 @@ def _init_custom_logger_compatible_class( _pointfive_logger: Final = PointFiveLogger() _in_memory_loggers.append(_pointfive_logger) return _pointfive_logger + elif logging_integration == "zerobus": + for callback in _in_memory_loggers: + if isinstance(callback, ZerobusLogger): + return callback + + _zerobus_logger: Final = ZerobusLogger() + _in_memory_loggers.append(_zerobus_logger) + return _zerobus_logger elif logging_integration == "aws_sqs": for callback in _in_memory_loggers: if isinstance(callback, SQSLogger): @@ -5352,6 +5354,10 @@ def get_custom_logger_compatible_class( for callback in _in_memory_loggers: if isinstance(callback, PointFiveLogger): return callback + elif logging_integration == "zerobus": + for callback in _in_memory_loggers: + if isinstance(callback, ZerobusLogger): + return callback elif logging_integration == "aws_sqs": for callback in _in_memory_loggers: if isinstance(callback, SQSLogger): diff --git a/litellm/litellm_core_utils/sentry_scrubbing.py b/litellm/litellm_core_utils/sentry_scrubbing.py new file mode 100644 index 00000000000..4c14cabc2ab --- /dev/null +++ b/litellm/litellm_core_utils/sentry_scrubbing.py @@ -0,0 +1,152 @@ +from __future__ import annotations + +import re +from collections.abc import Callable, Mapping, Sequence +from functools import reduce +from typing import TYPE_CHECKING, Final, TypeAlias, cast + +from pydantic import JsonValue +from sentry_sdk.scrubber import DEFAULT_DENYLIST, DEFAULT_PII_DENYLIST, EventScrubber +from typing_extensions import ReadOnly, TypedDict + +from litellm.constants import ( + LENGTH_OF_LITELLM_GENERATED_KEY, + MINIMUM_CUSTOM_KEY_LENGTH, + SENTRY_DENYLIST, + SENTRY_PII_DENYLIST, +) +from litellm.secret_managers.main import str_to_bool + +if TYPE_CHECKING: + from sentry_sdk.types import Event, Hint + +EventScrubFn: TypeAlias = "Callable[[Event, Hint], Event]" +JsonPath: TypeAlias = tuple[str, ...] + +FILTERED: Final = "[Filtered]" +SEND_DEFAULT_PII_ENV: Final = "SENTRY_SEND_DEFAULT_PII" +SECRET_FIELD_NAMES: Final = tuple(DEFAULT_DENYLIST) + tuple(SENTRY_DENYLIST) +PII_FIELD_NAMES: Final = tuple(DEFAULT_PII_DENYLIST) + tuple(SENTRY_PII_DENYLIST) + +KEY_PREFIX: Final = "sk-" + + +def build_key_pattern(custom_key_minimum: int, generated_key_bytes: int) -> re.Pattern[str]: + generated_suffix_length: Final = (generated_key_bytes * 4 + 2) // 3 + floor: Final = min(custom_key_minimum - len(KEY_PREFIX), generated_suffix_length) + return re.compile(rf"{KEY_PREFIX}[A-Za-z0-9_-]{{{floor},}}") + + +LITELLM_KEY_PATTERN: Final = build_key_pattern(MINIMUM_CUSTOM_KEY_LENGTH, LENGTH_OF_LITELLM_GENERATED_KEY) +SOURCE_CONTEXT_KEYS: Final = frozenset({"pre_context", "context_line", "post_context"}) +STACK_FRAME_PATHS: Final = frozenset( + { + ("exception", "values", "*", "stacktrace", "frames", "*"), + ("threads", "values", "*", "stacktrace", "frames", "*"), + ("stacktrace", "frames", "*"), + } +) +MAX_SCRUB_DEPTH: Final = 64 +EMAIL_PATTERN: Final = re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9-]+(?:\.[A-Za-z0-9-]+)*\.[A-Za-z]{2,}") +SHA256_HEX_PATTERN: Final = re.compile(r"(? re.Pattern[str]: + names: Final = "|".join(re.escape(name) for name in field_names) + return re.compile( + rf"(?P(?{QUOTED_VALUE}|{BRACKETED_VALUE}|{BARE_VALUE})", + re.IGNORECASE, + ) + + +def build_string_scrubber(send_default_pii: bool) -> Callable[[str], str]: + field_names: Final = SECRET_FIELD_NAMES if send_default_pii else SECRET_FIELD_NAMES + PII_FIELD_NAMES + field_pattern: Final = build_repr_field_pattern(field_names) + value_patterns: Final = ( + (LITELLM_KEY_PATTERN,) if send_default_pii else (LITELLM_KEY_PATTERN, EMAIL_PATTERN, SHA256_HEX_PATTERN) + ) + + def scrub(text: str) -> str: + fields_scrubbed: Final = field_pattern.sub(_filtered_field, text) + return _substitute_all(value_patterns, fields_scrubbed) + + return scrub + + +def _filtered_field(match: re.Match[str]) -> str: + quote: Final = '"' if match.group("value").startswith('"') else "'" + return f"{match.group('field')}{quote}{FILTERED}{quote}" + + +def _substitute_all(patterns: Sequence[re.Pattern[str]], text: str) -> str: + return reduce(lambda scrubbed, pattern: pattern.sub(FILTERED, scrubbed), patterns, text) + + +def scrub_json_strings(value: JsonValue, scrub: Callable[[str], str], path: JsonPath = ()) -> JsonValue: + if len(path) > MAX_SCRUB_DEPTH: + return FILTERED + if isinstance(value, str): + return scrub(value) + if isinstance(value, dict): + unscrubbed_keys: Final = SOURCE_CONTEXT_KEYS if path in STACK_FRAME_PATHS else frozenset[str]() + return { # mutable-ok: JSON object + key: item if key in unscrubbed_keys else scrub_json_strings(item, scrub, (*path, key)) + for key, item in value.items() + } + if isinstance(value, list): + return [scrub_json_strings(item, scrub, (*path, "*")) for item in value] # mutable-ok: JSON array + return value + + +def build_event_scrubber(send_default_pii: bool) -> EventScrubFn: + scrub: Final = build_string_scrubber(send_default_pii) + + def scrub_event(event: Event, _hint: Hint) -> Event: + json_event: Final = cast("JsonValue", event) # cast-ok: [LIT006] the SDK serialized the event to JSON already + return cast("Event", scrub_json_strings(json_event, scrub)) # cast-ok: [LIT006] same JSON shape going back + + return scrub_event + + +def send_default_pii_from_env(env: Mapping[str, str]) -> bool: + return str_to_bool(env.get(SEND_DEFAULT_PII_ENV)) is True + + +def build_sentry_init_options(env: Mapping[str, str]) -> SentryInitOptions: + send_default_pii: Final = send_default_pii_from_env(env) + scrub_event: Final = build_event_scrubber(send_default_pii) + return SentryInitOptions( + dsn=env.get("SENTRY_DSN"), + traces_sample_rate=float(env.get("SENTRY_API_TRACE_RATE") or "1.0"), + sample_rate=float(env.get("SENTRY_API_SAMPLE_RATE") or "1.0"), + send_default_pii=send_default_pii, + event_scrubber=EventScrubber( + denylist=list(SECRET_FIELD_NAMES), # mutable-ok: EventScrubber appends pii_denylist onto denylist in place + pii_denylist=list(PII_FIELD_NAMES), # mutable-ok: EventScrubber takes List[str] + recursive=True, + send_default_pii=send_default_pii, + ), + before_send=scrub_event, + before_send_transaction=scrub_event, + environment=env.get("SENTRY_ENVIRONMENT", "production"), + ) 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 20753afee5c..24d5b7f366e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -68,15 +68,11 @@ def _error_status_and_message(exc: Exception) -> tuple[int, str]: def _mid_stream_error_sse_event(exc: Exception) -> bytes: from litellm.anthropic_interface.exceptions.exception_mapping_utils import ( - AnthropicExceptionMapping, + anthropic_error_sse_frame, ) status_code, message = _error_status_and_message(exc) - error_response = AnthropicExceptionMapping.transform_to_anthropic_error( - status_code=status_code, - raw_message=message, - ) - return f"event: error\ndata: {json.dumps(error_response)}\n\n".encode() + return anthropic_error_sse_frame(status_code=status_code, raw_message=message).encode() def _delta_payload_field(delta_type: StreamingContentBlockDeltaType) -> str: diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 0bd46382fef..5550590d0c0 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -16,6 +16,7 @@ from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.anthropic.common_utils import ANTHROPIC_ERROR_STATUS_CODE_MAP +from litellm.llms.anthropic.experimental_pass_through.messages.utils import INCOMPLETE_STREAM_ERROR_MESSAGE from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) @@ -28,11 +29,6 @@ GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ: Final = PassThroughEndpointLogging() _UPSTREAM_PUMP_TASKS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: stdlib strong-ref set for pump tasks _DETACHED_STREAM_DRAINS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: bounded strong-ref set, detached drains -INCOMPLETE_STREAM_ERROR_MESSAGE: Final = ( - "Provider stream ended before emitting a message_stop event; " - "the response is incomplete and any partial content (e.g. tool_use input JSON) may be truncated." -) - def _is_message_stop_chunk(chunk: object) -> bool: if isinstance(chunk, dict): diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py index 89105c00428..fe8ac2cd7a2 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py @@ -15,6 +15,12 @@ if TYPE_CHECKING: from litellm.exceptions import ContentPolicyViolationError +INCOMPLETE_STREAM_ERROR_MESSAGE: Final = ( + "Provider stream ended before emitting a message_stop event; " + "the response is incomplete and any partial content (e.g. tool_use input JSON) may be truncated." +) + + def get_safeguard_refusal_stop_details(response: object) -> Mapping[str, Any] | None: """ Return the ``stop_details`` of an Anthropic Messages response refused by a diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index f753e87fee3..59ccde872fc 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -2,20 +2,25 @@ ## Translates OpenAI call to Anthropic `/v1/messages` format import asyncio import json -import traceback from collections import deque from collections.abc import AsyncIterator, Iterator, Mapping from typing import TYPE_CHECKING, Any, Final +from pydantic import BaseModel, ConfigDict, field_validator + from litellm import verbose_logger +from litellm._logging import redact_internal_details_from_client_message from litellm._uuid import uuid +from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.prompt_templates.common_utils import ( encrypted_reasoning_signature, ) from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( + INCOMPLETE_STREAM_ERROR_MESSAGE, refusal_stop_details, responses_output_refusal_text, ) +from litellm.responses.streaming_iterator import stream_error_status_and_message from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from .transformation import ( @@ -27,6 +32,72 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject +class _UpstreamFailure(BaseModel): + model_config = ConfigDict(frozen=True) + + status_code: int | None = None + message: str | None = None + + @field_validator("status_code", mode="before") + @classmethod + def http_error_status_or_none(cls, value: object) -> int | None: + candidate: Final = ( + value + if isinstance(value, int) and not isinstance(value, bool) + else int(value) + if isinstance(value, str) and value.isdecimal() + else None + ) + return candidate if candidate is not None and 400 <= candidate <= 599 else None + + @field_validator("message", mode="before") + @classmethod + def str_or_none(cls, value: object) -> str | None: + return value if isinstance(value, str) else None + + +class _FailedResponse(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + error: object | None = None + + +class _FailedResponseEvent(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + response: _FailedResponse | None = None + + +def _original_failure(exception: Exception) -> Exception: + failure = exception # rebind-ok: walks the MidStreamFallbackError chain down to the provider failure + while isinstance(failure, MidStreamFallbackError) and failure.original_exception is not None: + failure = failure.original_exception + return failure + + +def _failure_status_and_message(exception: Exception) -> tuple[int, str]: + original: Final = _original_failure(exception) + failure: Final = _UpstreamFailure.model_validate( + {"status_code": getattr(original, "status_code", None), "message": getattr(original, "message", None)} + ) + status_code: Final = failure.status_code if failure.status_code is not None else 500 + message: Final = failure.message or str(original) or INCOMPLETE_STREAM_ERROR_MESSAGE + return status_code, message + + +def _anthropic_error_chunk(status_code: int, message: str) -> dict[str, object]: + from litellm.anthropic_interface.exceptions.exception_mapping_utils import ( + AnthropicExceptionMapping, + ) + + return dict( + AnthropicExceptionMapping.transform_to_anthropic_error( + status_code=status_code, + raw_message=redact_internal_details_from_client_message(message), + ) + ) + + class AnthropicResponsesStreamWrapper: """ Wraps a Responses API streaming iterator and re-emits events in Anthropic SSE format. @@ -40,6 +111,7 @@ class AnthropicResponsesStreamWrapper: response.function_call_arguments.delta -> content_block_delta (input_json_delta) response.output_item.done -> content_block_delta (signature_delta) + content_block_stop response.completed -> message_delta + message_stop + response.failed -> error (the stream ends without message_stop) """ def __init__( @@ -60,6 +132,7 @@ class AnthropicResponsesStreamWrapper: self._pending_tool_ids: dict[str, str] = {} # item_id -> call_id / name accumulator self._sent_message_start = False self._sent_message_stop = False + self._stream_failed = False self._chunk_queue: deque[dict[str, object]] = deque() self._refusal_text: str = "" self._sync_responses_iterator: Iterator[object] | None = None @@ -293,10 +366,23 @@ class AnthropicResponsesStreamWrapper: ) return + if event_type == "response.failed": + failed: Final = _FailedResponseEvent.model_validate(event) + status_code, message = stream_error_status_and_message( + failed.response.error if failed.response is not None else None + ) + verbose_logger.error( + "AnthropicResponsesStreamWrapper: upstream Responses stream for %s failed (%s): %s", + self.model, + status_code, + message, + ) + self._fail_stream(status_code, message) + return + # ---- response completed -> message_delta + message_stop ---- if event_type in ( "response.completed", - "response.failed", "response.incomplete", ): response_obj: Final = getattr(event, "response", None) or ( @@ -350,21 +436,24 @@ class AnthropicResponsesStreamWrapper: self._sent_message_stop = True return + def _fail_stream(self, status_code: int, message: str) -> None: + self._stream_failed = True + self._chunk_queue.append(_anthropic_error_chunk(status_code, message)) + def __aiter__(self) -> "AnthropicResponsesStreamWrapper": return self async def __anext__(self) -> dict[str, object]: - # Return any queued chunks first if self._chunk_queue: return self._chunk_queue.popleft() + if self._stream_failed: + raise StopAsyncIteration - # Emit message_start if not yet done (fallback if response.created wasn't fired) if not self._sent_message_start: self._sent_message_start = True self._chunk_queue.append(self._make_message_start()) return self._chunk_queue.popleft() - # Consume the upstream stream try: if hasattr(self.responses_stream, "__aiter__"): async for event in self.responses_stream: @@ -382,10 +471,19 @@ class AnthropicResponsesStreamWrapper: return self._chunk_queue.popleft() except StopAsyncIteration: pass - except Exception as e: - verbose_logger.error("AnthropicResponsesStreamWrapper error: %s\n%s", e, traceback.format_exc()) + except Exception as e: # noqa: BLE001 # every upstream failure becomes a client error event + verbose_logger.exception( + "AnthropicResponsesStreamWrapper: upstream Responses stream for %s failed", self.model + ) + self._fail_stream(*_failure_status_and_message(e)) + + if not self._chunk_queue and not self._sent_message_stop and not self._stream_failed: + verbose_logger.error( + "AnthropicResponsesStreamWrapper: upstream Responses stream for %s ended without a terminal event", + self.model, + ) + self._fail_stream(500, INCOMPLETE_STREAM_ERROR_MESSAGE) - # Drain any remaining queued chunks if self._chunk_queue: return self._chunk_queue.popleft() diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index bd358805743..65a34f72167 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -69,7 +69,9 @@ def make_sync_call( completion_stream: Any = MockResponseIterator(model_response=model_response, json_mode=json_mode) else: decoder: Final = AWSEventStreamDecoder(model=model, json_mode=json_mode) - completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) + completion_stream = decoder.iter_bytes( + response.iter_bytes(chunk_size=stream_chunk_size), response_headers=response.headers + ) # LOGGING logging_obj.post_call( diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index c7b4018b80b..93804e20041 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -1,6 +1,6 @@ import types -from collections.abc import AsyncIterator, Iterator -from typing import Final, cast +from collections.abc import AsyncIterator, Iterator, Mapping +from typing import TYPE_CHECKING, Final, cast import httpx from pydantic import TypeAdapter @@ -51,7 +51,11 @@ from ..common_utils import ( bedrock_tool_name_mappings: Final[InMemoryCache] = InMemoryCache(max_size_in_memory=50, default_ttl=600) from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig +if TYPE_CHECKING: + from botocore.eventstream import EventStreamMessage + converse_config: Final = AmazonConverseConfig() +_STREAM_HEAD_BYTES: Final = 200 NOVA_INVOKE_STREAM_EVENT_TYPES: Final = ( "messageStart", "contentBlockStart", @@ -162,6 +166,22 @@ class AmazonCohereChatConfig: return optional_params +def _stream_decoder( + bedrock_invoke_provider: litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None, + *, + model: str, + json_mode: bool | None, + sync_stream: bool, +) -> "AWSEventStreamDecoder": + if bedrock_invoke_provider == "anthropic": + return AmazonAnthropicClaudeStreamDecoder(model=model, sync_stream=sync_stream, json_mode=json_mode) + if bedrock_invoke_provider == "deepseek_r1": + return AmazonDeepSeekR1StreamDecoder(model=model, sync_stream=sync_stream) + if bedrock_invoke_provider == "moonshot": + return AmazonOpenAICompatibleStreamDecoder(model=model, sync_stream=sync_stream) + return AWSEventStreamDecoder(model=model, json_mode=json_mode) + + async def make_call( client: AsyncHTTPHandler | None, api_base: str, @@ -218,28 +238,13 @@ async def make_call( completion_stream: MockResponseIterator | AsyncIterator[GChunk | ModelResponseStream | dict] = ( MockResponseIterator(model_response=model_response, json_mode=json_mode) ) - elif bedrock_invoke_provider == "anthropic": - decoder: AWSEventStreamDecoder = AmazonAnthropicClaudeStreamDecoder( - model=model, - sync_stream=False, - json_mode=json_mode, - ) - completion_stream = decoder.aiter_bytes(response.aiter_bytes(chunk_size=stream_chunk_size)) - elif bedrock_invoke_provider == "deepseek_r1": - decoder = AmazonDeepSeekR1StreamDecoder( - model=model, - sync_stream=False, - ) - completion_stream = decoder.aiter_bytes(response.aiter_bytes(chunk_size=stream_chunk_size)) - elif bedrock_invoke_provider == "moonshot": - decoder = AmazonOpenAICompatibleStreamDecoder( - model=model, - sync_stream=False, - ) - completion_stream = decoder.aiter_bytes(response.aiter_bytes(chunk_size=stream_chunk_size)) else: - decoder = AWSEventStreamDecoder(model=model, json_mode=json_mode) - completion_stream = decoder.aiter_bytes(response.aiter_bytes(chunk_size=stream_chunk_size)) + decoder: Final = _stream_decoder( + bedrock_invoke_provider, model=model, json_mode=json_mode, sync_stream=False + ) + completion_stream = decoder.aiter_bytes( + response.aiter_bytes(chunk_size=stream_chunk_size), response_headers=response.headers + ) # LOGGING logging_obj.post_call( @@ -322,28 +327,13 @@ def make_sync_call( completion_stream: MockResponseIterator | Iterator[GChunk | ModelResponseStream | dict] = ( MockResponseIterator(model_response=model_response, json_mode=json_mode) ) - elif bedrock_invoke_provider == "anthropic": - decoder: AWSEventStreamDecoder = AmazonAnthropicClaudeStreamDecoder( - model=model, - sync_stream=True, - json_mode=json_mode, - ) - completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) - elif bedrock_invoke_provider == "deepseek_r1": - decoder = AmazonDeepSeekR1StreamDecoder( - model=model, - sync_stream=True, - ) - completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) - elif bedrock_invoke_provider == "moonshot": - decoder = AmazonOpenAICompatibleStreamDecoder( - model=model, - sync_stream=True, - ) - completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) else: - decoder = AWSEventStreamDecoder(model=model, json_mode=json_mode) - completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) + decoder: Final = _stream_decoder( + bedrock_invoke_provider, model=model, json_mode=json_mode, sync_stream=True + ) + completion_stream = decoder.iter_bytes( + response.iter_bytes(chunk_size=stream_chunk_size), response_headers=response.headers + ) # LOGGING logging_obj.post_call( @@ -370,6 +360,49 @@ def make_sync_call( raise BedrockError(status_code=500, message=str(e)) +def _response_header(response_headers: Mapping[str, str] | None, name: str) -> str | None: + return None if response_headers is None else response_headers.get(name) + + +class _EventStreamTally: + def __init__(self) -> None: + self.bytes_received = 0 + self.bytes_decoded = 0 + self.events = 0 + self.head = b"" + + def add_chunk(self, chunk: bytes) -> None: + self.bytes_received += len(chunk) + if len(self.head) < _STREAM_HEAD_BYTES: + self.head = (self.head + chunk)[:_STREAM_HEAD_BYTES] + + def add_event(self, event: "EventStreamMessage") -> None: + self.events += 1 + self.bytes_decoded += event.prelude.total_length + + def undecoded_stream_error(self, response_headers: Mapping[str, str] | None) -> BedrockError | None: + undecoded: Final = self.bytes_received - self.bytes_decoded + if self.events and not undecoded: + return None + detail: Final = ( + f"content-type={_response_header(response_headers, 'content-type')!r}, " + f"x-amzn-requestid={_response_header(response_headers, 'x-amzn-requestid')!r}, " + f"{self.bytes_received} bytes received" + ) + if not self.events: + return BedrockError( + status_code=502, + message=( + "Bedrock answered the stream with HTTP 200 but its body decoded to no events " + f"({detail}, first bytes={self.head!r})" + ), + ) + return BedrockError( + status_code=502, + message=f"Bedrock stream ended with {undecoded} undecoded bytes after {self.events} events ({detail})", + ) + + class AWSEventStreamDecoder: def __init__(self, model: str, json_mode: bool | None = False) -> None: from botocore.parsers import EventStreamJSONParser @@ -709,32 +742,48 @@ class AWSEventStreamDecoder: tool_use=None, ) - def iter_bytes(self, iterator: Iterator[bytes]) -> Iterator[GChunk | ModelResponseStream | dict]: + def iter_bytes( + self, iterator: Iterator[bytes], *, response_headers: Mapping[str, str] | None = None + ) -> Iterator[GChunk | ModelResponseStream | dict]: """Given an iterator that yields lines, iterate over it & yield every event encountered""" from botocore.eventstream import EventStreamBuffer event_stream_buffer: Final = EventStreamBuffer() + tally: Final = _EventStreamTally() for chunk in iterator: event_stream_buffer.add_data(chunk) + tally.add_chunk(chunk) for event in event_stream_buffer: + tally.add_event(event) message = self._parse_message_from_event(event) if message: # sse_event = ServerSentEvent(data=message, event="completion") _data = json.loads(message) yield self._chunk_parser(chunk_data=_data) + undecoded_stream_error: Final = tally.undecoded_stream_error(response_headers) + if undecoded_stream_error is not None: + raise undecoded_stream_error - async def aiter_bytes(self, iterator: AsyncIterator[bytes]) -> AsyncIterator[GChunk | ModelResponseStream | dict]: + async def aiter_bytes( + self, iterator: AsyncIterator[bytes], *, response_headers: Mapping[str, str] | None = None + ) -> AsyncIterator[GChunk | ModelResponseStream | dict]: """Given an async iterator that yields lines, iterate over it & yield every event encountered""" from botocore.eventstream import EventStreamBuffer event_stream_buffer: Final = EventStreamBuffer() + tally: Final = _EventStreamTally() async for chunk in iterator: event_stream_buffer.add_data(chunk) + tally.add_chunk(chunk) for event in event_stream_buffer: + tally.add_event(event) message = self._parse_message_from_event(event) if message: _data = json.loads(message) yield self._chunk_parser(chunk_data=_data) + undecoded_stream_error: Final = tally.undecoded_stream_error(response_headers) + if undecoded_stream_error is not None: + raise undecoded_stream_error def _parse_message_from_event(self, event) -> str | None: response_stream_shape: Final = get_bedrock_response_stream_shape() diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 14bd2bee6cf..cefc8afed25 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -770,7 +770,9 @@ class AmazonAnthropicClaudeMessagesConfig( aws_decoder: Final = AmazonAnthropicClaudeMessagesStreamDecoder( model=model, ) - completion_stream: Final = aws_decoder.aiter_bytes(httpx_response.aiter_bytes()) + completion_stream: Final = aws_decoder.aiter_bytes( + httpx_response.aiter_bytes(), response_headers=httpx_response.headers + ) # Convert decoded Bedrock events to Server-Sent Events expected by Anthropic clients. return self.bedrock_sse_wrapper( completion_stream=completion_stream, diff --git a/litellm/llms/vertex_ai/files/handler.py b/litellm/llms/vertex_ai/files/handler.py index ac95d1348f9..f2da04a7db7 100644 --- a/litellm/llms/vertex_ai/files/handler.py +++ b/litellm/llms/vertex_ai/files/handler.py @@ -53,14 +53,18 @@ class VertexAIFilesHandler(GCSBucketBase): Sources them from the deployment's ``litellm_params`` (``gcs_bucket_name`` / ``bucket_name`` and ``vertex_credentials``), mirroring the write path in - ``VertexAIFilesConfig._get_configured_bucket_name``, and falls back to the global - ``GCS_BUCKET_NAME`` / ``GCS_PATH_SERVICE_ACCOUNT`` env vars. This lets Vertex batch - run entirely at the model-group level, so output written to a per-model bucket is - readable without setting the global env vars. + ``VertexAIFilesConfig._get_configured_bucket_name``, and falls back to the + ``GCS_BATCH_BUCKET_NAME`` then ``GCS_BUCKET_NAME`` / ``GCS_PATH_SERVICE_ACCOUNT`` + env vars. This lets Vertex batch run entirely at the model-group level, so output + written to a per-model bucket is readable without setting the global env vars. """ params: Final[Mapping[str, object]] = litellm_params or {} bucket_candidate: Final = params.get("gcs_bucket_name") or params.get("bucket_name") - configured_bucket_name = bucket_candidate if isinstance(bucket_candidate, str) else os.getenv("GCS_BUCKET_NAME") + configured_bucket_name = ( + bucket_candidate + if isinstance(bucket_candidate, str) + else os.getenv("GCS_BATCH_BUCKET_NAME") or os.getenv("GCS_BUCKET_NAME") + ) credentials: Final = params.get("vertex_credentials") or vertex_credentials if isinstance(credentials, dict): diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index dbb41b57348..2b0694697a4 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -961,7 +961,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _get_configured_bucket_name(self, litellm_params: dict) -> str: bucket_name: Final = ( - litellm_params.get("gcs_bucket_name") or litellm_params.get("bucket_name") or os.getenv("GCS_BUCKET_NAME") + litellm_params.get("gcs_bucket_name") + or litellm_params.get("bucket_name") + or os.getenv("GCS_BATCH_BUCKET_NAME") + or os.getenv("GCS_BUCKET_NAME") ) if not bucket_name: raise ValueError("GCS bucket_name is required") diff --git a/litellm/llms/vertex_ai/rag_engine/ingestion.py b/litellm/llms/vertex_ai/rag_engine/ingestion.py index d9916209a14..c10bac595b6 100644 --- a/litellm/llms/vertex_ai/rag_engine/ingestion.py +++ b/litellm/llms/vertex_ai/rag_engine/ingestion.py @@ -122,41 +122,26 @@ class VertexAIRAGIngestion(BaseRAGIngestion): """ import litellm - # Set GCS_BUCKET_NAME env var for litellm.files.create_file - # The handler uses this to determine where to upload - original_bucket: Final = os.environ.get("GCS_BUCKET_NAME") - if self.gcs_bucket: - os.environ["GCS_BUCKET_NAME"] = self.gcs_bucket + file_tuple: Final = (filename, file_content, content_type) - try: - # Create file tuple for litellm.files.acreate_file - file_tuple: Final = (filename, file_content, content_type) + verbose_logger.debug( + "Uploading file to GCS via litellm.files.acreate_file: %s (bucket: %s)", filename, self.gcs_bucket + ) - verbose_logger.debug( - "Uploading file to GCS via litellm.files.acreate_file: %s (bucket: %s)", filename, self.gcs_bucket - ) + response: Final = await litellm.acreate_file( + file=file_tuple, + purpose="assistants", + custom_llm_provider="vertex_ai", + gcs_bucket_name=self.gcs_bucket, + vertex_project=self.vertex_project, + vertex_location=self.vertex_location, + vertex_credentials=self.vertex_credentials, + ) - # Upload to GCS using LiteLLM's file upload - response: Final = await litellm.acreate_file( - file=file_tuple, - purpose="assistants", # Purpose for file storage - custom_llm_provider="vertex_ai", - vertex_project=self.vertex_project, - vertex_location=self.vertex_location, - vertex_credentials=self.vertex_credentials, - ) + gcs_uri: Final = response.id + verbose_logger.info("Uploaded file to GCS: %s", gcs_uri) - # The response.id should be the GCS URI - gcs_uri: Final = response.id - verbose_logger.info("Uploaded file to GCS: %s", gcs_uri) - - return gcs_uri - finally: - # Restore original env var - if original_bucket is not None: - os.environ["GCS_BUCKET_NAME"] = original_bucket - elif "GCS_BUCKET_NAME" in os.environ: - del os.environ["GCS_BUCKET_NAME"] + return gcs_uri async def _import_file_to_corpus_via_sdk( self, @@ -259,6 +244,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion): content_type: str | None, chunks: list[str], embeddings: list[list[float]] | None, + existing_file_id: str | None = None, ) -> tuple[str | None, str | None]: """ Store content in Vertex AI RAG corpus. @@ -274,6 +260,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion): content_type: MIME type chunks: Ignored - Vertex AI handles chunking embeddings: Ignored - Vertex AI handles embedding + existing_file_id: Existing provider file ID, unsupported for Vertex AI RAG Engine Returns: Tuple of (corpus_id, gcs_uri) diff --git a/tests/test_litellm/integrations/code_interpreter_interception/__init__.py b/litellm/llms/xai/batches/__init__.py similarity index 100% rename from tests/test_litellm/integrations/code_interpreter_interception/__init__.py rename to litellm/llms/xai/batches/__init__.py diff --git a/litellm/llms/xai/batches/handler.py b/litellm/llms/xai/batches/handler.py new file mode 100644 index 00000000000..62db1c4833a --- /dev/null +++ b/litellm/llms/xai/batches/handler.py @@ -0,0 +1,195 @@ +from collections.abc import Coroutine +from itertools import chain +from typing import Final + +import httpx +from typing_extensions import NotRequired, ReadOnly, TypedDict + +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + HTTPHandler, + get_async_httpx_client, +) +from litellm.types.llms.openai import CreateBatchRequest, HttpxBinaryResponseContent +from litellm.types.utils import LiteLLMBatch, LlmProviders + +from .transformation import ( + XAI_RESULTS_PAGE_SIZE, + OpenAIBatchListResponse, + XAIBatch, + XAIBatchList, + XAIBatchResult, + XAIBatchResultsPage, + get_xai_auth_headers, + raise_for_xai_status, + results_to_openai_jsonl, + to_create_batch_body, + to_litellm_batch, + to_openai_batch_list, + xai_batches_url, +) + +_JSONL_CONTENT_TYPE: Final = ("content-type", "application/jsonl") + + +class _PageParams(TypedDict): + limit: ReadOnly[int] + pagination_token: NotRequired[ReadOnly[str]] + + +def _results_params(after: str | None, limit: int | None) -> dict[str, object]: # mutable-ok: httpx params + if after is None: + return dict(_PageParams(limit=limit or XAI_RESULTS_PAGE_SIZE)) # mutable-ok: httpx params + return dict(_PageParams(limit=limit or XAI_RESULTS_PAGE_SIZE, pagination_token=after)) # mutable-ok: httpx params + + +def _flatten(pages: list[XAIBatchResultsPage]) -> tuple[XAIBatchResult, ...]: + return tuple(chain.from_iterable(page.results for page in pages)) + + +def _jsonl_response(url: str, results: tuple[XAIBatchResult, ...]) -> HttpxBinaryResponseContent: + return HttpxBinaryResponseContent( + response=httpx.Response( + status_code=200, + content=results_to_openai_jsonl(results), + headers=(_JSONL_CONTENT_TYPE,), + request=httpx.Request(method="GET", url=url), + ) + ) + + +class XAIBatchesHandler: + def __init__(self, sync_client: HTTPHandler | None = None, async_client: AsyncHTTPHandler | None = None) -> None: + self._sync_client = sync_client + self._async_client = async_client + + def _sync(self, timeout: float | httpx.Timeout) -> HTTPHandler: + return self._sync_client or HTTPHandler(timeout=timeout) + + def _async(self, timeout: float | httpx.Timeout) -> AsyncHTTPHandler: + return self._async_client or get_async_httpx_client( + llm_provider=LlmProviders.XAI, + params={"timeout": timeout}, # mutable-ok: get_async_httpx_client takes a dict + ) + + def create_batch( + self, + _is_async: bool, + create_batch_data: CreateBatchRequest, + api_base: str | None, + api_key: str | None, + timeout: float | httpx.Timeout, + ) -> LiteLLMBatch | Coroutine[None, None, LiteLLMBatch]: + url: Final = xai_batches_url(api_base) + headers: Final = get_xai_auth_headers(api_key=api_key) + body: Final = dict(to_create_batch_body(create_batch_data)) # mutable-ok: httpx json body + endpoint: Final = create_batch_data.get("endpoint") or "/v1/chat/completions" + if _is_async: + + async def _acreate() -> LiteLLMBatch: + response: Final = await self._async(timeout).post(url, json=body, headers=headers, timeout=timeout) + return to_litellm_batch(XAIBatch.model_validate(raise_for_xai_status(response).json()), endpoint) + + return _acreate() + response: Final = self._sync(timeout).post(url, json=body, headers=headers, timeout=timeout) + return to_litellm_batch(XAIBatch.model_validate(raise_for_xai_status(response).json()), endpoint) + + def retrieve_batch( + self, + _is_async: bool, + batch_id: str, + api_base: str | None, + api_key: str | None, + timeout: float | httpx.Timeout, + ) -> LiteLLMBatch | Coroutine[None, None, LiteLLMBatch]: + url: Final = xai_batches_url(api_base, batch_id) + headers: Final = get_xai_auth_headers(api_key=api_key) + if _is_async: + + async def _aretrieve() -> LiteLLMBatch: + response: Final = await self._async(timeout).get(url, headers=headers, timeout=timeout) + return to_litellm_batch(XAIBatch.model_validate(raise_for_xai_status(response).json())) + + return _aretrieve() + response: Final = self._sync(timeout).get(url, headers=headers, timeout=timeout) + return to_litellm_batch(XAIBatch.model_validate(raise_for_xai_status(response).json())) + + def cancel_batch( + self, + _is_async: bool, + batch_id: str, + api_base: str | None, + api_key: str | None, + timeout: float | httpx.Timeout, + ) -> LiteLLMBatch | Coroutine[None, None, LiteLLMBatch]: + url: Final = xai_batches_url(api_base, batch_id, suffix=":cancel") + headers: Final = get_xai_auth_headers(api_key=api_key) + if _is_async: + + async def _acancel() -> LiteLLMBatch: + response: Final = await self._async(timeout).post(url, headers=headers, timeout=timeout) + return to_litellm_batch(XAIBatch.model_validate(raise_for_xai_status(response).json())) + + return _acancel() + response: Final = self._sync(timeout).post(url, headers=headers, timeout=timeout) + return to_litellm_batch(XAIBatch.model_validate(raise_for_xai_status(response).json())) + + def list_batches( + self, + _is_async: bool, + api_base: str | None, + api_key: str | None, + timeout: float | httpx.Timeout, + after: str | None = None, + limit: int | None = None, + ) -> OpenAIBatchListResponse | Coroutine[None, None, OpenAIBatchListResponse]: + url: Final = xai_batches_url(api_base) + headers: Final = get_xai_auth_headers(api_key=api_key) + params: Final = _results_params(after, limit) + if _is_async: + + async def _alist() -> OpenAIBatchListResponse: + response: Final = await self._async(timeout).get(url, params=params, headers=headers, timeout=timeout) + return to_openai_batch_list(XAIBatchList.model_validate(raise_for_xai_status(response).json())) + + return _alist() + response: Final = self._sync(timeout).get(url, params=params, headers=headers, timeout=timeout) + return to_openai_batch_list(XAIBatchList.model_validate(raise_for_xai_status(response).json())) + + def batch_results_content( + self, + _is_async: bool, + batch_id: str, + api_base: str | None, + api_key: str | None, + timeout: float | httpx.Timeout, + ) -> HttpxBinaryResponseContent | Coroutine[None, None, HttpxBinaryResponseContent]: + url: Final = xai_batches_url(api_base, batch_id, suffix="/results") + headers: Final = get_xai_auth_headers(api_key=api_key) + if _is_async: + + async def _aresults() -> HttpxBinaryResponseContent: + client: Final = self._async(timeout) + + async def _page(after: str | None) -> XAIBatchResultsPage: + response: Final = await client.get( + url, params=_results_params(after, None), headers=headers, timeout=timeout + ) + return XAIBatchResultsPage.model_validate(raise_for_xai_status(response).json()) + + pages = [await _page(None)] # mutable-ok: page walk terminates on the cursor, not on a fixed count + while pages[-1].pagination_token and pages[-1].results: + pages.append(await _page(pages[-1].pagination_token)) + return _jsonl_response(url, _flatten(pages)) + + return _aresults() + client: Final = self._sync(timeout) + + def _page(after: str | None) -> XAIBatchResultsPage: + response: Final = client.get(url, params=_results_params(after, None), headers=headers, timeout=timeout) + return XAIBatchResultsPage.model_validate(raise_for_xai_status(response).json()) + + pages = [_page(None)] # mutable-ok: page walk terminates on the cursor, not on a fixed count + while pages[-1].pagination_token and pages[-1].results: + pages.append(_page(pages[-1].pagination_token)) + return _jsonl_response(url, _flatten(pages)) diff --git a/litellm/llms/xai/batches/transformation.py b/litellm/llms/xai/batches/transformation.py new file mode 100644 index 00000000000..8f305b8c203 --- /dev/null +++ b/litellm/llms/xai/batches/transformation.py @@ -0,0 +1,278 @@ +""" +xAI Batch API reference: https://docs.x.ai/developers/advanced-api-usage/batch-api + +xAI batches carry request counters, not a status, and no output file: results are paged from +``GET /v1/batches/{id}/results``, so LiteLLM hands back the batch id as ``output_file_id``. +""" + +import json +from collections.abc import Mapping, Sequence +from datetime import datetime, timezone +from types import MappingProxyType +from typing import Final, Literal, TypeAlias + +import httpx +from openai.types.batch import BatchRequestCounts +from openai.types.batch import Errors as BatchErrors +from openai.types.batch_error import BatchError +from pydantic import BaseModel, ConfigDict +from typing_extensions import NotRequired, ReadOnly, TypedDict + +from litellm.constants import XAI_API_BASE +from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.xai.common_utils import XAIModelInfo +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import CreateBatchRequest +from litellm.types.utils import LiteLLMBatch + +OpenAIBatchStatus: TypeAlias = Literal[ + "validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled" +] + +XAI_BATCH_ID_PREFIX: Final = "batch_" +XAI_RESULTS_PAGE_SIZE: Final = 1000 +DEFAULT_BATCH_NAME: Final = "litellm-batch" +DEFAULT_BATCH_ENDPOINT: Final = "/v1/chat/completions" +_EMPTY_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) + + +class XAIBatchesError(BaseLLMException): + pass + + +def xai_batches_error( + error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers +) -> XAIBatchesError: + return XAIBatchesError( + status_code=status_code, + message=error_message, + headers=headers if isinstance(headers, httpx.Headers) else httpx.Headers(tuple(headers.items())), + ) + + +def raise_for_xai_status(response: httpx.Response) -> httpx.Response: + if response.status_code >= 400: + raise xai_batches_error(response.text, response.status_code, response.headers) + return response + + +def get_xai_api_base(api_base: str | None) -> str: + resolved: Final = (api_base or get_secret_str("XAI_API_BASE") or XAI_API_BASE).rstrip("/") + return resolved.removesuffix("/v1") + + +def get_xai_auth_headers( + headers: Mapping[str, str] = _EMPTY_HEADERS, api_key: str | None = None +) -> dict[str, str]: # mutable-ok: BaseConfig.validate_environment contract returns dict + resolved_key: Final = XAIModelInfo.get_api_key(api_key) + if resolved_key is None: + raise xai_batches_error( + "Missing xAI API Key. Pass api_key, set litellm.xai_key or XAI_API_KEY", 401, _EMPTY_HEADERS + ) + return dict(headers, Authorization=f"Bearer {resolved_key}") # mutable-ok: BaseConfig contract returns dict + + +def xai_batches_url(api_base: str | None, batch_id: str | None = None, suffix: str = "") -> str: + base: Final = f"{get_xai_api_base(api_base)}/v1/batches" + if batch_id is None: + return base + return f"{base}/{encode_url_path_segment(batch_id, field_name='batch_id')}{suffix}" + + +def is_xai_batch_results_id(file_id: str) -> bool: + return file_id.startswith(XAI_BATCH_ID_PREFIX) + + +class XAICreateBatchRequest(TypedDict): + name: ReadOnly[str] + input_file_id: NotRequired[ReadOnly[str]] + + +class XAIBatchState(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + num_requests: int = 0 + num_pending: int = 0 + num_success: int = 0 + num_error: int = 0 + num_cancelled: int = 0 + + +class XAIBatch(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + batch_id: str + name: str = "" + create_time: str | None = None + expire_time: str | None = None + cancel_time: str | None = None + cancel_by_xai_message: str | None = None + state: XAIBatchState = XAIBatchState() + input_file_id: str | None = None + + +class XAIBatchList(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + batches: tuple[XAIBatch, ...] = () + pagination_token: str | None = None + + +class XAIBatchResultError(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + code: int | str | None = None + message: str = "" + + +class XAIBatchResultData(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + response: Mapping[str, Mapping[str, object]] | None = None + error: XAIBatchResultError | None = None + + +class XAIBatchResult(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + batch_request_id: str + batch_result: XAIBatchResultData = XAIBatchResultData() + + +class XAIBatchResultsPage(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + results: tuple[XAIBatchResult, ...] = () + pagination_token: str | None = None + + +def _to_unix_timestamp(value: str | None) -> int | None: + """xAI returns RFC 3339 timestamps over gRPC but a bare ``YYYY-MM-DD`` over REST.""" + if value is None: + return None + try: + parsed: Final = datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return None + return int((parsed if parsed.tzinfo is not None else parsed.replace(tzinfo=timezone.utc)).timestamp()) + + +def xai_batch_status(batch: XAIBatch) -> OpenAIBatchStatus: + """xAI exposes counters, not a status. A batch xAI itself cancelled (input validation failed) is a failure, + a caller-cancelled batch is cancelled, an empty batch is still validating its input file, and a batch + with nothing pending has completed.""" + if batch.cancel_time is not None: + return "failed" if batch.cancel_by_xai_message else "cancelled" + if batch.state.num_requests == 0: + return "validating" + if batch.state.num_pending > 0: + return "in_progress" + return "completed" + + +def to_litellm_batch(batch: XAIBatch, endpoint: str = DEFAULT_BATCH_ENDPOINT) -> LiteLLMBatch: + status: Final = xai_batch_status(batch) + created_at: Final = _to_unix_timestamp(batch.create_time) + cancelled_at: Final = _to_unix_timestamp(batch.cancel_time) + errors: Final = ( + BatchErrors(object="list", data=[BatchError(message=batch.cancel_by_xai_message)]) # mutable-ok: openai type + if batch.cancel_by_xai_message + else None + ) + return LiteLLMBatch( + id=batch.batch_id, + object="batch", + endpoint=endpoint, + input_file_id=batch.input_file_id or "", + completion_window="24h", + status=status, + created_at=created_at if created_at is not None else 0, + expires_at=_to_unix_timestamp(batch.expire_time), + failed_at=cancelled_at if status == "failed" else None, + cancelled_at=cancelled_at if status == "cancelled" else None, + output_file_id=batch.batch_id if status == "completed" else None, + errors=errors, + request_counts=BatchRequestCounts( + total=batch.state.num_requests, + completed=batch.state.num_success, + failed=batch.state.num_error + batch.state.num_cancelled, + ), + metadata={"name": batch.name} if batch.name else None, # mutable-ok: LiteLLMBatch.metadata is a dict + ) + + +class OpenAIBatchListResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + object: Literal["list"] = "list" + data: tuple[LiteLLMBatch, ...] + first_id: str | None + last_id: str | None + has_more: bool + next_page_token: str | None = None + + +def to_openai_batch_list(page: XAIBatchList) -> OpenAIBatchListResponse: + data: Final = tuple(to_litellm_batch(b) for b in page.batches) + return OpenAIBatchListResponse( + data=data, + first_id=data[0].id if data else None, + last_id=data[-1].id if data else None, + has_more=bool(page.pagination_token), + next_page_token=page.pagination_token or None, + ) + + +def to_create_batch_body(create_batch_data: CreateBatchRequest) -> XAICreateBatchRequest: + input_file_id: Final = create_batch_data.get("input_file_id") + if not input_file_id: + raise xai_batches_error("input_file_id is required to create an xAI batch", 400, _EMPTY_HEADERS) + metadata: Final = create_batch_data.get("metadata") + name: Final = metadata.get("name") if metadata else None + return XAICreateBatchRequest(name=name or DEFAULT_BATCH_NAME, input_file_id=input_file_id) + + +class OpenAIBatchOutputError(TypedDict): + code: ReadOnly[str] + message: ReadOnly[str] + + +class OpenAIBatchOutputResponse(TypedDict): + status_code: ReadOnly[int] + request_id: ReadOnly[object] + body: ReadOnly[Mapping[str, object]] + + +class OpenAIBatchOutputLine(TypedDict): + id: ReadOnly[str] + custom_id: ReadOnly[str] + response: ReadOnly[OpenAIBatchOutputResponse | None] + error: ReadOnly[OpenAIBatchOutputError | None] + + +def _result_to_openai_line(result: XAIBatchResult) -> OpenAIBatchOutputLine: + """One output JSONL line. xAI wraps the body in a one-key map named after the endpoint + (``chat_get_completion``, ``responses``, ``image_generation``, ...); the value is the OpenAI body.""" + error: Final = result.batch_result.error + response: Final = result.batch_result.response + body: Final = next(iter(response.values()), None) if response else None + if body is None: + message: Final = error.message if error is not None else "xAI returned no response for this request" + code: Final = str(error.code) if error is not None and error.code is not None else "request_failed" + return OpenAIBatchOutputLine( + id=f"batch_req_{result.batch_request_id}", + custom_id=result.batch_request_id, + response=None, + error=OpenAIBatchOutputError(code=code, message=message), + ) + return OpenAIBatchOutputLine( + id=f"batch_req_{result.batch_request_id}", + custom_id=result.batch_request_id, + response=OpenAIBatchOutputResponse(status_code=200, request_id=body.get("id"), body=body), + error=None, + ) + + +def results_to_openai_jsonl(results: Sequence[XAIBatchResult]) -> bytes: + return "".join(f"{json.dumps(_result_to_openai_line(r), ensure_ascii=False)}\n" for r in results).encode() diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index 33ee727dfab..e686d49e689 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -296,7 +296,7 @@ class XAIChatConfig(OpenAIGPTConfig): except Exception as e: verbose_logger.debug("Error extracting X.AI web search usage: %s", e) - self._fold_reasoning_tokens_into_completion(response) + self.fold_reasoning_tokens_into_completion(response) self._normalize_openai_compatible_usage_totals(getattr(response, "usage", None)) restated_usage: Final = _usage_restated_from_xai_ticks(getattr(response, "usage", None)) if restated_usage is not None: @@ -304,7 +304,7 @@ class XAIChatConfig(OpenAIGPTConfig): return response @staticmethod - def _fold_reasoning_tokens_into_completion( + def fold_reasoning_tokens_into_completion( target: ModelResponse | Usage | dict[str, Any] | None, ) -> None: """Reconcile xAI Usage to the OpenAI invariant. @@ -426,7 +426,7 @@ class XAIChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler): chunk["choices"] = [{"index": 0, "delta": {}, "finish_reason": None}] if "usage" in chunk and chunk["usage"] is not None: - XAIChatConfig._fold_reasoning_tokens_into_completion(chunk["usage"]) + XAIChatConfig.fold_reasoning_tokens_into_completion(chunk["usage"]) XAIChatConfig._normalize_openai_compatible_usage_totals(chunk["usage"]) parsed_chunk: Final = super().chunk_parser(chunk) diff --git a/tests/test_litellm/integrations/gitlab/__init__.py b/litellm/llms/xai/files/__init__.py similarity index 100% rename from tests/test_litellm/integrations/gitlab/__init__.py rename to litellm/llms/xai/files/__init__.py diff --git a/litellm/llms/xai/files/transformation.py b/litellm/llms/xai/files/transformation.py new file mode 100644 index 00000000000..dbccca47b25 --- /dev/null +++ b/litellm/llms/xai/files/transformation.py @@ -0,0 +1,247 @@ +""" +xAI Files API reference: https://docs.x.ai/developers/rest-api-reference/inference/files + +xAI stores ``purpose`` as an empty string; LiteLLM reports uploads as ``batch``, the only purpose xAI files serve. +""" + +import time +from collections.abc import Mapping, Sequence +from typing import Final + +import httpx +from openai.types.file_deleted import FileDeleted +from pydantic import BaseModel, ConfigDict +from typing_extensions import ReadOnly, TypedDict + +from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data +from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.files.transformation import BaseFilesConfig, LiteLLMLoggingObj +from litellm.types.llms.openai import ( + CreateFileRequest, + FileContentRequest, + HttpxBinaryResponseContent, + OpenAICreateFileRequestOptionalParams, + OpenAIFileObject, + OpenAIFilesPurpose, +) +from litellm.types.utils import LlmProviders + +from ..batches.transformation import ( + get_xai_api_base, + get_xai_auth_headers, + raise_for_xai_status, + xai_batches_error, +) + +_NO_QUERY_PARAMS: Final[dict[str, str]] = {} # mutable-ok: BaseFilesConfig request transforms return tuple[str, dict] +_DEFAULT_PURPOSE: Final[OpenAIFilesPurpose] = "batch" + + +class XAIMultipartUpload(TypedDict): + file: ReadOnly[tuple[str, object, str]] + purpose: ReadOnly[tuple[None, str]] + + +class XAIFile(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + id: str + bytes: int = 0 + created_at: int | None = None + filename: str = "" + purpose: str = "" + expires_at: int | None = None + + +class XAIFileList(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + data: tuple[XAIFile, ...] = () + pagination_token: str | None = None + + +class XAIFileDeleted(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + id: str + deleted: bool = True + + +def _to_openai_file_object(file: XAIFile) -> OpenAIFileObject: + return OpenAIFileObject( + id=file.id, + bytes=file.bytes, + created_at=file.created_at if file.created_at is not None else int(time.time()), + filename=file.filename, + object="file", + purpose=_DEFAULT_PURPOSE, + status="uploaded", + expires_at=file.expires_at, + ) + + +def _api_base_from(litellm_params: Mapping[str, object]) -> str: + api_base: Final = litellm_params.get("api_base") + return get_xai_api_base(api_base if isinstance(api_base, str) else None) + + +class XAIFilesConfig(BaseFilesConfig): + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.XAI + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + stream: bool | None = None, + ) -> str: + return f"{get_xai_api_base(api_base)}/v1/files" + + def _file_url(self, file_id: str, litellm_params: Mapping[str, object], suffix: str = "") -> str: + encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") + return f"{_api_base_from(litellm_params)}/v1/files/{encoded_file_id}{suffix}" + + def get_error_class( + self, error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers + ) -> BaseLLMException: + return xai_batches_error(error_message, status_code, headers) + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + messages: Sequence[object], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, str]: # mutable-ok: BaseFilesConfig signature + return get_xai_auth_headers(headers, api_key) + + def get_supported_openai_params( + self, model: str + ) -> list[OpenAICreateFileRequestOptionalParams]: # mutable-ok: BaseFilesConfig signature + return ["purpose"] # mutable-ok: BaseFilesConfig signature + + def map_openai_params( + self, + non_default_params: Mapping[str, object], + optional_params: dict[str, object], # mutable-ok: BaseConfig signature, returned as-is + model: str, + drop_params: bool, + ) -> dict[str, object]: # mutable-ok: BaseConfig signature + return optional_params + + def transform_create_file_request( + self, + model: str, + create_file_data: CreateFileRequest, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> dict[str, object]: # mutable-ok: BaseFilesConfig signature + if "file" not in create_file_data: + raise ValueError("File data is required") + extracted: Final = extract_file_data(create_file_data["file"]) + filename: Final = extracted["filename"] or f"file_{int(time.time())}.jsonl" + content_type: Final = extracted.get("content_type") or "application/octet-stream" + upload: Final = XAIMultipartUpload( + file=(filename, extracted["content"], content_type), + purpose=(None, create_file_data.get("purpose") or _DEFAULT_PURPOSE), + ) + return dict(upload) # mutable-ok: BaseFilesConfig signature + + def transform_create_file_response( + self, + model: str | None, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + ) -> OpenAIFileObject: + return _to_openai_file_object(XAIFile.model_validate(raise_for_xai_status(raw_response).json())) + + def transform_retrieve_file_request( + self, + file_id: str, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature + return self._file_url(file_id, litellm_params), _NO_QUERY_PARAMS + + def transform_retrieve_file_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + ) -> OpenAIFileObject: + return _to_openai_file_object(XAIFile.model_validate(raise_for_xai_status(raw_response).json())) + + def transform_delete_file_request( + self, + file_id: str, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature + return self._file_url(file_id, litellm_params), _NO_QUERY_PARAMS + + def transform_delete_file_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + ) -> FileDeleted: + deleted: Final = XAIFileDeleted.model_validate(raise_for_xai_status(raw_response).json()) + return FileDeleted(id=deleted.id, deleted=deleted.deleted, object="file") + + def transform_list_files_request( + self, + purpose: str | None, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature + return f"{_api_base_from(litellm_params)}/v1/files", _NO_QUERY_PARAMS + + def transform_list_files_next_request( + self, + raw_response: httpx.Response, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, str]] | None: # mutable-ok: BaseFilesConfig signature + page: Final = XAIFileList.model_validate(raw_response.json()) + if not page.pagination_token or not page.data: + return None + return f"{_api_base_from(litellm_params)}/v1/files", {"pagination_token": page.pagination_token} + + def transform_list_files_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + ) -> list[OpenAIFileObject]: # mutable-ok: BaseFilesConfig signature + return [ # mutable-ok: BaseFilesConfig signature + _to_openai_file_object(f) + for f in XAIFileList.model_validate(raise_for_xai_status(raw_response).json()).data + ] + + def transform_file_content_request( + self, + file_content_request: FileContentRequest, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature + file_id: Final = file_content_request.get("file_id") + if file_id is None: + raise ValueError("file_id is required to download file content") + return self._file_url(file_id, litellm_params, suffix="/content"), _NO_QUERY_PARAMS + + def transform_file_content_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + ) -> HttpxBinaryResponseContent: + return HttpxBinaryResponseContent(response=raw_response) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5a207dc4c02..a670200c132 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -51436,13 +51436,16 @@ }, "xai/grok-4.20-0309-reasoning": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, @@ -51450,8 +51453,11 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true @@ -51480,9 +51486,13 @@ "xai/grok-4.3": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_image_token": 1.25e-06, "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, @@ -51490,6 +51500,8 @@ "mode": "chat", "output_cost_per_token": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -51502,9 +51514,13 @@ "xai/grok-4.3-latest": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_image_token": 1.25e-06, "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, @@ -51512,6 +51528,8 @@ "mode": "chat", "output_cost_per_token": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -59483,13 +59501,16 @@ }, "xai/grok-4.20-0309-non-reasoning": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -59497,20 +59518,26 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-multi-agent-0309": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "responses", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": false, "supports_prompt_caching": true, @@ -59519,8 +59546,11 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true, "supported_endpoints": [ @@ -62787,13 +62817,16 @@ }, "xai/grok-4.20": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, @@ -62801,21 +62834,27 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true }, "xai/grok-4.20-reasoning": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, @@ -62823,21 +62862,27 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true }, "xai/grok-4.20-reasoning-latest": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, @@ -62845,8 +62890,11 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true @@ -63067,13 +63115,16 @@ }, "xai/grok-4.20-non-reasoning": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -63081,20 +63132,26 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-non-reasoning-latest": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -63102,20 +63159,26 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-multi-agent": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "responses", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supported_endpoints": [ "/v1/responses" @@ -63127,20 +63190,26 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-multi-agent-latest": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "responses", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supported_endpoints": [ "/v1/responses" @@ -63152,8 +63221,11 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, @@ -75459,13 +75531,16 @@ }, "xai/grok-4.20-0309": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, @@ -75473,8 +75548,11 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true diff --git a/litellm/passthrough/timeout_utils.py b/litellm/passthrough/timeout_utils.py index fc67aa8c553..f600c7817f2 100644 --- a/litellm/passthrough/timeout_utils.py +++ b/litellm/passthrough/timeout_utils.py @@ -5,6 +5,8 @@ from typing import Final from pydantic import TypeAdapter +from litellm.litellm_core_utils.request_timeout_resolver import get_configured_request_timeout + DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS: Final = 600.0 _SECONDS: Final = TypeAdapter(float) @@ -48,8 +50,8 @@ def resolve_llm_passthrough_timeout( Anthropic /v1/messages). Non-streaming precedence: kwargs timeout/request_timeout -> litellm_params - timeout/request_timeout -> router_timeout -> general_settings.pass_through_request_timeout - -> 600s default. + timeout/request_timeout -> router_timeout -> litellm.request_timeout (litellm_settings.request_timeout, + when explicitly set) -> general_settings.pass_through_request_timeout -> 600s default. Streaming (``kwargs["stream"]`` truthy) resolves ``stream_timeout`` at every level before any generic timeout, matching ``Router._get_stream_timeout`` on the completion route: @@ -73,6 +75,7 @@ def resolve_llm_passthrough_timeout( deployment.get("timeout"), deployment.get("request_timeout"), router_timeout, + get_configured_request_timeout(), ) winner: Final = next((val for val in candidates if val is not None), None) return resolve_pass_through_request_timeout() if winner is None else _SECONDS.validate_python(winner) diff --git a/litellm/proxy/_experimental/mcp_server/capabilities.py b/litellm/proxy/_experimental/mcp_server/capabilities.py new file mode 100644 index 00000000000..bfd00327eb4 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/capabilities.py @@ -0,0 +1,145 @@ +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from itertools import product +from types import MappingProxyType +from typing import Final, Literal + +from mcp.server.context import CallNext, HandlerResult, ServerRequestContext +from mcp.shared.exceptions import MCPError +from mcp.types import DiscoverResult, InitializeRequestParams, InitializeResult, ServerCapabilities +from mcp_types.methods import CLIENT_REQUESTS +from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS, LATEST_HANDSHAKE_VERSION +from pydantic import TypeAdapter + +from litellm.types.mcp import MCP_LEGACY_VERSIONS, MCPAdvertisedVersions, MCPLegacyVersion, MCPSpecVersion, MCPTransport + +GATEWAY_OPERATIONS: Final = frozenset( + { + "tools/list", + "tools/call", + "prompts/list", + "prompts/get", + "resources/list", + "resources/read", + "resources/templates/list", + } +) + + +@dataclass(frozen=True, slots=True) +class RevisionSupport: + transports: frozenset[MCPTransport] + operations: frozenset[str] + results: frozenset[Literal["complete", "input_required"]] + extensions: frozenset[str] + completed: bool + + +REVISION_SUPPORT: Final[Mapping[str, RevisionSupport]] = MappingProxyType( + { + version.value: RevisionSupport( + transports=frozenset(MCPTransport) + if version.value in HANDSHAKE_PROTOCOL_VERSIONS + else frozenset({MCPTransport.http, MCPTransport.stdio}), + operations=frozenset(method for method in GATEWAY_OPERATIONS if (method, version.value) in CLIENT_REQUESTS), + results=frozenset({"complete"}) + if version.value in HANDSHAKE_PROTOCOL_VERSIONS + else frozenset({"complete", "input_required"}), + extensions=frozenset(), + completed=version.value in HANDSHAKE_PROTOCOL_VERSIONS, + ) + for version in MCPSpecVersion + } +) +_COMPLETED_REVISIONS: Final = tuple(version for version, support in REVISION_SUPPORT.items() if support.completed) +TRANSLATION_PAIRS: Final = frozenset(product(_COMPLETED_REVISIONS, repeat=2)) +_ADVERTISED_VERSIONS: Final[TypeAdapter[tuple[MCPLegacyVersion, ...]]] = TypeAdapter(MCPAdvertisedVersions) + + +def configured_versions() -> tuple[str, ...]: + from litellm.proxy.proxy_server import general_settings_view + + configured: Final = general_settings_view().get("mcp_advertised_versions") + return _ADVERTISED_VERSIONS.validate_python(MCP_LEGACY_VERSIONS if configured is None else configured) + + +def build_discovery( + *, + configured: tuple[str, ...], + revision: str, + transport: MCPTransport, + authorized_operations: frozenset[str], + upstream_versions: frozenset[str], + capabilities: ServerCapabilities, + client_extensions: frozenset[str] = frozenset(), + upstream_extensions: frozenset[str] = frozenset(), + instructions: str | None = None, +) -> DiscoverResult: + supported: Final = tuple( + version + for version, support in REVISION_SUPPORT.items() + if version in configured and support.completed and transport in support.transports + ) + revision_support: Final = REVISION_SUPPORT.get(revision) + operations: Final[frozenset[str]] = ( + authorized_operations & revision_support.operations + if revision in supported + and revision_support is not None + and any((revision, upstream) in TRANSLATION_PAIRS for upstream in upstream_versions) + else frozenset() + ) + extensions: Final[frozenset[str]] = ( + revision_support.extensions & client_extensions & upstream_extensions + if operations and revision_support is not None + else frozenset() + ) + caller_capabilities: Final = capabilities.model_copy(deep=True) + return DiscoverResult( + supported_versions=list(supported), + capabilities=ServerCapabilities( + tools=caller_capabilities.tools if {"tools/list", "tools/call"} <= operations else None, + prompts=caller_capabilities.prompts if {"prompts/list", "prompts/get"} <= operations else None, + resources=caller_capabilities.resources if {"resources/list", "resources/read"} <= operations else None, + extensions={ + key: value for key, value in (caller_capabilities.extensions or {}).items() if key in extensions + } + or None, + ), + instructions=instructions, + cache_scope="private", + ttl_ms=0, + ) + + +class GatewayVersionPolicy: + def __init__(self, versions: Callable[[], tuple[str, ...]] = configured_versions) -> None: + self._versions = versions + + async def __call__(self, ctx: ServerRequestContext[object, object], call_next: CallNext) -> HandlerResult: + versions: Final = self._versions() + requested: Final = ( + InitializeRequestParams.model_validate(ctx.params or {}).protocol_version + if ctx.method == "initialize" + else ctx.protocol_version + ) + negotiated: Final = ( + (requested if requested in HANDSHAKE_PROTOCOL_VERSIONS else LATEST_HANDSHAKE_VERSION) + if ctx.method == "initialize" + else requested + ) + if negotiated not in versions: + raise MCPError(code=-32022, message="Unsupported MCP protocol version", data={"supported": list(versions)}) + result: Final = await call_next(ctx) + if ctx.method != "initialize": + return result + initialized: Final = InitializeResult.model_validate(result) + discovery: Final = build_discovery( + configured=versions, + revision=initialized.protocol_version, + transport=MCPTransport.http, + authorized_operations=GATEWAY_OPERATIONS, + upstream_versions=frozenset(HANDSHAKE_PROTOCOL_VERSIONS), + capabilities=initialized.capabilities, + instructions=initialized.instructions, + ) + return initialized.model_copy(update={"capabilities": discovery.capabilities}) diff --git a/litellm/proxy/_experimental/mcp_server/contracts.py b/litellm/proxy/_experimental/mcp_server/contracts.py index a88d400282c..a0a08dc08ce 100644 --- a/litellm/proxy/_experimental/mcp_server/contracts.py +++ b/litellm/proxy/_experimental/mcp_server/contracts.py @@ -28,6 +28,7 @@ class OperationContext: client_ip: str | None = None mcp_proxy_mode: bool = False wire_compat: WireCompat = WireCompat.LEGACY + protocol_version: str | None = None def __post_init__(self) -> None: object.__setattr__(self, "_caller", copy_caller(self._caller)) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 24cae976174..6baa695433c 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -193,6 +193,7 @@ from litellm.types.mcp import ( MCPAuth, MCPStdioConfig, MCPTokenEndpointAuthMethod, + MCPUpstreamProtocol, has_header, without_header, ) @@ -340,6 +341,7 @@ class MCPServerConfig(TypedDict, total=False): whatever the admin wrote, and each read applies its own default.""" server_id: ReadOnly[str] + protocol_version: ReadOnly[MCPUpstreamProtocol] alias: str description: str mcp_info: MCPInfo @@ -2549,6 +2551,9 @@ class MCPServerManager: new_server = MCPServer( server_id=server_id, name=name_for_prefix, + protocol_version=TypeAdapter(MCPUpstreamProtocol).validate_python( + server_config.get("protocol_version", mcp_info.get("protocol_version", "auto")) + ), alias=alias, server_name=server_name, spec_path=server_config.get("spec_path", None), @@ -3109,6 +3114,9 @@ class MCPServerManager: new_server: Final = MCPServer( server_id=mcp_server.server_id, name=name_for_prefix, + protocol_version=TypeAdapter(MCPUpstreamProtocol).validate_python( + _mcp_info.get("protocol_version", "auto") + ), alias=getattr(mcp_server, "alias", None), server_name=getattr(mcp_server, "server_name", None), url=mcp_server.url, @@ -4145,6 +4153,7 @@ class MCPServerManager: cred_provider: UpstreamCredentialProvider | None = None, raw_headers: Mapping[str, str] | None = None, client_ip: str | None = None, + protocol_version_override: MCPUpstreamProtocol | None = None, ) -> MCPClient: """ Create an MCPClient instance for the given server. @@ -4168,6 +4177,9 @@ class MCPServerManager: """ record_auth_resolution(server.server_id, AuthResolution.unresolved) resolved_server: Final = await self.ensure_oauth_metadata_discovered(server) + protocol_version: Final = ( + protocol_version_override if protocol_version_override is not None else resolved_server.protocol_version + ) transport: Final = resolved_server.transport or MCPTransport.sse spec = None if transport == MCPTransport.stdio else _to_server_spec_fail_closed(resolved_server) provider: Final = cred_provider or self._cred_provider @@ -4249,6 +4261,7 @@ class MCPServerManager: return MCPClient( server_url="", # Not used for stdio transport_type=transport, + protocol_version=protocol_version, auth_type=resolved_server.auth_type, auth_value=auth_value, timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT), @@ -4281,6 +4294,7 @@ class MCPServerManager: MCPClient( server_url=server_url, transport_type=transport, + protocol_version=protocol_version, auth_type=resolved_server.auth_type, timeout=( resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT @@ -4324,6 +4338,7 @@ class MCPServerManager: MCPClient( server_url=server_url, transport_type=transport, + protocol_version=protocol_version, auth_type=resolved_server.auth_type, auth_value=auth_value, auth_header_name=auth_header_name, diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index ebd26e4bf87..a19246b6e90 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -14,6 +14,8 @@ from mcp.types import ( CallToolRequest, CallToolRequestParams, CallToolResult, + DiscoverRequest, + DiscoverResult, GetPromptRequest, GetPromptRequestParams, GetPromptResult, @@ -28,10 +30,14 @@ from mcp.types import ( ListToolsResult, PaginatedRequestParams, Prompt, + PromptsCapability, ReadResourceRequest, ReadResourceRequestParams, + ResourcesCapability, ResourceTemplate, + ServerCapabilities, TextContent, + ToolsCapability, ) from mcp.types import Tool as MCPTool from pydantic import AnyUrl, ConfigDict, Field, TypeAdapter @@ -51,6 +57,11 @@ from litellm.proxy._experimental.mcp_server.byok_credential_cache import ( cache_byok_credential, get_cached_byok_credential, ) +from litellm.proxy._experimental.mcp_server.capabilities import ( + GATEWAY_OPERATIONS, + build_discovery, + configured_versions, +) from litellm.proxy._experimental.mcp_server.contracts import ( AuthorizedToolCall, OperationContext, @@ -122,7 +133,9 @@ from litellm.proxy.litellm_pre_call_utils import ( ) from litellm.types.mcp import ( DEFAULT_CREDENTIAL_HEADER, + MCP_LEGACY_VERSIONS, MCPAuth, + MCPTransport, without_header, ) from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer @@ -2657,7 +2670,11 @@ class _McpDeniedDetail(TypedDict): async def _execute_handle_list_tools( - context: OperationContext, params: PaginatedRequestParams, host_progress_callback: ProgressCallback | None = None + context: OperationContext, + params: PaginatedRequestParams, + host_progress_callback: ProgressCallback | None = None, + *, + log_list_tools_to_spendlogs: bool = True, ) -> ListToolsResult: try: ( @@ -2700,7 +2717,7 @@ async def _execute_handle_list_tools( mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, raw_headers=raw_headers, - log_list_tools_to_spendlogs=True, + log_list_tools_to_spendlogs=log_list_tools_to_spendlogs, list_tools_log_source="mcp_protocol", client_ip=_client_ip, ) @@ -3065,6 +3082,7 @@ def prepare_context( client_ip: str | None = None, mcp_proxy_mode: bool = False, wire_compat: WireCompat = WireCompat.LEGACY, + protocol_version: str | None = None, ) -> OperationContext: return OperationContext( _caller=user_api_key_auth, @@ -3076,11 +3094,13 @@ def prepare_context( client_ip=client_ip, mcp_proxy_mode=mcp_proxy_mode, wire_compat=wire_compat, + protocol_version=protocol_version, ) GatewayOperation: TypeAlias = ( AuthorizedToolCall + | DiscoverRequest | ListToolsRequest | CallToolRequest | ListPromptsRequest @@ -3090,7 +3110,8 @@ GatewayOperation: TypeAlias = ( | ReadResourceRequest ) GatewayResult: TypeAlias = ( - ListToolsResult + DiscoverResult + | ListToolsResult | CallToolResult | InputRequiredResult | ListPromptsResult @@ -3105,6 +3126,9 @@ class GatewayOperations: def __init__(self, host_progress_callback: ProgressCallback | None = None) -> None: self._host_progress_callback = host_progress_callback + @overload + async def execute(self, operation: DiscoverRequest, context: OperationContext) -> DiscoverResult: ... + @overload async def execute( self, operation: AuthorizedToolCall, context: OperationContext @@ -3137,6 +3161,51 @@ class GatewayOperations: async def execute(self, operation: GatewayOperation, context: OperationContext) -> GatewayResult: match operation: + case DiscoverRequest(): + listings: Final = ( + () + if context.mcp_proxy_mode + else (ListPromptsRequest(), ListResourcesRequest(), ListResourceTemplatesRequest()) + ) + tasks: Final = ( + asyncio.create_task( + _execute_handle_list_tools( + context, + PaginatedRequestParams(), + self._host_progress_callback, + log_list_tools_to_spendlogs=False, + ) + ), + *(asyncio.create_task(self.execute(listing, context)) for listing in listings), + ) + try: + results: Final = await asyncio.gather(*tasks) + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + return build_discovery( + configured=configured_versions(), + revision=context.protocol_version or "2025-11-25", + transport=MCPTransport.http, + authorized_operations=GATEWAY_OPERATIONS, + upstream_versions=frozenset(MCP_LEGACY_VERSIONS), + capabilities=ServerCapabilities( + tools=ToolsCapability() + if any(isinstance(result, ListToolsResult) and result.tools for result in results) + else None, + prompts=PromptsCapability() + if any(isinstance(result, ListPromptsResult) and result.prompts for result in results) + else None, + resources=ResourcesCapability() + if any( + (isinstance(result, ListResourcesResult) and result.resources) + or (isinstance(result, ListResourceTemplatesResult) and result.resource_templates) + for result in results + ) + else None, + ), + ) case AuthorizedToolCall(): auth, token, _servers, server_headers, oauth_headers, headers, _client_ip = context.legacy_auth() return await _execute_mcp_tool( diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 7f519e2c0d9..02694f110b1 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1375,7 +1375,16 @@ if MCP_AVAILABLE: and headers.get(MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_PRIMARY) else None ) - return _StagedServerTest(request=request, mcp_auth_header=mcp_auth_header, oauth2_headers=oauth2_headers) + preview_request: Final = ( + request.model_copy( + update={"mcp_info": {**(request.mcp_info or {}), "protocol_version": saved_server.protocol_version}} + ) + if saved_server is not None and "protocol_version" not in (request.mcp_info or {}) + else request + ) + return _StagedServerTest( + request=preview_request, mcp_auth_header=mcp_auth_header, oauth2_headers=oauth2_headers + ) async def _list_tools_within(client: MCPClient, deadline: float) -> list[MCPTool] | None: with anyio.move_on_after(deadline): @@ -1512,6 +1521,7 @@ if MCP_AVAILABLE: extra_headers=merged_headers, stdio_env=stdio_env, cred_provider=preview_cred_provider, + protocol_version_override=server_model.protocol_version, ) return await operation(client) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 1bd31d971b0..555aebc7434 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -125,12 +125,14 @@ def unsupported_protocol_version(scope: Scope) -> str | None: ``HANDSHAKE_PROTOCOL_VERSIONS`` to the modern single-exchange path, which bypasses litellm's session/auth model, so the ASGI entry rejects it. """ + from litellm.proxy._experimental.mcp_server.capabilities import configured_versions + headers: Final[Iterable[tuple[bytes, bytes]]] = scope.get("headers") or () values: Final = tuple( raw.decode("latin-1").strip() for key, raw in headers if key.lower() == _MCP_PROTOCOL_VERSION_HEADER ) for value in values: - if value and value not in HANDSHAKE_PROTOCOL_VERSIONS: + if value and value not in configured_versions(): return value return None @@ -149,7 +151,10 @@ try: from mcp.server.session import ServerSession as _McpServerSession from mcp.types import ( BlobResourceContents, + DiscoverRequest, + DiscoverResult, GetPromptResult, + RequestParams, ResourceTemplate, TextResourceContents, ) @@ -526,11 +531,11 @@ if MCP_AVAILABLE: PaginatedRequestParams, ReadResourceRequestParams, ) - from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import ( MCPAuthenticatedUser, ) + from litellm.proxy._experimental.mcp_server.capabilities import GatewayVersionPolicy, configured_versions from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, global_mcp_server_manager, @@ -585,6 +590,7 @@ if MCP_AVAILABLE: name=LITELLM_MCP_SERVER_NAME, version=LITELLM_MCP_SERVER_VERSION, ) + server.middleware.append(GatewayVersionPolicy()) server.create_initialization_options = types.MethodType(_gateway_create_initialization_options, server) sse: Final[SseServerTransport] = SseServerTransport("/sse/messages") @@ -830,6 +836,7 @@ if MCP_AVAILABLE: client_ip, _mcp_proxy_mode.get(), wire_compat_for(ctx.protocol_version), + ctx.protocol_version, ) async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams) -> ListToolsResult: @@ -948,6 +955,11 @@ if MCP_AVAILABLE: ReadResourceRequest(params=params), context ) + async def discover(ctx: ServerRequestContext, params: RequestParams) -> DiscoverResult: + async with _legacy_operation_context(ctx, trace=False) as context: + return await operations.GatewayOperations().execute(DiscoverRequest(params=params), context) + + server.add_request_handler("server/discover", RequestParams, discover) server.add_request_handler("tools/list", PaginatedRequestParams, handle_list_tools) server.add_request_handler("tools/call", CallToolRequestParams, mcp_server_tool_call) server.add_request_handler("prompts/list", PaginatedRequestParams, list_prompts) @@ -1954,7 +1966,7 @@ if MCP_AVAILABLE: reject_disallowed_mcp_origin(StarletteRequest(scope)) bad_version: Final = unsupported_protocol_version(scope) if bad_version is not None: - supported: Final = ", ".join(sorted(HANDSHAKE_PROTOCOL_VERSIONS)) + supported: Final = ", ".join(configured_versions()) await JSONResponse( status_code=400, content={ # mutable-ok: JSON-RPC error payload @@ -2299,7 +2311,7 @@ if MCP_AVAILABLE: reject_disallowed_mcp_origin(StarletteRequest(scope)) bad_version: Final = unsupported_protocol_version(scope) if bad_version is not None: - supported: Final = ", ".join(sorted(HANDSHAKE_PROTOCOL_VERSIONS)) + supported: Final = ", ".join(configured_versions()) await JSONResponse( status_code=400, content={ # mutable-ok: JSON-RPC error payload diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3da30070b9e..4597872d84e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -37,6 +37,7 @@ from litellm.types.llms.openai import ( ResponsesAPIResponse, ) from litellm.types.mcp import ( + MCPAdvertisedVersions, MCPAllowedClient, MCPAuth, MCPAuthType, @@ -2998,6 +2999,11 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="Custom CIDR ranges that define internal/private networks for MCP access control. When set, only these ranges are treated as internal. Defaults to RFC 1918 private ranges (10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16, 127.0.0.0/8).", ) + mcp_advertised_versions: MCPAdvertisedVersions | None = Field( + None, + description="MCP revisions enabled by the gateway. Defaults to all completed legacy revisions. " + "Modern protocol serving and Apps/Tasks remain disabled.", + ) mcp_allowed_clients: list[MCPAllowedClient] | None = Field( None, description="MCP client applications admitted by the gateway, each an {alias, value} pair where alias is the name shown in the dashboard and logs and value is the identity that must match exactly. When set, every MCP request must carry a client identity equal to one of the values: a JWT caller is identified by the claim named in litellm_jwtauth.mcp_client_id_jwt_field, any other caller by the header named in mcp_client_id_header. A request with no resolvable identity, or an unlisted one, is rejected with 403. Unset means every client is admitted.", @@ -4013,6 +4019,18 @@ class AllCallbacks(LiteLLMPydanticObjectBase): ], ) + zerobus: CallbackOnUI = CallbackOnUI( + litellm_callback_name="zerobus", + ui_callback_name="Databricks Zerobus", + litellm_callback_params=[ # mutable-ok: the registry field is typed list + "ZEROBUS_WORKSPACE_URL", + "ZEROBUS_SERVER_ENDPOINT", + "ZEROBUS_CLIENT_ID", + "ZEROBUS_CLIENT_SECRET", + "ZEROBUS_TABLE_NAME", + ], + ) + class HTTPExceptionErrorDetail(TypedDict): """The `{"error": }` shape most proxy endpoints raise as `HTTPException.detail`.""" @@ -5349,11 +5367,12 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): default=False, description=( "When True, users whose JWT contains no team claims are authenticated " - "using their database team memberships instead of receiving HTTP 403. " - "Usage is attributed to the user's first resolvable DB team, or to the " - "team specified via the x-litellm-team-id request header (validated " - "against DB membership). Requires user_id_upsert=True so that user " - "records exist before the fallback runs." + "using their database team memberships instead of receiving HTTP 403, " + "with usage attributed to the user's first resolvable DB team. Whether or " + "not the JWT carries team claims, the x-litellm-team-id request header may " + "select any team the user is a member of in the database (validated against " + "DB membership); without the header the JWT team stays the default. Requires " + "user_id_upsert=True so that user records exist before the fallback runs." ), ) issuers: list[JWTIssuerConfig] | None = Field( diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index e07b20fd5d5..4f41b283a33 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -1930,12 +1930,12 @@ class JWTAuthManager: ) -> HeaderTeam | None: """ The team named by x-litellm-team-id, which may carry a team id or a team - alias. A value that is already an allowed team id (or, under the DB - fallback, an existing team id) never costs an alias lookup; an alias is - accepted only when the team it names would have been accepted by id. - Under the DB fallback only a team row that is provably absent falls - through to the alias lookup; a read that failed for any other reason - keeps the membership denial the id path already gives. + alias. A value that is already an allowed team id never costs a lookup; + under the DB fallback any other value is accepted provisionally, by id + or alias, for the membership check auth_builder runs later. Under the + DB fallback only a team row that is provably absent falls through to + the alias lookup; a read that failed for any other reason keeps the + membership denial the id path already gives. Raises: HTTPException: 403 when neither the value nor the team it aliases is @@ -1948,7 +1948,11 @@ class JWTAuthManager: if not header_value: return None - if fallback_to_db_teams and not allowed_team_ids: + if header_value in allowed_team_ids: + verbose_proxy_logger.debug("Using team_id from x-litellm-team-id header: %s", header_value) + return HeaderTeam(header_value=header_value, team_id=header_value) + + if fallback_to_db_teams: try: await get_team_object( team_id=header_value, @@ -1969,10 +1973,6 @@ class JWTAuthManager: JWTAuthManager._raise_header_team_membership_denial(header_value) return HeaderTeam(header_value=header_value, team_id=header_value) - if header_value in allowed_team_ids: - verbose_proxy_logger.debug("Using team_id from x-litellm-team-id header: %s", header_value) - return HeaderTeam(header_value=header_value, team_id=header_value) - team_id_by_alias: Final = await JWTAuthManager._team_id_by_alias( header_value, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj ) @@ -2353,9 +2353,9 @@ class JWTAuthManager: header_value: str, ) -> None: """ - A provisional team_id from the x-litellm-team-id header (accepted without - JWT-team validation when the JWT carries no team claims) must exist in the - user's DB team memberships before it becomes request context. The denial + A provisional team_id from the x-litellm-team-id header (accepted under + fallback_to_db_teams because it is outside the JWT's teams) must exist in + the user's DB team memberships before it becomes request context. The denial names `header_value`, the id or alias the caller sent, not `team_id`. """ user_team_ids: Final = user_object.teams if user_object else [] @@ -2587,22 +2587,30 @@ class JWTAuthManager: if specific_team_id and not db_team_fallback: all_team_ids.add(specific_team_id) + header_db_fallback: Final = handler.litellm_jwtauth.fallback_to_db_teams and team_id is None + header_team: Final = await JWTAuthManager.resolve_team_from_header( request_headers=request_headers, allowed_team_ids=all_team_ids, - fallback_to_db_teams=db_team_fallback, + fallback_to_db_teams=header_db_fallback, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) + provisional_header_team: Final = ( + header_team + if header_team is not None and header_db_fallback and header_team.team_id not in all_team_ids + else None + ) if header_team: team_id = header_team.team_id - # A provisional header team (accepted only because the JWT carries no - # team claims) is validated against DB membership further down; never - # upsert it here or an attacker-supplied x-litellm-team-id would create - # an orphaned team row before that check runs. A genuine membership team - # already exists, so suppressing the upsert in that case costs nothing. + # A provisional header team (accepted because it is outside the + # JWT's teams under fallback_to_db_teams) is validated against DB + # membership further down; never upsert it here or an + # attacker-supplied x-litellm-team-id would create an orphaned team + # row before that check runs. A genuine membership team already + # exists, so suppressing the upsert in that case costs nothing. try: team_object = await get_team_object( team_id=team_id, @@ -2610,10 +2618,10 @@ class JWTAuthManager: user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, - team_id_upsert=(team_id_upsert and not db_team_fallback), + team_id_upsert=(team_id_upsert and provisional_header_team is None), ) except HTTPException: - if not db_team_fallback: + if provisional_header_team is None: raise JWTAuthManager._raise_header_team_membership_denial(header_team.header_value) elif not team_id and not db_team_fallback: @@ -2756,11 +2764,11 @@ class JWTAuthManager: proxy_logging_obj=proxy_logging_obj, team_id_upsert=team_id_upsert, ) - elif db_team_fallback and header_team is not None and team_id == header_team.team_id: + elif provisional_header_team is not None and team_id == provisional_header_team.team_id: JWTAuthManager._validate_header_team_in_db_membership( team_id=team_id, user_object=user_object, - header_value=header_team.header_value, + header_value=provisional_header_team.header_value, ) if not JWTAuthManager._is_team_route_allowed( route=route, @@ -2770,7 +2778,7 @@ class JWTAuthManager: raise HTTPException( status_code=403, detail=( - f"Team '{header_team.header_value}' (from x-litellm-team-id header) " + f"Team '{provisional_header_team.header_value}' (from x-litellm-team-id header) " f"is not allowed to access route '{route}'." ), ) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 7e1989b0aab..64b0c6c1967 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -32,6 +32,7 @@ from starlette.types import Receive, Scope, Send import litellm from litellm._logging import redact_internal_details_from_client_message, verbose_proxy_logger from litellm._uuid import uuid +from litellm.anthropic_interface.exceptions import AnthropicErrorSseFrame, anthropic_error_sse_frame from litellm.constants import ( DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE, DEFAULT_MAX_RECURSE_DEPTH, @@ -102,8 +103,11 @@ from litellm.proxy.common_utils.openai_error_payload import ( ) from litellm.proxy.common_utils.sse_keepalive import ( SSE_COMMENT_PING_BYTES, + SSE_STREAM_START_TAIL, + advance_sse_tail, coerce_keepalive_interval, resolve_ttft_keepalive_interval, + seal_open_sse_frame, wrap_sse_stream_with_keepalive_pings, ) from litellm.proxy.dd_span_tagger import DDSpanTagger @@ -999,6 +1003,17 @@ async def create_response( first_chunk_value = await _buffer_first_chunk_honoring_disconnect(generator, request) resolved_headers: Final = await _resolve_stream_headers(headers, refresh_headers) + if isinstance(first_chunk_value, AnthropicErrorSseFrame): + with contextlib.suppress(Exception): + await generator.aclose() + return JSONResponse( + status_code=first_chunk_value.status_code, + content=first_chunk_value.json_body( + error_body_call_id(general_settings, resolved_headers.get(LITELLM_CALL_ID_HEADER)) + ), + headers=resolved_headers, + ) + if first_chunk_value is not None: try: error_code_from_chunk: Final = await _parse_event_data_for_error(first_chunk_value) @@ -2781,15 +2796,6 @@ class ProxyBaseLLMRequestProcessing: request=request, ) if route_type == "aresponses": - # Streaming /v1/responses returns here without - # reaching the non-streaming ownership tail below. - # Wrap the SSE generator so container ownership is - # written once the upstream iterator finishes - # assembling ``completed_response`` — otherwise - # code-interpreter containers created during the - # stream stay unregistered and follow-up file API - # calls 403. Covers the background-polling path - # too, which loops ``body_iterator`` end-to-end. selected_data_generator = ( ProxyBaseLLMRequestProcessing._wrap_responses_stream_for_container_ownership( original_stream_response=response, @@ -3011,50 +3017,50 @@ class ProxyBaseLLMRequestProcessing: wrapped_generator: Any, user_api_key_dict: UserAPIKeyAuth, ): - """Forward SSE chunks, then record container ownership at stream end. + """Forward SSE chunks and record container ownership before the terminal chunk goes out. Streaming ``/v1/responses`` short-circuits out of ``base_process_llm_request`` before the non-streaming ownership - tail runs, so without this wrap the - ``LiteLLM_ManagedObjectTable`` row for any container created - during the stream is never written and follow-up file API calls - return 403. + tail runs. The OpenAI SDK closes the connection at ``data: [DONE]`` + and starlette cancels the body task on disconnect, so a write that + waits for the generator to finish never lands. The iterator sets + ``completed_response`` before it hands over its terminal chunk, so + the ``LiteLLM_ManagedObjectTable`` row is written the moment it + appears, ahead of the chunk carrying ``response.completed``. """ - try: - async for chunk in wrapped_generator: + async for chunk in wrapped_generator: + completed_obj = ProxyBaseLLMRequestProcessing._extract_completed_responses_response( + original_stream_response + ) + if completed_obj is None: yield chunk - finally: - try: - completed_obj: Final = ProxyBaseLLMRequestProcessing._extract_completed_responses_response( - original_stream_response - ) - if completed_obj is not None: - await ProxyBaseLLMRequestProcessing._record_container_owners_from_responses_if_needed( - response=completed_obj, - user_api_key_dict=user_api_key_dict, - ) - else: - # Silent skip caused #30210: the proxy's Router wrapper - # of the responses streaming iterator wasn't propagating - # ``completed_response``, so this hook recorded nothing - # and follow-up /v1/containers//files calls 403'd - # for non-admin keys with no proxy-side hint. Log a - # warning so future regressions of the same shape - # surface in operator logs. - verbose_proxy_logger.warning( - "Container ownership recording skipped on streaming " - "/v1/responses: no completed_response on stream " - "iterator %s. If this stream created any tool " - "container (e.g. code_interpreter), follow-up " - "/v1/containers//files calls will 403 for " - "non-admin keys.", - type(original_stream_response).__name__, - ) - except Exception as e: - verbose_proxy_logger.exception( - "Container ownership recording failed after streaming responses call: %s", - e, - ) + continue + await ProxyBaseLLMRequestProcessing._record_container_owners_from_responses_if_needed( + response=completed_obj, + user_api_key_dict=user_api_key_dict, + ) + yield chunk + async for remaining_chunk in wrapped_generator: + yield remaining_chunk + return + late_completed_obj: Final = ProxyBaseLLMRequestProcessing._extract_completed_responses_response( + original_stream_response + ) + if late_completed_obj is not None: + await ProxyBaseLLMRequestProcessing._record_container_owners_from_responses_if_needed( + response=late_completed_obj, + user_api_key_dict=user_api_key_dict, + ) + return + verbose_proxy_logger.warning( + "Container ownership recording skipped on streaming " + "/v1/responses: no completed_response on stream " + "iterator %s. If this stream created any tool " + "container (e.g. code_interpreter), follow-up " + "/v1/containers//files calls will 403 for " + "non-admin keys.", + type(original_stream_response).__name__, + ) async def base_passthrough_process_llm_request( self, @@ -3861,6 +3867,7 @@ class ProxyBaseLLMRequestProcessing: serialize_error: StreamErrorSerializer, request: Request | None = None, flush_tail: Callable[[], bytes] | None = None, + seal_open_frame: Callable[[bytes], str] | None = None, ) -> AsyncGenerator[str, None]: """ Shared streaming data generator: runs proxy iterator hook, per-chunk hook, @@ -3870,6 +3877,12 @@ class ProxyBaseLLMRequestProcessing: ``flush_tail`` runs once after the upstream iterator completes cleanly and its non-empty result is yielded, so a serializer that buffers bytes across chunks can emit anything still held at end of stream. + + ``seal_open_frame`` is given the tail of what has been yielded when the + error frame goes out, and what it returns is written first. A passthrough + relays raw upstream bytes, so an upstream that hangs up mid-frame leaves the + client inside an open frame, where an error frame would be swallowed or + misparsed instead of raised. """ verbose_proxy_logger.debug("inside generator") # Resolve per-stream (not per-chunk) whether the heavy per-chunk path @@ -3886,6 +3899,7 @@ class ProxyBaseLLMRequestProcessing: stream_completed = False client_disconnected = False delivered_chunk = False + recent_tail = SSE_STREAM_START_TAIL # rebind-ok: rolling window over the yielded bytes try: str_so_far = "" async for chunk in proxy_logging_obj.async_post_call_streaming_iterator_hook( @@ -3931,7 +3945,9 @@ class ProxyBaseLLMRequestProcessing: # False and refunds. A keepalive ping carries no provider output, # so it must not suppress that refund. delivered_chunk = delivered_chunk or chunk != STREAM_SSE_KEEPALIVE_PING_BYTES - yield serialize_chunk(chunk) + serialized = serialize_chunk(chunk) + recent_tail = advance_sse_tail(recent_tail, serialized) + yield serialized held_tail: Final = flush_tail() if flush_tail is not None else b"" if held_tail: yield serialize_chunk(held_tail) @@ -3979,7 +3995,9 @@ class ProxyBaseLLMRequestProcessing: code=stream_error_status, ) stream_completed = True - yield serialize_error(proxy_exception) + error_frame: Final = serialize_error(proxy_exception) + seal: Final = "" if seal_open_frame is None else seal_open_frame(recent_tail) + yield seal + error_frame if seal else error_frame finally: await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup( request=request, @@ -4001,7 +4019,7 @@ class ProxyBaseLLMRequestProcessing: restamp_model: str | None = None, ) -> AsyncGenerator[str, None]: """ - Anthropic /messages and Google /generateContent streaming data generator require SSE events. + Anthropic /messages streaming data generator, which requires SSE events. Returns the underlying ``async_streaming_data_generator`` configured with SSE serializers directly (rather than re-wrapping it in another @@ -4019,11 +4037,13 @@ class ProxyBaseLLMRequestProcessing: request_data=request_data, proxy_logging_obj=proxy_logging_obj, serialize_chunk=ProxyBaseLLMRequestProcessing._sse_chunk_serializer(restamper), - serialize_error=lambda proxy_exc: ( - f"{STREAM_SSE_DATA_PREFIX}{json.dumps({'error': proxy_exc.to_dict()})}\n\n" + serialize_error=lambda proxy_exc: anthropic_error_sse_frame( + status_code=error_status_code(proxy_exc, status.HTTP_500_INTERNAL_SERVER_ERROR), + raw_message=proxy_exc.message, ), request=request, flush_tail=None if restamper is None else restamper.flush, + seal_open_frame=seal_open_sse_frame, ) @overload diff --git a/litellm/proxy/common_utils/sse_keepalive.py b/litellm/proxy/common_utils/sse_keepalive.py index cf98a7e9224..d9685971f52 100644 --- a/litellm/proxy/common_utils/sse_keepalive.py +++ b/litellm/proxy/common_utils/sse_keepalive.py @@ -15,7 +15,7 @@ SSE_COMMENT_PING_BYTES: Final = SSE_COMMENT_PING.encode() # terminates a line with CRLF, LF or CR, so a blank line is any of these three. _SSE_FRAME_DELIMITERS: Final = (b"\r\n\r\n", b"\n\n", b"\r\r") _SSE_DELIMITER_LOOKBACK: Final = max(len(delimiter) for delimiter in _SSE_FRAME_DELIMITERS) -_STREAM_START_TAIL: Final = b"\n\n" +SSE_STREAM_START_TAIL: Final = b"\n\n" _SSE_MEDIA_TYPE: Final = "text/event-stream" @@ -128,7 +128,7 @@ async def _keepalive_ping_byte_stream( # Seeded as a delimiter because a stream starts at a frame boundary, and kept # across chunks because a delimiter can be split between two transport reads, # which testing only the latest chunk would miss for the rest of the stream. - recent_tail = _STREAM_START_TAIL # rebind-ok: rolling window over the relayed bytes + recent_tail = SSE_STREAM_START_TAIL # rebind-ok: rolling window over the relayed bytes try: while True: await asyncio.wait((pending,), timeout=ping_interval_seconds) @@ -155,6 +155,28 @@ async def _keepalive_ping_byte_stream( await stream.aclose() +def advance_sse_tail(recent_tail: bytes, chunk: object) -> bytes: + written: Final = _sse_tail_bytes(chunk) + if not written: + return recent_tail + return (recent_tail + written)[-_SSE_DELIMITER_LOOKBACK:] + + +def _sse_tail_bytes(chunk: object) -> bytes: + if isinstance(chunk, bytes): + return chunk[-_SSE_DELIMITER_LOOKBACK:] + if isinstance(chunk, str): + return chunk[-_SSE_DELIMITER_LOOKBACK:].encode() + return b"" + + +def seal_open_sse_frame(recent_tail: bytes) -> str: + if recent_tail.endswith(_SSE_FRAME_DELIMITERS): + return "" + line_break: Final = "" if recent_tail.endswith((b"\n", b"\r")) else "\n" + return f"{line_break}{ANTHROPIC_PING_SSE_CHUNK}" + + def resolve_ttft_keepalive_interval( deployments: Iterable[Mapping[str, object]], global_interval: float | str | None, diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 7859c678c07..618b200a14c 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -4618,6 +4618,13 @@ class GoogleSSOHandler: return result or {} +def _raise_if_sso_debug_disabled() -> None: + """The debug routes run the browser-redirect SSO flow, so they cannot carry a + bearer credential; an explicit opt-in flag is the only way to gate them.""" + if get_secret_bool("ENABLE_SSO_DEBUG") is not True: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Not Found") + + @router.get("/sso/debug/login", tags=["experimental"], include_in_schema=False) async def debug_sso_login(request: Request): """ @@ -4625,6 +4632,8 @@ async def debug_sso_login(request: Request): PROXY_BASE_URL should be the your deployed proxy endpoint, e.g. PROXY_BASE_URL="https://litellm-production-7002.up.railway.app/" Example: """ + _raise_if_sso_debug_disabled() + from litellm.proxy.proxy_server import premium_user microsoft_client_id: Final = os.getenv("MICROSOFT_CLIENT_ID", None) @@ -4670,6 +4679,8 @@ async def debug_sso_callback(request: Request): """ Returns the OpenID object returned by the SSO provider """ + _raise_if_sso_debug_disabled() + import json from fastapi.responses import HTMLResponse diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 70f2a7db6da..1ef39775bd3 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -10,7 +10,7 @@ from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from datetime import datetime from functools import lru_cache from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload, runtime_checkable +from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Protocol, overload, runtime_checkable import httpx from openai._streaming import SSEDecoder @@ -230,6 +230,11 @@ def _status_code_for_error_fields(error_type: str | None, error_code: str | None return next((status for status in map(_status_code_for_error_field, fields) if status is not None), 500) +def stream_error_status_and_message(error_obj: object) -> tuple[int, str]: + message, error_type, error_code = _error_event_fields(error_obj) + return _status_code_for_error_fields(error_type, error_code), message + + def _map_stream_error_to_exception(error_obj: object, model: str, custom_llm_provider: str) -> Exception: from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -260,6 +265,9 @@ def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool: return not isinstance(status_code, int) or status_code >= 500 or status_code == 429 +_PRE_OUTPUT_LIFECYCLE_EVENT_TYPES: Final = frozenset({"response.created", "response.in_progress", "response.queued"}) + + class BaseResponsesAPIStreamingIterator: """ Base class for streaming iterators that process responses from the Responses API. @@ -287,6 +295,7 @@ class BaseResponsesAPIStreamingIterator: self.start_time = getattr(logging_obj, "start_time", datetime.now()) self._failure_handled = False # Track if failure handler has been called self._yielded_first_chunk = False + self._output_started = False self._generated_content = "" self._generated_tool_arguments = "" self._completed_response_cached = False @@ -874,6 +883,46 @@ class BaseResponsesAPIStreamingIterator: except Exception: pass + def _note_yielded_event(self, event: ResponsesAPIStreamingResponse) -> None: + self._yielded_first_chunk = True + if event.type not in _PRE_OUTPUT_LIFECYCLE_EVENT_TYPES: + self._output_started = True + + def _fallback_error(self, original: Exception) -> MidStreamFallbackError: + return MidStreamFallbackError( + message=str(original), + model=self.model or "", + llm_provider=self.custom_llm_provider or "", + original_exception=original, + generated_content="", + is_pre_first_chunk=not self._yielded_first_chunk, + ) + + def _stream_ended_early_error(self) -> litellm.APIConnectionError: + return litellm.APIConnectionError( + message=( + f"{self.custom_llm_provider or 'provider'} closed the responses stream before any terminal event " + "(response.completed, response.incomplete or response.failed)" + ), + llm_provider=self.custom_llm_provider or "", + model=self.model or "", + ) + + def _raise_if_ended_without_terminal_event(self) -> None: + if self.completed_response is not None: + return + error: Final = self._stream_ended_early_error() + self._handle_failure(error) + if self._output_started: + raise error + raise self._fallback_error(error) from error + + def _raise_for_transport_error(self, error: httpx.ReadError | httpx.RemoteProtocolError) -> NoReturn: + self._handle_failure(error) + if self._output_started: + raise error + raise self._fallback_error(error) from error + async def call_post_streaming_hooks_for_testing( iterator: object, chunk: ResponsesAPIStreamingResponse @@ -929,12 +978,14 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): sse = await self.stream_iterator.__anext__() except StopAsyncIteration: self.finished = True + self._raise_if_ended_without_terminal_event() raise StopAsyncIteration self._check_max_streaming_duration() result = self._process_chunk(sse.data) if self.finished: + self._raise_if_ended_without_terminal_event() raise StopAsyncIteration elif result is not None: self._maybe_raise_for_error_event(result) @@ -943,7 +994,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): result = await self._call_post_streaming_deployment_hook( chunk=result, ) - self._yielded_first_chunk = True + self._note_yielded_event(result) return result # If result is None, continue the loop to get the next chunk @@ -952,10 +1003,9 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): raise except (httpx.ReadError, httpx.RemoteProtocolError) as e: self.finished = True - if self.completed_response is None: - self._handle_failure(e) - raise - raise StopAsyncIteration from e + if self.completed_response is not None: + raise StopAsyncIteration from e + self._raise_for_transport_error(e) except httpx.HTTPError as e: # Handle HTTP errors self.finished = True @@ -1011,12 +1061,14 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): sse = next(self.stream_iterator) except StopIteration: self.finished = True + self._raise_if_ended_without_terminal_event() raise StopIteration self._check_max_streaming_duration() result = self._process_chunk(sse.data) if self.finished: + self._raise_if_ended_without_terminal_event() raise StopIteration elif result is not None: self._maybe_raise_for_error_event(result) @@ -1025,7 +1077,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): async_function=self._call_post_streaming_deployment_hook, chunk=result, ) - self._yielded_first_chunk = True + self._note_yielded_event(result) return result # If result is None, continue the loop to get the next chunk @@ -1034,10 +1086,9 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): raise except (httpx.ReadError, httpx.RemoteProtocolError) as e: self.finished = True - if self.completed_response is None: - self._handle_failure(e) - raise - raise StopIteration from e + if self.completed_response is not None: + raise StopIteration from e + self._raise_for_transport_error(e) except httpx.HTTPError as e: # Handle HTTP errors self.finished = True diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 4e84bded9de..64252cbbfb3 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -502,12 +502,12 @@ class RouterBudgetLimiting(CustomLogger): response_cost: Final[float] = standard_logging_payload.get("response_cost", 0) model_id: Final[str] = str(standard_logging_payload.get("model_id", "")) - custom_llm_provider: Final[str] = kwargs.get("litellm_params", {}).get("custom_llm_provider", None) - if custom_llm_provider is None: - raise ValueError("custom_llm_provider is required") + custom_llm_provider: Final[str | None] = standard_logging_payload.get("custom_llm_provider") - budget_config: Final = self._get_budget_config_for_provider(custom_llm_provider) - if budget_config: + budget_config: Final = ( + self._get_budget_config_for_provider(custom_llm_provider) if custom_llm_provider is not None else None + ) + if custom_llm_provider is not None and budget_config is not None: # increment spend for provider spend_key: Final = f"provider_spend:{custom_llm_provider}:{budget_config.budget_duration}" start_time_key: Final = f"provider_budget_start_time:{custom_llm_provider}" diff --git a/litellm/types/integrations/zerobus.py b/litellm/types/integrations/zerobus.py new file mode 100644 index 00000000000..217002dbc65 --- /dev/null +++ b/litellm/types/integrations/zerobus.py @@ -0,0 +1,53 @@ +from dataclasses import dataclass, field +from typing import Final + +from pydantic import Field + +from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams + +RETRYABLE_INGEST_STATUS_CODES: Final = frozenset({408, 429, 500, 502, 503, 504}) + +TOKEN_REFRESH_LEEWAY_SECONDS: Final = 60 + + +class ZerobusInitParams(StandardCustomLoggerInitParams): + """ + Params for initializing a Databricks Zerobus logger on litellm. + + Every connection field falls back to its ``ZEROBUS_*`` environment variable, which is + what the proxy UI writes. ``table_name`` is the fully qualified ``catalog.schema.table``. + """ + + workspace_url: str | None = None + server_endpoint: str | None = None + client_id: str | None = None + client_secret: str | None = None + table_name: str | None = None + batch_size: int = Field(default=100, gt=0) + flush_interval: int = Field(default=10, gt=0) + + +@dataclass(frozen=True, slots=True) +class ZerobusConnection: + """Everything needed to mint a token for one table and post rows to it.""" + + workspace_url: str + workspace_id: str + server_endpoint: str + client_id: str + client_secret: str = field(repr=False) + table_name: str + + +@dataclass(frozen=True, slots=True) +class ZerobusAccessToken: + value: str = field(repr=False) + expires_at: float + + +@dataclass(frozen=True, slots=True) +class ZerobusIngestFailure: + """Why a batch could not be written, and whether a later attempt could still succeed.""" + + detail: str + retryable: bool diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 6e7e9da3498..99ab5920c4f 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -522,7 +522,19 @@ class CreateBatchRequest(TypedDict, total=False): """ completion_window: Literal["24h"] - endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses", "/v1/ocr"] + endpoint: Literal[ + "/v1/chat/completions", + "/v1/embeddings", + "/v1/completions", + "/v1/responses", + "/v1/ocr", + "/v1/images/generations", + "/v1/images/edits", + "/v1/videos/generations", + "/v1/videos", + "/v1/videos/edits", + "/v1/videos/extensions", + ] input_file_id: str metadata: dict[str, str] | None output_expires_after: FileExpiresAfter diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index 83e719810d5..fec5e84c8df 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -4,7 +4,7 @@ import enum import re from collections.abc import Awaitable, Callable, Mapping from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Annotated, Any, Final, Literal from urllib.parse import urlsplit import httpx @@ -34,6 +34,8 @@ class MCPSpecVersion(str, enum.Enum): nov_2024 = "2024-11-05" mar_2025 = "2025-03-26" jun_2025 = "2025-06-18" + nov_2025 = "2025-11-25" + jul_2026 = "2026-07-28" class MCPAuth(str, enum.Enum): @@ -59,7 +61,17 @@ DEFAULT_SUBJECT_TOKEN_TYPE: Final = "urn:ietf:params:oauth:token-type:access_tok # MCP Literals MCPTransportType = Literal[MCPTransport.sse, MCPTransport.http, MCPTransport.stdio] -MCPSpecVersionType = Literal[MCPSpecVersion.nov_2024, MCPSpecVersion.mar_2025, MCPSpecVersion.jun_2025] +MCPLegacyVersion = Literal["2024-11-05", "2025-03-26", "2025-06-18", "2025-11-25"] +MCP_LEGACY_VERSIONS: Final[tuple[MCPLegacyVersion, ...]] = ("2024-11-05", "2025-03-26", "2025-06-18", "2025-11-25") +MCPUpstreamProtocol = MCPLegacyVersion | Literal["auto"] +MCPAdvertisedVersions = Annotated[tuple[MCPLegacyVersion, ...], Field(min_length=1)] +MCPSpecVersionType = Literal[ + MCPSpecVersion.nov_2024, + MCPSpecVersion.mar_2025, + MCPSpecVersion.jun_2025, + MCPSpecVersion.nov_2025, + MCPSpecVersion.jul_2026, +] MCPAuthType = ( Literal[ MCPAuth.none, diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index cb32299b143..c3b106c11d5 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -1,7 +1,7 @@ from datetime import datetime -from typing import Any, Final, Literal +from typing import Annotated, Any, Final, Literal -from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from pydantic import AfterValidator, BaseModel, ConfigDict, Field, TypeAdapter, field_validator, model_validator from typing_extensions import Self from litellm.types.mcp import ( @@ -10,11 +10,19 @@ from litellm.types.mcp import ( MCPAuthType, MCPTokenEndpointAuthMethod, MCPTransportType, + MCPUpstreamProtocol, normalize_upstream_header_name, ) + # MCPInfo now allows arbitrary additional fields for custom metadata -MCPInfo = dict[str, Any] +def _validate_mcp_protocol_metadata(value: dict[str, object]) -> dict[str, object]: + if "protocol_version" in value: + TypeAdapter(MCPUpstreamProtocol).validate_python(value["protocol_version"]) + return value + + +MCPInfo = Annotated[dict[str, Any], AfterValidator(_validate_mcp_protocol_metadata)] class MCPOAuthMetadata(BaseModel): @@ -66,6 +74,7 @@ class MCPServer(BaseModel): server_name: str | None = None url: str | None = None transport: MCPTransportType + protocol_version: MCPUpstreamProtocol = "auto" spec_path: str | None = None auth_type: MCPAuthType | None = None authentication_token: str | None = None @@ -246,6 +255,14 @@ class MCPServer(BaseModel): """ return self.oauth2_flow == "client_credentials" + @model_validator(mode="after") + def resolve_protocol_version(self) -> Self: + if "protocol_version" not in self.model_fields_set and self.mcp_info is not None: + self.protocol_version = TypeAdapter(MCPUpstreamProtocol).validate_python( + self.mcp_info.get("protocol_version", "auto") + ) + return self + @model_validator(mode="after") def validate_identity_binding_mode(self) -> Self: binding: Final = self.oauth_identity_binding diff --git a/litellm/types/router.py b/litellm/types/router.py index b72809f625f..d545f7ae639 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -345,6 +345,7 @@ class CredentialLiteLLMParams(BaseModel): ## OBJECT STORAGE (files / batches) ## gcs_bucket_name: str | None = None + bucket_name: str | None = None ## AWS BEDROCK / SAGEMAKER ## aws_access_key_id: str | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 3e306b48887..cd336c9b989 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -299,6 +299,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): cache_read_input_token_cost_above_272k_tokens_flex: float | None cache_read_input_token_cost_above_512k_tokens: float | None cache_read_input_token_cost_batches: ReadOnly[float | None] + cache_read_input_token_cost_above_200k_tokens_batches: ReadOnly[float | None] cache_read_input_token_cost_above_272k_tokens_batches: ReadOnly[float | None] cache_creation_input_token_cost_batches: ReadOnly[float | None] cache_creation_input_token_cost_above_272k_tokens_batches: ReadOnly[float | None] @@ -327,8 +328,10 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): input_cost_per_second: float | None # for OpenAI Speech models input_cost_per_token_batches: float | None input_cost_per_video_token_batches: ReadOnly[float | None] + input_cost_per_token_above_200k_tokens_batches: ReadOnly[float | None] input_cost_per_token_above_272k_tokens_batches: ReadOnly[float | None] output_cost_per_token_batches: float | None + output_cost_per_token_above_200k_tokens_batches: ReadOnly[float | None] output_cost_per_token_above_272k_tokens_batches: ReadOnly[float | None] output_cost_per_token: Required[float | None] output_cost_per_token_flex: float | None # OpenAI flex service tier pricing @@ -3729,6 +3732,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): cache_read_input_token_cost_above_272k_tokens_priority: float | None = None cache_read_input_token_cost_above_272k_tokens_flex: float | None = None cache_read_input_token_cost_batches: float | None = None + cache_read_input_token_cost_above_200k_tokens_batches: float | None = None cache_read_input_token_cost_above_272k_tokens_batches: float | None = None cache_creation_input_token_cost_batches: float | None = None cache_creation_input_token_cost_above_272k_tokens_batches: float | None = None @@ -3742,6 +3746,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): input_cost_per_token_above_200k_tokens_priority: float | None = None input_cost_per_token_above_272k_tokens_priority: float | None = None input_cost_per_token_above_272k_tokens_flex: float | None = None + input_cost_per_token_above_200k_tokens_batches: float | None = None input_cost_per_token_above_272k_tokens_batches: float | None = None input_cost_per_query: float | None = None input_cost_per_image: float | None = None @@ -3766,6 +3771,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): output_cost_per_token_above_200k_tokens_priority: float | None = None output_cost_per_token_above_272k_tokens_priority: float | None = None output_cost_per_token_above_272k_tokens_flex: float | None = None + output_cost_per_token_above_200k_tokens_batches: float | None = None output_cost_per_token_above_272k_tokens_batches: float | None = None output_cost_per_character_above_128k_tokens: float | None = None output_cost_per_image: float | None = None @@ -4140,7 +4146,7 @@ FILE_CONTENT_STREAMING_PROVIDERS: Final[frozenset[str]] = frozenset( LITELLM_EXECUTED_BATCH_PROVIDERS: Final[frozenset[str]] = frozenset({LlmProviders.HOSTED_VLLM.value}) -ListBatchesSupportedProvider = Literal["openai", "azure", "hosted_vllm", "litellm_proxy", "vertex_ai"] +ListBatchesSupportedProvider = Literal["openai", "azure", "hosted_vllm", "litellm_proxy", "vertex_ai", "xai"] LIST_BATCHES_SUPPORTED_PROVIDERS: Final[frozenset[str]] = frozenset(get_args(ListBatchesSupportedProvider)) diff --git a/litellm/utils.py b/litellm/utils.py index 4ea0769ea11..be4388802f9 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6160,6 +6160,9 @@ def _get_model_info_helper( cache_read_input_token_cost_priority=_model_info.get("cache_read_input_token_cost_priority", None), cache_read_input_token_cost_ultrafast=_model_info.get("cache_read_input_token_cost_ultrafast", None), cache_read_input_token_cost_batches=_model_info.get("cache_read_input_token_cost_batches"), + cache_read_input_token_cost_above_200k_tokens_batches=_model_info.get( + "cache_read_input_token_cost_above_200k_tokens_batches" + ), cache_read_input_token_cost_above_272k_tokens_batches=_model_info.get( "cache_read_input_token_cost_above_272k_tokens_batches" ), @@ -6197,10 +6200,16 @@ def _get_model_info_helper( input_cost_per_video_per_second=_model_info.get("input_cost_per_video_per_second", None), input_cost_per_token_batches=_model_info.get("input_cost_per_token_batches"), input_cost_per_video_token_batches=_model_info.get("input_cost_per_video_token_batches", None), + input_cost_per_token_above_200k_tokens_batches=_model_info.get( + "input_cost_per_token_above_200k_tokens_batches" + ), input_cost_per_token_above_272k_tokens_batches=_model_info.get( "input_cost_per_token_above_272k_tokens_batches" ), output_cost_per_token_batches=_model_info.get("output_cost_per_token_batches"), + output_cost_per_token_above_200k_tokens_batches=_model_info.get( + "output_cost_per_token_above_200k_tokens_batches" + ), output_cost_per_token_above_272k_tokens_batches=_model_info.get( "output_cost_per_token_above_272k_tokens_batches" ), @@ -9357,6 +9366,10 @@ class ProviderConfigManager: from litellm.llms.mistral.files.transformation import MistralFilesConfig return MistralFilesConfig() + elif LlmProviders.XAI == provider: + from litellm.llms.xai.files.transformation import XAIFilesConfig + + return XAIFilesConfig() return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5a207dc4c02..a670200c132 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -51436,13 +51436,16 @@ }, "xai/grok-4.20-0309-reasoning": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, @@ -51450,8 +51453,11 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true @@ -51480,9 +51486,13 @@ "xai/grok-4.3": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_image_token": 1.25e-06, "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, @@ -51490,6 +51500,8 @@ "mode": "chat", "output_cost_per_token": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -51502,9 +51514,13 @@ "xai/grok-4.3-latest": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_image_token": 1.25e-06, "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, @@ -51512,6 +51528,8 @@ "mode": "chat", "output_cost_per_token": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -59483,13 +59501,16 @@ }, "xai/grok-4.20-0309-non-reasoning": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -59497,20 +59518,26 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-multi-agent-0309": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "responses", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": false, "supports_prompt_caching": true, @@ -59519,8 +59546,11 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true, "supported_endpoints": [ @@ -62787,13 +62817,16 @@ }, "xai/grok-4.20": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, @@ -62801,21 +62834,27 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true }, "xai/grok-4.20-reasoning": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, @@ -62823,21 +62862,27 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true }, "xai/grok-4.20-reasoning-latest": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, @@ -62845,8 +62890,11 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true @@ -63067,13 +63115,16 @@ }, "xai/grok-4.20-non-reasoning": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -63081,20 +63132,26 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-non-reasoning-latest": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, @@ -63102,20 +63159,26 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-multi-agent": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "responses", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supported_endpoints": [ "/v1/responses" @@ -63127,20 +63190,26 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-multi-agent-latest": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "responses", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supported_endpoints": [ "/v1/responses" @@ -63152,8 +63221,11 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, @@ -75459,13 +75531,16 @@ }, "xai/grok-4.20-0309": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1.6e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 2e-06, "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, @@ -75473,8 +75548,11 @@ "supports_vision": true, "supports_web_search": true, "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "output_cost_per_token_above_200k_tokens": 5e-06, + "output_cost_per_token_above_200k_tokens_batches": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_batches": 3.2e-07, "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 35624045fdf..fa1828c780a 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -170,6 +170,11 @@ "minimum": 0, "description": "Rate applied once the prompt exceeds the token threshold in the field name." }, + "cache_read_input_token_cost_above_200k_tokens_batches": { + "type": "number", + "minimum": 0, + "description": "Rate applied once the prompt exceeds the token threshold in the field name." + }, "cache_read_input_token_cost_above_200k_tokens_priority": { "type": "number", "minimum": 0, @@ -351,6 +356,11 @@ "minimum": 0, "description": "Rate applied once the prompt exceeds the token threshold in the field name." }, + "input_cost_per_token_above_200k_tokens_batches": { + "type": "number", + "minimum": 0, + "description": "Rate applied once the prompt exceeds the token threshold in the field name." + }, "input_cost_per_token_above_200k_tokens_priority": { "type": "number", "minimum": 0, @@ -708,6 +718,11 @@ "minimum": 0, "description": "Rate applied once the prompt exceeds the token threshold in the field name." }, + "output_cost_per_token_above_200k_tokens_batches": { + "type": "number", + "minimum": 0, + "description": "Rate applied once the prompt exceeds the token threshold in the field name." + }, "output_cost_per_token_above_200k_tokens_priority": { "type": "number", "minimum": 0, diff --git a/pyproject.toml b/pyproject.toml index f2364b5e77b..28b00379cc7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -249,6 +249,7 @@ proxy-dev = [ "prisma==0.11.0", "hypercorn==0.17.3", "prometheus-client==0.20.0", + "sentry-sdk==2.21.0", "opentelemetry-api==1.33.1", "opentelemetry-sdk==1.33.1", "opentelemetry-exporter-otlp==1.33.1", diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index dc1f8592612..659dc438f2d 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -72,6 +72,7 @@ IGNORE_FUNCTIONS = [ "_string_leaves", # bounded by the nesting depth of a safe_json_structure output (a finite JSON tree, no cycles possible). "_replace_string_leaves", # bounded by the nesting depth of a safe_json_structure output (a finite JSON tree, no cycles possible). "_sort_processed_sets", # bounded by the nesting depth of the log-record extra it walks (a finite JSON tree, no cycles possible). + "scrub_json_strings", # max depth set (MAX_SCRUB_DEPTH); fails closed by returning "[Filtered]" for anything nested past the cap. ] diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index bc19668e2c9..9cfe6e33ed6 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -85,6 +85,7 @@ - {id: llm.responses.azure_openai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "response_api_endpoints/endpoints.py:26", rationale: "Responses w/ Azure OpenAI (smoke)"} - {id: llm.responses.azure_openai.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: azure_openai, capability: tool_use, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Responses tool calls w/ Azure OpenAI"} - {id: llm.responses.azure_openai.code_interpreter.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: azure_openai, capability: code_interpreter, streaming: nonstream, assertions: [works], source: "llm_translation/test_containers_e2e.py", rationale: "An implicit code_interpreter container on an Azure deployment that carries its own api_base must serve GET /v1/containers/{id}/files/{fid}/content by its native cntr_ id to a team service-account key, the customer's shape (#27921, #28990)", fail_before_fix: proven} +- {id: llm.responses.azure_openai.code_interpreter.stream.works, module: llm, tier: P1, subject_endpoint: responses, route: azure_openai, capability: code_interpreter, streaming: stream, assertions: [works], source: "llm_translation/test_containers_e2e.py", rationale: "A container created by a streamed /v1/responses code_interpreter call must serve /v1/containers/{id}/files to the same service-account key right after the OpenAI SDK closes at [DONE]; the ownership row used to be written after the stream and the disconnect cancelled it (LIT-8612)", fail_before_fix: proven} - {id: llm.chat_completions.together_ai.thinking.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: thinking, streaming: nonstream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together reasoning surfaces as reasoning_content (LIT-5960)"} - {id: llm.chat_completions.together_ai.thinking.stream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: thinking, streaming: stream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together reasoning deltas stream as reasoning_content"} - {id: llm.chat_completions.together_ai.thinking.nonstream.template_kwargs_forwarded, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: thinking, streaming: nonstream, assertions: [template_kwargs_forwarded], source: "llm_translation/test_together_ai_e2e.py", rationale: "chat_template_kwargs reaches Together and turns thinking off"} @@ -102,6 +103,8 @@ - {id: llm.chat_completions.anthropic.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: chat_completions, route: anthropic, capability: basic, streaming: nonstream, assertions: [works, cost_logged], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic over /chat/completions: cost header and spend row agree"} - {id: llm.chat_completions.anthropic.multi_turn.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: anthropic, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic tool result round trip over /chat/completions"} - {id: llm.messages.anthropic.multi_turn.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: anthropic, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic tool result round trip over /v1/messages"} +- {id: llm.messages.anthropic.upstream_stream_failure.stream.error_event, module: llm, tier: P1, subject_endpoint: messages, route: anthropic, capability: upstream_stream_failure, streaming: stream, assertions: [error_event], source: "customer report", rationale: "An upstream that hangs up mid-stream must reach Anthropic clients as an event: error frame, not an OpenAI-shaped data-only error they silently drop"} +- {id: llm.messages.anthropic.upstream_stream_failure.stream.error_status, module: llm, tier: P1, subject_endpoint: messages, route: anthropic, capability: upstream_stream_failure, streaming: stream, assertions: [error_status], source: "customer report", rationale: "An upstream that hangs up before its first byte must answer as a JSON error carrying its status, so Anthropic clients raise the status-specific error and retry on it instead of reading a 200 stream that only carries an error event"} - {id: llm.messages.openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "OpenAI models served on the Anthropic Messages contract"} - {id: llm.messages.openai.basic.stream.works, module: llm, tier: P0, subject_endpoint: messages, route: openai, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "OpenAI over /v1/messages streams the Anthropic event grammar"} - {id: llm.messages.openai.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: messages, route: openai, capability: basic, streaming: nonstream, assertions: [works, cost_logged], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "OpenAI over /v1/messages: cost header and spend row agree"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index 5fd19212ab7..fec1934059c 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -86,6 +86,7 @@ LlmCapability = Literal[ "tool_search", "tool_search_history", "tool_use", + "upstream_stream_failure", "vision", "web_search", "web_search_server_tool", diff --git a/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md b/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md index a18c81fa01d..92330db0530 100644 --- a/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md +++ b/tests/e2e/llm_translation/LLM_TRANSLATION_COVERAGE_MATRIX.md @@ -48,7 +48,7 @@ most likely to silently break and the one a mock can't prove works. |----------|---------------|-----------|------------|-------------|--------| | Chat | live (spend suite) | live (spend suite) | gap | live | partial | | Embeddings | live (spend suite) | n/a | n/a | live | covered | -| Responses (Azure code_interpreter container files) | live | gap | live | gap | partial | +| Responses (Azure code_interpreter container files) | live | live | live | gap | partial | | Image / audio / rerank / realtime | - | - | - | - | gap | ## This suite's files @@ -63,6 +63,7 @@ most likely to silently break and the one a mock can't prove works. | `test_anthropic_passthrough_tool_call_logs_cost` | anthropic native, tool call, cost | | `test_vertex_passthrough_via_managed_model_logs_cost` | vertex_ai native, non-stream, cost | | `test_service_account_key_reads_container_file_by_native_id` | azure responses code_interpreter, non-stream, native container id, service-account key | +| `test_service_account_key_reads_container_file_created_by_a_streamed_response` | azure responses code_interpreter, stream, native container id, service-account key, upload right after `[DONE]` | Vertex keeps the credential on the proxy like gemini/anthropic, but the deployment is added at runtime instead of declared in the gateway config: the test POSTs `/model/new` diff --git a/tests/e2e/llm_translation/test_containers_e2e.py b/tests/e2e/llm_translation/test_containers_e2e.py index 887aecb8df1..3048a830810 100644 --- a/tests/e2e/llm_translation/test_containers_e2e.py +++ b/tests/e2e/llm_translation/test_containers_e2e.py @@ -31,10 +31,11 @@ A proxy whose env carries ``AZURE_API_BASE`` for the same resource masks the second regression, since the global-credential fallback then reaches the container anyway. -The streaming variant is not here: a streamed ``/v1/responses`` writes the -container ownership row only after the ``[DONE]`` frame, and the OpenAI SDK -closes the connection at ``[DONE]``, so the write is cancelled and every -follow-up container call 403s (LIT-8612). That cell comes with its fix. +The streaming cell repeats the flow with ``stream=True`` and uploads right +after the last event. The OpenAI SDK closes the connection at ``[DONE]``, so an +ownership row written after the stream is cancelled with the body task and every +follow-up container call 403s (LIT-8612); the row has to land before the +``response.completed`` frame goes out. """ from __future__ import annotations @@ -52,7 +53,7 @@ from lifecycle import ResourceManager from management.management_client import ManagementClient, build_client from models import KeyGenerateBody, KeyGenerateResponse, LiteLLMParamsBody, TeamNewBody, UserNewBody from openai import OpenAI -from openai.types.responses import Response, ResponseCodeInterpreterToolCall +from openai.types.responses import Response, ResponseCodeInterpreterToolCall, ResponseCompletedEvent from openai.types.responses.tool_param import CodeInterpreter from proxy_client import ProxyClient from sdk_clients import NO_PROXY_CACHE, SdkClients @@ -120,6 +121,24 @@ def _response_with_code_interpreter(client: OpenAI, model: str) -> Response: ) +def _streamed_response_with_code_interpreter(client: OpenAI, model: str) -> Response: + events: Final = tuple( + client.with_options(timeout=CODE_INTERPRETER_TIMEOUT).responses.create( + model=model, + input=PROMPT, + tools=[CODE_INTERPRETER], + tool_choice="required", + stream=True, + extra_body=NO_PROXY_CACHE, + ) + ) + assert events, "responses stream returned no events" + assert isinstance(events[-1], ResponseCompletedEvent), ( + f"responses stream did not terminate with response.completed: {events[-1].type}" + ) + return events[-1].response + + def _container_id(response: Response) -> str: calls: Final = tuple(item for item in response.output if isinstance(item, ResponseCodeInterpreterToolCall)) assert calls, f"no code_interpreter_call in the responses output: {response.output!r}" @@ -165,3 +184,17 @@ class TestAzureContainerFiles: f"container id is not the provider's own id: {native_id}" ) _assert_file_round_trip(client, native_id, marker) + + @pytest.mark.covers("llm.responses.azure_openai.code_interpreter.stream.works") + def test_service_account_key_reads_container_file_created_by_a_streamed_response( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + marker: Final = unique_marker() + model: Final = _register_two_azure_deployments(proxy, resources, marker) + key: Final = _service_account_key(proxy, resources, build_client(proxy), marker, model) + client: Final = sdk.openai(key) + native_id: Final = _native_container_id( + _container_id(_streamed_response_with_code_interpreter(client, model)) + ) + resources.defer(lambda: client.containers.delete(native_id, extra_query=AZURE_PROVIDER_QUERY)) + _assert_file_round_trip(client, native_id, marker) diff --git a/tests/e2e/llm_translation/test_messages_e2e.py b/tests/e2e/llm_translation/test_messages_e2e.py index d048d1343eb..871fd2f9aef 100644 --- a/tests/e2e/llm_translation/test_messages_e2e.py +++ b/tests/e2e/llm_translation/test_messages_e2e.py @@ -11,8 +11,11 @@ litellm-regression-tests/tests/test_inference_endpoints.py. from __future__ import annotations import time +from collections.abc import Callable +from types import MappingProxyType from typing import Final +import anthropic import pytest from anthropic import Anthropic from anthropic.types import ( @@ -30,12 +33,21 @@ from anthropic.types import ( ToolParam, ToolUseBlock, ) -from e2e_config import STREAM_MIN_LEAD_SECONDS, provider_edge_base, provider_paces_stream, unique_marker +from e2e_config import ( + PROVIDER_EDGE_ADVERTISE_HOST, + PROVIDER_EDGE_BIND_HOST, + STREAM_MIN_LEAD_SECONDS, + provider_edge_base, + provider_paces_stream, + unique_marker, +) from e2e_http import assert_client_error from lifecycle import ResourceManager -from models import ChatMessage, LiteLLMParamsBody, SpendLogRow +from models import AnthropicErrorEvent, AnthropicMessagesBody, ChatMessage, LiteLLMParamsBody, SpendLogRow +from provider_edge import EDGE_MOUNTS, LiveEdge, RunningEdge, StreamCut, start_provider_edge +from provider_edge_bedrock import bedrock_signer from proxy_client import ProxyClient -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError from sdk_clients import NO_PROXY_CACHE, SdkClients, response_header pytestmark = [pytest.mark.e2e, pytest.mark.replayable] @@ -385,3 +397,245 @@ class TestOpenAIMessagesToolContinuation: ) assert _text(continuation).strip() == receipt, "continuation did not consume the correlated tool result" assert all(not isinstance(block, ToolUseBlock) for block in continuation.content) + + +BEDROCK_BACKEND: Final = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" +BEDROCK_EDGE_REGION: Final = "us-east-1" +_STREAM_FAILURE_PROMPT: Final = "Count from 1 to 100, one number per line." +_FRAME_PAYLOAD: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_AT_FRAME_BOUNDARY: Final = StreamCut(after_content=True) +_MID_FRAME: Final = StreamCut(after_content=True, mid_chunk=True) +_BEFORE_FIRST_BYTE: Final = StreamCut(after_content=False) + +type _CutRegistration = Callable[[ProxyClient, ResourceManager, StreamCut], tuple[str, str]] + + +def _cut_edge(backend: LiveEdge, mount: str) -> RunningEdge: + return start_provider_edge( + backend, + mounts=MappingProxyType({mount: EDGE_MOUNTS[mount]}), + bind_host=PROVIDER_EDGE_BIND_HOST, + advertise_host=PROVIDER_EDGE_ADVERTISE_HOST, + ) + + +def _register_cut_bedrock(proxy: ProxyClient, resources: ResourceManager, cut: StreamCut) -> tuple[str, str]: + mount: Final = f"bedrock/{BEDROCK_EDGE_REGION}" + edge: Final = _cut_edge(LiveEdge(cut=cut, sign=bedrock_signer(BEDROCK_EDGE_REGION)), mount) + resources.defer(edge.shutdown) + return _register( + proxy, + resources, + LiteLLMParamsBody( + model=BEDROCK_BACKEND, + api_base=edge.edge.api_base(mount), + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name=BEDROCK_EDGE_REGION, + ), + prefix="e2e-messages-cut", + ) + + +def _register_cut_anthropic(proxy: ProxyClient, resources: ResourceManager, cut: StreamCut) -> tuple[str, str]: + edge: Final = _cut_edge(LiveEdge(cut=cut), "anthropic") + resources.defer(edge.shutdown) + return _register( + proxy, + resources, + LiteLLMParamsBody( + model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY", api_base=edge.edge.api_base("anthropic") + ), + prefix="e2e-messages-cut", + ) + + +_DROPPED_UPSTREAMS: Final[tuple[tuple[str, _CutRegistration, StreamCut], ...]] = ( + ("bedrock_at_a_frame_boundary", _register_cut_bedrock, _AT_FRAME_BOUNDARY), + ("anthropic_at_a_frame_boundary", _register_cut_anthropic, _AT_FRAME_BOUNDARY), + ("anthropic_mid_frame", _register_cut_anthropic, _MID_FRAME), +) +_DROPPED_BEFORE_FIRST_BYTE: Final[tuple[tuple[str, _CutRegistration, StreamCut], ...]] = ( + ("bedrock_before_the_first_byte", _register_cut_bedrock, _BEFORE_FIRST_BYTE), + ("anthropic_before_the_first_byte", _register_cut_anthropic, _BEFORE_FIRST_BYTE), +) + + +def _payload(frame: str) -> JsonValue | None: + try: + return _FRAME_PAYLOAD.validate_json(frame) + except ValidationError: + return None + + +def _bare_error_frame(frame: str) -> bool: + payload: Final = _payload(frame) + return isinstance(payload, dict) and "error" in payload and payload.get("type") != "error" + + +@pytest.mark.provider_edge_host +@pytest.mark.provider_live +class TestMessagesUpstreamStreamFailure: + @pytest.mark.covers("llm.messages.anthropic.upstream_stream_failure.stream.error_event") + @pytest.mark.parametrize( + ("register", "cut"), [case[1:] for case in _DROPPED_UPSTREAMS], ids=[case[0] for case in _DROPPED_UPSTREAMS] + ) + def test_interrupted_upstream_stream_raises_in_the_anthropic_sdk( + self, + proxy: ProxyClient, + resources: ResourceManager, + sdk: SdkClients, + register: _CutRegistration, + cut: StreamCut, + ) -> None: + model, key = register(proxy, resources, cut) + client: Final = sdk.anthropic(key) + + stream: Final = client.messages.create( + model=model, + max_tokens=300, + stream=True, + messages=[_user_turn(_STREAM_FAILURE_PROMPT)], + extra_body=NO_PROXY_CACHE, + ) + first: Final = next(stream) + assert first.type == "message_start", ( + f"the stream produced a first event that is not message_start, so this run proves a " + f"startup failure, not an interrupted stream: {first!r}" + ) + with pytest.raises(anthropic.APIStatusError) as raised: + for _ in stream: + pass + try: + AnthropicErrorEvent.model_validate(raised.value.body) + except ValidationError: + pytest.fail( + f"the SDK raised on the interrupted stream but without the Anthropic error envelope a " + f"client reads the failure from: body={raised.value.body!r} message={raised.value}" + ) + + @pytest.mark.covers("llm.messages.anthropic.upstream_stream_failure.stream.error_event") + @pytest.mark.parametrize( + ("register", "cut"), [case[1:] for case in _DROPPED_UPSTREAMS], ids=[case[0] for case in _DROPPED_UPSTREAMS] + ) + def test_interrupted_upstream_stream_is_an_anthropic_error_event( + self, proxy: ProxyClient, resources: ResourceManager, register: _CutRegistration, cut: StreamCut + ) -> None: + model, key = register(proxy, resources, cut) + + outcome: Final = proxy.messages_stream( + key, + AnthropicMessagesBody( + model=model, + max_tokens=300, + stream=True, + messages=[ChatMessage(role="user", content=_STREAM_FAILURE_PROMPT)], + ), + ) + frames: Final = outcome.stream_events + assert outcome.is_streaming, ( + f"/v1/messages did not answer with an SSE stream: status={outcome.status_code} body={outcome.body}" + ) + assert frames, ( + f"the proxy sent no SSE data frames although the upstream hung up; stream_error={outcome.stream_error!r}" + ) + assert outcome.stream_error == "event: error", ( + f"the interrupted stream was not announced by an 'event: error' line Anthropic clients read; " + f"stream_error={outcome.stream_error!r} frames={frames}" + ) + try: + AnthropicErrorEvent.model_validate_json(frames[-1]) + except ValidationError: + pytest.fail( + f'the last SSE frame was not an Anthropic {{"type": "error", "error": ...}} envelope; frames={frames}' + ) + torn: Final = tuple(index for index, frame in enumerate(frames) if _payload(frame) is None) + expected_torn: Final = 1 if cut.mid_chunk else 0 + assert len(torn) == expected_torn, ( + f"expected {expected_torn} data line(s) that are not JSON, since the edge tears one only when it " + f"cuts mid-frame, but the proxy relayed {[frames[index] for index in torn]}; all frames={frames}" + ) + for index in torn: + assert _payload(frames[index + 1]) == {"type": "ping"}, ( + f"the frame the upstream tore was not closed as a ping event before the error, so an " + f"Anthropic client parses the error inside it: after {frames[index]!r} came " + f"{frames[index + 1]!r}; all frames={frames}" + ) + bare: Final = tuple(frame for frame in frames if _bare_error_frame(frame)) + assert not bare, ( + f"the proxy emitted error frames without the Anthropic envelope, which Anthropic clients drop: " + f"{bare}; all frames={frames}" + ) + + @pytest.mark.covers("llm.messages.anthropic.upstream_stream_failure.stream.error_status") + @pytest.mark.parametrize( + ("register", "cut"), + [case[1:] for case in _DROPPED_BEFORE_FIRST_BYTE], + ids=[case[0] for case in _DROPPED_BEFORE_FIRST_BYTE], + ) + def test_upstream_that_hangs_up_before_the_first_byte_raises_with_its_status_in_the_anthropic_sdk( + self, + proxy: ProxyClient, + resources: ResourceManager, + sdk: SdkClients, + register: _CutRegistration, + cut: StreamCut, + ) -> None: + model, key = register(proxy, resources, cut) + client: Final = sdk.anthropic(key) + + with pytest.raises(anthropic.APIStatusError) as raised: + client.messages.create( + model=model, + max_tokens=300, + stream=True, + messages=[_user_turn(_STREAM_FAILURE_PROMPT)], + extra_body=NO_PROXY_CACHE, + ) + assert 500 <= raised.value.status_code < 600, ( + f"an upstream that hung up before sending anything must answer with a server error status the SDK " + f"retries on, not {raised.value.status_code}: {raised.value}" + ) + try: + AnthropicErrorEvent.model_validate(raised.value.body) + except ValidationError: + pytest.fail( + f"the SDK raised with the right status but without the Anthropic error envelope a client reads " + f"the failure from: body={raised.value.body!r} message={raised.value}" + ) + + @pytest.mark.covers("llm.messages.anthropic.upstream_stream_failure.stream.error_status") + @pytest.mark.parametrize( + ("register", "cut"), + [case[1:] for case in _DROPPED_BEFORE_FIRST_BYTE], + ids=[case[0] for case in _DROPPED_BEFORE_FIRST_BYTE], + ) + def test_upstream_that_hangs_up_before_the_first_byte_is_a_json_error_with_its_status( + self, proxy: ProxyClient, resources: ResourceManager, register: _CutRegistration, cut: StreamCut + ) -> None: + model, key = register(proxy, resources, cut) + + outcome: Final = proxy.messages_stream( + key, + AnthropicMessagesBody( + model=model, + max_tokens=300, + stream=True, + messages=[ChatMessage(role="user", content=_STREAM_FAILURE_PROMPT)], + ), + ) + assert not outcome.is_streaming, ( + f"nothing had been streamed when the upstream hung up, yet /v1/messages opened a 200 SSE stream " + f"instead of answering with the failure's status: stream_error={outcome.stream_error!r} " + f"frames={outcome.stream_events}" + ) + assert 500 <= outcome.status_code < 600, ( + f"/v1/messages answered {outcome.status_code} for an upstream that hung up before its first byte; " + f"body={outcome.body}" + ) + try: + AnthropicErrorEvent.model_validate_json(outcome.body) + except ValidationError: + pytest.fail( + f'the error body is not an Anthropic {{"type": "error", "error": ...}} envelope; body={outcome.body}' + ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index c96f4b0bef1..84399fd6155 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -598,6 +598,16 @@ class CountTokensResponse(BaseModel): input_tokens: int +class AnthropicErrorBody(BaseModel): + type: str + message: str + + +class AnthropicErrorEvent(BaseModel): + type: Literal["error"] + error: AnthropicErrorBody + + # ---------- mcp servers ---------- diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py index fc10dde2a77..3680375b6af 100644 --- a/tests/e2e/provider_edge.py +++ b/tests/e2e/provider_edge.py @@ -45,6 +45,7 @@ import hashlib import os import re import threading +import time from collections import deque from collections.abc import Callable, Generator, Mapping, Sequence from contextlib import closing, contextmanager @@ -56,6 +57,7 @@ from types import MappingProxyType from typing import Final, Literal, assert_never from urllib.parse import parse_qsl, urlsplit +from botocore.eventstream import EventStreamBuffer from e2e_http import ( NetworkError, StreamChunk, @@ -96,16 +98,18 @@ from fixture_mode import ( ) from fixture_profile import IneligibleRequest, MatchProfile, match_profile, strict_identity from provider_cache import ( + JSON_VALUE, SIGNATURE_HEADERS, CacheEdge, MountPolicy, RequestSigner, + invoke_chunk_value, is_bedrock, scoped_edge_base, split_test_segment, ) from provider_cache_routing import LIVE_PROVIDER_REQUIRED -from pydantic import JsonValue, TypeAdapter +from pydantic import JsonValue, TypeAdapter, ValidationError BEDROCK_REGIONS: Final[tuple[str, ...]] = ("us-east-1",) @@ -537,10 +541,33 @@ class ReplayEdge: source: ReplaySource +@dataclass(frozen=True, slots=True) +class StreamCut: + """Where a live edge hangs up on a streamed upstream body: before its first byte, or with + ``after_content`` set, right after the first transfer chunk carrying assistant output (a + ``content_block_delta``). That frame is what commits the proxy's mid-stream fallback + wrapper to the client: it holds the lifecycle frames before it back and drops them when + the transport fails first, so a cut after a fixed number of chunks landed on either side + of that commit depending on how the provider batched its frames. With ``mid_chunk`` set + the hang-up comes part way through the next ``data:`` line the provider sends after that, + so the client is left inside an SSE frame the way a dropped transport leaves it. + + Whatever was relayed sits on the wire for ``_CUT_SETTLE_SECONDS`` before the hang-up, so + the client has read it by then instead of receiving the data and the close in one burst, + where its reader can surface the close before what it buffered.""" + + after_content: bool + mid_chunk: bool = False + + +_CUT_SETTLE_SECONDS: Final = 1.0 + + @dataclass(frozen=True, slots=True) class LiveEdge: observe_request: Callable[[str, Mapping[str, str], bytes | None], None] | None = None sign: RequestSigner | None = None + cut: StreamCut | None = None type EdgeBackend = RecordEdge | ReplayEdge | LiveEdge | CacheEdge @@ -786,11 +813,128 @@ def _handle_record( assert_never(head) +def _data_line_start(data: bytes) -> int: + if data.startswith(b"data:"): + return 0 + at_line_start: Final = data.find(b"\ndata:") + return -1 if at_line_start < 0 else at_line_start + 1 + + +def _torn_prefix(data: bytes) -> bytes: + start: Final = _data_line_start(data) + line_end: Final = data.find(b"\n", start) + end: Final = len(data) if line_end < 0 else line_end + return data[: start + (end - start) // 2] + + +class _DataLineTearer: + __slots__ = ("_unfinished_line",) + + _unfinished_line: bytes + + def __init__(self) -> None: + self._unfinished_line = b"" + + def observe(self, data: bytes) -> None: + self._unfinished_line = (self._unfinished_line + data).rsplit(b"\n", 1)[-1] + + def tear(self, data: bytes) -> bytes | None: + buffered: Final = self._unfinished_line + data + if _data_line_start(buffered) < 0: + self.observe(data) + return None + return _torn_prefix(buffered)[len(self._unfinished_line):] + + +def _is_content_delta(value: JsonValue | None) -> bool: + return isinstance(value, dict) and value.get("type") == "content_block_delta" + + +def _sse_data_carries_content(line: bytes) -> bool: + if not line.startswith(b"data:"): + return False + try: + return _is_content_delta(JSON_VALUE.validate_json(line[len(b"data:"):].strip())) + except ValidationError: + return False + + +class _AnthropicContentDetector: + __slots__ = ("_unfinished_line",) + + _unfinished_line: bytes + + def __init__(self) -> None: + self._unfinished_line = b"" + + def __call__(self, data: bytes) -> bool: + lines: Final = (self._unfinished_line + data).split(b"\n") + self._unfinished_line = lines[-1] + return any(_sse_data_carries_content(line.rstrip(b"\r")) for line in lines[:-1]) + + +def _invoke_frame_carries_content(payload: bytes) -> bool: + try: + return _is_content_delta(invoke_chunk_value(JSON_VALUE.validate_json(payload))) + except ValidationError: + return False + + +def _bedrock_content_detector() -> Callable[[bytes], bool]: + """Bedrock's invoke stream wraps each Anthropic event in an eventstream frame that a + transfer chunk can split, so the frames are reassembled across chunks before being read.""" + frames: Final = EventStreamBuffer() + + def carries_content(data: bytes) -> bool: + frames.add_data(data) + return any(_invoke_frame_carries_content(frame.payload) for frame in frames) + + return carries_content + + +def _content_detector(mount: str) -> Callable[[bytes], bool]: + return _bedrock_content_detector() if is_bedrock(mount) else _AnthropicContentDetector() + + +def _cut_steps( + steps: Generator[StreamStep, None, None], cut: StreamCut, carries_content: Callable[[bytes], bool] +) -> Generator[StreamStep, None, None]: + with closing(steps) as source: + tearer: Final = _DataLineTearer() + if cut.after_content: + for step in source: + yield step + if isinstance(step, StreamTruncation): + return + tearer.observe(step.data) + if carries_content(step.data): + break + else: + return + if cut.mid_chunk: + for step in source: + if isinstance(step, StreamTruncation): + yield step + return + if (torn := tearer.tear(step.data)) is None: + yield step + continue + if torn: + yield StreamChunk(data=torn) + break + else: + return + if cut.after_content or cut.mid_chunk: + time.sleep(_CUT_SETTLE_SECONDS) + yield StreamTruncation(reason=f"edge cut the upstream stream: {cut!r}") + + def _handle_live( method: str, url: str, headers: Mapping[str, str], body: bytes | None, timeout: float, cache: CacheEdge | None = None, mount: str = "", test_key: str | None = None, observe_request: Callable[[str, Mapping[str, str], bytes | None], None] | None = None, sign: RequestSigner | None = None, + cut: StreamCut | None = None, ) -> EdgeOutcome: forwarded: Final = { name: value for name, value in headers.items() if name.lower() not in _REQUEST_DROPPED_HEADERS @@ -805,6 +949,8 @@ def _handle_live( match head: case NetworkError(message=message): return _recorded_outcome(_network_error_response(message)) + case StreamHead() if cut is not None: + return EdgeStream(head.status_code, _filtered_response_headers(head.headers), _cut_steps(head.steps, cut, _content_detector(mount))) case StreamHead() if _is_streamed(head.headers): return EdgeStream(head.status_code, _filtered_response_headers(head.headers), head.steps) case StreamHead(): @@ -875,10 +1021,10 @@ def handle_edge_request( method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout, backend, mount, test_key, ) - case LiveEdge(observe_request=observe_request, sign=sign): + case LiveEdge(observe_request=observe_request, sign=sign, cut=cut): return _handle_live( method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout, - observe_request=observe_request, sign=sign, + mount=mount, observe_request=observe_request, sign=sign, cut=cut, ) case RecordEdge(): return _handle_record( diff --git a/tests/e2e/test_provider_edge.py b/tests/e2e/test_provider_edge.py index d776c338ef7..978c3671a77 100644 --- a/tests/e2e/test_provider_edge.py +++ b/tests/e2e/test_provider_edge.py @@ -54,11 +54,13 @@ from provider_edge import ( EdgeBackend, EdgeReply, EdgeStream, + LiveEdge, ProviderEdge, ProviderRequestObservation, RecordEdge, ReplayEdge, ReplaySource, + StreamCut, edge_request, handle_edge_request, observed_provider_edge, @@ -1000,6 +1002,36 @@ def stream_chunks(response: RecordedStreamedResponse) -> list[bytes]: return [base64.b64decode(chunk) for chunk in response.chunks_b64] +SECOND_DATA_LINE: Final = b'data: {"type":"content_block_delta","delta":{"text":" two"}}' +SPLIT_MARKER_CHUNKS: tuple[bytes, ...] = ( + b'data: {"type":"content_block_delta","delta":{"text":"one"}}\n\nda', + b"ta" + SECOND_DATA_LINE[4:] + b"\n\nda", + b'ta: {"type":"message_delta","usage":{"output_tokens":7}}\n\nda', + b"ta: [DONE]\n\n", +) + + +class TestStreamCut: + def test_a_mid_frame_cut_tears_a_data_line_whose_marker_is_split_across_chunks(self) -> None: + """Every ``data:`` marker after the first content delta straddles a transfer + chunk boundary, so a tearer that inspects each chunk on its own never finds + one and lets the stream finish cleanly instead of cutting it.""" + backend: Final = LiveEdge(cut=StreamCut(after_content=True, mid_chunk=True)) + with chunked_provider(chunks=SPLIT_MARKER_CHUNKS) as provider: + with running_edge(backend, {"openai": provider_url(provider)}) as edge: + head, chunks, ending = raw_stream_post(edge.port, STREAM_PATH, STREAM_BODY) + + assert head.startswith("HTTP/1.1 200 OK") + assert ending == "truncated" + relayed: Final = b"".join(chunks) + whole: Final = b"".join(SPLIT_MARKER_CHUNKS) + assert whole.startswith(relayed) and relayed != whole + assert relayed.startswith(SPLIT_MARKER_CHUNKS[0]) + torn_line: Final = relayed.rsplit(b"\n", 1)[-1] + assert torn_line and SECOND_DATA_LINE.startswith(torn_line) and torn_line != SECOND_DATA_LINE + assert b"[DONE]" not in relayed + + class TestStreamingFidelity: """LIT-5742: a streamed response records and replays as the chunk sequence the provider actually sent, not as one coalesced body. The unit of fidelity is the diff --git a/tests/integration/mcp/test_mcp_protocol_errors.py b/tests/integration/mcp/test_mcp_protocol_errors.py index bb06d8c6068..257e90e9d0e 100644 --- a/tests/integration/mcp/test_mcp_protocol_errors.py +++ b/tests/integration/mcp/test_mcp_protocol_errors.py @@ -84,3 +84,42 @@ def test_jsonrpc_error_and_malformed_tool_result_remain_errors(gateway: Gateway) control: Final = call_tool(gateway, key, identity, names["add"], {"a": 3, "b": 5}) assert control.status_code == 200 and control.json()["isError"] is False, control.text assert control.json()["content"][0]["text"] == "8" + + +@pytest.mark.parametrize("ingress", ("http", "sse")) +def test_configured_revision_blocks_unadvertised_handshake_and_keeps_allowed_control( + gateway: Gateway, tmp_path, ingress: str +) -> None: + import asyncio + from pathlib import Path + + import yaml + from integration._support.mcp import mcp_peer + from integration._support.process import owned_proxy + from litellm.experimental_mcp_client.client import MCPClient + from litellm.types.mcp import MCPTransport + from mcp import MCPError + from mcp.types import CallToolRequestParams + + with mcp_peer() as upstream, gateway.scenario() as scenario: + alias: Final = "restricted" + uuid.uuid4().hex[:8] + identity: Final = register_mcp(scenario, upstream, alias) + key: Final = scenario.key(object_permission={"mcp_servers": [identity]}) + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["general_settings"]["mcp_advertised_versions"] = ["2024-11-05"] + config_path: Final = tmp_path / "restricted.yaml" + config_path.write_text(yaml.safe_dump(config)) + with owned_proxy(gateway, tmp_path, {"DISABLE_SCHEMA_UPDATE": "true"}, config=config_path) as restricted: + endpoint: Final = str(restricted.client.base_url).rstrip("/") + ("/mcp/sse" if ingress == "sse" else "/mcp") + headers: Final = {"Authorization": f"Bearer {key}", "x-mcp-servers": identity} + denied: Final = MCPClient(server_url=endpoint, transport_type=MCPTransport(ingress), protocol_version="2025-11-25", extra_headers=headers) + allowed: Final = MCPClient(server_url=endpoint, transport_type=MCPTransport(ingress), protocol_version="2024-11-05", extra_headers=headers) + + async def exercise() -> None: + with pytest.raises(MCPError, match="Unsupported MCP protocol version"): + await denied.list_tools(raise_on_error=True) + assert f"{alias}-add" in tuple(tool.name for tool in await allowed.list_tools(raise_on_error=True)) + result: Final = await allowed.call_tool(CallToolRequestParams(name=f"{alias}-add", arguments={"a": 2, "b": 5})) + assert result.is_error is False and result.content[0].text == "7" + + asyncio.run(exercise()) diff --git a/tests/integration/mcp/test_mcp_transports.py b/tests/integration/mcp/test_mcp_transports.py index 1862f11d07e..16004cdd501 100644 --- a/tests/integration/mcp/test_mcp_transports.py +++ b/tests/integration/mcp/test_mcp_transports.py @@ -155,3 +155,43 @@ def test_server_initiated_sampling_and_elicitation_surface_as_errors_not_success assert outcome.error is not None, outcome.raw assert outcome.text is None or not outcome.text.startswith(("sampled:", "elicited:")), outcome.raw assert len(tool_calls(peer.drain())) == 1 + + +@pytest.mark.parametrize("downstream", ("2024-11-05", "2025-03-26", "2025-06-18", "2025-11-25")) +@pytest.mark.parametrize("upstream", ("2024-11-05", "2025-03-26", "2025-06-18", "2025-11-25")) +@pytest.mark.parametrize("peer_kind", ("http", "sse", "stdio")) +@pytest.mark.parametrize("ingress", ("http", "sse")) +def test_pinned_revision_pairs_list_and_call_through_gateway( + gateway: Gateway, downstream: str, upstream: str, peer_kind: PeerKind, ingress: str +) -> None: + import asyncio + + from mcp.types import CallToolRequestParams + + from litellm.experimental_mcp_client.client import MCPClient + from litellm.types.mcp import MCPTransport + + with peer_of(peer_kind) as peer, gateway.scenario() as scenario: + alias: Final = "versions" + uuid.uuid4().hex[:8] + identity: Final = register_mcp(scenario, peer, alias, mcp_info={"protocol_version": upstream}) + key: Final = scenario.key(object_permission={"mcp_servers": [identity]}) + endpoint: Final = str(gateway.client.base_url).rstrip("/") + ("/mcp/sse" if ingress == "sse" else "/mcp") + client: Final = MCPClient( + server_url=endpoint, transport_type=MCPTransport(ingress), protocol_version=downstream, + extra_headers={"Authorization": f"Bearer {key}", "x-mcp-servers": identity}, timeout=15, + ) + + async def exercise() -> None: + tools: Final = await client.list_tools(raise_on_error=True) + assert f"{alias}-add" in tuple(tool.name for tool in tools) + result: Final = await client.call_tool(CallToolRequestParams(name=f"{alias}-add", arguments={"a": 3, "b": 4})) + assert result.is_error is False + assert result.content[0].text == "7" + + peer.drain() + asyncio.run(exercise()) + observed: Final = peer.drain() + negotiations: Final = tuple(item["body"] for item in observed if item["body"].get("method") == "initialize") + assert negotiations, "The operation must reach the upstream negotiation" + assert all(request["params"]["protocolVersion"] == upstream for request in negotiations), negotiations + assert len(tool_calls(observed)) == 1 diff --git a/tests/llm_responses_api_testing/test_base_responses_api_streaming_iterator.py b/tests/llm_responses_api_testing/test_base_responses_api_streaming_iterator.py index 47b377dc9a4..da37803b64a 100644 --- a/tests/llm_responses_api_testing/test_base_responses_api_streaming_iterator.py +++ b/tests/llm_responses_api_testing/test_base_responses_api_streaming_iterator.py @@ -381,6 +381,36 @@ class TestBaseResponsesAPIStreamingIterator: ) raise + @staticmethod + def _config_completing_after_one_delta() -> Mock: + mock_config = Mock(spec=BaseResponsesAPIConfig) + completed_response = ResponsesAPIResponse( + id="resp_123", + created_at=0, + status="completed", + model="gpt-5.5", + object="response", + output=[], + usage=ResponseAPIUsage(input_tokens=1, output_tokens=1, total_tokens=2), + ) + + def _transform(model, parsed_chunk, logging_obj): + if parsed_chunk.get("type") == "response.completed": + return ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=completed_response, + ) + return OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_123", + output_index=0, + content_index=0, + delta=parsed_chunk["delta"], + ) + + mock_config.transform_streaming_response.side_effect = _transform + return mock_config + @pytest.mark.asyncio async def test_stop_async_iteration_not_logged_as_failure(self): """ @@ -399,6 +429,7 @@ class TestBaseResponsesAPIStreamingIterator: async def mock_aiter_bytes(): yield b'data: {"type": "response.output_text.delta", "delta": "test"}\n\n' + yield b'data: {"type": "response.completed", "response": {"id": "resp_123"}}\n\n' mock_response.aiter_bytes = mock_aiter_bytes @@ -408,11 +439,7 @@ class TestBaseResponsesAPIStreamingIterator: mock_logging_obj.async_failure_handler = Mock() mock_logging_obj.failure_handler = Mock() - mock_config = Mock(spec=BaseResponsesAPIConfig) - mock_delta_event = Mock() - mock_delta_event.type = ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA - mock_delta_event.delta = "test" - mock_config.transform_streaming_response.return_value = mock_delta_event + mock_config = self._config_completing_after_one_delta() # Create the iterator instance iterator = ResponsesAPIStreamingIterator( @@ -432,8 +459,9 @@ class TestBaseResponsesAPIStreamingIterator: except StopAsyncIteration: pass # This is expected - # Verify we got the chunk - assert len(chunks_received) == 1 + # Verify we got the delta and the terminal event + assert len(chunks_received) == 2 + assert iterator.completed_response is not None # CRITICAL: Verify that failure handlers were NOT called # StopAsyncIteration is a normal end of stream, not a failure @@ -460,6 +488,7 @@ class TestBaseResponsesAPIStreamingIterator: def mock_iter_bytes(): yield b'data: {"type": "response.output_text.delta", "delta": "test"}\n\n' + yield b'data: {"type": "response.completed", "response": {"id": "resp_123"}}\n\n' mock_response.iter_bytes = mock_iter_bytes @@ -469,11 +498,7 @@ class TestBaseResponsesAPIStreamingIterator: mock_logging_obj.async_failure_handler = Mock() mock_logging_obj.failure_handler = Mock() - mock_config = Mock(spec=BaseResponsesAPIConfig) - mock_delta_event = Mock() - mock_delta_event.type = ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA - mock_delta_event.delta = "test" - mock_config.transform_streaming_response.return_value = mock_delta_event + mock_config = self._config_completing_after_one_delta() # Create the iterator instance iterator = SyncResponsesAPIStreamingIterator( @@ -493,8 +518,9 @@ class TestBaseResponsesAPIStreamingIterator: except StopIteration: pass # This is expected - # Verify we got the chunk - assert len(chunks_received) == 1 + # Verify we got the delta and the terminal event + assert len(chunks_received) == 2 + assert iterator.completed_response is not None # CRITICAL: Verify that failure handlers were NOT called # StopIteration is a normal end of stream, not a failure diff --git a/tests/llm_translation/test_bedrock_gpt_oss.py b/tests/llm_translation/test_bedrock_gpt_oss.py index 4af81ee81f7..b264c16601f 100644 --- a/tests/llm_translation/test_bedrock_gpt_oss.py +++ b/tests/llm_translation/test_bedrock_gpt_oss.py @@ -22,7 +22,7 @@ class TestBedrockGPTOSS(BaseLLMChatTest): """Bedrock GPT-OSS intermittently emits truncated toolUse.input deltas on the live endpoint, which makes the inherited live integration test flaky. The accumulation side is covered deterministically by - tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py::test_transform_tool_calls_index; + tests/unit/llms/bedrock/chat/test_invoke_handler.py::test_transform_tool_calls_index; the GPT-OSS-specific request-body transformation is covered by test_function_calling_request_body_gpt_oss below. """ diff --git a/tests/local_testing/test_function_calling.py b/tests/local_testing/test_function_calling.py index 3a5e2209f1e..2d79f8a6af6 100644 --- a/tests/local_testing/test_function_calling.py +++ b/tests/local_testing/test_function_calling.py @@ -324,7 +324,7 @@ def test_parallel_function_call_anthropic_error_msg(model, messages): Anthropic (and Bedrock Invoke via ``AnthropicConfig.transform_request``) inject a dummy tool so CLIs work with ``modify_params`` left off. Bedrock Converse's no-raise behavior is covered offline in - ``tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py`` + ``tests/unit/llms/bedrock/chat/test_converse_transformation.py`` (see #24158, #27138), which needs no live credentials. """ # Force modify_params off as a clean baseline: it exercises the Anthropic diff --git a/tests/local_testing/test_handler_gc_does_not_close_client.py b/tests/local_testing/test_handler_gc_does_not_close_client.py index 1a6ab1b1827..63c5694dd89 100644 --- a/tests/local_testing/test_handler_gc_does_not_close_client.py +++ b/tests/local_testing/test_handler_gc_does_not_close_client.py @@ -17,7 +17,7 @@ body can still arrive, released once the caller is done with the response. Nothing here re-tests the shapes ``_handler_may_close_client`` covers -- a borrowed ``handler.client``, a caller-supplied client, an evicted-but-held -client. Those are pinned in ``tests/test_litellm/llms/custom_httpx/ +client. Those are pinned in ``tests/unit/llms/custom_httpx/ test_http_handler.py``. What is uncovered there is the in-flight response, so no test here may keep the client in a local: that inflates the very refcount under test, and the test then passes on a broken handler. They hold weak references diff --git a/tests/local_testing/test_sagemaker_nova_integration.py b/tests/local_testing/test_sagemaker_nova_integration.py index beeb1fa2db3..95f28fe9892 100644 --- a/tests/local_testing/test_sagemaker_nova_integration.py +++ b/tests/local_testing/test_sagemaker_nova_integration.py @@ -4,7 +4,7 @@ Integration tests for SageMaker Nova provider. These tests require a live SageMaker Nova endpoint and AWS credentials. They are skipped by default — run manually with: - pytest tests/test_litellm/llms/sagemaker/test_sagemaker_nova_integration.py -v --no-header -rN + pytest tests/local_testing/test_sagemaker_nova_integration.py -v --no-header -rN Prerequisites: export AWS_PROFILE= # or set AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY @@ -251,7 +251,7 @@ class TestSagemakerNova2LiteIntegration: Run with: export SAGEMAKER_NOVA2_LITE_ENDPOINT= - pytest tests/test_litellm/llms/sagemaker/test_sagemaker_nova_integration.py::TestSagemakerNova2LiteIntegration -v + pytest tests/local_testing/test_sagemaker_nova_integration.py::TestSagemakerNova2LiteIntegration -v """ def test_should_accept_reasoning_effort_low(self): diff --git a/tests/search_tests/test_bing_grounding_search.py b/tests/search_tests/test_bing_grounding_search.py index 3d1737477a1..f532158e462 100644 --- a/tests/search_tests/test_bing_grounding_search.py +++ b/tests/search_tests/test_bing_grounding_search.py @@ -85,7 +85,7 @@ class TestBingGroundingSearch(BaseSearchTest): class TestBingGroundingSearchTransformation: """ Full-stack tests through `litellm.search` / `litellm.asearch` with the HTTP layer mocked. - Transformation details are unit-tested in tests/test_litellm/llms/azure/search/. + Transformation details are unit-tested in tests/unit/llms/azure/search/. """ @pytest.fixture(autouse=True) diff --git a/tests/search_tests/test_nimble_search.py b/tests/search_tests/test_nimble_search.py index df432f8ae84..3426fc712f4 100644 --- a/tests/search_tests/test_nimble_search.py +++ b/tests/search_tests/test_nimble_search.py @@ -58,7 +58,7 @@ class TestNimbleSearch(BaseSearchTest): class TestNimbleSearchTransformation: """ Full-stack tests through `litellm.search` / `litellm.asearch` with the HTTP layer mocked. - Transformation details are unit-tested in tests/test_litellm/llms/nimble/search/. + Transformation details are unit-tested in tests/unit/llms/nimble/search/. """ @pytest.fixture(autouse=True) diff --git a/tests/test_litellm/integrations/levo/__init__.py b/tests/test_litellm/integrations/levo/__init__.py deleted file mode 100644 index 1560e78b7b9..00000000000 --- a/tests/test_litellm/integrations/levo/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Levo integration tests diff --git a/tests/test_litellm/integrations/zerobus/test_zerobus_client.py b/tests/test_litellm/integrations/zerobus/test_zerobus_client.py new file mode 100644 index 00000000000..ae7610536f3 --- /dev/null +++ b/tests/test_litellm/integrations/zerobus/test_zerobus_client.py @@ -0,0 +1,258 @@ +import base64 +import json +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from itertools import chain, repeat + +import httpx +import pytest + +from litellm.integrations.zerobus.client import ZerobusIngestClient +from litellm.types.integrations.zerobus import ZerobusAccessToken, ZerobusConnection, ZerobusIngestFailure + +CONNECTION = ZerobusConnection( + workspace_url="https://dbc-a1b2c3d4-e5f6.cloud.databricks.com/", + workspace_id="1234567890123456", + server_endpoint="https://1234567890123456.zerobus.us-west-2.cloud.databricks.com", + client_id="sp-client-id", + client_secret="sp-client-secret", + table_name="main.litellm.traces", +) +ROWS = ({"id": "a", "model": "gpt-4o"}, {"id": "b", "model": "gpt-4o"}) + + +def _token(value: str = "tok-1", expires_in: float = 3600) -> httpx.Response: + return httpx.Response(200, text=json.dumps({"access_token": value, "expires_in": expires_in})) + + +def _accepted() -> httpx.Response: + return httpx.Response(200, text="{}") + + +@dataclass(frozen=True, slots=True) +class TokenCall: + url: str + data: Mapping[str, str] + headers: Mapping[str, str] + + +@dataclass(frozen=True, slots=True) +class InsertCall: + url: str + content: bytes + headers: Mapping[str, str] + + +def _results(results: Sequence[httpx.Response | Exception]) -> Iterator[httpx.Response | Exception]: + """Results are served in order, and the last one repeats.""" + return chain(results[:-1], repeat(results[-1])) + + +class FakeHTTPClient: + """Stands in for AsyncHTTPHandler, including its habit of raising on error statuses.""" + + def __init__( + self, + token: Sequence[httpx.Response | Exception] = (), + insert: Sequence[httpx.Response | Exception] = (), + ) -> None: + self.token_results = _results(token or (_token(),)) + self.insert_results = _results(insert or (_accepted(),)) + self.token_calls: tuple[TokenCall, ...] = () + self.insert_calls: tuple[InsertCall, ...] = () + + async def post( + self, + url: str, + data: Mapping[str, str] | None = None, + content: bytes | None = None, + headers: Mapping[str, str] | None = None, + ) -> httpx.Response: + if url.endswith("/oidc/v1/token"): + self.token_calls = (*self.token_calls, TokenCall(url, data or {}, headers or {})) + return _raise_like_the_handler(next(self.token_results), url) + self.insert_calls = (*self.insert_calls, InsertCall(url, content or b"", headers or {})) + return _raise_like_the_handler(next(self.insert_results), url) + + +def _raise_like_the_handler(result: httpx.Response | Exception, url: str) -> httpx.Response: + if isinstance(result, Exception): + raise result + if result.status_code >= 300: + raise httpx.HTTPStatusError( + "boom", + request=httpx.Request("POST", url), + response=httpx.Response(result.status_code, text=result.text), + ) + return result + + +class FakeClock: + def __init__(self, now: float = 1_000.0) -> None: + self.now = now + + def __call__(self) -> float: + return self.now + + +def _client(http_client: FakeHTTPClient, clock: FakeClock | None = None) -> ZerobusIngestClient: + return ZerobusIngestClient(connection=CONNECTION, http_client=http_client, clock=clock or FakeClock()) + + +@pytest.mark.asyncio +async def test_rows_are_posted_as_one_json_list_to_the_table_insert_endpoint(): + http_client = FakeHTTPClient() + + outcome = await _client(http_client).insert(ROWS) + + assert outcome is None + (call,) = http_client.insert_calls + # Insert endpoint per the Zerobus Ingest docs, read 2026-09-19: + # https://docs.databricks.com/aws/en/ingestion/lakeflow-connect/zerobus-ingest + assert call.url == ( + "https://1234567890123456.zerobus.us-west-2.cloud.databricks.com/zerobus/v1/tables/main.litellm.traces/insert" + ) + assert json.loads(call.content) == [{"id": "a", "model": "gpt-4o"}, {"id": "b", "model": "gpt-4o"}] + assert call.headers["Content-Type"] == "application/json" + assert call.headers["Authorization"] == "Bearer tok-1" + + +@pytest.mark.asyncio +async def test_the_token_is_minted_for_the_zerobus_resource_with_the_table_privileges(): + """Zerobus refuses a plain workspace token: it must name its own resource and the table's UC privileges.""" + http_client = FakeHTTPClient() + + await _client(http_client).insert(ROWS) + + (call,) = http_client.token_calls + # Token form per the Zerobus Ingest docs (REST API authentication), read 2026-09-19: + # https://docs.databricks.com/aws/en/ingestion/lakeflow-connect/zerobus-ingest + assert call.url == "https://dbc-a1b2c3d4-e5f6.cloud.databricks.com/oidc/v1/token" + assert call.data["grant_type"] == "client_credentials" + assert call.data["scope"] == "all-apis" + assert call.data["resource"] == "api://databricks/workspaces/1234567890123456/zerobusDirectWriteApi" + details = json.loads(call.data["authorization_details"]) + assert [(d["object_type"], d["object_full_path"], d["privileges"]) for d in details] == [ + ("CATALOG", "main", ["USE CATALOG"]), + ("SCHEMA", "main.litellm", ["USE SCHEMA"]), + ("TABLE", "main.litellm.traces", ["SELECT", "MODIFY"]), + ] + assert all(d["type"] == "unity_catalog_privileges" for d in details) + + +@pytest.mark.asyncio +async def test_the_service_principal_authenticates_with_http_basic(): + http_client = FakeHTTPClient() + + await _client(http_client).insert(ROWS) + + scheme, credentials = http_client.token_calls[0].headers["Authorization"].split(" ") + assert scheme == "Basic" + assert base64.b64decode(credentials).decode() == "sp-client-id:sp-client-secret" + + +def test_the_client_secret_and_minted_token_stay_out_of_reprs_and_tracebacks(): + token = ZerobusAccessToken(value="tok-secret", expires_at=1.0) + + assert "sp-client-secret" not in repr(CONNECTION) + assert "sp-client-id" in repr(CONNECTION) + assert "tok-secret" not in repr(token) + assert "expires_at=1.0" in repr(token) + + +@pytest.mark.asyncio +async def test_the_token_is_reused_across_inserts_until_it_nears_expiry(): + clock = FakeClock(now=1_000.0) + http_client = FakeHTTPClient(token=[_token("tok-1", expires_in=600), _token("tok-2")]) + client = _client(http_client, clock) + + await client.insert(ROWS) + clock.now = 1_000.0 + 600 - 61 + await client.insert(ROWS) + clock.now = 1_000.0 + 600 - 59 + await client.insert(ROWS) + + assert len(http_client.token_calls) == 2 + assert [call.headers["Authorization"] for call in http_client.insert_calls] == [ + "Bearer tok-1", + "Bearer tok-1", + "Bearer tok-2", + ] + + +@pytest.mark.asyncio +async def test_a_401_discards_the_token_so_the_next_insert_mints_a_fresh_one(): + http_client = FakeHTTPClient( + token=[_token("tok-1"), _token("tok-2")], + insert=[httpx.Response(401, text="expired"), _accepted()], + ) + client = _client(http_client) + + first = await client.insert(ROWS) + second = await client.insert(ROWS) + + assert first == ZerobusIngestFailure(detail="insert returned 401, token discarded", retryable=True) + assert second is None + assert http_client.insert_calls[1].headers["Authorization"] == "Bearer tok-2" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [429, 500, 503]) +async def test_a_transient_insert_status_is_retryable(status: int): + http_client = FakeHTTPClient(insert=[httpx.Response(status, text="later")]) + + outcome = await _client(http_client).insert(ROWS) + + assert isinstance(outcome, ZerobusIngestFailure) + assert outcome.retryable is True + assert str(status) in outcome.detail + + +@pytest.mark.asyncio +async def test_a_schema_rejection_is_not_retryable_and_says_why(): + http_client = FakeHTTPClient(insert=[httpx.Response(400, text="unknown column foo")]) + + outcome = await _client(http_client).insert(ROWS) + + assert outcome == ZerobusIngestFailure(detail="insert returned 400: unknown column foo", retryable=False) + + +@pytest.mark.asyncio +async def test_a_network_failure_on_insert_is_retryable(): + http_client = FakeHTTPClient(insert=[httpx.ConnectError("connection refused")]) + + outcome = await _client(http_client).insert(ROWS) + + assert isinstance(outcome, ZerobusIngestFailure) + assert outcome.retryable is True + + +@pytest.mark.asyncio +async def test_bad_credentials_fail_the_insert_without_posting_rows(): + http_client = FakeHTTPClient(token=[httpx.Response(401, text="invalid_client")]) + + outcome = await _client(http_client).insert(ROWS) + + assert outcome == ZerobusIngestFailure(detail="token request returned 401: invalid_client", retryable=False) + assert http_client.insert_calls == () + + +@pytest.mark.asyncio +async def test_a_token_endpoint_outage_is_retryable(): + http_client = FakeHTTPClient(token=[httpx.Response(503, text="try later")]) + + outcome = await _client(http_client).insert(ROWS) + + assert isinstance(outcome, ZerobusIngestFailure) + assert outcome.retryable is True + + +@pytest.mark.asyncio +async def test_a_token_response_without_a_token_is_reported_not_raised(): + http_client = FakeHTTPClient(token=[httpx.Response(200, text='{"token_type": "Bearer"}')]) + + outcome = await _client(http_client).insert(ROWS) + + assert isinstance(outcome, ZerobusIngestFailure) + assert outcome.retryable is False + assert "token response" in outcome.detail diff --git a/tests/test_litellm/integrations/zerobus/test_zerobus_logger.py b/tests/test_litellm/integrations/zerobus/test_zerobus_logger.py new file mode 100644 index 00000000000..a85a9e2e6a0 --- /dev/null +++ b/tests/test_litellm/integrations/zerobus/test_zerobus_logger.py @@ -0,0 +1,392 @@ +import asyncio +from collections.abc import Callable, Iterator, Mapping, Sequence +from itertools import chain, repeat + +import pytest + +import litellm +from litellm.integrations.zerobus.client import ZerobusIngestError +from litellm.integrations.zerobus.logger import ZerobusLogger, connection_for +from litellm.types.integrations.zerobus import ZerobusIngestFailure, ZerobusInitParams + +WORKSPACE_URL = "https://dbc-a1b2c3d4-e5f6.cloud.databricks.com" +SERVER_ENDPOINT = "https://1234567890123456.zerobus.us-west-2.cloud.databricks.com" + + +Row = Mapping[str, object] + + +class FakeIngestClient: + """Records the rows each flush would have written; outcomes are served in order and the last one repeats.""" + + def __init__( + self, + outcomes: Sequence[ZerobusIngestFailure | None] = (None,), + on_insert: Callable[[], None] | None = None, + ) -> None: + self.outcomes: Iterator[ZerobusIngestFailure | None] = chain(outcomes[:-1], repeat(outcomes[-1])) + self.on_insert = on_insert + self.batches: tuple[tuple[Row, ...], ...] = () + + async def insert(self, rows: Sequence[Row]) -> ZerobusIngestFailure | None: + if self.on_insert is not None: + self.on_insert() + self.batches = (*self.batches, tuple(rows)) + return next(self.outcomes) + + def ids(self) -> tuple[object, ...]: + return tuple(row["id"] for batch in self.batches for row in batch) + + +def _logger(client: FakeIngestClient, **params: object) -> ZerobusLogger: + return ZerobusLogger(params=ZerobusInitParams.model_validate(params), client=client) + + +def _event(request_id: str, **payload: object) -> dict[str, object]: + return { + "standard_logging_object": { + "id": request_id, + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "response": {"choices": []}, + **payload, + } + } + + +async def _settle(logger: ZerobusLogger) -> None: + for _ in range(200): + await asyncio.sleep(0.001) + task = logger._batch_flush_task + if (task is None or task.done()) and not logger._flushing: + return + + +@pytest.mark.asyncio +async def test_a_full_batch_is_written_as_one_insert_of_table_rows(): + client = FakeIngestClient() + logger = _logger(client, batch_size=3) + + for request_id in ("a", "b", "c"): + await logger.async_log_success_event(_event(request_id), None, None, None) + + await _settle(logger) + assert len(client.batches) == 1 + assert client.ids() == ("a", "b", "c") + assert client.batches[0][0]["model"] == "gpt-4o" + assert logger.log_queue == [] + + +@pytest.mark.asyncio +async def test_rows_are_held_until_the_batch_is_full(): + client = FakeIngestClient() + logger = _logger(client, batch_size=3) + + await logger.async_log_success_event(_event("a"), None, None, None) + + assert client.batches == () + assert len(logger.log_queue) == 1 + + +@pytest.mark.asyncio +async def test_failed_requests_are_written_too(): + client = FakeIngestClient() + logger = _logger(client, batch_size=1) + + await logger.async_log_failure_event(_event("failed", status="failure", error_str="boom"), None, None, None) + + await _settle(logger) + assert client.ids() == ("failed",) + assert client.batches[0][0]["status"] == "failure" + assert client.batches[0][0]["error_str"] == "boom" + + +@pytest.mark.asyncio +async def test_an_event_without_a_standard_payload_is_skipped(): + client = FakeIngestClient() + logger = _logger(client, batch_size=1) + + await logger.async_log_success_event({"kwargs": "but no payload"}, None, None, None) + + assert client.batches == () + assert logger.log_queue == [] + + +@pytest.mark.asyncio +async def test_a_retryable_failure_keeps_the_rows_for_the_next_flush(): + client = FakeIngestClient([ZerobusIngestFailure("zerobus is down", retryable=True)]) + logger = _logger(client, batch_size=2) + + for request_id in ("a", "b"): + await logger.async_log_success_event(_event(request_id), None, None, None) + + await _settle(logger) + assert [row["id"] for row in logger.log_queue] == ["a", "b"] + + +@pytest.mark.asyncio +async def test_a_retryable_failure_surfaces_so_the_base_logger_can_preserve_it(): + client = FakeIngestClient([ZerobusIngestFailure("zerobus is down", retryable=True)]) + logger = _logger(client, batch_size=99) + logger.log_queue.append({"id": "a"}) + + with pytest.raises(ZerobusIngestError, match="zerobus is down"): + await logger.async_send_batch() + + +@pytest.mark.asyncio +async def test_a_rejected_batch_is_dropped_rather_than_blocking_the_queue(): + client = FakeIngestClient([ZerobusIngestFailure("unknown column", retryable=False)]) + logger = _logger(client, batch_size=2) + + for request_id in ("a", "b"): + await logger.async_log_success_event(_event(request_id), None, None, None) + + await _settle(logger) + assert logger.log_queue == [] + + +@pytest.mark.asyncio +async def test_a_row_that_arrives_mid_flush_is_kept_for_the_next_one(): + client = FakeIngestClient() + logger = _logger(client, batch_size=1) + client.on_insert = lambda: logger.log_queue.append({"id": "late"}) + + await logger.async_log_success_event(_event("first"), None, None, None) + + await _settle(logger) + assert client.ids() == ("first",) + assert [row["id"] for row in logger.log_queue] == ["late"] + + +@pytest.mark.asyncio +async def test_the_queue_cap_holds_while_an_insert_is_in_flight(): + """A slow insert must not let the queue grow past max_queue_size, nor disturb the in-flight head.""" + insert_started = asyncio.Event() + finish_insert = asyncio.Event() + + class SlowClient: + batches: tuple[tuple[Row, ...], ...] = () + + async def insert(self, rows: Sequence[Row]) -> None: + insert_started.set() + await finish_insert.wait() + self.batches = (*self.batches, tuple(rows)) + + client = SlowClient() + logger = ZerobusLogger(params=ZerobusInitParams(batch_size=2), client=client) + logger.max_queue_size = 3 + + for request_id in ("a", "b"): + await logger.async_log_success_event(_event(request_id), None, None, None) + await insert_started.wait() + for request_id in ("c", "d", "e"): + await logger.async_log_success_event(_event(request_id), None, None, None) + finish_insert.set() + await _settle(logger) + + assert [[row["id"] for row in batch] for batch in client.batches] == [["a", "b"]] + assert [row["id"] for row in logger.log_queue] == ["c"] + + +@pytest.mark.asyncio +async def test_a_client_error_does_not_break_the_request_path(): + class ExplodingClient: + async def insert(self, rows: Sequence[Row]) -> None: + raise RuntimeError("bug") + + logger = ZerobusLogger(params=ZerobusInitParams(batch_size=1), client=ExplodingClient()) + + await logger.async_log_success_event(_event("a"), None, None, None) + await _settle(logger) + + assert [row["id"] for row in logger.log_queue] == ["a"] + + +@pytest.mark.asyncio +async def test_turn_off_message_logging_redacts_prompts_and_responses_but_keeps_the_rest(): + client = FakeIngestClient() + logger = _logger(client, batch_size=1, turn_off_message_logging=True) + + await logger.async_log_success_event( + _event("a", prompt_tokens=10, response={"choices": [{"message": {"content": "the secret answer"}}]}), + None, + None, + None, + ) + + await _settle(logger) + (row,) = client.batches[0] + assert row["id"] == "a" + assert row["prompt_tokens"] == 10 + assert '"hi"' not in str(row["messages"]) + assert "the secret answer" not in str(row["response"]) + + +def test_connection_comes_from_the_environment_the_proxy_ui_writes(monkeypatch): + monkeypatch.setenv("ZEROBUS_WORKSPACE_URL", WORKSPACE_URL) + monkeypatch.setenv("ZEROBUS_SERVER_ENDPOINT", SERVER_ENDPOINT) + monkeypatch.setenv("ZEROBUS_CLIENT_ID", "sp-id") + monkeypatch.setenv("ZEROBUS_CLIENT_SECRET", "sp-secret") + monkeypatch.setenv("ZEROBUS_TABLE_NAME", "main.litellm.traces") + + connection = connection_for(ZerobusInitParams()) + + assert connection.workspace_url == WORKSPACE_URL + assert connection.server_endpoint == SERVER_ENDPOINT + assert connection.workspace_id == "1234567890123456" + assert connection.client_id == "sp-id" + assert connection.client_secret == "sp-secret" + assert connection.table_name == "main.litellm.traces" + + +def test_config_yaml_params_win_over_the_environment(monkeypatch): + monkeypatch.setenv("ZEROBUS_TABLE_NAME", "env.schema.table") + monkeypatch.setenv("ZEROBUS_CLIENT_SECRET", "from-env") + + connection = connection_for( + ZerobusInitParams( + workspace_url=WORKSPACE_URL, + server_endpoint=SERVER_ENDPOINT, + client_id="sp-id", + client_secret="from-config", + table_name="cfg.schema.table", + ) + ) + + assert connection.table_name == "cfg.schema.table" + assert connection.client_secret == "from-config" + + +def test_a_secret_reference_in_config_yaml_is_resolved(monkeypatch): + monkeypatch.setenv("MY_SP_SECRET", "resolved-secret") + + connection = connection_for( + ZerobusInitParams( + workspace_url=WORKSPACE_URL, + server_endpoint=SERVER_ENDPOINT, + client_id="sp-id", + client_secret="os.environ/MY_SP_SECRET", + table_name="main.litellm.traces", + ) + ) + + assert connection.client_secret == "resolved-secret" + + +def test_a_missing_setting_names_the_env_var_to_set(monkeypatch): + monkeypatch.delenv("ZEROBUS_CLIENT_SECRET", raising=False) + + with pytest.raises(ValueError, match="ZEROBUS_CLIENT_SECRET"): + connection_for( + ZerobusInitParams( + workspace_url=WORKSPACE_URL, + server_endpoint=SERVER_ENDPOINT, + client_id="sp-id", + table_name="main.litellm.traces", + ) + ) + + +def test_a_table_that_is_not_fully_qualified_is_refused(): + with pytest.raises(ValueError, match=r"catalog\.schema\.table"): + connection_for( + ZerobusInitParams( + workspace_url=WORKSPACE_URL, + server_endpoint=SERVER_ENDPOINT, + client_id="sp-id", + client_secret="sp-secret", + table_name="traces", + ) + ) + + +def test_an_endpoint_without_a_workspace_id_is_refused(): + """The token's resource needs the numeric workspace id, which only the Zerobus hostname carries.""" + with pytest.raises(ValueError, match="ZEROBUS_SERVER_ENDPOINT"): + connection_for( + ZerobusInitParams( + workspace_url=WORKSPACE_URL, + server_endpoint=WORKSPACE_URL, + client_id="sp-id", + client_secret="sp-secret", + table_name="main.litellm.traces", + ) + ) + + +def test_a_misconfigured_logger_fails_at_startup_not_at_first_flush(monkeypatch): + for name in ("WORKSPACE_URL", "SERVER_ENDPOINT", "CLIENT_ID", "CLIENT_SECRET", "TABLE_NAME"): + monkeypatch.delenv(f"ZEROBUS_{name}", raising=False) + monkeypatch.setattr(litellm, "zerobus_params", None) + + with pytest.raises(ValueError, match="ZEROBUS_"): + ZerobusLogger() + + +def test_litellm_zerobus_params_configure_the_logger(monkeypatch): + monkeypatch.setattr( + litellm, + "zerobus_params", + { + "workspace_url": WORKSPACE_URL, + "server_endpoint": SERVER_ENDPOINT, + "client_id": "sp-id", + "client_secret": "sp-secret", + "table_name": "main.litellm.traces", + "batch_size": 7, + "flush_interval": 3, + }, + ) + + logger = ZerobusLogger() + + assert logger.batch_size == 7 + assert logger.flush_interval == 3 + assert logger.client.connection.table_name == "main.litellm.traces" + + +def test_the_client_is_kept_while_the_connection_is_unchanged_and_rebuilt_when_it_changes(monkeypatch): + """The client caches its token, so it must survive across flushes, yet a UI edit must take effect.""" + monkeypatch.setenv("ZEROBUS_WORKSPACE_URL", WORKSPACE_URL) + monkeypatch.setenv("ZEROBUS_SERVER_ENDPOINT", SERVER_ENDPOINT) + monkeypatch.setenv("ZEROBUS_CLIENT_ID", "sp-id") + monkeypatch.setenv("ZEROBUS_CLIENT_SECRET", "sp-secret") + monkeypatch.setenv("ZEROBUS_TABLE_NAME", "main.litellm.traces") + monkeypatch.setattr(litellm, "zerobus_params", None) + logger = ZerobusLogger() + + first = logger.client + unchanged = logger.client + monkeypatch.setenv("ZEROBUS_TABLE_NAME", "main.litellm.traces_v2") + rebuilt = logger.client + + assert unchanged is first + assert rebuilt is not first + assert rebuilt.connection.table_name == "main.litellm.traces_v2" + + +def test_callbacks_zerobus_builds_one_logger_and_reuses_it(monkeypatch): + """`litellm_settings.callbacks: ["zerobus"]` goes through litellm_logging, which must hand back one instance.""" + from litellm.litellm_core_utils import litellm_logging as logging_module + + monkeypatch.setenv("ZEROBUS_WORKSPACE_URL", WORKSPACE_URL) + monkeypatch.setenv("ZEROBUS_SERVER_ENDPOINT", SERVER_ENDPOINT) + monkeypatch.setenv("ZEROBUS_CLIENT_ID", "sp-id") + monkeypatch.setenv("ZEROBUS_CLIENT_SECRET", "sp-secret") + monkeypatch.setenv("ZEROBUS_TABLE_NAME", "main.litellm.traces") + monkeypatch.setattr(litellm, "zerobus_params", None) + monkeypatch.setattr(logging_module, "_in_memory_loggers", []) + + assert logging_module.get_custom_logger_compatible_class("zerobus") is None + + first = logging_module._init_custom_logger_compatible_class( + logging_integration="zerobus", internal_usage_cache=None, llm_router=None, custom_logger_init_args={} + ) + second = logging_module._init_custom_logger_compatible_class( + logging_integration="zerobus", internal_usage_cache=None, llm_router=None, custom_logger_init_args={} + ) + + assert isinstance(first, ZerobusLogger) + assert second is first + assert logging_module.get_custom_logger_compatible_class("zerobus") is first diff --git a/tests/test_litellm/integrations/zerobus/test_zerobus_row.py b/tests/test_litellm/integrations/zerobus/test_zerobus_row.py new file mode 100644 index 00000000000..b73c3bae48f --- /dev/null +++ b/tests/test_litellm/integrations/zerobus/test_zerobus_row.py @@ -0,0 +1,139 @@ +import json + +from litellm.integrations.zerobus.row import TRACE_TABLE_COLUMNS, create_table_sql, trace_row + + +def _payload() -> dict[str, object]: + return { + "id": "chatcmpl-1", + "trace_id": "trace-1", + "session_id": "session-1", + "litellm_call_id": "call-1", + "call_type": "acompletion", + "status": "success", + "model": "gpt-4o", + "model_group": "gpt-4o-group", + "custom_llm_provider": "openai", + "api_base": "https://api.openai.com", + "stream": False, + "cache_hit": None, + "startTime": 1_700_000_000.25, + "endTime": 1_700_000_001.5, + "completionStartTime": 1_700_000_000.75, + "response_time": 1.25, + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + "response_cost": 0.0015, + "saved_cache_cost": 0.0, + "end_user": "end-user-1", + "requester_ip_address": "10.0.0.1", + "user_agent": "curl/8", + "request_tags": ["prod"], + "messages": [{"role": "user", "content": "hi"}], + "response": {"choices": [{"message": {"role": "assistant", "content": "hello"}}]}, + "error_str": None, + "error_information": None, + "metadata": { + "user_api_key_hash": "hash-1", + "user_api_key_alias": "alias-1", + "user_api_key_team_id": "team-1", + "user_api_key_team_alias": "team-alias-1", + "user_api_key_user_id": "user-1", + "user_api_key_org_id": "org-1", + }, + "model_parameters": {"temperature": 0.2}, + "hidden_params": {"response_cost": 0.0015}, + "guardrail_information": None, + "cost_breakdown": {"input_cost": 0.001, "output_cost": 0.0005}, + } + + +def test_every_row_has_exactly_the_documented_columns(): + """Zerobus rejects a record naming a column the table lacks, so the row and the DDL must agree.""" + assert tuple(trace_row(_payload())) == tuple(TRACE_TABLE_COLUMNS) + assert tuple(trace_row({})) == tuple(TRACE_TABLE_COLUMNS) + + +def test_scalars_land_in_their_columns(): + row = trace_row(_payload()) + + assert row["id"] == "chatcmpl-1" + assert row["trace_id"] == "trace-1" + assert row["status"] == "success" + assert row["model"] == "gpt-4o" + assert row["stream"] is False + assert row["prompt_tokens"] == 10 + assert row["total_tokens"] == 15 + assert row["response_cost"] == 0.0015 + assert row["end_user"] == "end-user-1" + + +def test_key_and_team_identity_is_lifted_out_of_metadata(): + """Filtering spend by team or key is the main query, so those live in their own columns.""" + row = trace_row(_payload()) + + assert row["api_key_hash"] == "hash-1" + assert row["api_key_alias"] == "alias-1" + assert row["team_id"] == "team-1" + assert row["team_alias"] == "team-alias-1" + assert row["user_id"] == "user-1" + assert row["org_id"] == "org-1" + + +def test_timestamps_become_epoch_microseconds(): + row = trace_row(_payload()) + + assert row["start_time"] == 1_700_000_000_250_000 + assert row["end_time"] == 1_700_000_001_500_000 + assert row["completion_start_time"] == 1_700_000_000_750_000 + + +def test_a_zero_timestamp_is_null_rather_than_1970(): + """LiteLLM leaves completionStartTime at 0 when there is no first token, which is not a real time.""" + row = trace_row({**_payload(), "completionStartTime": 0}) + + assert row["completion_start_time"] is None + + +def test_nested_fields_are_json_text_for_the_variant_columns(): + row = trace_row(_payload()) + + assert json.loads(str(row["messages"])) == [{"role": "user", "content": "hi"}] + assert json.loads(str(row["metadata"]))["user_api_key_team_id"] == "team-1" + assert json.loads(str(row["request_tags"])) == ["prod"] + assert json.loads(str(row["cost_breakdown"])) == {"input_cost": 0.001, "output_cost": 0.0005} + + +def test_missing_and_null_fields_are_null(): + row = trace_row({**_payload(), "messages": None, "guardrail_information": None}) + + assert row["messages"] is None + assert row["guardrail_information"] is None + assert row["error_str"] is None + assert row["cache_hit"] is None + + +def test_a_wrongly_typed_field_is_null_instead_of_a_rejected_record(): + """One odd payload must not poison the whole batch: the table type wins.""" + row = trace_row({**_payload(), "prompt_tokens": "ten", "stream": "yes", "startTime": "now"}) + + assert row["prompt_tokens"] is None + assert row["stream"] is None + assert row["start_time"] is None + + +def test_the_row_survives_a_json_round_trip_unchanged(): + row = trace_row(_payload()) + + assert json.loads(json.dumps(dict(row))) == dict(row) + + +def test_create_table_sql_declares_every_column_with_its_type(): + sql = create_table_sql("main.litellm.traces") + + assert sql.startswith("CREATE TABLE main.litellm.traces (") + assert " start_time TIMESTAMP," in sql + assert " messages VARIANT," in sql + assert " cost_breakdown VARIANT\n);" in sql + assert sql.count(",") == len(TRACE_TABLE_COLUMNS) - 1 diff --git a/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py b/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py index 1cc96cb1256..c3478c0d5eb 100644 --- a/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py @@ -364,6 +364,26 @@ async def test_batch_health_check_uses_alist_batches_for_supported_providers(): mock_alist.assert_called_once() +@pytest.mark.asyncio +async def test_batch_health_check_hands_the_resolved_provider_to_alist_batches(): + filtered_model_params: Final = { + "model": "xai/grok-4.3", + "api_key": "sk-test", + "litellm_metadata": {"tags": [LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME]}, + } + + with patch("litellm.alist_batches", new_callable=AsyncMock, return_value={}) as mock_alist: + await HealthCheckHelpers._batch_health_check( + custom_llm_provider="xai", + model_params={**filtered_model_params, "messages": []}, + filtered_model_params=filtered_model_params, + ) + + assert mock_alist.call_args.kwargs["custom_llm_provider"] == "xai" + assert mock_alist.call_args.kwargs["model"] == "xai/grok-4.3" + assert mock_alist.call_args.kwargs["api_key"] == "sk-test" + + @pytest.mark.asyncio async def test_batch_health_check_falls_back_to_acompletion_for_unsupported(): """Providers not in LIST_BATCHES_SUPPORTED_PROVIDERS fall back to acompletion.""" 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 3afa31cc801..265fdb50836 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -20,7 +20,7 @@ from openai._legacy_response import HttpxBinaryResponseContent import litellm from litellm._logging import session_id_var, trace_id_var -from litellm.constants import SENTRY_DENYLIST, SENTRY_PII_DENYLIST +from litellm.constants import SENTRY_PII_DENYLIST from litellm.cost_calculator import ocr_batch_cost from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging @@ -357,108 +357,23 @@ def test_post_call_serializes_dict_with_datetime(logging_obj): assert "2026-05-11" in serialized -def test_sentry_sample_rate(monkeypatch): - existing_sample_rate = os.getenv("SENTRY_API_SAMPLE_RATE") - try: - # test with default value by removing the environment variable - if existing_sample_rate: - del os.environ["SENTRY_API_SAMPLE_RATE"] - - set_callbacks(["sentry"]) - # Check if the default sample rate is set to 1.0 - assert os.environ.get("SENTRY_API_SAMPLE_RATE") == "1.0" - - # test with custom value - monkeypatch.setenv("SENTRY_API_SAMPLE_RATE", "0.5") - - set_callbacks(["sentry"]) - # Check if the custom sample rate is set correctly - assert os.environ.get("SENTRY_API_SAMPLE_RATE") == "0.5" - except Exception as e: - print(f"Error: {e}") - finally: - # Restore the original environment variable - if existing_sample_rate: - monkeypatch.setenv("SENTRY_API_SAMPLE_RATE", existing_sample_rate) - else: - if "SENTRY_API_SAMPLE_RATE" in os.environ: - del os.environ["SENTRY_API_SAMPLE_RATE"] - - def test_sentry_environment(monkeypatch): - """Test that SENTRY_ENVIRONMENT is properly handled during Sentry initialization""" - existing_environment = os.getenv("SENTRY_ENVIRONMENT") - existing_dsn = os.getenv("SENTRY_DSN") + import sentry_sdk - # Create mock sentry_sdk module - mock_event_scrubber_instance = MagicMock() - mock_event_scrubber_cls = MagicMock(return_value=mock_event_scrubber_instance) - - mock_scrubber_module = MagicMock() - mock_scrubber_module.EventScrubber = mock_event_scrubber_cls - - mock_sentry_sdk = MagicMock() - mock_sentry_sdk.scrubber = mock_scrubber_module mock_init = MagicMock() - mock_sentry_sdk.init = mock_init + monkeypatch.setattr(sentry_sdk, "init", mock_init) + monkeypatch.setenv("SENTRY_DSN", "https://test@sentry.io/123456") + monkeypatch.delenv("SENTRY_ENVIRONMENT", raising=False) - # Inject mocks into sys.modules - sys.modules["sentry_sdk"] = mock_sentry_sdk - sys.modules["sentry_sdk.scrubber"] = mock_scrubber_module - - try: - # Set a mock DSN to allow Sentry initialization - monkeypatch.setenv("SENTRY_DSN", "https://test@sentry.io/123456") - - # Test with default value (no environment set) - if existing_environment: - del os.environ["SENTRY_ENVIRONMENT"] + set_callbacks(["sentry"]) + assert mock_init.call_args[1]["environment"] == "production" + for environment in ("development", "staging"): + monkeypatch.setenv("SENTRY_ENVIRONMENT", environment) mock_init.reset_mock() set_callbacks(["sentry"]) - # Check that init was called with default environment "production" mock_init.assert_called_once() - call_kwargs = mock_init.call_args[1] - assert call_kwargs["environment"] == "production" - - # Test with custom environment value - monkeypatch.setenv("SENTRY_ENVIRONMENT", "development") - - mock_init.reset_mock() - set_callbacks(["sentry"]) - # Check that init was called with custom environment "development" - mock_init.assert_called_once() - call_kwargs = mock_init.call_args[1] - assert call_kwargs["environment"] == "development" - - # Test with staging environment - monkeypatch.setenv("SENTRY_ENVIRONMENT", "staging") - - mock_init.reset_mock() - set_callbacks(["sentry"]) - # Check that init was called with custom environment "staging" - mock_init.assert_called_once() - call_kwargs = mock_init.call_args[1] - assert call_kwargs["environment"] == "staging" - - except Exception as e: - print(f"Error: {e}") - raise - finally: - # Restore the original environment variables - if existing_environment: - monkeypatch.setenv("SENTRY_ENVIRONMENT", existing_environment) - else: - if "SENTRY_ENVIRONMENT" in os.environ: - del os.environ["SENTRY_ENVIRONMENT"] - - if existing_dsn: - monkeypatch.setenv("SENTRY_DSN", existing_dsn) - else: - if "SENTRY_DSN" in os.environ: - del os.environ["SENTRY_DSN"] - - + assert mock_init.call_args[1]["environment"] == environment def test_use_custom_pricing_for_model(): from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model @@ -3100,37 +3015,34 @@ def test_speech_call_is_still_priced_from_input_characters(call_type): def test_sentry_event_scrubber_initialization(monkeypatch): - # Step 1: Create a fake sentry_sdk.scrubber module - mock_event_scrubber_instance = MagicMock() - mock_event_scrubber_cls = MagicMock(return_value=mock_event_scrubber_instance) + import sentry_sdk - mock_scrubber_module = MagicMock() - mock_scrubber_module.EventScrubber = mock_event_scrubber_cls - - # Step 2: Create a fake sentry_sdk module and insert into sys.modules - mock_sentry_sdk = MagicMock() - mock_sentry_sdk.scrubber = mock_scrubber_module mock_init = MagicMock() - mock_sentry_sdk.init = mock_init + monkeypatch.setattr(sentry_sdk, "init", mock_init) + monkeypatch.delenv("SENTRY_SEND_DEFAULT_PII", raising=False) - # Step 3: Inject both into sys.modules BEFORE import occurs - sys.modules["sentry_sdk"] = mock_sentry_sdk - sys.modules["sentry_sdk.scrubber"] = mock_scrubber_module - - # Step 4: Run the actual sentry setup code set_callbacks(["sentry"]) - # Step 5: Assert the EventScrubber was constructed correctly - mock_event_scrubber_cls.assert_called_once_with( - denylist=SENTRY_DENYLIST, - pii_denylist=SENTRY_PII_DENYLIST, - ) - - # Step 6: Assert the event_scrubber and PII args were passed mock_init.assert_called_once() call_args = mock_init.call_args[1] - assert call_args["event_scrubber"] == mock_event_scrubber_instance assert call_args["send_default_pii"] is False + assert call_args["event_scrubber"].recursive is True + assert {name.lower() for name in SENTRY_PII_DENYLIST} <= {name.lower() for name in call_args["event_scrubber"].denylist} + assert call_args["before_send"] is call_args["before_send_transaction"] + + +def test_sentry_send_default_pii_opt_in(monkeypatch): + import sentry_sdk + + mock_init = MagicMock() + monkeypatch.setattr(sentry_sdk, "init", mock_init) + monkeypatch.setenv("SENTRY_SEND_DEFAULT_PII", "true") + + set_callbacks(["sentry"]) + + call_args = mock_init.call_args[1] + assert call_args["send_default_pii"] is True + assert not {name.lower() for name in SENTRY_PII_DENYLIST} & {name.lower() for name in call_args["event_scrubber"].denylist} def test_get_masked_values(): @@ -5396,6 +5308,19 @@ def test_handle_anthropic_messages_response_logging_passes_model_response_throug assert logging_obj._handle_anthropic_messages_response_logging(result=model_response) is model_response +def test_anthropic_messages_logged_response_tolerates_a_stream_that_assembled_nothing(): + """A /v1/messages stream whose upstream yielded no chunks assembles to None; the spend + row must still land under the message id the caller was served instead of crashing.""" + logging_obj = _anthropic_messages_logging_obj() + logging_obj.record_streamed_anthropic_message_id("msg_served") + + result = logging_obj._anthropic_messages_logged_response(result=None) + + assert isinstance(result, ModelResponse) + assert result.id == "msg_served" + assert result.model == "openai/my-local" + + def test_handle_anthropic_messages_response_logging_degrades_on_unparseable_responses_payload(): """If the Responses translation raises (eg. empty output on an incomplete response), the row must still land: a minimal ModelResponse with model + usage is returned.""" @@ -7856,6 +7781,9 @@ _PUBLISHED_BATCH_RATES: Final = MappingProxyType( "output_cost_per_token_batches": 4.1e-6, "cache_read_input_token_cost_batches": 1.2e-7, "cache_creation_input_token_cost_batches": 1.3e-6, + "input_cost_per_token_above_200k_tokens_batches": 2.1e-6, + "output_cost_per_token_above_200k_tokens_batches": 5.1e-6, + "cache_read_input_token_cost_above_200k_tokens_batches": 2.2e-7, "input_cost_per_token_above_272k_tokens_batches": 3.1e-6, "output_cost_per_token_above_272k_tokens_batches": 7.1e-6, "cache_read_input_token_cost_above_272k_tokens_batches": 3.2e-7, @@ -7864,14 +7792,17 @@ _PUBLISHED_BATCH_RATES: Final = MappingProxyType( ) _PUBLISHED_INPUT_BATCH_KEYS: Final = ( "input_cost_per_token_batches", + "input_cost_per_token_above_200k_tokens_batches", "input_cost_per_token_above_272k_tokens_batches", "cache_read_input_token_cost_batches", + "cache_read_input_token_cost_above_200k_tokens_batches", "cache_read_input_token_cost_above_272k_tokens_batches", "cache_creation_input_token_cost_batches", "cache_creation_input_token_cost_above_272k_tokens_batches", ) _PUBLISHED_OUTPUT_BATCH_KEYS: Final = ( "output_cost_per_token_batches", + "output_cost_per_token_above_200k_tokens_batches", "output_cost_per_token_above_272k_tokens_batches", ) @@ -7960,22 +7891,22 @@ def test_batch_cost_calculator_bills_the_carried_output_tier_when_the_deployment ) +@pytest.mark.parametrize( + "tier_key", + ["input_cost_per_token_above_200k_tokens_batches", "input_cost_per_token_above_272k_tokens_batches"], +) def test_deployment_pricing_model_info_honors_a_tier_only_batch_override_over_the_published_flat_rates( - _published_batch_model: None, + _published_batch_model: None, tier_key: str ) -> None: from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info - info: Final = deployment_pricing_model_info( - _batch_deployment_id({"input_cost_per_token_above_272k_tokens_batches": 1e-3}), _PUBLISHED_BATCH_DEPLOYMENT - ) + info: Final = deployment_pricing_model_info(_batch_deployment_id({tier_key: 1e-3}), _PUBLISHED_BATCH_DEPLOYMENT) carried_keys: Final = tuple( - key - for key in (*_PUBLISHED_INPUT_BATCH_KEYS, *_PUBLISHED_OUTPUT_BATCH_KEYS) - if key != "input_cost_per_token_above_272k_tokens_batches" + key for key in (*_PUBLISHED_INPUT_BATCH_KEYS, *_PUBLISHED_OUTPUT_BATCH_KEYS) if key != tier_key ) assert info is not None - assert info["input_cost_per_token_above_272k_tokens_batches"] == 1e-3 + assert info[tier_key] == 1e-3 assert {key: info[key] for key in carried_keys} == {key: _PUBLISHED_BATCH_RATES[key] for key in carried_keys} diff --git a/tests/test_litellm/litellm_core_utils/test_sentry_scrubbing.py b/tests/test_litellm/litellm_core_utils/test_sentry_scrubbing.py new file mode 100644 index 00000000000..9aae3999129 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_sentry_scrubbing.py @@ -0,0 +1,278 @@ +import hashlib +import json +import secrets +from collections.abc import Callable, Mapping +from functools import reduce +from typing import Final, cast + +import pytest +import sentry_sdk +from pydantic import JsonValue +from sentry_sdk.envelope import Envelope +from sentry_sdk.transport import Transport +from sentry_sdk.utils import event_from_exception + +from litellm.constants import LENGTH_OF_LITELLM_GENERATED_KEY, MINIMUM_CUSTOM_KEY_LENGTH +from litellm.litellm_core_utils.sentry_scrubbing import ( + FILTERED, + MAX_SCRUB_DEPTH, + build_key_pattern, + build_sentry_init_options, + build_string_scrubber, + scrub_json_strings, +) +from litellm.proxy._types import LiteLLM_UserTable, UserAPIKeyAuth + +EMAIL: Final = "qa.user@example.com" +VIRTUAL_KEY: Final = "sk-virtual-key-under-test" +KEY_HASH: Final = hashlib.sha256(VIRTUAL_KEY.encode()).hexdigest() +MASTER_KEY: Final = "sk-master-key-under-test" +DATABASE_URL: Final = "postgresql://litellm:db-password-under-test@db.internal:5432/litellm" +PII_ON: Final = {"SENTRY_DSN": "https://key@sentry.example/1", "SENTRY_SEND_DEFAULT_PII": "true"} +PII_OFF: Final = {"SENTRY_DSN": "https://key@sentry.example/1"} + + +class RecordingTransport(Transport): + def __init__(self) -> None: + super().__init__() + self.last_envelope: Envelope | None = None + + def capture_envelope(self, envelope: Envelope) -> None: + self.last_envelope = envelope + + +def reject_request( + valid_token: UserAPIKeyAuth, + user_obj: LiteLLM_UserTable, + general_settings: Mapping[str, str], + data: Mapping[str, Mapping[str, str]], + raw_headers: Mapping[str, str], +) -> None: + raise RuntimeError(f"key {valid_token.token} owned by {user_obj.user_email} was rejected") + + +def raise_with_identity_locals() -> None: + reject_request( + valid_token=UserAPIKeyAuth(token=KEY_HASH, key_name="sk-...test", user_id=EMAIL, user_email=EMAIL), + user_obj=LiteLLM_UserTable(user_id=EMAIL, user_email=EMAIL, user_role="internal_user"), + general_settings={"master_key": MASTER_KEY, "database_url": DATABASE_URL}, + data={"metadata": {"user_api_key_hash": KEY_HASH, "user_api_key_user_email": EMAIL}}, + raw_headers={"authorization": f"Bearer {VIRTUAL_KEY}", "x-api-key": VIRTUAL_KEY, "content-type": "application/json"}, + ) + + +def raise_with_source_context_named_locals() -> None: + metadata: Final = {"context_line": f"Bearer {VIRTUAL_KEY}", "pre_context": [EMAIL], "post_context": [KEY_HASH]} + stacktrace: Final = {"frames": [{"context_line": MASTER_KEY, "pre_context": [EMAIL]}]} + raise RuntimeError(f"rejected with {len(metadata)} metadata fields and {len(stacktrace)} stack fields") + + +def capture_serialized_event(env: Mapping[str, str], raiser: Callable[[], None] = raise_with_identity_locals) -> str: + transport: Final = RecordingTransport() + client: Final = sentry_sdk.Client(transport=transport, **build_sentry_init_options(env)) + try: + raiser() + except RuntimeError as error: + event, hint = event_from_exception(error, client_options=client.options) + client.capture_event(event, hint=hint) + assert transport.last_envelope is not None + return json.dumps(transport.last_envelope.items[0].payload.json) + + +def innermost_frame_vars(serialized: str) -> dict[str, JsonValue]: + event: Final = json.loads(serialized) + frames: Final = event["exception"]["values"][0]["stacktrace"]["frames"] + return frames[-1]["vars"] + + +def test_default_event_carries_no_email_hash_or_secret_anywhere() -> None: + serialized: Final = capture_serialized_event(PII_OFF) + assert EMAIL not in serialized + assert KEY_HASH not in serialized + assert MASTER_KEY not in serialized + assert VIRTUAL_KEY not in serialized + assert "db-password-under-test" not in serialized + frame_vars: Final = innermost_frame_vars(serialized) + assert frame_vars["raw_headers"] == {"authorization": FILTERED, "x-api-key": FILTERED, "content-type": "'application/json'"} + assert f"token='{FILTERED}'" in frame_vars["valid_token"] + assert f"user_id='{FILTERED}'" in frame_vars["valid_token"] + assert f"user_email='{FILTERED}'" in frame_vars["user_obj"] + assert frame_vars["general_settings"] == {"master_key": FILTERED, "database_url": FILTERED} + assert frame_vars["data"] == {"metadata": {"user_api_key_hash": FILTERED, "user_api_key_user_email": FILTERED}} + assert "key_name='sk-...test'" in frame_vars["valid_token"] + assert "user_role='internal_user'" in frame_vars["user_obj"] + + +def test_source_context_lines_are_left_readable() -> None: + frames: Final = json.loads(capture_serialized_event(PII_OFF))["exception"]["values"][0]["stacktrace"]["frames"] + source_lines: Final = tuple( + line + for frame in frames + for line in (*frame.get("pre_context", []), frame.get("context_line", ""), *frame.get("post_context", [])) + ) + assert any("token=KEY_HASH" in line for line in source_lines) + assert not any(FILTERED in line for line in source_lines) + + +def test_source_context_names_outside_stack_frames_are_scrubbed() -> None: + serialized: Final = capture_serialized_event(PII_OFF, raise_with_source_context_named_locals) + assert VIRTUAL_KEY not in serialized + assert MASTER_KEY not in serialized + assert EMAIL not in serialized + assert KEY_HASH not in serialized + frame_vars: Final = innermost_frame_vars(serialized) + assert frame_vars["metadata"] == { + "context_line": f"'Bearer {FILTERED}'", + "pre_context": [f"'{FILTERED}'"], + "post_context": [f"'{FILTERED}'"], + } + assert frame_vars["stacktrace"] == {"frames": [{"context_line": f"'{FILTERED}'", "pre_context": [f"'{FILTERED}'"]}]} + innermost_frame: Final = json.loads(serialized)["exception"]["values"][0]["stacktrace"]["frames"][-1] + assert "raise RuntimeError" in innermost_frame["context_line"] + assert FILTERED not in json.dumps(innermost_frame["pre_context"]) + + +def test_default_event_keeps_the_exception_message_shape() -> None: + serialized: Final = capture_serialized_event(PII_OFF) + message: Final = json.loads(serialized)["exception"]["values"][0]["value"] + assert message == f"key {FILTERED} owned by {FILTERED} was rejected" + + +def test_pii_opt_in_keeps_identifiers_and_still_scrubs_secrets() -> None: + serialized: Final = capture_serialized_event(PII_ON) + frame_vars: Final = innermost_frame_vars(serialized) + assert f"user_id='{EMAIL}'" in frame_vars["valid_token"] + assert f"user_email='{EMAIL}'" in frame_vars["user_obj"] + assert frame_vars["data"] == { + "metadata": {"user_api_key_hash": f"'{KEY_HASH}'", "user_api_key_user_email": f"'{EMAIL}'"} + } + assert f"token='{FILTERED}'" in frame_vars["valid_token"] + assert frame_vars["general_settings"] == {"master_key": FILTERED, "database_url": FILTERED} + assert frame_vars["raw_headers"] == {"authorization": FILTERED, "x-api-key": FILTERED, "content-type": "'application/json'"} + assert MASTER_KEY not in serialized + assert VIRTUAL_KEY not in serialized + assert "db-password-under-test" not in serialized + + +def test_transaction_events_are_scrubbed_too() -> None: + transport: Final = RecordingTransport() + client: Final = sentry_sdk.Client(transport=transport, **build_sentry_init_options(PII_OFF)) + client.capture_event( + { + "type": "transaction", + "transaction": "/user/info", + "contexts": {"trace": {"trace_id": "a" * 32, "span_id": "b" * 16}}, + "spans": [{"description": f"lookup {EMAIL} by {KEY_HASH}", "span_id": "c" * 16, "trace_id": "a" * 32}], + } + ) + assert transport.last_envelope is not None + serialized: Final = json.dumps(transport.last_envelope.items[0].payload.json) + assert EMAIL not in serialized + assert KEY_HASH not in serialized + assert f"lookup {FILTERED} by {FILTERED}" in serialized + + +@pytest.mark.parametrize( + ("text", "expected"), + [ + ( + "UserAPIKeyAuth(token='abc', key_alias='team-a', user_id=None)", + f"UserAPIKeyAuth(token='{FILTERED}', key_alias='team-a', user_id=None)", + ), + ('{"api_key": "sk-1", "model": "gpt-5"}', f'{{"api_key": "{FILTERED}", "model": "gpt-5"}}'), + ("{'user_id': 'u-1', 'max_budget': 5}", f"{{'user_id': '{FILTERED}', 'max_budget': 5}}"), + ("Config(OPENAI_API_KEY=sk-live, timeout=10)", f"Config(OPENAI_API_KEY='{FILTERED}', timeout=10)"), + ("lookup for somebody@example.com failed", f"lookup for {FILTERED} failed"), + (f"hashed key {KEY_HASH} not found", f"hashed key {FILTERED} not found"), + ("request id 0123456789abcdef0123456789abcdef stays", "request id 0123456789abcdef0123456789abcdef stays"), + ("monkey=banana", "monkey=banana"), + ( + "{'x-api-key': 'k-1', 'cookie': 'session=abc', 'content-type': 'application/json'}", + f"{{'x-api-key': '{FILTERED}', 'cookie': '{FILTERED}', 'content-type': 'application/json'}}", + ), + ( + "headers={'x-tenant-key': 'sk-custom-header-key-0123456789'} key_name='sk-...6789'", + f"headers={{'x-tenant-key': '{FILTERED}'}} key_name='sk-...6789'", + ), + ( + "master_key={'value': 'not-a-litellm-key'} timeout=10", + f"master_key='{FILTERED}' timeout=10", + ), + ( + "credentials=[{'value': ('deep', 'secret')}], model='gpt-5'", + f"credentials='{FILTERED}', model='gpt-5'", + ), + ], +) +def test_string_scrubber_rewrites_field_and_value_forms(text: str, expected: str) -> None: + assert build_string_scrubber(send_default_pii=False)(text) == expected + + +def test_bare_key_floor_follows_the_custom_key_minimum() -> None: + scrub: Final = build_string_scrubber(send_default_pii=False) + shortest_key: Final = "sk-" + "a" * (MINIMUM_CUSTOM_KEY_LENGTH - len("sk-")) + assert scrub(f"label={shortest_key} model=gpt-5") == f"label={FILTERED} model=gpt-5" + assert scrub(f"label={shortest_key[:-1]} model=gpt-5") == f"label={shortest_key[:-1]} model=gpt-5" + + +def test_key_pattern_floor_never_exceeds_a_generated_key() -> None: + generated_key: Final = "sk-" + secrets.token_urlsafe(LENGTH_OF_LITELLM_GENERATED_KEY) + stricter_custom_minimum: Final = len(generated_key) + 10 + assert build_key_pattern(stricter_custom_minimum, LENGTH_OF_LITELLM_GENERATED_KEY).fullmatch(generated_key) + assert build_key_pattern(stricter_custom_minimum, LENGTH_OF_LITELLM_GENERATED_KEY).fullmatch(generated_key[:-1]) is None + + +def test_json_walk_fails_closed_past_the_depth_cap() -> None: + scrub: Final = build_string_scrubber(send_default_pii=False) + nested: Final = reduce(lambda inner, _: [inner], range(MAX_SCRUB_DEPTH + 1), cast("JsonValue", "api_key=sk-1")) + assert FILTERED in json.dumps(scrub_json_strings(nested, scrub)) + assert "sk-1" not in json.dumps(scrub_json_strings(nested, scrub)) + assert scrub_json_strings([["api_key=sk-1"]], scrub) == [[f"api_key='{FILTERED}'"]] + + +def test_string_scrubber_with_pii_on_only_scrubs_secrets() -> None: + scrub: Final = build_string_scrubber(send_default_pii=True) + assert scrub(f"user_id='{EMAIL}', token='{KEY_HASH}', email {EMAIL} hash {KEY_HASH}") == ( + f"user_id='{EMAIL}', token='{FILTERED}', email {EMAIL} hash {KEY_HASH}" + ) + assert scrub(f"headers={{'authorization': 'Bearer {VIRTUAL_KEY}'}} sent {VIRTUAL_KEY}") == ( + f"headers={{'authorization': '{FILTERED}'}} sent {FILTERED}" + ) + + +@pytest.mark.parametrize( + ("env", "expected"), + [ + ({}, False), + ({"SENTRY_SEND_DEFAULT_PII": "true"}, True), + ({"SENTRY_SEND_DEFAULT_PII": "True"}, True), + ({"SENTRY_SEND_DEFAULT_PII": "false"}, False), + ({"SENTRY_SEND_DEFAULT_PII": "yes please"}, False), + ], +) +def test_send_default_pii_comes_from_the_environment(env: Mapping[str, str], expected: bool) -> None: + assert build_sentry_init_options(env)["send_default_pii"] is expected + + +def test_init_options_read_dsn_rates_and_environment() -> None: + options: Final = build_sentry_init_options( + { + "SENTRY_DSN": "https://key@sentry.example/7", + "SENTRY_API_TRACE_RATE": "0.25", + "SENTRY_API_SAMPLE_RATE": "0.5", + "SENTRY_ENVIRONMENT": "staging", + } + ) + assert options["dsn"] == "https://key@sentry.example/7" + assert options["traces_sample_rate"] == 0.25 + assert options["sample_rate"] == 0.5 + assert options["environment"] == "staging" + assert options["event_scrubber"].recursive is True + + +def test_init_options_defaults() -> None: + options: Final = build_sentry_init_options({}) + assert options["dsn"] is None + assert options["traces_sample_rate"] == 1.0 + assert options["sample_rate"] == 1.0 + assert options["environment"] == "production" diff --git a/tests/test_litellm/llms/cometapi/chat/test_cometapi_chat_transformation.py b/tests/test_litellm/llms/cometapi/chat/test_cometapi_chat_transformation.py index 7a69b676667..f692259db2e 100644 --- a/tests/test_litellm/llms/cometapi/chat/test_cometapi_chat_transformation.py +++ b/tests/test_litellm/llms/cometapi/chat/test_cometapi_chat_transformation.py @@ -9,171 +9,6 @@ import os import pytest -from litellm.llms.cometapi.chat.transformation import ( - CometAPIChatCompletionStreamingHandler, - CometAPIConfig, -) -from litellm.llms.cometapi.common_utils import CometAPIException - - -class TestCometAPIChatCompletionStreamingHandler: - def test_chunk_parser_successful(self): - handler = CometAPIChatCompletionStreamingHandler( - streaming_response=None, sync_stream=True - ) - - # Test input chunk - chunk = { - "id": "test_id", - "created": 1234567890, - "model": "gpt-3.5-turbo", - "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, - "choices": [ - {"delta": {"content": "test content", "reasoning": "test reasoning"}} - ], - } - - # Parse chunk - result = handler.chunk_parser(chunk) - - # Verify response - assert result.id == "test_id" - assert result.object == "chat.completion.chunk" - assert result.created == 1234567890 - assert result.model == "gpt-3.5-turbo" - assert result.usage.prompt_tokens == chunk["usage"]["prompt_tokens"] - assert result.usage.completion_tokens == chunk["usage"]["completion_tokens"] - assert result.usage.total_tokens == chunk["usage"]["total_tokens"] - assert len(result.choices) == 1 - assert result.choices[0]["delta"]["reasoning_content"] == "test reasoning" - - def test_chunk_parser_error_response(self): - handler = CometAPIChatCompletionStreamingHandler( - streaming_response=None, sync_stream=True - ) - - # Test error chunk - error_chunk = { - "error": { - "message": "test error", - "code": 400, - } - } - - # Verify error handling - with pytest.raises(CometAPIException) as exc_info: - handler.chunk_parser(error_chunk) - - assert "CometAPI Error: test error" in str(exc_info.value) - assert exc_info.value.status_code == 400 - - def test_chunk_parser_key_error(self): - handler = CometAPIChatCompletionStreamingHandler( - streaming_response=None, sync_stream=True - ) - - # Test invalid chunk missing required fields - invalid_chunk = {"incomplete": "data"} - - # Verify KeyError handling - with pytest.raises(CometAPIException) as exc_info: - handler.chunk_parser(invalid_chunk) - - assert "KeyError" in str(exc_info.value) - assert exc_info.value.status_code == 400 - - -class TestCometAPIConfig: - def test_transform_request_basic(self): - """Test basic request transformation""" - config = CometAPIConfig() - - transformed_request = config.transform_request( - model="cometapi/gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hello, world!"}], - optional_params={}, - litellm_params={}, - headers={}, - ) - - assert transformed_request["model"] == "cometapi/gpt-3.5-turbo" - assert transformed_request["messages"] == [ - {"role": "user", "content": "Hello, world!"} - ] - - def test_transform_request_with_extra_body(self): - """Test request transformation with extra_body parameters""" - config = CometAPIConfig() - - transformed_request = config.transform_request( - model="cometapi/gpt-4", - messages=[{"role": "user", "content": "Hello, world!"}], - optional_params={"extra_body": {"custom_param": "custom_value"}}, - litellm_params={}, - headers={}, - ) - - # Validate that extra_body parameters are merged into the request - assert transformed_request["custom_param"] == "custom_value" - assert transformed_request["messages"] == [ - {"role": "user", "content": "Hello, world!"} - ] - - def test_cache_control_flag_removal(self): - """Test cache control flag removal from messages""" - config = CometAPIConfig() - - transformed_request = config.transform_request( - model="cometapi/gpt-3.5-turbo", - messages=[ - { - "role": "user", - "content": "Hello, world!", - "cache_control": {"type": "ephemeral"}, - } - ], - optional_params={}, - litellm_params={}, - headers={}, - ) - - # CometAPI should remove cache_control flags by default - assert transformed_request["messages"][0].get("cache_control") is None - - def test_map_openai_params(self): - """Test OpenAI parameter mapping""" - config = CometAPIConfig() - - non_default_params = { - "temperature": 0.7, - "max_tokens": 100, - "top_p": 0.9, - } - - mapped_params = config.map_openai_params( - non_default_params=non_default_params, - optional_params={}, - model="cometapi/gpt-3.5-turbo", - drop_params=False, - ) - - assert mapped_params["temperature"] == 0.7 - assert mapped_params["max_tokens"] == 100 - assert mapped_params["top_p"] == 0.9 - - def test_get_error_class(self): - """Test error class creation""" - config = CometAPIConfig() - - error = config.get_error_class( - error_message="Test error", - status_code=400, - headers={"Content-Type": "application/json"}, - ) - - assert isinstance(error, CometAPIException) - assert error.message == "Test error" - assert error.status_code == 400 # Integration test example (requires real API key) diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py deleted file mode 100644 index a3391a2c585..00000000000 --- a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py +++ /dev/null @@ -1,79 +0,0 @@ -import json -from typing import Final - -import httpx -import respx - -import litellm - - -def test_completion_merges_leading_system_and_developer_messages_for_chat_template_models( - respx_mock: respx.MockRouter, -): - upstream: Final = respx_mock.post("https://example.databricks.test/serving-endpoints/chat/completions").mock( - return_value=httpx.Response( - status_code=200, - json={ - "id": "chatcmpl-123", - "object": "chat.completion", - "created": 1677652288, - "model": "my-custom-model", - "choices": [{"index": 0, "message": {"role": "assistant", "content": "Answer"}, "finish_reason": "stop"}], - "usage": {"prompt_tokens": 9, "completion_tokens": 1, "total_tokens": 10}, - }, - ) - ) - - response: Final = litellm.completion( - model="databricks/my-custom-model", - messages=[ - {"role": "system", "content": "You are terse."}, - {"role": "developer", "content": "Skills: none."}, - {"role": "user", "content": "Hello"}, - ], - api_base="https://example.databricks.test/serving-endpoints", - api_key="fake-databricks-api-key", - num_retries=0, - ) - - assert upstream.call_count == 1 - request_body: Final = json.loads(upstream.calls[0].request.read()) - assert request_body["messages"] == [ - {"role": "system", "content": "You are terse.\n\nSkills: none."}, - {"role": "user", "content": "Hello"}, - ] - assert response.choices[0].message.content == "Answer" - - -def test_completion_merges_system_messages_when_one_has_empty_content(respx_mock: respx.MockRouter): - upstream: Final = respx_mock.post("https://example.databricks.test/serving-endpoints/chat/completions").mock( - return_value=httpx.Response( - status_code=200, - json={ - "id": "chatcmpl-123", - "object": "chat.completion", - "created": 1677652288, - "model": "my-custom-model", - "choices": [{"index": 0, "message": {"role": "assistant", "content": "Answer"}, "finish_reason": "stop"}], - "usage": {"prompt_tokens": 9, "completion_tokens": 1, "total_tokens": 10}, - }, - ) - ) - - litellm.completion( - model="databricks/my-custom-model", - messages=[ - {"role": "system", "content": "You are terse."}, - {"role": "system", "content": ""}, - {"role": "user", "content": "Hello"}, - ], - api_base="https://example.databricks.test/serving-endpoints", - api_key="fake-databricks-api-key", - num_retries=0, - ) - - request_body: Final = json.loads(upstream.calls[0].request.read()) - assert request_body["messages"] == [ - {"role": "system", "content": "You are terse."}, - {"role": "user", "content": "Hello"}, - ] diff --git a/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_integration.py b/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_integration.py deleted file mode 100644 index 5b013681864..00000000000 --- a/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_integration.py +++ /dev/null @@ -1,433 +0,0 @@ -""" -Integration tests for DeepInfra rerank functionality. -Tests the full rerank flow following the repository patterns. -""" - -import asyncio -import json -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -import litellm - - -def assert_response_shape(response, custom_llm_provider): - """Helper function to validate response structure specific to DeepInfra.""" - assert hasattr(response, "id") - assert hasattr(response, "results") - assert hasattr(response, "meta") - assert isinstance(response.results, list) - - for result in response.results: - assert "index" in result - assert "relevance_score" in result - assert isinstance(result["index"], int) - assert isinstance(result["relevance_score"], (int, float)) - - # Check meta structure - assert "tokens" in response.meta - assert "billed_units" in response.meta - assert "input_tokens" in response.meta["tokens"] - assert "total_tokens" in response.meta["billed_units"] - - -@pytest.mark.parametrize("sync_mode", [True, False]) -@patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post") -@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") -def test_basic_rerank_deepinfra(mock_sync_post, mock_async_post, sync_mode): - """Test basic DeepInfra rerank functionality.""" - # Mock response data that matches DeepInfra API format - mock_response_data = { - "scores": [0.9, 0.1], - "input_tokens": 25, - "request_id": "deepinfra-request-123", - "inference_status": { - "status": "success", - "runtime_ms": 150, - "cost": 0.0001, - "tokens_generated": 0, - "tokens_input": 25, - }, - } - - def return_val(): - return mock_response_data - - api_key = "test_deepinfra_api_key" - api_base = "https://api.deepinfra.com" - - if sync_mode: - # Create mock response object for sync - mock_response = MagicMock() - mock_response.json = return_val - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.text = json.dumps(mock_response_data) - mock_sync_post.return_value = mock_response - - response = litellm.rerank( - model="deepinfra/Qwen/Qwen3-Reranker-0.6B", - query="hello", - documents=["hello", "world"], - top_n=2, - custom_llm_provider="deepinfra", - api_key=api_key, - api_base=api_base, - ) - mock_sync_post.assert_called_once() - else: - # Create mock response object for async - mock_response = AsyncMock() - - def return_val(): - return mock_response_data - - mock_response.json = return_val - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.text = json.dumps(mock_response_data) - mock_async_post.return_value = mock_response - - response = asyncio.run( - litellm.arerank( - model="deepinfra/Qwen/Qwen3-Reranker-0.6B", - query="hello", - documents=["hello", "world"], - top_n=2, - custom_llm_provider="deepinfra", - api_key=api_key, - api_base=api_base, - ) - ) - mock_async_post.assert_called_once() - - # Verify response structure - assert response.id == "deepinfra-request-123" - assert response.results is not None - assert len(response.results) == 2 - assert response.results[0]["index"] == 0 - assert response.results[0]["relevance_score"] == 0.9 - assert response.results[1]["index"] == 1 - assert response.results[1]["relevance_score"] == 0.1 - - # Verify metadata - assert response.meta["tokens"]["input_tokens"] == 25 - assert response.meta["billed_units"]["total_tokens"] == 25 - - # Verify hidden params specific to DeepInfra - assert response._hidden_params["status"] == "success" - assert response._hidden_params["runtime_ms"] == 150 - assert response._hidden_params["cost"] == 0.0001 - # Note: The model name is processed and the 'deepinfra/' prefix is removed - assert response._hidden_params["model"] == "Qwen/Qwen3-Reranker-0.6B" - - assert_response_shape(response, custom_llm_provider="deepinfra") - - -@pytest.mark.parametrize("sync_mode", [True, False]) -@patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post") -@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") -def test_deepinfra_rerank_with_queries_param( - mock_sync_post, mock_async_post, sync_mode -): - """Test DeepInfra rerank with multiple queries parameter.""" - mock_response_data = { - "scores": [0.8, 0.6, 0.2], - "input_tokens": 35, - "request_id": "deepinfra-multi-query-123", - "inference_status": {"status": "success", "runtime_ms": 200}, - } - - def return_val(): - return mock_response_data - - if sync_mode: - mock_response = MagicMock() - mock_response.json = return_val - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.text = json.dumps(mock_response_data) - mock_sync_post.return_value = mock_response - - response = litellm.rerank( - model="deepinfra/Qwen/Qwen3-Reranker-4B", - query="hello", - documents=["hello", "world", "test"], - queries=["hello", "hi there"], # DeepInfra specific param - custom_llm_provider="deepinfra", - api_key="test_key", - api_base="https://api.deepinfra.com", - ) - - mock_sync_post.assert_called_once() - # Verify that queries parameter was passed in request - call_data = json.loads(mock_sync_post.call_args.kwargs["data"]) - assert "queries" in call_data - assert call_data["queries"] == ["hello", "hi there"] - else: - mock_response = AsyncMock() - mock_response.json = return_val - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.text = json.dumps(mock_response_data) - mock_async_post.return_value = mock_response - - response = asyncio.run( - litellm.arerank( - model="deepinfra/Qwen/Qwen3-Reranker-4B", - query="hello", - documents=["hello", "world", "test"], - queries=["hello", "hi there"], - custom_llm_provider="deepinfra", - api_key="test_key", - api_base="https://api.deepinfra.com", - ) - ) - - mock_async_post.assert_called_once() - call_data = json.loads(mock_async_post.call_args.kwargs["data"]) - assert "queries" in call_data - assert call_data["queries"] == ["hello", "hi there"] - - assert response.results is not None - assert len(response.results) == 3 - - -@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") -def test_deepinfra_rerank_with_service_tier(mock_post): - """Test DeepInfra rerank with service_tier parameter.""" - mock_response_data = { - "scores": [0.95, 0.75], - "input_tokens": 30, - "request_id": "deepinfra-premium-123", - } - - def return_val(): - return mock_response_data - - mock_response = MagicMock() - mock_response.json = return_val - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.text = json.dumps(mock_response_data) - mock_post.return_value = mock_response - - response = litellm.rerank( - model="deepinfra/Qwen/Qwen3-Reranker-8B", - query="premium search", - documents=["doc1", "doc2"], - service_tier="premium", # DeepInfra specific param - custom_llm_provider="deepinfra", - api_key="test_key", - api_base="https://api.deepinfra.com", - ) - - mock_post.assert_called_once() - - # Verify URL - call_url = mock_post.call_args.kwargs["url"] - assert "api.deepinfra.com/inference/Qwen/Qwen3-Reranker-8B" in call_url - - # Verify request contains service_tier - call_data = json.loads(mock_post.call_args.kwargs["data"]) - assert call_data["service_tier"] == "premium" - - assert response.results is not None - - -@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") -def test_deepinfra_rerank_with_env_vars(mock_post, monkeypatch): - """Test DeepInfra rerank with environment variable configuration.""" - monkeypatch.setenv("DEEPINFRA_API_KEY", "env_test_key") - monkeypatch.setenv("DEEPINFRA_API_BASE", "https://custom-deepinfra.com") - - mock_response_data = { - "scores": [0.88, 0.22], - "input_tokens": 28, - "request_id": "env-test-123", - } - - def return_val(): - return mock_response_data - - mock_response = MagicMock() - mock_response.json = return_val - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.text = json.dumps(mock_response_data) - mock_post.return_value = mock_response - - response = litellm.rerank( - model="deepinfra/Qwen/Qwen3-Reranker-0.6B", - query="hello", - documents=["hello", "world"], - custom_llm_provider="deepinfra", - ) - - mock_post.assert_called_once() - - # Verify headers contain env API key - headers = mock_post.call_args.kwargs.get("headers", {}) - assert "Bearer env_test_key" in headers.get("Authorization", "") - - assert response.results is not None - - -@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") -def test_deepinfra_rerank_error_handling(mock_post): - """Test DeepInfra rerank error handling.""" - error_response = {"detail": {"error": "Invalid API key"}} - - def return_val(): - return error_response - - mock_response = MagicMock() - mock_response.status_code = 401 - mock_response.json = return_val - mock_response.text = json.dumps(error_response) - mock_response.headers = {"content-type": "application/json"} - mock_post.return_value = mock_response - - # The current implementation handles errors gracefully, so we expect a successful response - # with the error information in the hidden params - response = litellm.rerank( - model="deepinfra/Qwen/Qwen3-Reranker-0.6B", - query="hello", - documents=["hello", "world"], - custom_llm_provider="deepinfra", - api_key="invalid_key", - api_base="https://api.deepinfra.com", - ) - - # Verify that the response contains error information - assert ( - response._hidden_params["status"] == "unknown" - ) # Default status when error occurs - - -@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") -def test_deepinfra_rerank_defaults_api_base_when_missing(mock_post, monkeypatch): - """With no api_base anywhere, the call still goes out against DeepInfra's own base.""" - monkeypatch.delenv("DEEPINFRA_API_BASE", raising=False) - - mock_response = MagicMock() - mock_response.json = lambda: {"scores": [0.9, 0.1], "input_tokens": 20} - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_post.return_value = mock_response - - response = litellm.rerank( - model="deepinfra/Qwen/Qwen3-Reranker-0.6B", - query="hello", - documents=["hello", "world"], - custom_llm_provider="deepinfra", - api_key="test_key", - # api_base is intentionally missing - ) - - assert "api.deepinfra.com" in mock_post.call_args.kwargs["url"] - assert [result["relevance_score"] for result in response.results] == [0.9, 0.1] - - -@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") -def test_deepinfra_rerank_request_format(mock_post): - """Test that the request is properly formatted for DeepInfra API.""" - mock_response_data = {"scores": [0.9, 0.1], "input_tokens": 20} - - def return_val(): - return mock_response_data - - mock_response = MagicMock() - mock_response.json = return_val - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.text = json.dumps(mock_response_data) - mock_post.return_value = mock_response - - response = litellm.rerank( - model="deepinfra/Qwen/Qwen3-Reranker-0.6B", - query="test query", - documents=["doc1", "doc2"], - custom_llm_provider="deepinfra", - api_key="test_key", - api_base="https://api.deepinfra.com", - instruction="custom instruction", - webhook="https://webhook.example.com", - ) - - mock_post.assert_called_once() - - # Verify URL format - call_url = mock_post.call_args.kwargs["url"] - assert call_url == "https://api.deepinfra.com/inference/Qwen/Qwen3-Reranker-0.6B" - - # Verify headers - headers = mock_post.call_args.kwargs["headers"] - assert headers["Authorization"] == "Bearer test_key" - assert headers["accept"] == "application/json" - assert headers["content-type"] == "application/json" - - # Verify request body format - request_data = json.loads(mock_post.call_args.kwargs["data"]) - assert request_data["queries"] == [ - "test query", - "test query", - ] # DeepInfra requires queries to match documents length - assert request_data["documents"] == ["doc1", "doc2"] - assert request_data["instruction"] == "custom instruction" - assert request_data["webhook"] == "https://webhook.example.com" - - assert response.results is not None - - -def test_deepinfra_rerank_models(): - """Test that DeepInfra Qwen rerank models are recognized.""" - # These should not raise errors during model validation - models = [ - "deepinfra/Qwen/Qwen3-Reranker-0.6B", - "deepinfra/Qwen/Qwen3-Reranker-4B", - "deepinfra/Qwen/Qwen3-Reranker-8B", - ] - - for model in models: - resolved_model, provider, _, api_base = litellm.get_llm_provider(model=model) - assert provider == "deepinfra" - assert resolved_model == model.removeprefix("deepinfra/") - assert api_base == "https://api.deepinfra.com/v1/openai" - - -@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") -def test_deepinfra_rerank_minimal_response(mock_post): - """Test handling of minimal DeepInfra response.""" - # Minimal response with just scores - mock_response_data = {"scores": [0.7, 0.3]} - - def return_val(): - return mock_response_data - - mock_response = MagicMock() - mock_response.json = return_val - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.text = json.dumps(mock_response_data) - mock_post.return_value = mock_response - - response = litellm.rerank( - model="deepinfra/Qwen/Qwen3-Reranker-0.6B", - query="hello", - documents=["hello", "world"], - custom_llm_provider="deepinfra", - api_key="test_key", - api_base="https://api.deepinfra.com", - ) - - # Should handle minimal response gracefully - assert response.results is not None - assert len(response.results) == 2 - assert response.results[0]["relevance_score"] == 0.7 - assert response.results[1]["relevance_score"] == 0.3 - - # Should have default values for missing fields - assert response.meta["tokens"]["input_tokens"] == 0 # Default when missing - assert response._hidden_params["status"] == "unknown" # Default when missing diff --git a/tests/test_litellm/llms/gemini/files/__init__.py b/tests/test_litellm/llms/gemini/files/__init__.py deleted file mode 100644 index f48fe7dbe2b..00000000000 --- a/tests/test_litellm/llms/gemini/files/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Tests for Gemini files functionality""" diff --git a/tests/test_litellm/llms/gemini/videos/__init__.py b/tests/test_litellm/llms/gemini/videos/__init__.py deleted file mode 100644 index e0780c08321..00000000000 --- a/tests/test_litellm/llms/gemini/videos/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Gemini Video Generation Tests diff --git a/tests/test_litellm/llms/manus/__init__.py b/tests/test_litellm/llms/manus/__init__.py deleted file mode 100644 index c9121a7b2a4..00000000000 --- a/tests/test_litellm/llms/manus/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Manus provider tests diff --git a/tests/test_litellm/llms/manus/responses/__init__.py b/tests/test_litellm/llms/manus/responses/__init__.py deleted file mode 100644 index ea7ebb64d55..00000000000 --- a/tests/test_litellm/llms/manus/responses/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Manus Responses API tests diff --git a/tests/test_litellm/llms/minimax/__init__.py b/tests/test_litellm/llms/minimax/__init__.py deleted file mode 100644 index 451f542f4ad..00000000000 --- a/tests/test_litellm/llms/minimax/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# MiniMax tests diff --git a/tests/test_litellm/llms/minimax/chat/__init__.py b/tests/test_litellm/llms/minimax/chat/__init__.py deleted file mode 100644 index 4a7916ae6cf..00000000000 --- a/tests/test_litellm/llms/minimax/chat/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# MiniMax chat tests diff --git a/tests/test_litellm/llms/minimax/messages/__init__.py b/tests/test_litellm/llms/minimax/messages/__init__.py deleted file mode 100644 index de5a80602ea..00000000000 --- a/tests/test_litellm/llms/minimax/messages/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# MiniMax messages tests diff --git a/tests/test_litellm/llms/mistral/audio_transcription/test_mistral_audio_transcription_transformation.py b/tests/test_litellm/llms/mistral/audio_transcription/test_mistral_audio_transcription_transformation.py index d1eb6241ceb..db77eabba23 100644 --- a/tests/test_litellm/llms/mistral/audio_transcription/test_mistral_audio_transcription_transformation.py +++ b/tests/test_litellm/llms/mistral/audio_transcription/test_mistral_audio_transcription_transformation.py @@ -1,19 +1,9 @@ import os from typing import Dict -from unittest.mock import MagicMock -import httpx import litellm import pytest -from litellm.llms.base_llm.audio_transcription.transformation import ( - BaseAudioTranscriptionConfig, -) -from litellm.llms.mistral.audio_transcription.transformation import ( - MistralAudioTranscriptionConfig, -) -from litellm.types.utils import TranscriptionResponse -from litellm.utils import ProviderConfigManager from tests.llm_translation.base_audio_transcription_unit_tests import ( BaseLLMAudioTranscriptionTest, ) @@ -37,184 +27,3 @@ class TestMistralAudioTranscription(BaseLLMAudioTranscriptionTest): "Async audio transcription test for Mistral is skipped in this suite; " "async test plugins (e.g. pytest-asyncio/anyio) are not configured here." ) - - -def test_mistral_audio_transcription_config_installed(): - """Ensure Mistral audio transcription config is registered with ProviderConfigManager.""" - config = ProviderConfigManager.get_provider_audio_transcription_config( - model="mistral/voxtral-mini-latest", - provider=litellm.LlmProviders.MISTRAL, - ) - assert config is not None - assert isinstance(config, BaseAudioTranscriptionConfig) - assert isinstance(config, MistralAudioTranscriptionConfig) - - -def test_mistral_audio_transcription_get_complete_url(): - config = MistralAudioTranscriptionConfig() - url = config.get_complete_url( - api_base=None, - api_key="fake-key", - model="voxtral-mini-latest", - optional_params={}, - litellm_params={}, - ) - assert url == "https://api.mistral.ai/v1/audio/transcriptions" - - -def test_mistral_audio_transcription_get_complete_url_custom_base(): - config = MistralAudioTranscriptionConfig() - url = config.get_complete_url( - api_base="https://custom.api.example.com/v1/", - api_key="fake-key", - model="voxtral-mini-latest", - optional_params={}, - litellm_params={}, - ) - assert url == "https://custom.api.example.com/v1/audio/transcriptions" - - -def test_mistral_audio_transcription_validate_environment(): - config = MistralAudioTranscriptionConfig() - headers = config.validate_environment( - headers={}, - model="voxtral-mini-latest", - messages=[], - optional_params={}, - litellm_params={}, - api_key="test-key-123", - ) - assert headers["Authorization"] == "Bearer test-key-123" - assert headers["accept"] == "application/json" - - -def test_mistral_audio_transcription_supported_params(): - config = MistralAudioTranscriptionConfig() - params = config.get_supported_openai_params("voxtral-mini-latest") - assert "language" in params - assert "temperature" in params - assert "response_format" in params - assert "timestamp_granularities" in params - - -def test_mistral_audio_transcription_request_transform(): - config = MistralAudioTranscriptionConfig() - - wav_path = os.path.join( - os.path.dirname(__file__), - "../../../../..", - "tests", - "llm_translation", - "gettysburg.wav", - ) - audio_file = open(wav_path, "rb") - - result = config.transform_audio_transcription_request( - model="voxtral-mini-latest", - audio_file=audio_file, - optional_params={"language": "en", "temperature": 0.0}, - litellm_params={}, - ) - - audio_file.close() - - assert isinstance(result.data, dict) - assert result.data["model"] == "voxtral-mini-latest" - assert result.data["language"] == "en" - assert result.data["temperature"] == 0.0 - assert result.files is not None - assert "file" in result.files - - -def test_mistral_audio_transcription_request_with_diarize(): - """Test that Mistral-specific params like diarize are passed through.""" - config = MistralAudioTranscriptionConfig() - - wav_path = os.path.join( - os.path.dirname(__file__), - "../../../../..", - "tests", - "llm_translation", - "gettysburg.wav", - ) - audio_file = open(wav_path, "rb") - - result = config.transform_audio_transcription_request( - model="voxtral-mini-latest", - audio_file=audio_file, - optional_params={"diarize": True}, - litellm_params={}, - ) - - audio_file.close() - - assert isinstance(result.data, dict) - assert result.data["diarize"] == "true" - - -def test_mistral_audio_transcription_response_transform(): - config = MistralAudioTranscriptionConfig() - - mock_response = MagicMock(spec=httpx.Response) - mock_response.json.return_value = {"text": "Four score and seven years ago..."} - - response = config.transform_audio_transcription_response(mock_response) - - assert isinstance(response, TranscriptionResponse) - assert response.text == "Four score and seven years ago..." - - -def test_mistral_audio_transcription_response_transform_diarized(): - """Test that diarized responses preserve segments and language.""" - config = MistralAudioTranscriptionConfig() - - mock_response = MagicMock(spec=httpx.Response) - mock_response.json.return_value = { - "model": "voxtral-mini-latest", - "text": "Hello, how are you? I am fine.", - "language": None, - "segments": [ - { - "text": "Hello, how are you?", - "start": 0.3, - "end": 2.1, - "speaker_id": "speaker_1", - "type": "transcription_segment", - }, - { - "text": "I am fine.", - "start": 2.5, - "end": 3.8, - "speaker_id": "speaker_2", - "type": "transcription_segment", - }, - ], - "usage": { - "prompt_audio_seconds": 4, - "prompt_tokens": 5, - "total_tokens": 50, - "completion_tokens": 20, - }, - } - - response = config.transform_audio_transcription_response(mock_response) - - assert isinstance(response, TranscriptionResponse) - assert response.text == "Hello, how are you? I am fine." - assert response["segments"] is not None - assert len(response["segments"]) == 2 - assert response["segments"][0]["speaker_id"] == "speaker_1" - assert response["segments"][1]["speaker_id"] == "speaker_2" - assert response["language"] is None - - -def test_mistral_audio_transcription_response_transform_empty(): - config = MistralAudioTranscriptionConfig() - - mock_response = MagicMock(spec=httpx.Response) - mock_response.json.return_value = {} - - response = config.transform_audio_transcription_response(mock_response) - - assert isinstance(response, TranscriptionResponse) - assert response.text == "" diff --git a/tests/test_litellm/llms/openai_like/test_json_providers.py b/tests/test_litellm/llms/openai_like/test_json_providers.py index d84cc8d3237..55703063fae 100644 --- a/tests/test_litellm/llms/openai_like/test_json_providers.py +++ b/tests/test_litellm/llms/openai_like/test_json_providers.py @@ -3,321 +3,12 @@ Tests for JSON-based provider configuration system. """ import os -import sys -from unittest.mock import patch -try: - import pytest -except ImportError: - # pytest not available, will run as standalone script - pytest = None - -# Add workspace to path -workspace_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) -sys.path.insert(0, workspace_path) +import pytest import litellm -class TestJSONProviderLoader: - """Test JSON provider loading and configuration""" - - def test_load_json_providers(self): - """Test that JSON providers load correctly""" - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - # Verify publicai is loaded - assert JSONProviderRegistry.exists("publicai") - - # Get publicai config - publicai = JSONProviderRegistry.get("publicai") - assert publicai is not None - assert publicai.base_url == "https://api.publicai.co/v1" - assert publicai.api_key_env == "PUBLICAI_API_KEY" - assert publicai.api_base_env == "PUBLICAI_API_BASE" - assert publicai.param_mappings.get("max_completion_tokens") == "max_tokens" - - def test_dynamic_config_generation(self): - """Test dynamic config class creation""" - from litellm.llms.openai_like.dynamic_config import create_config_class - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - provider = JSONProviderRegistry.get("publicai") - config_class = create_config_class(provider) - config = config_class() - - # Test API info resolution - api_base, api_key = config._get_openai_compatible_provider_info(None, None) - assert api_base == "https://api.publicai.co/v1" - - # Test with custom base - api_base, api_key = config._get_openai_compatible_provider_info( - "https://custom.api.com", "test-key" - ) - assert api_base == "https://custom.api.com" - assert api_key == "test-key" - - def test_parameter_mapping(self): - """Test parameter mapping works""" - from litellm.llms.openai_like.dynamic_config import create_config_class - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - provider = JSONProviderRegistry.get("publicai") - config_class = create_config_class(provider) - config = config_class() - - # Test parameter mapping - optional_params = {} - non_default_params = {"max_completion_tokens": 100, "temperature": 0.7} - result = config.map_openai_params( - non_default_params, optional_params, "gpt-4", False - ) - - # max_completion_tokens should be mapped to max_tokens - assert "max_tokens" in result - assert result["max_tokens"] == 100 - assert "max_completion_tokens" not in result - - # temperature should be passed through - assert result["temperature"] == 0.7 - - def test_supported_params(self): - """Test that config returns supported params""" - from litellm.llms.openai_like.dynamic_config import create_config_class - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - provider = JSONProviderRegistry.get("publicai") - config_class = create_config_class(provider) - config = config_class() - - # Get supported params - supported = config.get_supported_openai_params("gpt-4") - - # Should have standard OpenAI params - assert isinstance(supported, list) - assert len(supported) > 0 - - def test_tool_params_excluded_when_function_calling_not_supported(self): - """Test that tool-related params are excluded for models that don't support - function calling. Regression test for https://github.com/BerriAI/litellm/issues/21125 - """ - from litellm.llms.openai_like.dynamic_config import create_config_class - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - provider = JSONProviderRegistry.get("publicai") - config_class = create_config_class(provider) - config = config_class() - - # Mock supports_function_calling to return False - with patch("litellm.utils.supports_function_calling", return_value=False): - supported = config.get_supported_openai_params("some-model-without-fc") - - tool_params = [ - "tools", - "tool_choice", - "function_call", - "functions", - "parallel_tool_calls", - ] - for param in tool_params: - assert ( - param not in supported - ), f"'{param}' should not be in supported params when function calling is not supported" - - # Non-tool params should still be present - assert "temperature" in supported - assert "max_tokens" in supported - assert "stop" in supported - - def test_tool_params_included_when_function_calling_supported(self): - """Test that tool-related params are included for models that support function calling.""" - from litellm.llms.openai_like.dynamic_config import create_config_class - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - provider = JSONProviderRegistry.get("publicai") - config_class = create_config_class(provider) - config = config_class() - - # Mock supports_function_calling to return True - with patch("litellm.utils.supports_function_calling", return_value=True): - supported = config.get_supported_openai_params("some-model-with-fc") - - assert "tools" in supported - assert "tool_choice" in supported - - def test_provider_resolution(self): - """Test that provider resolution finds JSON providers""" - from litellm.litellm_core_utils.get_llm_provider_logic import ( - get_llm_provider, - ) - - model, provider, api_key, api_base = get_llm_provider( - model="publicai/gpt-4", - custom_llm_provider=None, - api_base=None, - api_key=None, - ) - - assert model == "gpt-4" - assert provider == "publicai" - assert api_base == "https://api.publicai.co/v1" - - def test_provider_config_manager(self): - """Test that ProviderConfigManager returns JSON-based configs""" - from litellm import LlmProviders - from litellm.utils import ProviderConfigManager - - config = ProviderConfigManager.get_provider_chat_config( - model="gpt-4", provider=LlmProviders.PUBLICAI - ) - - assert config is not None - assert config.custom_llm_provider == "publicai" - - -class TestPinstripes: - """Tests for Pinstripes JSON-configured provider""" - - def test_pinstripes_json_config_exists(self): - """Test that pinstripes is configured in providers.json""" - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - assert JSONProviderRegistry.exists("pinstripes") - - pinstripes = JSONProviderRegistry.get("pinstripes") - assert pinstripes is not None - assert pinstripes.base_url == "https://pinstripes.io/v1" - assert pinstripes.api_key_env == "PINSTRIPES_API_KEY" - assert pinstripes.param_mappings.get("max_completion_tokens") == "max_tokens" - - def test_pinstripes_provider_resolution(self): - """Test that provider resolution finds pinstripes and returns the default base URL""" - from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider - - model, provider, api_key, api_base = get_llm_provider( - model="pinstripes/ps/glm-4.5-air", - custom_llm_provider=None, - api_base=None, - api_key=None, - ) - - assert model == "ps/glm-4.5-air" - assert provider == "pinstripes" - assert api_base == "https://pinstripes.io/v1" - - def test_pinstripes_dynamic_config(self): - """Test dynamic config class creation for pinstripes""" - from litellm.llms.openai_like.dynamic_config import create_config_class - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - provider = JSONProviderRegistry.get("pinstripes") - config_class = create_config_class(provider) - config = config_class() - - api_base, api_key = config._get_openai_compatible_provider_info(None, None) - assert api_base == "https://pinstripes.io/v1" - - api_base, api_key = config._get_openai_compatible_provider_info( - "https://custom.pinstripes.io/v1", "test-key" - ) - assert api_base == "https://custom.pinstripes.io/v1" - assert api_key == "test-key" - - def test_pinstripes_parameter_mapping(self): - """Test that max_completion_tokens is mapped to max_tokens for pinstripes""" - from litellm.llms.openai_like.dynamic_config import create_config_class - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - provider = JSONProviderRegistry.get("pinstripes") - config_class = create_config_class(provider) - config = config_class() - - optional_params = {} - non_default_params = {"max_completion_tokens": 100, "temperature": 0.7} - result = config.map_openai_params( - non_default_params, optional_params, "ps/glm-4.5-air", False - ) - - assert "max_tokens" in result - assert result["max_tokens"] == 100 - assert "max_completion_tokens" not in result - assert result["temperature"] == 0.7 - - -class TestDarkbloom: - def test_darkbloom_json_config_exists(self): - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - darkbloom = JSONProviderRegistry.get("darkbloom") - assert darkbloom is not None - assert darkbloom.base_url == "https://api.darkbloom.dev/v1" - assert darkbloom.api_key_env == "DARKBLOOM_API_KEY" - assert darkbloom.api_base_env == "DARKBLOOM_API_BASE" - assert darkbloom.param_mappings.get("max_completion_tokens") == "max_tokens" - - def test_darkbloom_provider_resolution(self): - from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider - - model, provider, api_key, api_base = get_llm_provider( - model="darkbloom/gemma-4-26b", - custom_llm_provider=None, - api_base=None, - api_key=None, - ) - - assert model == "gemma-4-26b" - assert provider == "darkbloom" - assert api_key is None - assert api_base == "https://api.darkbloom.dev/v1" - - def test_darkbloom_dynamic_config(self): - from litellm.llms.openai_like.dynamic_config import create_config_class - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - provider = JSONProviderRegistry.get("darkbloom") - config_class = create_config_class(provider) - config = config_class() - - api_base, api_key = config._get_openai_compatible_provider_info(None, None) - assert api_base == "https://api.darkbloom.dev/v1" - - api_base, api_key = config._get_openai_compatible_provider_info( - "https://custom.darkbloom.dev/v1", "test-key" - ) - assert api_base == "https://custom.darkbloom.dev/v1" - assert api_key == "test-key" - - def test_darkbloom_complete_url_appends_endpoint(self): - from litellm.llms.openai_like.dynamic_config import create_config_class - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - provider = JSONProviderRegistry.get("darkbloom") - config_class = create_config_class(provider) - config = config_class() - - url = config.get_complete_url( - api_base="https://api.darkbloom.dev/v1", - api_key="test-key", - model="darkbloom/gemma-4-26b", - optional_params={}, - litellm_params={}, - stream=True, - ) - - assert url == "https://api.darkbloom.dev/v1/chat/completions" - - def test_darkbloom_provider_config_manager(self): - from litellm import LlmProviders - from litellm.utils import ProviderConfigManager - - config = ProviderConfigManager.get_provider_chat_config( - model="gemma-4-26b", provider=LlmProviders.DARKBLOOM - ) - - assert config is not None - assert config.custom_llm_provider == "darkbloom" - - class TestPublicAIIntegration: """Integration tests for PublicAI provider""" @@ -457,55 +148,3 @@ class TestPublicAIIntegration: pytest.fail(f"Content list conversion test failed: {str(e)}") else: raise - - -if __name__ == "__main__": - # Run basic tests - print("Testing JSON Provider System...") - - test_loader = TestJSONProviderLoader() - print("\n1. Testing JSON provider loading...") - test_loader.test_load_json_providers() - print(" ✓ JSON providers loaded") - - print("\n2. Testing dynamic config generation...") - test_loader.test_dynamic_config_generation() - print(" ✓ Dynamic config works") - - print("\n3. Testing parameter mapping...") - test_loader.test_parameter_mapping() - print(" ✓ Parameter mapping works") - - print("\n4. Testing excluded params...") - test_loader.test_excluded_params() - print(" ✓ Excluded params work") - - print("\n5. Testing provider resolution...") - test_loader.test_provider_resolution() - print(" ✓ Provider resolution works") - - print("\n6. Testing provider config manager...") - test_loader.test_provider_config_manager() - print(" ✓ Config manager works") - - print("\n" + "=" * 50) - print("PublicAI Integration Tests...") - print("=" * 50) - - test_integration = TestPublicAIIntegration() - - print("\n7. Testing basic completion...") - test_integration.test_publicai_completion_basic() - - print("\n8. Testing streaming...") - test_integration.test_publicai_completion_with_streaming() - - print("\n9. Testing parameter mapping...") - test_integration.test_publicai_parameter_mapping() - - print("\n10. Testing content list conversion...") - test_integration.test_publicai_content_list_conversion() - - print("\n" + "=" * 50) - print("✓ All tests passed!") - print("=" * 50) diff --git a/tests/test_litellm/llms/openai_like/test_xiaomi_mimo.py b/tests/test_litellm/llms/openai_like/test_xiaomi_mimo.py index 8104fb12943..580994f60b8 100644 --- a/tests/test_litellm/llms/openai_like/test_xiaomi_mimo.py +++ b/tests/test_litellm/llms/openai_like/test_xiaomi_mimo.py @@ -4,86 +4,12 @@ Related to issue #18794 """ import os -import sys -from unittest.mock import MagicMock, patch -try: - import pytest -except ImportError: - pytest = None - -# Add workspace to path -workspace_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) -sys.path.insert(0, workspace_path) +import pytest import litellm -class TestXiaomiMiMoProviderConfig: - """Test Xiaomi MiMo provider configuration""" - - def test_xiaomi_mimo_in_provider_list(self): - """Test that xiaomi_mimo is in the provider list (fixes #18794)""" - from litellm import LlmProviders - - # Verify xiaomi_mimo is in the enum - assert hasattr(LlmProviders, "XIAOMI_MIMO") - assert LlmProviders.XIAOMI_MIMO.value == "xiaomi_mimo" - - # Verify it's in the provider list - assert "xiaomi_mimo" in litellm.provider_list - - def test_xiaomi_mimo_json_config_exists(self): - """Test that xiaomi_mimo is configured in providers.json""" - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - # Verify xiaomi_mimo is loaded - assert JSONProviderRegistry.exists("xiaomi_mimo") - - # Get xiaomi_mimo config - xiaomi_mimo = JSONProviderRegistry.get("xiaomi_mimo") - assert xiaomi_mimo is not None - assert xiaomi_mimo.base_url == "https://api.xiaomimimo.com/v1" - assert xiaomi_mimo.api_key_env == "XIAOMI_MIMO_API_KEY" - assert xiaomi_mimo.param_mappings.get("max_completion_tokens") == "max_tokens" - - def test_xiaomi_mimo_provider_resolution(self): - """Test that provider resolution finds xiaomi_mimo""" - from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider - - model, provider, api_key, api_base = get_llm_provider( - model="xiaomi_mimo/mimo-v2-flash", - custom_llm_provider=None, - api_base=None, - api_key=None, - ) - - assert model == "mimo-v2-flash" - assert provider == "xiaomi_mimo" - assert api_base == "https://api.xiaomimimo.com/v1" - - def test_xiaomi_mimo_router_config(self): - """Test that xiaomi_mimo can be used in Router configuration (fixes #18794)""" - from litellm import Router - - # This should not raise "Unsupported provider - xiaomi_mimo" - router = Router( - model_list=[ - { - "model_name": "mimo-v2-flash", - "litellm_params": { - "model": "xiaomi_mimo/mimo-v2-flash", - "api_key": "test-key", - }, - } - ] - ) - - # Verify the deployment was created successfully - assert len(router.model_list) == 1 - assert router.model_list[0]["model_name"] == "mimo-v2-flash" - - class TestXiaomiMiMoIntegration: """Integration tests for Xiaomi MiMo provider""" @@ -128,30 +54,3 @@ class TestXiaomiMiMoIntegration: pytest.fail(f"Xiaomi MiMo completion failed: {str(e)}") else: raise - - -if __name__ == "__main__": - # Run basic tests - print("Testing Xiaomi MiMo Provider...") - - test_config = TestXiaomiMiMoProviderConfig() - - print("\n1. Testing provider in list...") - test_config.test_xiaomi_mimo_in_provider_list() - print(" ✓ xiaomi_mimo in provider list") - - print("\n2. Testing JSON config...") - test_config.test_xiaomi_mimo_json_config_exists() - print(" ✓ xiaomi_mimo JSON config loaded") - - print("\n3. Testing provider resolution...") - test_config.test_xiaomi_mimo_provider_resolution() - print(" ✓ Provider resolution works") - - print("\n4. Testing router configuration...") - test_config.test_xiaomi_mimo_router_config() - print(" ✓ Router configuration works (issue #18794 fixed)") - - print("\n" + "=" * 50) - print("✓ All configuration tests passed!") - print("=" * 50) diff --git a/tests/test_litellm/llms/ovhcloud/test_ovhcloud_audio_transcription_transformation.py b/tests/test_litellm/llms/ovhcloud/test_ovhcloud_audio_transcription_transformation.py index c8751fb2d95..8cc46dc98d0 100644 --- a/tests/test_litellm/llms/ovhcloud/test_ovhcloud_audio_transcription_transformation.py +++ b/tests/test_litellm/llms/ovhcloud/test_ovhcloud_audio_transcription_transformation.py @@ -54,61 +54,3 @@ def test_ovhcloud_audio_transcription_config_installed(): assert config is not None assert isinstance(config, BaseAudioTranscriptionConfig) - - - -class TestOVHCloudDurationFieldMigration: - """Tests for OVHCloud duration -> seconds field migration.""" - - def test_seconds_field_mapped_to_duration(self): - """New `seconds` field should be normalized to `duration`.""" - from litellm.llms.ovhcloud.audio_transcription.transformation import ( - OVHCloudAudioTranscriptionConfig, - ) - from unittest.mock import MagicMock - - config = OVHCloudAudioTranscriptionConfig() - mock_response = MagicMock() - mock_response.json.return_value = { - "text": "Hello world", - "seconds": 3.14, - } - - result = config.transform_audio_transcription_response(mock_response) - - assert result.text == "Hello world" - assert result._hidden_params["duration"] == 3.14 - - def test_legacy_duration_field_still_works(self): - """Legacy `duration` field should still be accepted.""" - from litellm.llms.ovhcloud.audio_transcription.transformation import ( - OVHCloudAudioTranscriptionConfig, - ) - from unittest.mock import MagicMock - - config = OVHCloudAudioTranscriptionConfig() - mock_response = MagicMock() - mock_response.json.return_value = { - "text": "Hello world", - "duration": 2.71, - } - - result = config.transform_audio_transcription_response(mock_response) - - assert result.text == "Hello world" - assert result._hidden_params["duration"] == 2.71 - - - - def test_seconds_zero_mapped_to_duration(self): - """seconds=0.0 must not be treated as falsy and lost.""" - from litellm.llms.ovhcloud.audio_transcription.transformation import ( - OVHCloudAudioTranscriptionConfig, - ) - from unittest.mock import MagicMock - - config = OVHCloudAudioTranscriptionConfig() - mock_response = MagicMock() - mock_response.json.return_value = {"text": "silence", "seconds": 0.0} - result = config.transform_audio_transcription_response(mock_response) - assert result._hidden_params["duration"] == 0.0 \ No newline at end of file diff --git a/tests/test_litellm/llms/ovhcloud/test_ovhcloud_chat_transformation.py b/tests/test_litellm/llms/ovhcloud/test_ovhcloud_chat_transformation.py index 057ab9ede9a..34954587ed0 100644 --- a/tests/test_litellm/llms/ovhcloud/test_ovhcloud_chat_transformation.py +++ b/tests/test_litellm/llms/ovhcloud/test_ovhcloud_chat_transformation.py @@ -6,174 +6,12 @@ import os import pytest -from litellm.llms.ovhcloud.utils import OVHCloudException -from litellm.utils import get_optional_params -from litellm.llms.ovhcloud.chat.transformation import ( - OVHCloudChatCompletionStreamingHandler, - OVHCloudChatConfig, -) -config = OVHCloudChatConfig() model = "ovhcloud/Mistral-7B-Instruct-v0.3" -class TestOvhCloudChatCompletionStreamingHandler: - def test_chunk_parser_successful(self): - handler = OVHCloudChatCompletionStreamingHandler( - streaming_response=None, sync_stream=True - ) - - chunk = { - "id": "test_id", - "created": 1234567890, - "model": "gpt-oss-20b", - "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, - "choices": [ - {"delta": {"content": "test content", "reasoning": "test reasoning"}} - ], - } - - result = handler.chunk_parser(chunk) - - assert result.id == "test_id" - assert result.object == "chat.completion.chunk" - assert result.created == 1234567890 - assert result.model == "gpt-oss-20b" - assert result.usage.prompt_tokens == chunk["usage"]["prompt_tokens"] - assert result.usage.completion_tokens == chunk["usage"]["completion_tokens"] - assert result.usage.total_tokens == chunk["usage"]["total_tokens"] - assert len(result.choices) == 1 - assert result.choices[0]["delta"]["reasoning_content"] == "test reasoning" - - def test_chunk_parser_error_response(self): - handler = OVHCloudChatCompletionStreamingHandler( - streaming_response=None, sync_stream=True - ) - - error_chunk = { - "error": { - "message": "test error", - "code": 400, - } - } - - with pytest.raises(OVHCloudException) as exc_info: - handler.chunk_parser(error_chunk) - - assert "OVHCloud Error: test error" in str(exc_info.value) - assert exc_info.value.status_code == 400 - - def test_chunk_parser_key_error(self): - handler = OVHCloudChatCompletionStreamingHandler( - streaming_response=None, sync_stream=True - ) - - invalid_chunk = {"incomplete": "data"} - - with pytest.raises(OVHCloudException) as exc_info: - handler.chunk_parser(invalid_chunk) - - assert "KeyError" in str(exc_info.value) - assert exc_info.value.status_code == 400 - - -class TestOVHCloudConfig: - def test_transform_request_basic(self): - """Test basic request transformation""" - transformed_request = config.transform_request( - model, - messages=[{"role": "user", "content": "Hello, world!"}], - optional_params={}, - litellm_params={}, - headers={}, - ) - - assert transformed_request["model"] == model - assert transformed_request["messages"] == [ - {"role": "user", "content": "Hello, world!"} - ] - - def test_transform_request_with_extra_body(self): - """Test request transformation with extra_body parameters""" - transformed_request = config.transform_request( - model, - messages=[{"role": "user", "content": "Hello, world!"}], - optional_params={"extra_body": {"custom_param": "custom_value"}}, - litellm_params={}, - headers={}, - ) - - assert transformed_request["custom_param"] == "custom_value" - assert transformed_request["messages"] == [ - {"role": "user", "content": "Hello, world!"} - ] - - def test_map_openai_params(self): - """Test OpenAI parameter mapping""" - non_default_params = { - "temperature": 0.7, - "max_tokens": 100, - "top_p": 0.9, - } - - mapped_params = config.map_openai_params( - non_default_params=non_default_params, - optional_params={}, - model=model, - drop_params=False, - ) - - assert mapped_params["temperature"] == 0.7 - assert mapped_params["max_tokens"] == 100 - assert mapped_params["top_p"] == 0.9 - - def test_get_error_class(self): - """Test error class creation""" - error = config.get_error_class( - error_message="Test error", - status_code=400, - headers={"Content-Type": "application/json"}, - ) - - assert isinstance(error, OVHCloudException) - assert error.message == "Test error" - assert error.status_code == 400 - - @pytest.mark.parametrize( - "model", - [ - "Meta-Llama-3_3-70B-Instruct", - "Meta-Llama-3_1-70B-Instruct", - "Mixtral-8x7B-Instruct-v0.1", - "gpt-oss-120b", - "some-model-not-in-the-cost-map", - ], - ) - def test_tools_not_filtered_by_static_model_map(self, model): - """ - OVHCloud AI Endpoints are OpenAI-compatible; tools/tool_choice must pass - through for any model. The server is responsible for rejecting unsupported - tool calls — LiteLLM must not strip them based on a stale static catalog. - """ - - params = get_optional_params( - model=model, - custom_llm_provider="ovhcloud", - tools=[ - { - "type": "function", - "function": {"name": "x", "parameters": {}}, - } - ], - tool_choice="auto", - ) - - assert "tools" in params - assert "tool_choice" in params - - def test_ovhcloud_integration(): from litellm import completion @@ -285,78 +123,3 @@ def test_ovhcloud_with_custom_base_url(): if __name__ == "__main__": pytest.main([__file__, "-v"]) - - -class TestOVHCloudReasoningFieldMigration: - """Tests for OVHCloud reasoning_content -> reasoning field migration.""" - - def test_streaming_new_reasoning_field(self): - """New `reasoning` field should be mapped to `reasoning_content`.""" - handler = OVHCloudChatCompletionStreamingHandler( - streaming_response=iter([]), - sync_stream=True, - ) - chunk = { - "id": "test-id", - "created": 1234567890, - "model": "test-model", - "choices": [ - { - "delta": { - "role": "assistant", - "reasoning": "Let me think...", - }, - "index": 0, - } - ], - } - result = handler.chunk_parser(chunk) - assert result.choices[0]["delta"]["reasoning_content"] == "Let me think..." - - def test_streaming_legacy_reasoning_content_unchanged(self): - """Legacy `reasoning_content` field should pass through untouched.""" - handler = OVHCloudChatCompletionStreamingHandler( - streaming_response=iter([]), - sync_stream=True, - ) - chunk = { - "id": "test-id", - "created": 1234567890, - "model": "test-model", - "choices": [ - { - "delta": { - "role": "assistant", - "reasoning_content": "Already correct field.", - }, - "index": 0, - } - ], - } - result = handler.chunk_parser(chunk) - assert result.choices[0]["delta"]["reasoning_content"] == "Already correct field." - - def test_streaming_both_fields_legacy_wins(self): - """When both fields present, existing `reasoning_content` is not overwritten.""" - handler = OVHCloudChatCompletionStreamingHandler( - streaming_response=iter([]), - sync_stream=True, - ) - chunk = { - "id": "test-id", - "created": 1234567890, - "model": "test-model", - "choices": [ - { - "delta": { - "reasoning": "new field", - "reasoning_content": "legacy field", - }, - "index": 0, - } - ], - } - result = handler.chunk_parser(chunk) - assert result.choices[0]["delta"]["reasoning_content"] == "legacy field" - - diff --git a/tests/test_litellm/llms/reducto/__init__.py b/tests/test_litellm/llms/reducto/__init__.py deleted file mode 100644 index 8b137891791..00000000000 --- a/tests/test_litellm/llms/reducto/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/tests/test_litellm/llms/s3_vectors/__init__.py b/tests/test_litellm/llms/s3_vectors/__init__.py deleted file mode 100644 index d4b0c4d8550..00000000000 --- a/tests/test_litellm/llms/s3_vectors/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# S3 Vectors tests diff --git a/tests/test_litellm/llms/s3_vectors/vector_stores/__init__.py b/tests/test_litellm/llms/s3_vectors/vector_stores/__init__.py deleted file mode 100644 index 231735c1de7..00000000000 --- a/tests/test_litellm/llms/s3_vectors/vector_stores/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# S3 Vectors vector store tests diff --git a/tests/test_litellm/llms/soniox/__init__.py b/tests/test_litellm/llms/soniox/__init__.py deleted file mode 100644 index b2cd496d66a..00000000000 --- a/tests/test_litellm/llms/soniox/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Soniox provider tests.""" diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index 4679b978f78..d3a7ba7a1bd 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -1,1681 +1,13 @@ -import base64 - import pytest from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_gemini_tool_call_result, ) -from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - _transform_request_body, - check_if_part_exists_in_parts, - _get_highest_media_resolution, - _extract_max_media_resolution_from_messages, -) from litellm.types.llms.vertex_ai import BlobType -from litellm.types.utils import Message - - -def test_check_if_part_exists_in_parts(): - parts = [ - {"text": "Hello", "thought": True}, - {"text": "World", "thought": False}, - ] - part = {"text": "Hello", "thought": True} - new_part = {"text": "Hello World", "thought": True} - assert check_if_part_exists_in_parts(parts, part) - assert not check_if_part_exists_in_parts(parts, new_part, ["thought"]) - assert check_if_part_exists_in_parts(parts, new_part, ["text"]) - - -def test_check_if_part_exists_in_parts_camel_case_snake_case(): - """Test that function handles both camelCase and snake_case key variations""" - # Test snake_case to camelCase matching - parts_with_snake_case = [ - { - "function_call": { - "name": "get_current_weather", - "args": {"location": "San Francisco, CA"}, - } - }, - {"text": "Some other content"}, - ] - - part_with_camel_case = { - "functionCall": { - "name": "get_current_weather", - "args": {"location": "San Francisco, CA"}, - } - } - - # Should find match between function_call and functionCall - assert check_if_part_exists_in_parts(parts_with_snake_case, part_with_camel_case) - - # Test camelCase to snake_case matching - parts_with_camel_case = [ - {"functionCall": {"name": "calculate_sum", "args": {"a": 1, "b": 2}}} - ] - - part_with_snake_case = { - "function_call": {"name": "calculate_sum", "args": {"a": 1, "b": 2}} - } - - # Should find match between functionCall and function_call - assert check_if_part_exists_in_parts(parts_with_camel_case, part_with_snake_case) - - # Test no match when values differ - part_with_different_values = { - "function_call": {"name": "different_function", "args": {"x": 5}} - } - - assert not check_if_part_exists_in_parts( - parts_with_snake_case, part_with_different_values - ) - - # Test multiple keys with mixed casing - parts_mixed = [ - { - "function_call": {"name": "test"}, - "thoughtSignature": "reasoning", - "text": "content", - } - ] - - part_mixed_casing = { - "functionCall": {"name": "test"}, - "thought_signature": "reasoning", - "text": "content", - } - - assert check_if_part_exists_in_parts(parts_mixed, part_mixed_casing) - - -def test_cached_content_respects_modify_params_for_cache_incompatible_fields(): - """Regression: cachedContent drops system/tools/toolConfig only when modify_params=True.""" - import litellm - - cache_name = "projects/p/locations/us-central1/cachedContents/abc123" - messages = [ - {"role": "system", "content": "You are helpful"}, - {"role": "user", "content": "hi"}, - ] - optional_params = { - "tools": [ - { - "functionDeclarations": [ - {"name": "get_weather", "description": "Get weather"}, - ] - } - ], - "tool_choice": {"functionCallingConfig": {"mode": "AUTO"}}, - } - - original_modify_params = litellm.modify_params - try: - # With modify_params=False (default), keep fields even with cachedContent. - litellm.modify_params = False - result = _transform_request_body( - messages=list(messages), - model="gemini-2.5-pro", - optional_params=dict(optional_params), - custom_llm_provider="vertex_ai", - litellm_params={}, - cached_content=cache_name, - ) - assert result.get("cachedContent") == cache_name - assert "system_instruction" in result - assert "tools" in result - assert "toolConfig" in result - assert "contents" in result - - # With modify_params=True, drop cache-incompatible fields. - litellm.modify_params = True - result_modify_true = _transform_request_body( - messages=list(messages), - model="gemini-2.5-pro", - optional_params=dict(optional_params), - custom_llm_provider="vertex_ai", - litellm_params={}, - cached_content=cache_name, - ) - assert result_modify_true.get("cachedContent") == cache_name - assert "system_instruction" not in result_modify_true - assert "tools" not in result_modify_true - assert "toolConfig" not in result_modify_true - assert "contents" in result_modify_true - - # Without cache, fields are always included. - result_no_cache = _transform_request_body( - messages=list(messages), - model="gemini-2.5-pro", - optional_params=dict(optional_params), - custom_llm_provider="vertex_ai", - litellm_params={}, - cached_content=None, - ) - assert "system_instruction" in result_no_cache - assert "tools" in result_no_cache - assert "toolConfig" in result_no_cache - finally: - litellm.modify_params = original_modify_params - - -# Tests for issue #14556: Labels field provider-aware filtering -def test_google_genai_excludes_labels(): - """Test that Google GenAI/AI Studio endpoints exclude labels when custom_llm_provider='gemini'""" - messages = [{"role": "user", "content": "test"}] - optional_params = {"labels": {"project": "test", "team": "ai"}} - litellm_params = {} - - result = _transform_request_body( - messages=messages, - model="gemini-2.5-pro", - optional_params=optional_params, - custom_llm_provider="gemini", - litellm_params=litellm_params, - cached_content=None, - ) - - # Google GenAI/AI Studio should NOT include labels - assert "labels" not in result - assert "contents" in result - - -def test_vertex_ai_includes_labels(): - """Test that Vertex AI endpoints include labels when custom_llm_provider='vertex_ai'""" - messages = [{"role": "user", "content": "test"}] - optional_params = {"labels": {"project": "test", "team": "ai"}} - litellm_params = {} - - result = _transform_request_body( - messages=messages, - model="gemini-2.5-pro", - optional_params=optional_params, - custom_llm_provider="vertex_ai", - litellm_params=litellm_params, - cached_content=None, - ) - - # Vertex AI SHOULD include labels - assert "labels" in result - assert result["labels"] == {"project": "test", "team": "ai"} - - -def test_service_tier_forwarded_to_vertex_ai(): - """Test that service_tier in optional_params is mapped to serviceTier in request body.""" - messages = [{"role": "user", "content": "test"}] - optional_params = {"service_tier": "flex"} - litellm_params = {} - - result = _transform_request_body( - messages=messages, - model="gemini-2.5-pro", - optional_params=optional_params, - custom_llm_provider="vertex_ai", - litellm_params=litellm_params, - cached_content=None, - ) - - assert "serviceTier" in result - assert result["serviceTier"] == "flex" - - -def test_extra_body_cache_not_forwarded_to_vertex_ai(): - """ - 'cache' inside extra_body is a LiteLLM-internal proxy caching control. - It must NOT be forwarded to the Vertex AI request body. - - Regression test for: "Invalid JSON payload received. Unknown name \"cache\": Cannot find field." - Vertex AI enforces a strict JSON schema and rejects any unknown field. - """ - messages = [{"role": "user", "content": "test"}] - optional_params = { - "extra_body": { - "cache": {"use-cache": True, "ttl": 86400}, # LiteLLM-internal - "some_vertex_param": "value", # legitimate provider extra - }, - } - litellm_params = {} - - result = _transform_request_body( - messages=messages, - model="gemini-2.5-pro", - optional_params=optional_params, - custom_llm_provider="vertex_ai", - litellm_params=litellm_params, - cached_content=None, - ) - - # 'cache' must be stripped — Vertex AI has no such field - assert "cache" not in result, ( - "extra_body.cache must not be forwarded to Vertex AI. " - 'Vertex AI rejects it with 400: Unknown name "cache": Cannot find field.' - ) - - # Other legitimate extra_body keys should still pass through - assert "some_vertex_param" in result - assert result["some_vertex_param"] == "value" - - # Core request fields must be present - assert "contents" in result - - -def test_extra_body_tags_not_forwarded_to_vertex_ai(): - """ - 'tags' inside extra_body is a LiteLLM-internal param for logging/tracking. - It must NOT be forwarded to the Vertex AI request body. - Documented in litellm_proxy.md: "Send tags by including them in the extra_body parameter" - """ - messages = [{"role": "user", "content": "test"}] - optional_params = { - "extra_body": { - "tags": ["user:alice", "env:prod"], - "custom_param": "allowed", - }, - } - litellm_params = {} - - result = _transform_request_body( - messages=messages, - model="gemini-2.5-pro", - optional_params=optional_params, - custom_llm_provider="vertex_ai", - litellm_params=litellm_params, - cached_content=None, - ) - - assert "tags" not in result - assert "custom_param" in result - assert result["custom_param"] == "allowed" - - -def test_extra_body_google_maps_rewrites_json_response_format(): - messages = [{"role": "user", "content": "test"}] - optional_params = { - "response_mime_type": "application/json", - "response_schema": { - "type": "object", - "properties": {"answer": {"type": "string"}}, - }, - "extra_body": { - "tools": [{"googleMaps": {}}], - }, - } - - result = _transform_request_body( - messages=messages, - model="gemini-2.5-pro", - optional_params=optional_params, - custom_llm_provider="vertex_ai", - litellm_params={}, - cached_content=None, - ) - - generation_config = result["generationConfig"] - assert "response_mime_type" not in generation_config - assert generation_config["responseFormat"] == { - "text": { - "mimeType": "APPLICATION_JSON", - "schema": { - "type": "object", - "properties": {"answer": {"type": "string"}}, - }, - } - } - - -def test_extra_body_generation_config_cannot_restore_google_maps_json_mime_type(): - messages = [{"role": "user", "content": "test"}] - optional_params = { - "tools": [{"googleMaps": {}}], - "response_mime_type": "application/json", - "extra_body": { - "generationConfig": { - "response_mime_type": "application/json", - "response_json_schema": { - "type": "object", - "properties": {"answer": {"type": "string"}}, - }, - }, - }, - } - - result = _transform_request_body( - messages=messages, - model="gemini-2.5-pro", - optional_params=optional_params, - custom_llm_provider="vertex_ai", - litellm_params={}, - cached_content=None, - ) - - generation_config = result["generationConfig"] - assert "response_mime_type" not in generation_config - assert "response_json_schema" not in generation_config - assert generation_config["responseFormat"] == { - "text": { - "mimeType": "APPLICATION_JSON", - "schema": { - "type": "object", - "properties": {"answer": {"type": "string"}}, - }, - } - } - - -def test_metadata_to_labels_vertex_only(): - """Test that metadata->labels conversion only happens for Vertex AI""" - messages = [{"role": "user", "content": "test"}] - optional_params = {} - litellm_params = { - "metadata": { - "requester_metadata": {"user": "john_doe", "project": "test-project"} - } - } - - # Google GenAI/AI Studio should not include labels from metadata - result = _transform_request_body( - messages=messages, - model="gemini-2.5-pro", - optional_params=optional_params.copy(), - custom_llm_provider="gemini", - litellm_params=litellm_params.copy(), - cached_content=None, - ) - assert "labels" not in result - - # Vertex AI should include labels from metadata - result = _transform_request_body( - messages=messages, - model="gemini-2.5-pro", - optional_params=optional_params.copy(), - custom_llm_provider="vertex_ai", - litellm_params=litellm_params.copy(), - cached_content=None, - ) - assert "labels" in result - assert result["labels"] == {"user": "john_doe", "project": "test-project"} - - -def test_empty_content_handling(): - """Test that empty content strings are properly handled in Gemini message transformation""" - # Test with empty content in user message - messages = [{"content": "", "role": "user"}] - - contents = _gemini_convert_messages_with_history(messages=messages) - - # Verify that the content was properly transformed - assert len(contents) == 1 - assert contents[0]["role"] == "user" - assert len(contents[0]["parts"]) == 1 - assert "text" in contents[0]["parts"][0] - assert contents[0]["parts"][0]["text"] == "" - - -def test_thought_signature_extraction_from_response(): - """Test that thought signatures are extracted from Gemini response parts and stored in provider_specific_fields""" - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexGeminiConfig, - ) - from litellm.types.llms.vertex_ai import HttpxPartType - - # Test case: Single function call with thought signature - test_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" - - parts_with_signature = [ - HttpxPartType( - functionCall={ - "name": "get_current_temperature", - "args": {"location": "Paris"}, - }, - thoughtSignature=test_signature, - ) - ] - - function, tools, _ = VertexGeminiConfig._transform_parts( - parts=parts_with_signature, - cumulative_tool_call_idx=0, - is_function_call=False, - ) - - # Verify thought signature is stored in provider_specific_fields - assert tools is not None - assert len(tools) == 1 - assert "provider_specific_fields" in tools[0] - assert tools[0]["provider_specific_fields"]["thought_signature"] == test_signature - - -def test_thought_signature_parallel_function_calls(): - """Test that only the first function call in parallel calls has thought signature""" - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexGeminiConfig, - ) - from litellm.types.llms.vertex_ai import HttpxPartType - - test_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" - - # Parallel function calls - only first has signature - parts_parallel = [ - HttpxPartType( - functionCall={ - "name": "get_current_temperature", - "args": {"location": "Paris"}, - }, - thoughtSignature=test_signature, # First FC has signature - ), - HttpxPartType( - functionCall={ - "name": "get_current_temperature", - "args": {"location": "London"}, - }, - # Second FC has no signature (parallel call) - ), - ] - - function, tools, _ = VertexGeminiConfig._transform_parts( - parts=parts_parallel, - cumulative_tool_call_idx=0, - is_function_call=False, - ) - - # Verify only first tool call has thought signature - assert tools is not None - assert len(tools) == 2 - assert "provider_specific_fields" in tools[0] - assert tools[0]["provider_specific_fields"]["thought_signature"] == test_signature - # Second tool call should not have thought signature - assert "provider_specific_fields" not in tools[ - 1 - ] or "thought_signature" not in tools[1].get("provider_specific_fields", {}) - - -def test_thought_signature_preservation_in_conversion(): - """Test that thought signatures are preserved when converting assistant messages back to Gemini format""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - test_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" - - # Assistant message with tool calls containing thought signatures - assistant_message = { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_abc123", - "type": "function", - "function": { - "name": "get_current_temperature", - "arguments": '{"location": "Paris"}', - }, - "index": 0, - "provider_specific_fields": { - "thought_signature": test_signature, - }, - }, - { - "id": "call_def456", - "type": "function", - "function": { - "name": "get_current_temperature", - "arguments": '{"location": "London"}', - }, - "index": 1, - # No thought signature for parallel call - }, - ], - } - - gemini_parts = convert_to_gemini_tool_call_invoke(assistant_message) - - # Verify thought signature is preserved in first function call part - assert len(gemini_parts) == 2 - assert "function_call" in gemini_parts[0] - assert "thoughtSignature" in gemini_parts[0] - assert gemini_parts[0]["thoughtSignature"] == test_signature - - # Verify second function call part does not have thought signature - assert "function_call" in gemini_parts[1] - assert "thoughtSignature" not in gemini_parts[1] - - -def test_thought_signature_sequential_function_calls(): - """Test that each sequential function call preserves its own thought signature""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - signature_1 = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" - signature_2 = "DifferentSignatureForSecondCall1234567890ABCDEFGHIJKLMNOPQRSTUVWXYZ" - - # Sequential function calls - each has its own signature - # This simulates a multi-step conversation where each step has a signature - assistant_message_step1 = { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_step1", - "type": "function", - "function": { - "name": "check_flight", - "arguments": '{"flight": "AA100"}', - }, - "index": 0, - "provider_specific_fields": { - "thought_signature": signature_1, - }, - }, - ], - } - - assistant_message_step2 = { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_step2", - "type": "function", - "function": { - "name": "book_taxi", - "arguments": '{"destination": "airport"}', - }, - "index": 0, - "provider_specific_fields": { - "thought_signature": signature_2, - }, - }, - ], - } - - gemini_parts_step1 = convert_to_gemini_tool_call_invoke(assistant_message_step1) - gemini_parts_step2 = convert_to_gemini_tool_call_invoke(assistant_message_step2) - - # Verify each step preserves its own signature - assert len(gemini_parts_step1) == 1 - assert gemini_parts_step1[0]["thoughtSignature"] == signature_1 - - assert len(gemini_parts_step2) == 1 - assert gemini_parts_step2[0]["thoughtSignature"] == signature_2 - - -def test_thought_signature_with_function_call_mode(): - """Test thought signature extraction in function_call mode (is_function_call=True)""" - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexGeminiConfig, - ) - from litellm.types.llms.vertex_ai import HttpxPartType - - test_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" - - parts_with_signature = [ - HttpxPartType( - functionCall={ - "name": "get_current_weather", - "args": {"location": "Tokyo"}, - }, - thoughtSignature=test_signature, - ) - ] - - function, tools, _ = VertexGeminiConfig._transform_parts( - parts=parts_with_signature, - cumulative_tool_call_idx=0, - is_function_call=True, - ) - - # Verify thought signature is stored in function's provider_specific_fields - assert function is not None - # Function should be dict-like (TypedDict or dict) - assert hasattr(function, "__getitem__") or isinstance(function, dict) - assert "provider_specific_fields" in function - assert function["provider_specific_fields"]["thought_signature"] == test_signature - assert tools is None - - -def test_dummy_signature_added_for_gemini_3_conversation_history(): - """Test that dummy signatures are added when transferring conversation history from older models (like gemini-2.5-flash) to gemini-3.""" - import base64 - - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - # Simulate conversation history from gemini-2.5-flash (no thought signature) - assistant_message_from_older_model = { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_abc123", - "type": "function", - "function": { - "name": "get_current_temperature", - "arguments": '{"location": "Paris"}', - }, - "index": 0, - # No provider_specific_fields - older model doesn't provide signatures - }, - ], - } - - # Convert to Gemini format for gemini-3-pro-preview (should add dummy signature) - gemini_parts = convert_to_gemini_tool_call_invoke( - assistant_message_from_older_model, model="gemini-3-pro-preview" - ) - - # Verify dummy signature is added - assert len(gemini_parts) == 1 - assert "function_call" in gemini_parts[0] - assert "thoughtSignature" in gemini_parts[0] - - # Verify it's the expected dummy signature (base64 encoded "skip_thought_signature_validator") - expected_dummy = base64.b64encode(b"skip_thought_signature_validator").decode( - "utf-8" - ) - assert gemini_parts[0]["thoughtSignature"] == expected_dummy - - -def test_dummy_signature_not_added_for_gemini_2_5(): - """Test that dummy signatures are NOT added when target model is not gemini-3.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - # Simulate conversation history from gemini-2.5-flash (no thought signature) - assistant_message = { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_abc123", - "type": "function", - "function": { - "name": "get_current_temperature", - "arguments": '{"location": "Paris"}', - }, - "index": 0, - # No provider_specific_fields - }, - ], - } - - # Convert to Gemini format for gemini-2.5-flash (should NOT add dummy signature) - gemini_parts = convert_to_gemini_tool_call_invoke( - assistant_message, model="gemini-2.5-flash" - ) - - # Verify no dummy signature is added for non-gemini-3 models - assert len(gemini_parts) == 1 - assert "function_call" in gemini_parts[0] - assert "thoughtSignature" not in gemini_parts[0] - - -def test_dummy_signature_not_added_when_signature_exists(): - """Test that dummy signatures are NOT added when a real signature already exists.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - real_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" - - # Assistant message with existing thought signature - assistant_message_with_signature = { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_abc123", - "type": "function", - "function": { - "name": "get_current_temperature", - "arguments": '{"location": "Paris"}', - "provider_specific_fields": { - "thought_signature": real_signature, - }, - }, - "index": 0, - }, - ], - } - - # Convert to Gemini format for gemini-3-pro-preview - gemini_parts = convert_to_gemini_tool_call_invoke( - assistant_message_with_signature, model="gemini-3-pro-preview" - ) - - # Verify real signature is preserved, not replaced with dummy - assert len(gemini_parts) == 1 - assert "function_call" in gemini_parts[0] - assert "thoughtSignature" in gemini_parts[0] - assert gemini_parts[0]["thoughtSignature"] == real_signature - - -def test_dummy_signature_with_function_call_mode(): - """Test that dummy signatures are added for function_call mode when converting to gemini-3.""" - import base64 - - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - # Assistant message with function_call (not tool_calls) and no signature - assistant_message_function_call = { - "role": "assistant", - "content": None, - "function_call": { - "name": "get_current_temperature", - "arguments": '{"location": "Paris"}', - # No provider_specific_fields - }, - } - - # Convert to Gemini format for gemini-3-pro-preview - gemini_parts = convert_to_gemini_tool_call_invoke( - assistant_message_function_call, model="gemini-3-pro-preview" - ) - - # Verify dummy signature is added - assert len(gemini_parts) == 1 - assert "function_call" in gemini_parts[0] - assert "thoughtSignature" in gemini_parts[0] - - # Verify it's the expected dummy signature - expected_dummy = base64.b64encode(b"skip_thought_signature_validator").decode( - "utf-8" - ) - assert gemini_parts[0]["thoughtSignature"] == expected_dummy - - -def _parallel_tool_calls(*signatures): - return [ - { - "id": f"call_{idx}", - "type": "function", - "function": { - "name": f"tool_{idx}", - "arguments": '{"location": "Paris"}', - **( - {"provider_specific_fields": {"thought_signature": signature}} - if signature is not None - else {} - ), - }, - "index": idx, - } - for idx, signature in enumerate(signatures) - ] - - -def _parallel_tool_calls_signed_via_id(*signatures): - """Parallel tool calls in the shape LiteLLM actually hands back to clients. - - The signature rides in the tool call id behind __thought__, which is what an - OpenAI-format client echoes back on the next turn. - """ - from litellm.litellm_core_utils.prompt_templates.factory import ( - _encode_tool_call_id_with_signature, - ) - - return [ - { - "id": _encode_tool_call_id_with_signature(f"call_{idx}", signature), - "type": "function", - "function": {"name": f"tool_{idx}", "arguments": '{"location": "Paris"}'}, - "index": idx, - } - for idx, signature in enumerate(signatures) - ] - - -REAL_THOUGHT_SIGNATURE = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n" -PLACEHOLDER_SIGNATURE = base64.b64encode(b"skip_thought_signature_validator").decode( - "utf-8" -) - - -def test_dummy_signature_only_on_first_parallel_tool_call(): - """Google documents the placeholder as a last resort that degrades quality, so an unsigned - parallel turn replayed to gemini-3 gets a budget of exactly one.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - gemini_parts = convert_to_gemini_tool_call_invoke( - { - "role": "assistant", - "content": None, - "tool_calls": _parallel_tool_calls(None, None, None), - }, - model="gemini-3-pro-preview", - ) - - assert len(gemini_parts) == 3 - assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE - assert "thoughtSignature" not in gemini_parts[1] - assert "thoughtSignature" not in gemini_parts[2] - - -def test_real_signature_on_first_parallel_tool_call_leaves_siblings_empty(): - """Gemini signs only the first of N parallel function calls, so a faithful replay has - nothing to attach to the siblings.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - gemini_parts = convert_to_gemini_tool_call_invoke( - { - "role": "assistant", - "content": None, - "tool_calls": _parallel_tool_calls(REAL_THOUGHT_SIGNATURE, None, None), - }, - model="gemini-3-pro-preview", - ) - - assert len(gemini_parts) == 3 - assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE - assert "thoughtSignature" not in gemini_parts[1] - assert "thoughtSignature" not in gemini_parts[2] - - -def test_real_signature_on_later_parallel_tool_call_is_preserved(): - """Clients may reorder or drop calls, so a signature that lands on a non-first call is - still the model's own and must survive the round trip.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - gemini_parts = convert_to_gemini_tool_call_invoke( - { - "role": "assistant", - "content": None, - "tool_calls": _parallel_tool_calls(None, REAL_THOUGHT_SIGNATURE), - }, - model="gemini-3-pro-preview", - ) - - assert len(gemini_parts) == 2 - assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE - assert gemini_parts[1]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE - - -def test_no_signatures_on_parallel_tool_calls_for_gemini_2_5(): - """Non-gemini-3 models never get a placeholder signature, on any call.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - gemini_parts = convert_to_gemini_tool_call_invoke( - { - "role": "assistant", - "content": None, - "tool_calls": _parallel_tool_calls(None, None), - }, - model="gemini-2.5-flash", - ) - - assert len(gemini_parts) == 2 - assert all("thoughtSignature" not in part for part in gemini_parts) - - -def test_signature_embedded_in_tool_call_id_only_on_first_parallel_call(): - """The production shape: the signature arrives inside the first call's id, siblings have bare ids.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - gemini_parts = convert_to_gemini_tool_call_invoke( - { - "role": "assistant", - "content": None, - "tool_calls": _parallel_tool_calls_signed_via_id( - REAL_THOUGHT_SIGNATURE, None, None - ), - }, - model="gemini-3-pro-preview", - ) - - assert len(gemini_parts) == 3 - assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE - assert "thoughtSignature" not in gemini_parts[1] - assert "thoughtSignature" not in gemini_parts[2] - - -def test_tool_level_provider_specific_fields_signature_leaves_siblings_empty(): - """A signature on the tool call itself, rather than on its function, behaves the same way.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - tool_calls = _parallel_tool_calls(None, None) - tool_calls[0]["provider_specific_fields"] = { - "thought_signature": REAL_THOUGHT_SIGNATURE - } - - gemini_parts = convert_to_gemini_tool_call_invoke( - {"role": "assistant", "content": None, "tool_calls": tool_calls}, - model="gemini-3-pro-preview", - ) - - assert len(gemini_parts) == 2 - assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE - assert "thoughtSignature" not in gemini_parts[1] - - -def test_placeholder_lands_on_first_emitted_part_not_first_tool_call_entry(): - """A non-function entry (e.g. an OpenAI custom tool call) emits no part, so it must not - consume the one placeholder slot and leave the real first function call bare.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - tool_calls = [ - {"id": "call_custom", "type": "custom", "custom": {"name": "noop", "input": ""}} - ] + _parallel_tool_calls(None, None) - - gemini_parts = convert_to_gemini_tool_call_invoke( - {"role": "assistant", "content": None, "tool_calls": tool_calls}, - model="gemini-3-pro-preview", - ) - - assert len(gemini_parts) == 2 - assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE - assert "thoughtSignature" not in gemini_parts[1] - - -def test_no_placeholder_when_model_is_unknown(): - """Without a model there is nothing to prove the target needs a placeholder, so none is added.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - gemini_parts = convert_to_gemini_tool_call_invoke( - { - "role": "assistant", - "content": None, - "tool_calls": _parallel_tool_calls(None, None), - }, - ) - - assert len(gemini_parts) == 2 - assert all("thoughtSignature" not in part for part in gemini_parts) - - -def test_real_signature_forwarded_to_gemini_2_5_without_placeholder_siblings(): - """Older models still receive a real signature that a client replays, and still get no placeholder.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - gemini_parts = convert_to_gemini_tool_call_invoke( - { - "role": "assistant", - "content": None, - "tool_calls": _parallel_tool_calls(REAL_THOUGHT_SIGNATURE, None), - }, - model="gemini-2.5-flash", - ) - - assert len(gemini_parts) == 2 - assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE - assert "thoughtSignature" not in gemini_parts[1] - - -def test_parallel_tool_call_history_replayed_through_full_message_conversion(): - """End to end through the message-history converter, the path a real /chat/completions replay takes.""" - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - ) - - messages = [ - {"role": "user", "content": "Weather in Paris, London and Tokyo?"}, - { - "role": "assistant", - "content": None, - "tool_calls": _parallel_tool_calls_signed_via_id( - REAL_THOUGHT_SIGNATURE, None, None - ), - }, - ] - - contents = _gemini_convert_messages_with_history( - messages=messages, model="gemini-3-pro-preview" - ) - - model_parts = contents[1]["parts"] - assert len(model_parts) == 3 - assert model_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE - assert "thoughtSignature" not in model_parts[1] - assert "thoughtSignature" not in model_parts[2] - - -@pytest.mark.parametrize( - "model", - ["gemini-3.5-flash", "vertex_ai/gemini-3.5-flash", "gemini/gemini-3.5-flash"], -) -def test_natively_signed_parallel_turn_never_carries_a_placeholder(model): - """A native gemini-3.5 parallel turn replays with zero skip_thought_signature_validator parts. - - Fabricating the placeholder alongside a real signature is what produced empty text responses - on gemini-3.5 parallel function calling, so the whole payload has to stay placeholder-free. - """ - import json - - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - ) - - messages = [ - {"role": "user", "content": "Weather in Paris, London and Tokyo?"}, - { - "role": "assistant", - "content": None, - "tool_calls": _parallel_tool_calls_signed_via_id( - REAL_THOUGHT_SIGNATURE, None, None - ), - }, - ] - - contents = _gemini_convert_messages_with_history(messages=messages, model=model) - - model_parts = contents[1]["parts"] - assert len(model_parts) == 3 - assert model_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE - assert "thoughtSignature" not in model_parts[1] - assert "thoughtSignature" not in model_parts[2] - assert PLACEHOLDER_SIGNATURE not in json.dumps(contents) - - -@pytest.mark.parametrize( - "model", - [ - "gemini-3-pro-preview", - "gemini-3-flash-preview", - "gemini-3.1-pro-preview", - "gemini-3.5-flash", - "gemini-3.6-flash", - "gemini-3.7-flash", - "gemini-3.8-flash", - "vertex_ai/gemini-3.5-flash", - "vertex_ai/gemini-3.7-flash", - "vertex_ai/gemini-3.8-flash", - "gemini/gemini-3.5-flash", - "gemini/gemini-3.7-flash", - "gemini/gemini-3.8-flash", - ], -) -def test_placeholder_scoped_to_first_call_across_gemini_3_variants(model): - """The gemini-3 gate is a substring match, so every family member and prefix form has to - land on the same one-placeholder budget rather than only the versions we happened to try.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_gemini_tool_call_invoke, - ) - - gemini_parts = convert_to_gemini_tool_call_invoke( - { - "role": "assistant", - "content": None, - "tool_calls": _parallel_tool_calls(None, None, None), - }, - model=model, - ) - - assert len(gemini_parts) == 3 - assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE - assert "thoughtSignature" not in gemini_parts[1] - assert "thoughtSignature" not in gemini_parts[2] - - -def test_signed_text_part_survives_alongside_unsigned_parallel_tool_calls(): - """Text-part and function-call signatures are collected by separate code paths, so scoping the - placeholder must not disturb a real signature that arrived on the text part.""" - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - ) - - msg = { - "role": "assistant", - "content": "Checking all three cities.", - "provider_specific_fields": {"thought_signatures": ["real_25_signature"]}, - "tool_calls": _parallel_tool_calls(None, None, None), - } - - parts = _gemini_convert_messages_with_history( - messages=[msg], model="gemini-3-pro-preview" - )[0]["parts"] - - assert parts[0]["text"] == "Checking all three cities." - assert parts[0]["thoughtSignature"] == "real_25_signature" - assert parts[1]["thoughtSignature"] == PLACEHOLDER_SIGNATURE - assert "thoughtSignature" not in parts[2] - assert "thoughtSignature" not in parts[3] - - -# Tests for media_resolution (detail parameter) handling - Issue #17084 -class TestMediaResolution: - """Tests for media_resolution handling in Gemini 2.x models""" - - def test_get_highest_media_resolution_high_wins(self): - """Test that 'high' resolution takes precedence over 'low'""" - assert _get_highest_media_resolution("low", "high") == "high" - assert _get_highest_media_resolution("high", "low") == "high" - assert _get_highest_media_resolution(None, "high") == "high" - assert _get_highest_media_resolution("high", None) == "high" - - def test_get_highest_media_resolution_low_over_none(self): - """Test that 'low' resolution takes precedence over None""" - assert _get_highest_media_resolution(None, "low") == "low" - assert _get_highest_media_resolution("low", None) == "low" - - def test_get_highest_media_resolution_same_values(self): - """Test handling of same resolution values""" - assert _get_highest_media_resolution("high", "high") == "high" - assert _get_highest_media_resolution("low", "low") == "low" - assert _get_highest_media_resolution(None, None) is None - - def test_get_highest_media_resolution_medium(self): - """Test that 'medium' resolution is correctly ranked between 'low' and 'high'""" - assert _get_highest_media_resolution("low", "medium") == "medium" - assert _get_highest_media_resolution("medium", "low") == "medium" - assert _get_highest_media_resolution("medium", "high") == "high" - assert _get_highest_media_resolution("high", "medium") == "high" - assert _get_highest_media_resolution(None, "medium") == "medium" - assert _get_highest_media_resolution("medium", None) == "medium" - - def test_get_highest_media_resolution_ultra_high(self): - """Test that 'ultra_high' resolution takes precedence over all others""" - assert _get_highest_media_resolution("high", "ultra_high") == "ultra_high" - assert _get_highest_media_resolution("ultra_high", "high") == "ultra_high" - assert _get_highest_media_resolution("medium", "ultra_high") == "ultra_high" - assert _get_highest_media_resolution("low", "ultra_high") == "ultra_high" - assert _get_highest_media_resolution(None, "ultra_high") == "ultra_high" - assert _get_highest_media_resolution("ultra_high", None) == "ultra_high" - - def test_extract_max_media_resolution_single_image_high(self): - """Test extraction of media resolution from single image with detail=high""" - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What is this?"}, - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,abc123", - "detail": "high", - }, - }, - ], - } - ] - assert _extract_max_media_resolution_from_messages(messages) == "high" - - def test_extract_max_media_resolution_single_image_low(self): - """Test extraction of media resolution from single image with detail=low""" - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What is this?"}, - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,abc123", - "detail": "low", - }, - }, - ], - } - ] - assert _extract_max_media_resolution_from_messages(messages) == "low" - - def test_extract_max_media_resolution_no_detail(self): - """Test extraction when no detail parameter is provided""" - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What is this?"}, - { - "type": "image_url", - "image_url": {"url": "data:image/png;base64,abc123"}, - }, - ], - } - ] - assert _extract_max_media_resolution_from_messages(messages) is None - - def test_extract_max_media_resolution_multiple_images_mixed(self): - """Test that highest resolution is returned when multiple images have different details""" - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "Compare these images"}, - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,abc123", - "detail": "low", - }, - }, - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,def456", - "detail": "high", - }, - }, - ], - } - ] - assert _extract_max_media_resolution_from_messages(messages) == "high" - - def test_extract_max_media_resolution_text_only(self): - """Test extraction from messages with no images""" - messages = [ - {"role": "user", "content": "Hello, how are you?"}, - {"role": "assistant", "content": "I'm doing well!"}, - ] - assert _extract_max_media_resolution_from_messages(messages) is None - - def test_transform_request_body_gemini_2x_adds_media_resolution(self): - """Test that media_resolution is added to generationConfig for Gemini 2.x models""" - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What is this?"}, - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,iVBORw0KGgo=", - "detail": "high", - }, - }, - ], - } - ] - - result = _transform_request_body( - messages=messages, - model="gemini-2.5-flash", - optional_params={}, - custom_llm_provider="gemini", - litellm_params={}, - cached_content=None, - ) - - assert "generationConfig" in result - assert "mediaResolution" in result["generationConfig"] - assert result["generationConfig"]["mediaResolution"] == "MEDIA_RESOLUTION_HIGH" - - def test_transform_request_body_gemini_2x_low_resolution(self): - """Test that low media_resolution is correctly added for Gemini 2.x""" - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What is this?"}, - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,iVBORw0KGgo=", - "detail": "low", - }, - }, - ], - } - ] - - result = _transform_request_body( - messages=messages, - model="gemini-2.5-flash", - optional_params={}, - custom_llm_provider="gemini", - litellm_params={}, - cached_content=None, - ) - - assert "generationConfig" in result - assert "mediaResolution" in result["generationConfig"] - assert result["generationConfig"]["mediaResolution"] == "MEDIA_RESOLUTION_LOW" - - def test_transform_request_body_gemini_3_no_global_media_resolution(self): - """Test that Gemini 3 models don't add media_resolution to generationConfig (they use per-part)""" - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What is this?"}, - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,iVBORw0KGgo=", - "detail": "high", - }, - }, - ], - } - ] - - result = _transform_request_body( - messages=messages, - model="gemini-3-pro-preview", - optional_params={}, - custom_llm_provider="gemini", - litellm_params={}, - cached_content=None, - ) - - # Gemini 3 should NOT have mediaResolution in generationConfig - # (it's handled per-part in the content transformation) - if "generationConfig" in result: - assert "mediaResolution" not in result["generationConfig"] - - def test_transform_request_body_no_detail_no_media_resolution(self): - """Test that no mediaResolution is added when detail is not specified""" - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What is this?"}, - { - "type": "image_url", - "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}, - }, - ], - } - ] - - result = _transform_request_body( - messages=messages, - model="gemini-2.5-flash", - optional_params={}, - custom_llm_provider="gemini", - litellm_params={}, - cached_content=None, - ) - - # When no detail is specified, mediaResolution should not be in generationConfig - if "generationConfig" in result: - assert "mediaResolution" not in result["generationConfig"] - - def test_extract_max_media_resolution_file_type_with_detail(self): - """Test that detail is extracted from file content type, not just image_url""" - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What is in this file?"}, - { - "type": "file", - "file": { - "url": "data:image/png;base64,abc123", - "detail": "high", - }, - }, - ], - } - ] - assert _extract_max_media_resolution_from_messages(messages) == "high" - - def test_extract_max_media_resolution_mixed_image_and_file(self): - """Test that highest detail is returned across both image_url and file types""" - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "Compare these"}, - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,abc123", - "detail": "low", - }, - }, - { - "type": "file", - "file": { - "url": "data:image/png;base64,def456", - "detail": "high", - }, - }, - ], - } - ] - assert _extract_max_media_resolution_from_messages(messages) == "high" - - def test_transform_request_body_gemini_1x_no_media_resolution(self): - """Test that Gemini 1.x models don't get mediaResolution in generationConfig""" - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "What is this?"}, - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,iVBORw0KGgo=", - "detail": "high", - }, - }, - ], - } - ] - - result = _transform_request_body( - messages=messages, - model="gemini-1.5-pro", - optional_params={}, - custom_llm_provider="gemini", - litellm_params={}, - cached_content=None, - ) - - # Gemini 1.x should NOT have mediaResolution (not supported) - if "generationConfig" in result: - assert "mediaResolution" not in result["generationConfig"] - - -# Tests for VideoMetadata support across all Gemini models (Issue #25474) -class TestVideoMetadataAllGeminiModels: - """Tests that video_metadata (fps, start_offset, end_offset) works for all Gemini models""" - - def _make_video_messages(self, video_metadata: dict) -> list: - return [ - { - "role": "user", - "content": [ - {"type": "text", "text": "Analyze this video"}, - { - "type": "file", - "file": { - "file_id": "gs://bucket/video.mp4", - "format": "video/mp4", - "video_metadata": video_metadata, - }, - }, - ], - } - ] - - def _get_file_part(self, contents: list) -> dict: - for part in contents[0]["parts"]: - if "file_data" in part: - return part - raise AssertionError("No file part found in contents") - - def test_video_metadata_fps_gemini_2_5_flash(self): - """Gemini 2.5 Flash: fps in video_metadata should be forwarded (Issue #25474)""" - messages = self._make_video_messages({"fps": 5}) - contents = _gemini_convert_messages_with_history( - messages=messages, model="gemini-2.5-flash" - ) - file_part = self._get_file_part(contents) - assert "video_metadata" in file_part - assert file_part["video_metadata"]["fps"] == 5 - - def test_video_metadata_fps_gemini_2_5_pro(self): - """Gemini 2.5 Pro: fps in video_metadata should be forwarded (Issue #25474)""" - messages = self._make_video_messages({"fps": 10}) - contents = _gemini_convert_messages_with_history( - messages=messages, model="gemini-2.5-pro" - ) - file_part = self._get_file_part(contents) - assert "video_metadata" in file_part - assert file_part["video_metadata"]["fps"] == 10 - - def test_video_metadata_offsets_gemini_2_5_flash(self): - """Gemini 2.5 Flash: start_offset/end_offset converted to camelCase (Issue #25474)""" - messages = self._make_video_messages( - {"start_offset": "5s", "end_offset": "30s"} - ) - contents = _gemini_convert_messages_with_history( - messages=messages, model="gemini-2.5-flash" - ) - file_part = self._get_file_part(contents) - assert "video_metadata" in file_part - vm = file_part["video_metadata"] - assert vm["startOffset"] == "5s" - assert vm["endOffset"] == "30s" - - def test_video_metadata_all_fields_gemini_2_5_flash(self): - """Gemini 2.5 Flash: all video_metadata fields forwarded correctly (Issue #25474)""" - messages = self._make_video_messages( - {"fps": 5, "start_offset": "10s", "end_offset": "60s"} - ) - contents = _gemini_convert_messages_with_history( - messages=messages, model="gemini-2.5-flash" - ) - file_part = self._get_file_part(contents) - assert "video_metadata" in file_part - vm = file_part["video_metadata"] - assert vm["fps"] == 5 - assert vm["startOffset"] == "10s" - assert vm["endOffset"] == "60s" - - def test_video_metadata_gemini_1_5_pro(self): - """Gemini 1.5 Pro: video_metadata should also be forwarded (Issue #25474)""" - messages = self._make_video_messages({"fps": 2}) - contents = _gemini_convert_messages_with_history( - messages=messages, model="gemini-1.5-pro" - ) - file_part = self._get_file_part(contents) - assert "video_metadata" in file_part - assert file_part["video_metadata"]["fps"] == 2 - - -def test_convert_tool_response_with_base64_image(): - """Test tool response with base64 data URI image.""" - # Create a small test image (1x1 red pixel PNG) - test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" - image_data_uri = f"data:image/png;base64,{test_image_base64}" - - # Create tool message with image - tool_message = { - "role": "tool", - "tool_call_id": "call_test123", - "content": [ - { - "type": "text", - "text": '{"url": "https://example.com", "status": "success"}', - }, - {"type": "input_image", "image_url": image_data_uri}, - ], - } - - # Mock last message with tool calls - last_message_with_tool_calls = { - "tool_calls": [ - { - "id": "call_test123", - "function": {"name": "click_at", "arguments": '{"x": 100, "y": 200}'}, - } - ] - } - - # Convert tool response with nested multimodal functionResponse.parts. - result = convert_to_gemini_tool_call_result( - tool_message, last_message_with_tool_calls - ) - - assert isinstance(result, list), "Should return a parts list when media is present" - assert len(result) == 1, "Should return one function_response part" - result_part = result[0] - assert "function_response" in result_part - assert "inline_data" not in result_part - function_response = result_part["function_response"] - assert function_response["name"] == "click_at" - assert "response" in function_response - # Verify JSON response is parsed correctly - assert "url" in function_response["response"] - assert function_response["response"]["url"] == "https://example.com" - - # Check inline_data is nested under functionResponse.parts. - assert "parts" in function_response - assert len(function_response["parts"]) == 1 - inline_data: BlobType = function_response["parts"][0]["inline_data"] - assert "data" in inline_data - assert "mime_type" in inline_data - assert inline_data["mime_type"] == "image/png" - assert inline_data["data"] == test_image_base64 - - -def test_gemini_history_nests_multimodal_tool_response_parts(): - """Full history conversion should not emit sibling inline_data tool result parts.""" - test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" - messages = [ - {"role": "user", "content": "Get me an image"}, - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_get_image", - "type": "function", - "function": {"name": "get_image", "arguments": "{}"}, - } - ], - }, - { - "role": "tool", - "tool_call_id": "call_get_image", - "content": [ - {"type": "text", "text": '{"image_ref": "inline"}'}, - { - "type": "image", - "source": { - "type": "base64", - "media_type": "image/png", - "data": test_image_base64, - }, - }, - ], - }, - ] - - contents = _gemini_convert_messages_with_history(messages=messages) - - tool_response_parts = contents[-1]["parts"] - assert len(tool_response_parts) == 1 - assert "inline_data" not in tool_response_parts[0] - function_response = tool_response_parts[0]["function_response"] - assert function_response["parts"] == [ - { - "inline_data": { - "data": test_image_base64, - "mime_type": "image/png", - } - } - ] def test_convert_tool_response_with_url_image(): """Test tool response with HTTP URL image (will download and convert).""" - import pytest - # Use a publicly accessible test image URL test_image_url = "https://via.placeholder.com/1x1.png" @@ -1701,13 +33,9 @@ def test_convert_tool_response_with_url_image(): } try: - result = convert_to_gemini_tool_call_result( - tool_message, last_message_with_tool_calls - ) + result = convert_to_gemini_tool_call_result(tool_message, last_message_with_tool_calls) - assert isinstance( - result, list - ), "Should return a parts list when media is present" + assert isinstance(result, list), "Should return a parts list when media is present" assert len(result) == 1, "Should return one function_response part" result_part = result[0] assert "function_response" in result_part @@ -1724,1060 +52,3 @@ def test_convert_tool_response_with_url_image(): except Exception as e: # Skip test if URL download fails (no internet connection, etc.) pytest.skip(f"Failed to download image from URL: {e}") - - -def test_convert_tool_response_text_only(): - """Test tool response with only text (no image).""" - tool_message = { - "role": "tool", - "tool_call_id": "call_test789", - "content": [ - {"type": "text", "text": '{"status": "completed", "result": "success"}'} - ], - } - - last_message_with_tool_calls = { - "tool_calls": [ - { - "id": "call_test789", - "function": {"name": "wait_5_seconds", "arguments": "{}"}, - } - ] - } - - result = convert_to_gemini_tool_call_result( - tool_message, last_message_with_tool_calls - ) - - # Should be a single part (no list) when no image - assert not isinstance(result, list), "Should return single part when no image" - - # Check function_response exists - assert "function_response" in result - function_response = result["function_response"] - assert function_response["name"] == "wait_5_seconds" - # Verify JSON response is parsed correctly - assert "status" in function_response["response"] - assert function_response["response"]["status"] == "completed" - - # Check inline_data does NOT exist (no image provided) - assert "inline_data" not in result - - -def test_file_data_field_order(): - """ - Test that file_data fields are in the correct order (mime_type before file_uri). - - The Gemini API is sensitive to field order in the file_data object. - This test verifies that mime_type comes before file_uri in both: - 1. Dictionary key order - 2. JSON serialization - - Related issue: Gemini API returns 400 INVALID_ARGUMENT when fields are in wrong order. - """ - import json - - from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media - - # Test with HTTPS URL and explicit format (audio file) - file_url = "https://generativelanguage.googleapis.com/v1beta/files/test123" - format = "audio/mpeg" - - result = _process_gemini_media(image_url=file_url, format=format) - - # Verify the result has file_data - assert "file_data" in result - file_data = result["file_data"] - - # Verify both fields are present - assert "mime_type" in file_data - assert "file_uri" in file_data - assert file_data["mime_type"] == "audio/mpeg" - assert file_data["file_uri"] == file_url - - # Verify field order by checking dictionary keys - # In Python 3.7+, dict maintains insertion order - file_data_keys = list(file_data.keys()) - assert file_data_keys.index("mime_type") < file_data_keys.index( - "file_uri" - ), "mime_type must come before file_uri in the file_data dict" - - # Also verify by serializing to JSON string - json_str = json.dumps(file_data) - mime_type_pos = json_str.find('"mime_type"') - file_uri_pos = json_str.find('"file_uri"') - assert ( - mime_type_pos < file_uri_pos - ), "mime_type must appear before file_uri in JSON serialization" - - -def test_file_data_field_order_gcs_urls(): - """Test that GCS URLs also maintain correct field order.""" - import json - - from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media - - # Test with GCS URL - gcs_url = "gs://bucket/audio.mp3" - - result = _process_gemini_media(image_url=gcs_url) - - # Verify the result has file_data - assert "file_data" in result - file_data = result["file_data"] - - # Verify both fields are present - assert "mime_type" in file_data - assert "file_uri" in file_data - - # Verify field order - file_data_keys = list(file_data.keys()) - assert file_data_keys.index("mime_type") < file_data_keys.index( - "file_uri" - ), "mime_type must come before file_uri in the file_data dict" - - -def test_gemini_files_api_uri_without_format(): - """ - Test that Gemini Files API URIs work WITHOUT an explicit format/mime_type. - - When a user uploads a file via the Gemini Files API and then references it - by URI (https://generativelanguage.googleapis.com/v1beta/files/...), - the file is already on Google's servers. These URLs return 403 when - fetched directly, so _process_gemini_media must NOT try to resolve the - MIME type via HTTP. Instead it should pass the URI through as file_data - and let the Gemini API resolve the type from its stored metadata. - - Related issue: https://github.com/BerriAI/litellm/issues/24907 - """ - from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media - - file_url = "https://generativelanguage.googleapis.com/v1beta/files/37eh7rsw1vfe" - - # Should NOT raise — previously this hit the generic https:// handler - # which called _get_image_mime_type_from_url() and got a 403. - result = _process_gemini_media(image_url=file_url) - - assert "file_data" in result - file_data = result["file_data"] - assert file_data["file_uri"] == file_url - # When no format is provided, mime_type should be absent so the - # Gemini API infers it from the stored file metadata. - assert "mime_type" not in file_data - - -def test_gemini_files_api_uri_with_format(): - """ - Test that Gemini Files API URIs correctly forward an explicit format. - - Related issue: https://github.com/BerriAI/litellm/issues/24907 - """ - from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media - - file_url = "https://generativelanguage.googleapis.com/v1beta/files/n1vhxa28lyaw" - - result = _process_gemini_media(image_url=file_url, format="text/plain") - - assert "file_data" in result - file_data = result["file_data"] - assert file_data["file_uri"] == file_url - assert file_data["mime_type"] == "text/plain" - - -def test_extract_file_data_with_path_object(): - """ - Test that filename is correctly extracted from Path objects for MIME type detection. - - When uploading files using Path objects (e.g., Path("speech.mp3")), the filename - must be extracted to enable proper MIME type detection. Without this, files get - uploaded with 'application/octet-stream' instead of the correct MIME type. - - Related issue: Files uploaded with wrong MIME type cause Gemini API to reject - requests where the specified format doesn't match the uploaded file's MIME type. - """ - import os - import tempfile - from pathlib import Path - - from litellm.litellm_core_utils.prompt_templates.common_utils import ( - extract_file_data, - ) - - # Create a temporary MP3 file - with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as tmp: - tmp.write(b"fake mp3 content") - tmp_path = tmp.name - - try: - # Test with Path object - path_obj = Path(tmp_path) - extracted = extract_file_data(path_obj) - - # Verify filename was extracted - assert extracted["filename"] is not None - assert extracted["filename"].endswith(".mp3") - - # Verify MIME type was correctly detected - assert ( - extracted["content_type"] == "audio/mpeg" - ), f"Expected 'audio/mpeg' but got '{extracted['content_type']}'" - - # Verify content was read - assert extracted["content"] == b"fake mp3 content" - - finally: - # Clean up temporary file - os.unlink(tmp_path) - - -def test_extract_file_data_with_pathlib_path(): - """Test that filename is correctly extracted from pathlib.Path inputs. - Bare str paths are rejected — when this runs in a proxy request handler - the value is attacker-controlled and opening it as a path is an LFI.""" - import os - import tempfile - from pathlib import Path - - from litellm.litellm_core_utils.prompt_templates.common_utils import ( - extract_file_data, - ) - - with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: - tmp.write(b"fake wav content") - tmp_path = Path(tmp.name) - - try: - extracted = extract_file_data(tmp_path) - - assert extracted["filename"] is not None - assert extracted["filename"].endswith(".wav") - assert extracted["content_type"] in [ - "audio/wav", - "audio/x-wav", - ], f"Expected 'audio/wav' or 'audio/x-wav' but got '{extracted['content_type']}'" - assert extracted["content"] == b"fake wav content" - finally: - os.unlink(str(tmp_path)) - - -def test_extract_file_data_with_tuple_format(): - """Test that tuple format (with explicit content_type) still works correctly.""" - from litellm.litellm_core_utils.prompt_templates.common_utils import ( - extract_file_data, - ) - - # Test with tuple format: (filename, content, content_type) - filename = "test_audio.mp3" - content = b"test audio content" - content_type = "audio/mpeg" - - extracted = extract_file_data((filename, content, content_type)) - - # Verify all fields are correct - assert extracted["filename"] == filename - assert extracted["content"] == content - assert extracted["content_type"] == content_type - - -def test_extract_file_data_fallback_to_octet_stream(): - """Unknown file types fall back to application/octet-stream.""" - import os - import tempfile - from pathlib import Path - - from litellm.litellm_core_utils.prompt_templates.common_utils import ( - extract_file_data, - ) - - with tempfile.NamedTemporaryFile(suffix=".xyz123", delete=False) as tmp: - tmp.write(b"unknown content") - tmp_path = Path(tmp.name) - - try: - extracted = extract_file_data(tmp_path) - - assert extracted["filename"] is not None - assert extracted["filename"].endswith(".xyz123") - assert ( - extracted["content_type"] == "application/octet-stream" - ), f"Expected 'application/octet-stream' for unknown type, got '{extracted['content_type']}'" - finally: - os.unlink(str(tmp_path)) - - -def test_convert_tool_response_with_pdf_file(): - """Test tool response with PDF file content using file_data field.""" - # Create a minimal test PDF (base64 encoded) - test_pdf_base64 = "JVBERi0xLjQKJeLjz9MKMSAwIG9iago8PC9UeXBlL0NhdGFsb2cvUGFnZXMgMiAwIFI+PgplbmRvYmoKdHJhaWxlcgo8PC9TaXplIDQvUm9vdCAxIDAgUj4+CnN0YXJ0eHJlZgoyMTYKJSVFT0Y=" - file_data_uri = f"data:application/pdf;base64,{test_pdf_base64}" - - # Create tool message with file - tool_message = { - "role": "tool", - "tool_call_id": "call_pdf_test", - "content": [ - {"type": "text", "text": '{"status": "success", "pages": 1}'}, - {"type": "file", "file_data": file_data_uri}, - ], - } - - # Mock last message with tool calls - last_message_with_tool_calls = { - "tool_calls": [ - { - "id": "call_pdf_test", - "function": { - "name": "analyze_document", - "arguments": '{"path": "/tmp/doc.pdf"}', - }, - } - ] - } - - # Convert tool response with nested multimodal functionResponse.parts. - result = convert_to_gemini_tool_call_result( - tool_message, last_message_with_tool_calls - ) - - assert isinstance(result, list), "Should return a parts list when media is present" - assert len(result) == 1, "Should return one function_response part" - result_part = result[0] - assert "function_response" in result_part - assert "inline_data" not in result_part - function_response = result_part["function_response"] - assert function_response["name"] == "analyze_document" - assert "response" in function_response - # Verify JSON response is parsed correctly - assert "status" in function_response["response"] - assert function_response["response"]["status"] == "success" - - # Check inline_data is nested under functionResponse.parts. - assert "parts" in function_response - assert len(function_response["parts"]) == 1 - inline_data: BlobType = function_response["parts"][0]["inline_data"] - assert "data" in inline_data - assert "mime_type" in inline_data - assert inline_data["mime_type"] == "application/pdf" - assert inline_data["data"] == test_pdf_base64 - - -def test_convert_tool_response_with_input_file_type(): - """Test tool response with input_file content type (Responses API format).""" - # Create a minimal test PDF (base64 encoded) - test_pdf_base64 = "JVBERi0xLjQKJeLjz9MKMSAwIG9iago8PC9UeXBlL0NhdGFsb2cvUGFnZXMgMiAwIFI+PgplbmRvYmoKdHJhaWxlcgo8PC9TaXplIDQvUm9vdCAxIDAgUj4+CnN0YXJ0eHJlZgoyMTYKJSVFT0Y=" - file_data_uri = f"data:application/pdf;base64,{test_pdf_base64}" - - # Create tool message with input_file type - tool_message = { - "role": "tool", - "tool_call_id": "call_input_file_test", - "content": [{"type": "input_file", "file_data": file_data_uri}], - } - - # Mock last message with tool calls - last_message_with_tool_calls = { - "tool_calls": [ - { - "id": "call_input_file_test", - "function": {"name": "read_file", "arguments": "{}"}, - } - ] - } - - # Convert tool response - result = convert_to_gemini_tool_call_result( - tool_message, last_message_with_tool_calls - ) - - # Check inline_data is nested under functionResponse.parts. - assert isinstance(result, list), "Should return a parts list when media is present" - assert len(result) == 1, "Should return one function_response part" - function_response = result[0]["function_response"] - assert ( - function_response["parts"][0]["inline_data"]["mime_type"] == "application/pdf" - ) - - -def test_convert_tool_response_with_nested_file_object(): - """Test tool response with file content using nested file object format.""" - # Create a minimal test PDF (base64 encoded) - test_pdf_base64 = "JVBERi0xLjQKJeLjz9MKMSAwIG9iago8PC9UeXBlL0NhdGFsb2cvUGFnZXMgMiAwIFI+PgplbmRvYmoKdHJhaWxlcgo8PC9TaXplIDQvUm9vdCAxIDAgUj4+CnN0YXJ0eHJlZgoyMTYKJSVFT0Y=" - file_data_uri = f"data:application/pdf;base64,{test_pdf_base64}" - - # Create tool message with nested file object (OpenAI Agents SDK format) - tool_message = { - "role": "tool", - "tool_call_id": "call_nested_test", - "content": [{"type": "file", "file": {"file_data": file_data_uri}}], - } - - # Mock last message with tool calls - last_message_with_tool_calls = { - "tool_calls": [ - { - "id": "call_nested_test", - "function": {"name": "process_document", "arguments": "{}"}, - } - ] - } - - # Convert tool response - result = convert_to_gemini_tool_call_result( - tool_message, last_message_with_tool_calls - ) - - # Check inline_data is nested under functionResponse.parts. - assert isinstance(result, list), "Should return a parts list when media is present" - assert len(result) == 1, "Should return one function_response part" - function_response = result[0]["function_response"] - inline_data: BlobType = function_response["parts"][0]["inline_data"] - assert "data" in inline_data - assert "mime_type" in inline_data - assert inline_data["mime_type"] == "application/pdf" - assert inline_data["data"] == test_pdf_base64 - - -def test_assistant_message_with_images_field(): - """ - Test that assistant messages with images field are properly converted to Gemini format. - - This handles the case where an assistant message contains generated images in the - `images` field (e.g., from image generation models like gemini-2.5-flash-image). - The images should be converted to inline_data parts in the Gemini format. - """ - # Create a small test image (1x1 red pixel PNG) - test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" - image_data_uri = f"data:image/png;base64,{test_image_base64}" - - # Create messages with assistant message containing images field - messages = [ - { - "role": "user", - "content": "Generate an image of a banana wearing a costume that says LiteLLM", - }, - { - "role": "assistant", - "content": "Here's your banana in a LiteLLM costume!", - "images": [ - { - "image_url": {"url": image_data_uri, "detail": "auto"}, - "index": 0, - "type": "image_url", - } - ], - }, - ] - - # Convert messages to Gemini format - contents = _gemini_convert_messages_with_history(messages=messages) - - # Verify structure - assert len(contents) == 2, f"Expected 2 content blocks, got {len(contents)}" - - # Verify user message - assert contents[0]["role"] == "user" - assert len(contents[0]["parts"]) == 1 - assert ( - contents[0]["parts"][0]["text"] - == "Generate an image of a banana wearing a costume that says LiteLLM" - ) - - # Verify assistant message - assert contents[1]["role"] == "model" - assert ( - len(contents[1]["parts"]) == 2 - ), f"Expected 2 parts (text + image), got {len(contents[1]['parts'])}" - - # Find text part and inline_data part - text_part = None - inline_data_part = None - for part in contents[1]["parts"]: - if "text" in part: - text_part = part - elif "inline_data" in part: - inline_data_part = part - - # Verify text part - assert text_part is not None, "Missing text part in assistant message" - assert text_part["text"] == "Here's your banana in a LiteLLM costume!" - - # Verify inline_data part (image) - assert inline_data_part is not None, "Missing inline_data part in assistant message" - inline_data: BlobType = inline_data_part["inline_data"] - assert "data" in inline_data - assert "mime_type" in inline_data - assert inline_data["mime_type"] == "image/png" - assert inline_data["data"] == test_image_base64 - - -def test_assistant_message_with_multiple_images(): - """Test that assistant messages with multiple images are properly converted.""" - # Create two test images - test_image1_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" - test_image2_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg==" - image1_data_uri = f"data:image/png;base64,{test_image1_base64}" - image2_data_uri = f"data:image/jpeg;base64,{test_image2_base64}" - - messages = [ - {"role": "user", "content": "Generate two images"}, - { - "role": "assistant", - "content": "Here are your images:", - "images": [ - { - "image_url": {"url": image1_data_uri, "detail": "auto"}, - "index": 0, - "type": "image_url", - }, - { - "image_url": {"url": image2_data_uri, "detail": "high"}, - "index": 1, - "type": "image_url", - }, - ], - }, - ] - - # Convert messages to Gemini format - contents = _gemini_convert_messages_with_history(messages=messages) - - # Verify assistant message has 3 parts (1 text + 2 images) - assert contents[1]["role"] == "model" - assert ( - len(contents[1]["parts"]) == 3 - ), f"Expected 3 parts (text + 2 images), got {len(contents[1]['parts'])}" - - # Count inline_data parts - inline_data_parts = [part for part in contents[1]["parts"] if "inline_data" in part] - assert ( - len(inline_data_parts) == 2 - ), f"Expected 2 inline_data parts, got {len(inline_data_parts)}" - - # Verify first image - assert inline_data_parts[0]["inline_data"]["mime_type"] == "image/png" - assert inline_data_parts[0]["inline_data"]["data"] == test_image1_base64 - - # Verify second image - assert inline_data_parts[1]["inline_data"]["mime_type"] == "image/jpeg" - assert inline_data_parts[1]["inline_data"]["data"] == test_image2_base64 - - -def test_assistant_message_with_images_using_message_object(): - """Test that Message objects with images field are properly converted.""" - # Create a small test image - test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" - image_data_uri = f"data:image/png;base64,{test_image_base64}" - - # Create messages using Message object (as returned by LiteLLM) - user_message = {"role": "user", "content": "Generate an image"} - - assistant_message = Message( - content="Here's your image!", - role="assistant", - tool_calls=None, - function_call=None, - images=[ - { - "image_url": {"url": image_data_uri, "detail": "auto"}, - "index": 0, - "type": "image_url", - } - ], - ) - - messages = [user_message, assistant_message] - - # Convert messages to Gemini format - contents = _gemini_convert_messages_with_history(messages=messages) - - # Verify assistant message has both text and image - assert contents[1]["role"] == "model" - assert len(contents[1]["parts"]) == 2 - - # Verify image was converted - inline_data_parts = [part for part in contents[1]["parts"] if "inline_data" in part] - assert len(inline_data_parts) == 1 - assert inline_data_parts[0]["inline_data"]["mime_type"] == "image/png" - assert inline_data_parts[0]["inline_data"]["data"] == test_image_base64 - - -def test_assistant_message_with_images_in_conversation_history(): - """ - Test multi-turn conversation where assistant message with images is in history. - - This simulates the real use case where: - 1. User asks for image generation - 2. Assistant generates image (with images field) - 3. User asks follow-up question about the image - """ - test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" - image_data_uri = f"data:image/png;base64,{test_image_base64}" - - messages = [ - {"role": "user", "content": "Generate an image of a cat"}, - { - "role": "assistant", - "content": "Here's a cat image:", - "images": [ - { - "image_url": {"url": image_data_uri, "detail": "auto"}, - "index": 0, - "type": "image_url", - } - ], - }, - {"role": "user", "content": "Can you make it more colorful?"}, - ] - - # Convert messages to Gemini format - contents = _gemini_convert_messages_with_history(messages=messages) - - # Verify structure: user -> model (with image) -> user - assert len(contents) == 3 - assert contents[0]["role"] == "user" - assert contents[1]["role"] == "model" - assert contents[2]["role"] == "user" - - # Verify assistant message has image in history - inline_data_parts = [part for part in contents[1]["parts"] if "inline_data" in part] - assert len(inline_data_parts) == 1 - assert inline_data_parts[0]["inline_data"]["mime_type"] == "image/png" - - -def test_function_response_has_user_role(): - """ - Test that function response ContentType blocks include role="user". - - Gemini API only accepts two roles: "user" and "model". Function responses - must be sent with role="user". Previously, LiteLLM omitted the role field - entirely, causing 400 errors from the Gemini API. - - Fixes: https://github.com/BerriAI/litellm/issues/22003 - Fixes: https://github.com/BerriAI/litellm/issues/20690 - """ - messages = [ - {"role": "user", "content": "What is the weather in Berlin?"}, - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_abc123", - "type": "function", - "function": { - "name": "get_weather", - "arguments": '{"city": "Berlin"}', - }, - } - ], - }, - { - "role": "tool", - "tool_call_id": "call_abc123", - "content": '{"temperature": "15°C", "condition": "Cloudy"}', - }, - ] - - contents = _gemini_convert_messages_with_history(messages=messages) - - # Expect: user -> model (functionCall) -> user (functionResponse) - assert len(contents) == 3 - - assert contents[0]["role"] == "user" - assert contents[1]["role"] == "model" - assert "function_call" in contents[1]["parts"][0] - - # The critical assertion: function response must have role="user" - assert contents[2]["role"] == "user" - assert "function_response" in contents[2]["parts"][0] - - -def test_multi_turn_function_calling_roles(): - """ - Test a full multi-turn function calling conversation produces correct roles. - - Simulates: user asks → model calls tool → tool responds → model answers → user asks again. - Every content block must have an explicit role of "user" or "model". - - Fixes: https://github.com/BerriAI/litellm/issues/22003 - """ - messages = [ - {"role": "user", "content": "What is the weather in Berlin?"}, - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_001", - "type": "function", - "function": { - "name": "get_weather", - "arguments": '{"city": "Berlin"}', - }, - } - ], - }, - { - "role": "tool", - "tool_call_id": "call_001", - "content": '{"temperature": "15°C"}', - }, - { - "role": "assistant", - "content": "The weather in Berlin is 15°C.", - }, - {"role": "user", "content": "And in Paris?"}, - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_002", - "type": "function", - "function": { - "name": "get_weather", - "arguments": '{"city": "Paris"}', - }, - } - ], - }, - { - "role": "tool", - "tool_call_id": "call_002", - "content": '{"temperature": "18°C"}', - }, - ] - - contents = _gemini_convert_messages_with_history(messages=messages) - - # Every content block must have a valid role - for i, content in enumerate(contents): - assert "role" in content, f"Content block {i} missing 'role' field" - assert content["role"] in ( - "user", - "model", - ), f"Content block {i} has invalid role: {content.get('role')}" - - # Verify the function response blocks specifically have role="user" - for i, content in enumerate(contents): - for part in content["parts"]: - if "function_response" in part: - assert ( - content["role"] == "user" - ), f"Content block {i} with function_response has role='{content['role']}', expected 'user'" - - -def test_gemini_thought_signature_preservation_real_response(): - """Test that thought signatures are preserved on the text part if originally there, without dropping or duplicating (real response case).""" - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexGeminiConfig, - ) - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - ) - - real_candidate = { - "content": { - "parts": [ - { - "text": "I will explain and then list files.", - "thoughtSignature": "mock_signature_from_text_part", - }, - { - "functionCall": { - "name": "list_files", - "args": {}, - } - }, - ] - } - } - - parts = real_candidate["content"]["parts"] - - content, reasoning_content = ( - VertexGeminiConfig().get_assistant_content_message(parts=parts) - ) - thought_signatures = ( - VertexGeminiConfig()._extract_thought_signatures_from_parts( - parts=parts - ) - ) - functions, tools, _ = VertexGeminiConfig._transform_parts( - parts=parts, - cumulative_tool_call_idx=0, - is_function_call=False, - ) - - msg: dict = {"role": "assistant"} - if content is not None: - msg["content"] = content - if tools: - msg["tool_calls"] = tools - if functions is not None: - msg["function_call"] = functions - if thought_signatures is not None: - msg["provider_specific_fields"] = { - "thought_signatures": thought_signatures - } - - converted_real = _gemini_convert_messages_with_history( - messages=[msg], - model="gemini-2.5-pro", - ) - - assert len(converted_real) == 1 - assert "parts" in converted_real[0] - parts_out = converted_real[0]["parts"] - assert len(parts_out) == 2 - assert "text" in parts_out[0] - assert ( - parts_out[0]["thoughtSignature"] == "mock_signature_from_text_part" - ) - assert "function_call" in parts_out[1] - assert "thoughtSignature" not in parts_out[1] - - -def test_gemini_thought_signature_deduplication_assumed_response(): - """Test that thought signatures are deduplicated and not attached to the text part if already present in the tool call (assumed response case).""" - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - ) - - pr_assumed_msg = { - "role": "assistant", - "content": "I will list the directory.", - "provider_specific_fields": { - "thought_signatures": ["mock_signature_63k"] - }, - "tool_calls": [ - { - "id": "call_1", - "type": "function", - "function": {"name": "list_files", "arguments": "{}"}, - "provider_specific_fields": { - "thought_signature": "mock_signature_63k" - }, - } - ], - } - - converted_pr = _gemini_convert_messages_with_history( - messages=[pr_assumed_msg], - model="gemini-2.5-pro", - ) - - assert len(converted_pr) == 1 - assert "parts" in converted_pr[0] - parts_out = converted_pr[0]["parts"] - assert len(parts_out) == 2 - assert "text" in parts_out[0] - assert "thoughtSignature" not in parts_out[0] - assert "function_call" in parts_out[1] - assert parts_out[1]["thoughtSignature"] == "mock_signature_63k" - - -def test_gemini_thought_signature_pure_text(): - """Test that thought signatures are preserved on the text part for responses with no tool calls.""" - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - ) - - msg = { - "role": "assistant", - "content": "Hello, I am a model.", - "provider_specific_fields": { - "thought_signatures": ["pure_text_signature"] - }, - } - - converted = _gemini_convert_messages_with_history( - messages=[msg], - model="gemini-2.5-pro", - ) - - assert len(converted) == 1 - assert "parts" in converted[0] - parts_out = converted[0]["parts"] - assert len(parts_out) == 1 - assert "text" in parts_out[0] - assert parts_out[0]["thoughtSignature"] == "pure_text_signature" - - -def test_gemini_thought_signature_pure_tool_call(): - """Test that thought signatures are preserved on the tool call for responses with no intermediate text.""" - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - ) - - msg = { - "role": "assistant", - "content": None, - "provider_specific_fields": { - "thought_signatures": ["pure_tool_signature"] - }, - "tool_calls": [ - { - "id": "call_1", - "type": "function", - "function": {"name": "list_files", "arguments": "{}"}, - "provider_specific_fields": { - "thought_signature": "pure_tool_signature" - }, - } - ], - } - - converted = _gemini_convert_messages_with_history( - messages=[msg], - model="gemini-2.5-pro", - ) - - assert len(converted) == 1 - assert "parts" in converted[0] - parts_out = converted[0]["parts"] - assert len(parts_out) == 1 - assert "function_call" in parts_out[0] - assert parts_out[0]["thoughtSignature"] == "pure_tool_signature" - - -def test_gemini_distinct_text_and_tool_signatures_are_both_preserved(): - """A text-part signature that differs from the tool-call signature must stay on the text part.""" - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - ) - - msg = { - "role": "assistant", - "content": "Some analysis.", - "provider_specific_fields": { - "thought_signatures": ["text_signature", "tool_signature"] - }, - "tool_calls": [ - { - "id": "call_1", - "type": "function", - "function": {"name": "list_files", "arguments": "{}"}, - "provider_specific_fields": {"thought_signature": "tool_signature"}, - } - ], - } - - parts = _gemini_convert_messages_with_history( - messages=[msg], model="gemini-2.5-pro" - )[0]["parts"] - - assert parts[0]["text"] == "Some analysis." - assert parts[0]["thoughtSignature"] == "text_signature" - assert "function_call" in parts[1] - assert parts[1]["thoughtSignature"] == "tool_signature" - - -def test_gemini_25_text_signature_survives_replay_to_gemini_3(): - """gemini-2.5 history (signed text, unsigned tool call) replayed to gemini-3 keeps the real - text signature; the dummy signature synthesized for the unsigned tool call must not suppress it.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - _get_dummy_thought_signature, - ) - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - ) - - msg = { - "role": "assistant", - "content": "I will list the directory.", - "provider_specific_fields": {"thought_signatures": ["real_25_signature"]}, - "tool_calls": [ - { - "id": "call_1", - "type": "function", - "function": {"name": "list_files", "arguments": "{}"}, - } - ], - } - - parts = _gemini_convert_messages_with_history(messages=[msg], model="gemini-3-pro")[ - 0 - ]["parts"] - - assert parts[0]["text"] == "I will list the directory." - assert parts[0]["thoughtSignature"] == "real_25_signature" - assert "function_call" in parts[1] - assert parts[1]["thoughtSignature"] == _get_dummy_thought_signature() - - -def test_gemini_function_call_signature_round_trip_no_duplicate(): - """End to end: a gemini-3-style response (unsigned text + signed functionCall) parsed and - re-serialized sends the signature exactly once, on the function-call part.""" - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - ) - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexGeminiConfig, - ) - - response_parts = [ - {"text": "I will calculate the result for you."}, - { - "functionCall": {"name": "add_numbers", "args": {"a": 17, "b": 25}}, - "thoughtSignature": "signature_from_function_call", - }, - ] - - config = VertexGeminiConfig() - content, _ = config.get_assistant_content_message(parts=response_parts) - thought_signatures = config._extract_thought_signatures_from_parts( - parts=response_parts - ) - _, tools, _ = VertexGeminiConfig._transform_parts( - parts=response_parts, cumulative_tool_call_idx=0, is_function_call=False - ) - - msg = { - "role": "assistant", - "content": content, - "tool_calls": tools, - "provider_specific_fields": {"thought_signatures": thought_signatures}, - } - - parts = _gemini_convert_messages_with_history(messages=[msg], model="gemini-3-pro")[ - 0 - ]["parts"] - - signatures = [p["thoughtSignature"] for p in parts if "thoughtSignature" in p] - assert signatures == ["signature_from_function_call"] - assert "thoughtSignature" not in parts[0] - assert "function_call" in parts[1] - - -def test_gemini_server_side_tool_signature_not_duplicated_on_text(): - """A signature already re-injected on a server-side toolCall part is not attached to the text part again.""" - from litellm.llms.vertex_ai.gemini.transformation import ( - _gemini_convert_messages_with_history, - ) - - msg = { - "role": "assistant", - "content": "The weather in Buenos Aires is sunny.", - "provider_specific_fields": { - "thought_signatures": ["server_side_signature"], - "server_side_tool_invocations": [ - { - "tool_type": "GOOGLE_SEARCH_WEB", - "id": "abc123", - "args": {"queries": ["weather Buenos Aires"]}, - "response": {"weather": "Sunny"}, - "thought_signature": "server_side_signature", - } - ], - }, - } - - parts = _gemini_convert_messages_with_history( - messages=[msg], model="gemini-2.5-pro" - )[0]["parts"] - - text_part = next(p for p in parts if "text" in p) - assert "thoughtSignature" not in text_part - tool_call_part = next(p for p in parts if "toolCall" in p) - assert tool_call_part["thoughtSignature"] == "server_side_signature" diff --git a/tests/test_litellm/llms/vertex_ai/image_edit/__init__.py b/tests/test_litellm/llms/vertex_ai/image_edit/__init__.py deleted file mode 100644 index 50135ba1f92..00000000000 --- a/tests/test_litellm/llms/vertex_ai/image_edit/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Vertex AI Image Edit Tests diff --git a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py index 54607cc5284..aeba9f0fa3c 100644 --- a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py @@ -1,13 +1,9 @@ import os -from unittest.mock import MagicMock, patch +from unittest.mock import patch -import httpx import pytest -from litellm.llms.vertex_ai.image_generation import ( - get_vertex_ai_image_generation_config, -) from litellm.llms.vertex_ai.image_generation.vertex_gemini_transformation import ( VertexAIGeminiImageGenerationConfig, ) @@ -16,588 +12,6 @@ from litellm.llms.vertex_ai.image_generation.vertex_imagen_transformation import ) -class TestVertexAIGeminiImageGenerationConfig: - def setup_method(self): - """Set up test fixtures""" - self.config = VertexAIGeminiImageGenerationConfig() - - def test_get_supported_openai_params(self): - """Test get_supported_openai_params returns correct params""" - supported = self.config.get_supported_openai_params("gemini-2.5-flash-image") - assert "n" in supported - assert "size" in supported - - def test_map_openai_params_n(self): - """Test mapping n parameter to candidate_count""" - non_default_params = {"n": 3} - optional_params = {} - result = self.config.map_openai_params(non_default_params, optional_params, "gemini-2.5-flash-image", False) - assert result.get("candidate_count") == 3 - - def test_map_openai_params_size(self): - """Test mapping size parameter to aspectRatio""" - non_default_params = {"size": "1024x1024"} - optional_params = {} - result = self.config.map_openai_params(non_default_params, optional_params, "gemini-2.5-flash-image", False) - assert result.get("aspectRatio") == "1:1" - - def test_map_openai_params_size_16_9(self): - """Test mapping 16:9 size""" - non_default_params = {"size": "1792x1024"} - optional_params = {} - result = self.config.map_openai_params(non_default_params, optional_params, "gemini-2.5-flash-image", False) - assert result.get("aspectRatio") == "16:9" - - def test_map_size_to_aspect_ratio(self): - """Test size to aspect ratio mapping""" - assert self.config._map_size_to_aspect_ratio("1024x1024") == "1:1" - assert self.config._map_size_to_aspect_ratio("1792x1024") == "16:9" - assert self.config._map_size_to_aspect_ratio("1024x1792") == "9:16" - assert self.config._map_size_to_aspect_ratio("1280x896") == "4:3" - assert self.config._map_size_to_aspect_ratio("896x1280") == "3:4" - assert self.config._map_size_to_aspect_ratio("unknown") == "1:1" # default - - def test_get_supported_openai_params_includes_native_gemini_params(self): - """Test that native Gemini imageConfig params are supported""" - supported = self.config.get_supported_openai_params("gemini-3-pro-image-preview") - assert "aspectRatio" in supported - assert "aspect_ratio" in supported - assert "imageSize" in supported - assert "image_size" in supported - assert "imageConfig" in supported - - def test_map_openai_params_aspect_ratio_camel_case(self): - """Test mapping native aspectRatio parameter""" - result = self.config.map_openai_params({"aspectRatio": "9:16"}, {}, "gemini-3-pro-image-preview", False) - assert result["aspectRatio"] == "9:16" - - def test_map_openai_params_aspect_ratio_snake_case(self): - """Test mapping native aspect_ratio parameter""" - result = self.config.map_openai_params({"aspect_ratio": "16:9"}, {}, "gemini-3-pro-image-preview", False) - assert result["aspectRatio"] == "16:9" - - def test_map_openai_params_image_size_camel_case(self): - """Test mapping native imageSize parameter""" - result = self.config.map_openai_params({"imageSize": "4K"}, {}, "gemini-3-pro-image-preview", False) - assert result["imageSize"] == "4K" - - def test_map_openai_params_image_size_snake_case(self): - """Test mapping native image_size parameter""" - result = self.config.map_openai_params({"image_size": "2K"}, {}, "gemini-3-pro-image-preview", False) - assert result["imageSize"] == "2K" - - def test_map_openai_params_image_config_dict_stored_whole(self): - """imageConfig dict is stored as-is so all fields survive""" - result = self.config.map_openai_params( - {"imageConfig": {"aspectRatio": "16:9", "imageSize": "2K"}}, - {}, - "gemini-3.1-flash-image", - False, - ) - assert result["imageConfig"] == {"aspectRatio": "16:9", "imageSize": "2K"} - - def test_map_openai_params_image_config_all_fields(self): - """All ImageConfig fields (personGeneration, imageOutputOptions) pass through""" - payload = { - "imageConfig": { - "aspectRatio": "9:16", - "imageSize": "4K", - "personGeneration": "DONT_ALLOW", - "imageOutputOptions": { - "mimeType": "image/jpeg", - "compressionQuality": 80, - }, - } - } - result = self.config.map_openai_params(payload, {}, "gemini-3.1-flash-image", False) - assert result["imageConfig"] == payload["imageConfig"] - - def test_map_openai_params_image_config_non_dict_warns_and_drops(self): - """Non-dict imageConfig is dropped with a warning, not silently discarded""" - with patch("litellm.llms.vertex_ai.image_generation.vertex_gemini_transformation.verbose_logger") as mock_log: - result = self.config.map_openai_params( - {"imageConfig": "bad-string-value"}, {}, "gemini-3.1-flash-image", False - ) - assert "imageConfig" not in result - mock_log.warning.assert_called_once() - - def test_transform_image_generation_request_from_image_config(self): - """Full imageConfig dict is forwarded verbatim into generationConfig""" - full_config = { - "aspectRatio": "16:9", - "imageSize": "2K", - "personGeneration": "DONT_ALLOW", - "imageOutputOptions": {"mimeType": "image/jpeg", "compressionQuality": 85}, - } - mapped = self.config.map_openai_params( - {"imageConfig": full_config}, - {}, - "gemini-3.1-flash-image", - False, - ) - request = self.config.transform_image_generation_request( - model="gemini-3.1-flash-image", - prompt="A nano banana on a desk", - optional_params=mapped, - litellm_params={}, - headers={}, - ) - assert request["generationConfig"]["imageConfig"] == full_config - - def test_transform_image_generation_flat_params_override_image_config(self): - """Explicit flat params win over the same key inside imageConfig""" - request = self.config.transform_image_generation_request( - model="gemini-3.1-flash-image", - prompt="A nano banana", - optional_params={ - "imageConfig": {"aspectRatio": "1:1", "personGeneration": "DONT_ALLOW"}, - "aspectRatio": "16:9", # should win - }, - litellm_params={}, - headers={}, - ) - assert request["generationConfig"]["imageConfig"]["aspectRatio"] == "16:9" - assert request["generationConfig"]["imageConfig"]["personGeneration"] == "DONT_ALLOW" - - def test_transform_image_generation_request_basic(self): - """Test basic request transformation""" - request = self.config.transform_image_generation_request( - model="gemini-2.5-flash-image", - prompt="A nano banana", - optional_params={}, - litellm_params={}, - headers={}, - ) - assert "contents" in request - assert "generationConfig" in request - assert request["generationConfig"]["responseModalities"] == ["IMAGE"] - assert request["contents"][0]["parts"][0]["text"] == "A nano banana" - - def test_transform_image_generation_request_with_aspect_ratio(self): - """Test request transformation with aspectRatio""" - request = self.config.transform_image_generation_request( - model="gemini-2.5-flash-image", - prompt="A nano banana", - optional_params={"aspectRatio": "16:9"}, - litellm_params={}, - headers={}, - ) - assert request["generationConfig"]["imageConfig"]["aspectRatio"] == "16:9" - - def test_transform_image_generation_request_with_image_size(self): - """Test request transformation with imageSize (Gemini 3 Pro)""" - request = self.config.transform_image_generation_request( - model="gemini-3-pro-image-preview", - prompt="A nano banana", - optional_params={"imageSize": "4K"}, - litellm_params={}, - headers={}, - ) - assert request["generationConfig"]["imageConfig"]["imageSize"] == "4K" - - def test_map_openai_params_web_search_options(self): - """Test web_search_options maps to googleSearch tool""" - result = self.config.map_openai_params({"web_search_options": {}}, {}, "gemini-3.1-flash-image-preview", False) - assert result["tools"] == [{"googleSearch": {}}] - - def test_transform_image_generation_request_with_web_search_tools(self): - """Test request transformation includes googleSearch tools""" - request = self.config.transform_image_generation_request( - model="gemini-3.1-flash-image-preview", - prompt="Generate an image of the latest iPhone", - optional_params={"tools": [{"googleSearch": {}}]}, - litellm_params={}, - headers={}, - ) - assert request["tools"] == [{"googleSearch": {}}] - - def test_transform_image_generation_request_forwards_tool_config(self): - """Test request transformation forwards toolConfig side-effects from tool mapping""" - mapped = self.config.map_openai_params( - {"tools": [{"googleMaps": {"latitude": 37.7, "longitude": -122.4}}]}, - {}, - "gemini-3.1-flash-image-preview", - False, - ) - request = self.config.transform_image_generation_request( - model="gemini-3.1-flash-image-preview", - prompt="Generate an image of a coffee shop nearby", - optional_params=mapped, - litellm_params={}, - headers={}, - ) - assert request["tools"] == [{"googleMaps": {}}] - assert request["toolConfig"] == {"retrievalConfig": {"latLng": {"latitude": 37.7, "longitude": -122.4}}} - - def test_transform_image_generation_request_with_candidate_count(self): - """Test request transformation with candidate_count""" - request = self.config.transform_image_generation_request( - model="gemini-2.5-flash-image", - prompt="A nano banana", - optional_params={"candidate_count": 2}, - litellm_params={}, - headers={}, - ) - assert request["generationConfig"]["candidateCount"] == 2 - - def test_transform_image_generation_request_with_n(self): - """Test request transformation with n parameter""" - request = self.config.transform_image_generation_request( - model="gemini-2.5-flash-image", - prompt="A nano banana", - optional_params={"n": 2}, - litellm_params={}, - headers={}, - ) - assert request["generationConfig"]["candidateCount"] == 2 - - def test_transform_image_generation_response(self): - """Test response transformation""" - mock_response = MagicMock(spec=httpx.Response) - mock_response.status_code = 200 - mock_response.json.return_value = { - "candidates": [ - { - "content": { - "parts": [ - { - "inlineData": { - "mimeType": "image/png", - "data": "base64_encoded_image_data", - } - } - ] - } - } - ], - "usageMetadata": { - "promptTokenCount": 93, - "promptTokensDetails": [ - { - "modality": "TEXT", - "tokenCount": 54, - }, - { - "modality": "IMAGE", - "tokenCount": 39, - }, - ], - "candidatesTokenCount": 17, - "totalTokenCount": 110, - }, - } - mock_response.headers = {} - - from litellm.types.utils import ImageResponse - - model_response = ImageResponse() - result = self.config.transform_image_generation_response( - model="gemini-2.5-flash-image", - raw_response=mock_response, - model_response=model_response, - logging_obj=MagicMock(), - request_data={}, - optional_params={}, - litellm_params={}, - encoding=None, - ) - - assert len(result.data) == 1 - assert result.data[0].b64_json == "base64_encoded_image_data" - assert result.data[0].url is None - assert result.usage.input_tokens == 93 - assert result.usage.input_tokens_details.text_tokens == 54 - assert result.usage.input_tokens_details.image_tokens == 39 - assert result.usage.output_tokens == 17 - assert result.usage.total_tokens == 110 - - def test_transform_image_generation_response_multiple_images(self): - """Test response transformation with multiple images""" - mock_response = MagicMock(spec=httpx.Response) - mock_response.status_code = 200 - mock_response.json.return_value = { - "candidates": [ - { - "content": { - "parts": [ - { - "inlineData": { - "mimeType": "image/png", - "data": "image1", - } - }, - { - "inlineData": { - "mimeType": "image/png", - "data": "image2", - } - }, - ] - } - } - ] - } - mock_response.headers = {} - - from litellm.types.utils import ImageResponse - - model_response = ImageResponse() - result = self.config.transform_image_generation_response( - model="gemini-2.5-flash-image", - raw_response=mock_response, - model_response=model_response, - logging_obj=MagicMock(), - request_data={}, - optional_params={}, - litellm_params={}, - encoding=None, - ) - - assert len(result.data) == 2 - assert result.data[0].b64_json == "image1" - assert result.data[1].b64_json == "image2" - - def test_transform_image_generation_response_signature(self): - """Test response transformation includes thoughtSignature for Gemini 3 Pro""" - mock_response = MagicMock(spec=httpx.Response) - mock_response.status_code = 200 - mock_response.json.return_value = { - "candidates": [ - { - "content": { - "parts": [ - { - "inlineData": { - "mimeType": "image/png", - "data": "base64_encoded_image_data", - }, - "thoughtSignature": "test_signature_abc123", - } - ] - } - } - ] - } - mock_response.headers = {} - - from litellm.types.utils import ImageResponse - - model_response = ImageResponse() - result = self.config.transform_image_generation_response( - model="gemini-3-pro-image-preview", - raw_response=mock_response, - model_response=model_response, - logging_obj=MagicMock(), - request_data={}, - optional_params={}, - litellm_params={}, - encoding=None, - ) - - assert len(result.data) == 1 - assert result.data[0].b64_json == "base64_encoded_image_data" - assert result.data[0].provider_specific_fields["thought_signature"] == "test_signature_abc123" - - def test_transform_image_generation_response_tracks_web_search_requests(self): - """Grounding queries are carried onto usage so search spend can be billed""" - mock_response = MagicMock(spec=httpx.Response) - mock_response.status_code = 200 - mock_response.json.return_value = { - "candidates": [ - { - "content": { - "parts": [ - { - "inlineData": { - "mimeType": "image/png", - "data": "base64_encoded_image_data", - } - } - ] - }, - "groundingMetadata": {"webSearchQueries": ["eiffel tower", "paris skyline"]}, - } - ], - "usageMetadata": { - "promptTokenCount": 93, - "candidatesTokenCount": 17, - "totalTokenCount": 110, - }, - } - mock_response.headers = {} - - from litellm.types.utils import ImageResponse - - result = self.config.transform_image_generation_response( - model="gemini-2.5-flash-image", - raw_response=mock_response, - model_response=ImageResponse(), - logging_obj=MagicMock(), - request_data={}, - optional_params={}, - litellm_params={}, - encoding=None, - ) - - assert result.usage.web_search_requests == 2 - - -class TestVertexAIImagenImageGenerationConfig: - def setup_method(self): - """Set up test fixtures""" - self.config = VertexAIImagenImageGenerationConfig() - - def test_get_supported_openai_params(self): - """Test get_supported_openai_params returns correct params""" - supported = self.config.get_supported_openai_params("imagegeneration@006") - assert "n" in supported - assert "size" in supported - - def test_map_openai_params_n(self): - """Test mapping n parameter to sampleCount""" - non_default_params = {"n": 3} - optional_params = {} - result = self.config.map_openai_params(non_default_params, optional_params, "imagegeneration@006", False) - assert result.get("sampleCount") == 3 - - def test_map_openai_params_size(self): - """Test mapping size parameter to aspectRatio""" - non_default_params = {"size": "1024x1024"} - optional_params = {} - result = self.config.map_openai_params(non_default_params, optional_params, "imagegeneration@006", False) - assert result.get("aspectRatio") == "1:1" - - def test_map_size_to_aspect_ratio(self): - """Test size to aspect ratio mapping""" - assert self.config._map_size_to_aspect_ratio("1024x1024") == "1:1" - assert self.config._map_size_to_aspect_ratio("1792x1024") == "16:9" - assert self.config._map_size_to_aspect_ratio("unknown") == "1:1" # default - - def test_transform_image_generation_request_basic(self): - """Test basic request transformation""" - request = self.config.transform_image_generation_request( - model="imagegeneration@006", - prompt="A cat", - optional_params={}, - litellm_params={}, - headers={}, - ) - assert "instances" in request - assert "parameters" in request - assert request["instances"][0]["prompt"] == "A cat" - assert request["parameters"]["sampleCount"] == 1 - - def test_transform_image_generation_request_with_params(self): - """Test request transformation with parameters""" - request = self.config.transform_image_generation_request( - model="imagegeneration@006", - prompt="A cat", - optional_params={"sampleCount": 2, "aspectRatio": "16:9"}, - litellm_params={}, - headers={}, - ) - assert request["parameters"]["sampleCount"] == 2 - assert request["parameters"]["aspectRatio"] == "16:9" - - def test_transform_image_generation_request_labels_from_metadata(self): - """Billing labels from litellm_params.metadata.requester_metadata on predict body.""" - request = self.config.transform_image_generation_request( - model="imagegeneration@006", - prompt="A cat", - optional_params={}, - litellm_params={"metadata": {"requester_metadata": {"team": "platform", "env": "prod"}}}, - headers={}, - ) - assert request["labels"] == {"team": "platform", "env": "prod"} - assert "labels" not in request["parameters"] - - def test_transform_image_generation_response(self): - """Test response transformation""" - mock_response = MagicMock(spec=httpx.Response) - mock_response.status_code = 200 - mock_response.json.return_value = {"predictions": [{"bytesBase64Encoded": "base64_encoded_image_data"}]} - mock_response.headers = {} - - from litellm.types.utils import ImageResponse - - model_response = ImageResponse() - result = self.config.transform_image_generation_response( - model="imagegeneration@006", - raw_response=mock_response, - model_response=model_response, - logging_obj=MagicMock(), - request_data={}, - optional_params={}, - litellm_params={}, - encoding=None, - ) - - assert len(result.data) == 1 - assert result.data[0].b64_json == "base64_encoded_image_data" - assert result.data[0].url is None - - def test_transform_image_generation_response_multiple_images(self): - """Test response transformation with multiple images""" - mock_response = MagicMock(spec=httpx.Response) - mock_response.status_code = 200 - mock_response.json.return_value = { - "predictions": [ - {"bytesBase64Encoded": "image1"}, - {"bytesBase64Encoded": "image2"}, - ] - } - mock_response.headers = {} - - from litellm.types.utils import ImageResponse - - model_response = ImageResponse() - result = self.config.transform_image_generation_response( - model="imagegeneration@006", - raw_response=mock_response, - model_response=model_response, - logging_obj=MagicMock(), - request_data={}, - optional_params={}, - litellm_params={}, - encoding=None, - ) - - assert len(result.data) == 2 - assert result.data[0].b64_json == "image1" - assert result.data[1].b64_json == "image2" - - -class TestGetVertexAIImageGenerationConfig: - """Test the router function that selects the correct config""" - - def test_get_gemini_model_config(self): - """Test that Gemini models return Gemini config""" - config = get_vertex_ai_image_generation_config("gemini-2.5-flash-image") - assert isinstance(config, VertexAIGeminiImageGenerationConfig) - - config = get_vertex_ai_image_generation_config("gemini-3-pro-image-preview") - assert isinstance(config, VertexAIGeminiImageGenerationConfig) - - config = get_vertex_ai_image_generation_config("vertex_ai/gemini-2.5-flash-image") - assert isinstance(config, VertexAIGeminiImageGenerationConfig) - - def test_get_imagen_model_config(self): - """Test that Imagen models return Imagen config""" - config = get_vertex_ai_image_generation_config("imagegeneration@006") - assert isinstance(config, VertexAIImagenImageGenerationConfig) - - config = get_vertex_ai_image_generation_config("imagen-4.0-generate-001") - assert isinstance(config, VertexAIImagenImageGenerationConfig) - - config = get_vertex_ai_image_generation_config("vertex_ai/imagegeneration@006") - assert isinstance(config, VertexAIImagenImageGenerationConfig) - - def test_get_non_gemini_model_config(self): - """Test that non-Gemini models default to Imagen config""" - config = get_vertex_ai_image_generation_config("some-other-model") - assert isinstance(config, VertexAIImagenImageGenerationConfig) - - class TestVertexAIImageGenerationIntegration: """Integration tests for Vertex AI image generation""" @@ -642,39 +56,3 @@ class TestVertexAIImageGenerationIntegration: litellm_params={}, ) assert "Authorization" in headers - - def test_gemini_get_complete_url(self): - """Test Gemini config URL generation""" - config = VertexAIGeminiImageGenerationConfig() - url = config.get_complete_url( - api_base=None, - api_key=None, - model="gemini-2.5-flash-image", - optional_params={}, - litellm_params={ - "vertex_project": "test-project", - "vertex_location": "us-central1", - }, - ) - assert "test-project" in url - assert "us-central1" in url - assert "gemini-2.5-flash-image" in url - assert "generateContent" in url - - def test_imagen_get_complete_url(self): - """Test Imagen config URL generation""" - config = VertexAIImagenImageGenerationConfig() - url = config.get_complete_url( - api_base=None, - api_key=None, - model="imagegeneration@006", - optional_params={}, - litellm_params={ - "vertex_project": "test-project", - "vertex_location": "us-central1", - }, - ) - assert "test-project" in url - assert "us-central1" in url - assert "imagegeneration@006" in url - assert "predict" in url diff --git a/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/__init__.py b/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/__init__.py deleted file mode 100644 index 8b41c5ab3f8..00000000000 --- a/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Tests for Vertex AI Gemma-AI models""" diff --git a/tests/test_litellm/llms/vertex_ai/videos/__init__.py b/tests/test_litellm/llms/vertex_ai/videos/__init__.py deleted file mode 100644 index f29c2a16fd5..00000000000 --- a/tests/test_litellm/llms/vertex_ai/videos/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -""" -Tests for Vertex AI video generation. -""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_capabilities.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_capabilities.py new file mode 100644 index 00000000000..f104e655704 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_capabilities.py @@ -0,0 +1,106 @@ +from typing import Final + +import pytest +from mcp import Client +from mcp.server import Server +from mcp.shared.exceptions import MCPError +from mcp.types import PromptsCapability, ResourcesCapability, ServerCapabilities, ToolsCapability +from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS + +from litellm.proxy._experimental.mcp_server.capabilities import ( + GATEWAY_OPERATIONS, + REVISION_SUPPORT, + TRANSLATION_PAIRS, + GatewayVersionPolicy, + build_discovery, +) +from litellm.types.mcp import MCPTransport + + +@pytest.mark.parametrize("revision", HANDSHAKE_PROTOCOL_VERSIONS) +@pytest.mark.parametrize("transport", tuple(MCPTransport)) +def test_discovery_only_exposes_authorized_completed_support(revision, transport): + result = build_discovery( + configured=(revision, "2026-07-28", "unknown"), + revision=revision, + transport=transport, + authorized_operations=frozenset({"tools/list", "tools/call"}), + upstream_versions=frozenset(HANDSHAKE_PROTOCOL_VERSIONS), + capabilities=ServerCapabilities( + tools=ToolsCapability(), prompts=PromptsCapability(), resources=ResourcesCapability(), + extensions={"io.modelcontextprotocol/ui": {}}, + ), + client_extensions=frozenset({"io.modelcontextprotocol/ui"}), + upstream_extensions=frozenset({"io.modelcontextprotocol/ui"}), + ) + assert result.supported_versions == [revision] + assert result.capabilities.tools is not None + assert result.capabilities.prompts is None + assert result.capabilities.resources is None + assert result.capabilities.extensions is None + assert result.capabilities.tasks is None + assert result.cache_scope == "private" + assert result.ttl_ms == 0 + + +@pytest.mark.parametrize("upstream", [frozenset(), frozenset({"unknown"}), frozenset({"2026-07-28"})]) +def test_unproven_translation_never_advertises_operations(upstream): + result = build_discovery( + configured=HANDSHAKE_PROTOCOL_VERSIONS, + revision="2025-11-25", + transport=MCPTransport.http, + authorized_operations=GATEWAY_OPERATIONS, + upstream_versions=upstream, + capabilities=ServerCapabilities(tools=ToolsCapability()), + ) + assert result.capabilities.tools is None + + +@pytest.mark.parametrize("revision", ["2026-07-28", "unknown", "2024-11-05"]) +def test_unadvertised_revision_never_gains_capabilities(revision): + result = build_discovery( + configured=("2025-11-25",), revision=revision, transport=MCPTransport.http, + authorized_operations=GATEWAY_OPERATIONS, upstream_versions=frozenset(HANDSHAKE_PROTOCOL_VERSIONS), + capabilities=ServerCapabilities(tools=ToolsCapability()), + ) + assert result.capabilities.tools is None + + +def test_discovery_results_do_not_share_mutable_capabilities(): + capabilities = ServerCapabilities(tools=ToolsCapability(), prompts=PromptsCapability(), resources=ResourcesCapability()) + args = dict( + configured=HANDSHAKE_PROTOCOL_VERSIONS, revision="2025-11-25", transport=MCPTransport.http, + upstream_versions=frozenset(HANDSHAKE_PROTOCOL_VERSIONS), capabilities=capabilities, + ) + allowed = build_discovery(**args, authorized_operations=GATEWAY_OPERATIONS) + denied = build_discovery(**args, authorized_operations=frozenset()) + assert allowed.capabilities.prompts is not None + assert allowed.capabilities.resources is not None + assert denied.capabilities.model_dump(exclude_none=True) == {} + assert allowed.capabilities.tools is not None + allowed.capabilities.tools.list_changed = True + assert capabilities.tools.list_changed is not True + + +def test_modern_candidates_do_not_enable_public_serving(): + modern = REVISION_SUPPORT["2026-07-28"] + assert modern.completed is False + assert "input_required" in modern.results + assert MCPTransport.sse not in modern.transports + assert not any("2026-07-28" in pair for pair in TRANSLATION_PAIRS) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("versions,accepted", [(("2025-11-25",), True), (("2025-06-18",), False)]) +async def test_version_policy_gates_the_actual_sdk_handshake(versions, accepted): + server: Final = Server("test-gateway", version="1") + server.middleware.append(GatewayVersionPolicy(lambda: versions)) + if accepted: + async with Client(server, mode="legacy") as client: + assert client.protocol_version == "2025-11-25" + result = await client.session.send_ping() + assert result is not None + else: + with pytest.RaisesGroup(pytest.RaisesExc(MCPError, match="Unsupported MCP protocol version"), flatten_subgroups=True): + async with Client(server, mode="legacy"): + pytest.fail("The excluded revision must not initialize") diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 97b242831a2..ab00ec4da1e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -10444,11 +10444,12 @@ async def test_active_request_ctx_var_feeds_auth_resolution_recording(_mcp_reque ) @pytest.mark.parametrize("handler", ("handle_streamable_http_mcp", "handle_sse_mcp")) async def test_streamable_http_rejects_modern_protocol_version( - header_value: str, expected_rejected: bool, handler: str + header_value: str, expected_rejected: bool, handler: str, monkeypatch: pytest.MonkeyPatch ) -> None: from litellm.proxy._experimental.mcp_server import server as mcp_module from litellm.proxy._experimental.mcp_server.server import unsupported_protocol_version + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) scope: Scope = { "type": "http", "method": "POST", @@ -10659,3 +10660,32 @@ async def test_legacy_sse_mount_emits_message_endpoint( await incoming.put({"type": "http.disconnect"}) await asyncio.wait_for(task, 2) assert await post(initialization) == 404 + + +@pytest.mark.parametrize("revision,rejected", [("2024-11-05", False), ("2025-11-25", True), ("2026-07-28", True)]) +def test_protocol_header_respects_configured_advertisement(revision, rejected): + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.server import unsupported_protocol_version + + with patch.dict(proxy_server.general_settings, {"mcp_advertised_versions": ["2024-11-05"]}): + result = unsupported_protocol_version({"headers": [(b"mcp-protocol-version", revision.encode())]}) + assert result == (revision if rejected else None) + + +@pytest.mark.asyncio +async def test_discovery_adapter_preserves_authenticated_context(_mcp_request_ctx): + from mcp.types import DiscoverResult, RequestParams, ServerCapabilities + from litellm.proxy._experimental.mcp_server import server + + expected = DiscoverResult(supported_versions=["2025-11-25"], capabilities=ServerCapabilities()) + dispatched = AsyncMock(return_value=expected) + auth = UserAPIKeyAuth(user_id="discover-caller") + with ( + patch.object(server, "get_or_extract_auth_context", AsyncMock(return_value=(auth, None, ["allowed"], None, None, None, None))), + patch.object(server.operations.GatewayOperations, "execute", dispatched), + ): + result = await server.discover(_mcp_request_ctx(), RequestParams()) + assert result is expected + context = dispatched.await_args.args[1] + assert context.user_api_key_auth.user_id == "discover-caller" + assert context.mcp_servers == ("allowed",) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 64d94065674..16bffa1a356 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -63,7 +63,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.types.llms.custom_http import httpxSpecialProvider -from litellm.types.mcp import MCPAuth, MCPAuthType +from litellm.types.mcp import MCPAuth, MCPAuthType, MCPUpstreamProtocol from litellm.types.mcp_server.mcp_server_manager import MCPOAuthMetadata, MCPServer from litellm.caching.caching import DualCache from litellm.caching.llm_caching_handler import LLMClientCache @@ -14637,3 +14637,27 @@ class TestSharedIdentifierPrefixWarning: assert "srv-b" in shared_warnings[0] assert "srv-c" not in shared_warnings[0] assert "'shared'" in shared_warnings[0] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("revision", ["auto", "2024-11-05", "2025-03-26", "2025-06-18", "2025-11-25"]) +async def test_configured_protocol_reaches_the_upstream_client(config_only_mcp_manager_factory, revision): + manager = config_only_mcp_manager_factory() + await manager.load_servers_from_config({"versions": {"url": "http://127.0.0.1:9/mcp", "transport": "http", "protocol_version": revision}}) + server = next(iter(manager.config_mcp_servers.values())) + client = await manager._create_mcp_client(server) + assert server.protocol_version == revision + assert client.protocol_version == revision + + +@pytest.mark.parametrize("revision", ("auto", "2024-11-05", "2025-06-18")) +@pytest.mark.parametrize("explicit", (None, "auto", "2025-11-25")) +def test_runtime_protocol_metadata_preserves_explicit_precedence( + revision: MCPUpstreamProtocol, explicit: MCPUpstreamProtocol | None +) -> None: + server: Final = MCPServer.model_validate({ + "server_id": "preview", "name": "preview", "transport": "http", + "mcp_info": {"protocol_version": revision}, + **({"protocol_version": explicit} if explicit is not None else {}), + }) + assert server.protocol_version == (explicit if explicit is not None else revision) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py index 81f81045740..bb900de4f98 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py @@ -542,3 +542,126 @@ async def test_local_tool_json_array_is_converted_once_for_the_caller_revision(c assert [block.text for block in result.content] == [body] assert result.structured_content == (["a", "b"] if compat == "modern" else None) + + +@pytest.mark.asyncio +async def test_discovery_preserves_caller_scope_and_proxy_restrictions(): + from mcp.types import DiscoverRequest, ListToolsResult, Tool + + listed = AsyncMock(return_value=ListToolsResult(tools=[Tool(name="allowed", input_schema={"type": "object"})])) + context = prepare_context(UserAPIKeyAuth(user_id="scoped"), mcp_servers=["only-this"], mcp_proxy_mode=True, protocol_version="2025-06-18") + with patch("litellm.proxy._experimental.mcp_server.operations._execute_handle_list_tools", listed): + result = await GatewayOperations().execute(DiscoverRequest(), context) + assert result.capabilities.tools is not None + assert result.capabilities.resources is None + assert result.capabilities.prompts is None + assert listed.await_args.args[0] is context + assert listed.await_args.args[0].user_api_key_auth.user_id == "scoped" + assert listed.await_args.args[0].mcp_servers == ("only-this",) + + +@pytest.mark.asyncio +async def test_discovery_denial_cannot_advertise_tools(): + from mcp.types import DiscoverRequest + from fastapi import HTTPException + + denied = AsyncMock(side_effect=HTTPException(status_code=403, detail="Forbidden")) + with patch("litellm.proxy._experimental.mcp_server.operations._execute_handle_list_tools", denied): + with pytest.raises(HTTPException) as error: + await GatewayOperations().execute(DiscoverRequest(), prepare_context(UserAPIKeyAuth(user_id="denied"))) + assert error.value.status_code == 403 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("available", ["none", "resources", "templates", "prompts"]) +async def test_discovery_lists_each_capability_with_the_same_caller(available): + from mcp.types import ( + DiscoverRequest, ListToolsResult, ListPromptsResult, ListResourcesResult, + ListResourceTemplatesResult, Prompt, Resource, ResourceTemplate, + ) + from litellm.proxy._experimental.mcp_server import operations + + context = prepare_context(UserAPIKeyAuth(user_id="scoped"), mcp_servers=["authorized"]) + tools = AsyncMock(return_value=ListToolsResult(tools=[])) + prompts = AsyncMock(return_value=ListPromptsResult(prompts=[Prompt(name="allowed")] if available == "prompts" else [])) + resources = AsyncMock(return_value=ListResourcesResult(resources=[Resource(name="allowed", uri="test://allowed")] if available == "resources" else [])) + templates = AsyncMock(return_value=ListResourceTemplatesResult(resource_templates=[ResourceTemplate(name="allowed", uri_template="test://{id}")] if available == "templates" else [])) + with ( + patch.object(operations, "_execute_handle_list_tools", tools), + patch.object(operations, "_execute_list_prompts", prompts), + patch.object(operations, "_execute_list_resources", resources), + patch.object(operations, "_execute_list_resource_templates", templates), + ): + result = await GatewayOperations().execute(DiscoverRequest(), context) + assert result.capabilities.tools is None + assert (result.capabilities.prompts is not None) == (available == "prompts") + assert (result.capabilities.resources is not None) == (available in {"resources", "templates"}) + for listing in (tools, prompts, resources, templates): + assert listing.await_args.args[0] is context + + +@pytest.mark.asyncio +@pytest.mark.parametrize("outcome", ["success", "failure", "cancel"]) +async def test_discovery_concurrent_listings_drain_on_failure_and_cancellation(outcome): + from mcp.types import DiscoverRequest, ListToolsResult, ListPromptsResult, ListResourcesResult, ListResourceTemplatesResult + from litellm.proxy._experimental.mcp_server import operations + + ready = [asyncio.Event() for _ in range(4)] + closed = [asyncio.Event() for _ in range(4)] + release = asyncio.Event() + responses = (ListToolsResult(tools=[]), ListPromptsResult(prompts=[]), ListResourcesResult(resources=[]), ListResourceTemplatesResult(resource_templates=[])) + + def listing(index): + async def run(*args, **kwargs): + ready[index].set() + try: + await release.wait() + if index == 0 and outcome == "failure": + raise ValueError("discovery failed") + if outcome != "success": + await asyncio.Event().wait() + return responses[index] + finally: + closed[index].set() + return run + + with ( + patch.object(operations, "_execute_handle_list_tools", side_effect=listing(0)) as tools, + patch.object(operations, "_execute_list_prompts", side_effect=listing(1)), + patch.object(operations, "_execute_list_resources", side_effect=listing(2)), + patch.object(operations, "_execute_list_resource_templates", side_effect=listing(3)), + ): + task = asyncio.create_task(GatewayOperations().execute(DiscoverRequest(), prepare_context(UserAPIKeyAuth(user_id="scoped")))) + try: + await asyncio.wait_for(asyncio.gather(*(event.wait() for event in ready)), 1) + if outcome == "cancel": + task.cancel() + else: + release.set() + if outcome == "success": + result = await asyncio.wait_for(task, 1) + assert result.capabilities.model_dump(exclude_none=True) == {} + else: + with pytest.raises(asyncio.CancelledError if outcome == "cancel" else ValueError): + await asyncio.wait_for(task, 1) + assert all(event.is_set() for event in closed) + assert tools.call_args.kwargs["log_list_tools_to_spendlogs"] is False + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("log_enabled", [False, True]) +async def test_tools_listing_preserves_explicit_spend_log_policy(log_enabled): + from mcp.types import PaginatedRequestParams + from litellm.proxy._experimental.mcp_server import operations + + listing = AsyncMock(return_value=operations.AggregateToolListing(tools=[], outcomes={})) + with patch.object(operations, "_list_mcp_tools", listing): + result = await operations._execute_handle_list_tools( + prepare_context(UserAPIKeyAuth(user_id="caller")), PaginatedRequestParams(), + log_list_tools_to_spendlogs=log_enabled, + ) + assert result.tools == [] + assert listing.await_args.kwargs["log_list_tools_to_spendlogs"] is log_enabled 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 4da120cb26f..e82ab28bb4c 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 @@ -28,7 +28,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.types.mcp import MCPAuth, MCPTransport +from litellm.types.mcp import MCPAuth, MCPTransport, MCPUpstreamProtocol from litellm.types.mcp_server.mcp_server_manager import MCPServer _OK_TOOL_RESULT: Final = CallToolResult(content=[TextContent(type="text", text='{"result": "ok"}')], is_error=False) @@ -1476,7 +1476,7 @@ class TestListToolsRestAPI: monkeypatch, ): """The REST tools/list path should include tools beyond the upstream first page.""" - from mcp.types import ListToolsResult, PaginatedRequestParams + from mcp.types import Implementation, InitializeResult, ListToolsResult, PaginatedRequestParams, ServerCapabilities from mcp.types import Tool as MCPTool import litellm.experimental_mcp_client.client as mcp_client_module @@ -1512,7 +1512,11 @@ class TestListToolsRestAPI: mock_session_ctx = AsyncMock() mock_session_instance = AsyncMock() - mock_session_instance.initialize = AsyncMock(return_value=None) + mock_session_instance.initialize = AsyncMock(return_value=InitializeResult( + protocol_version="2025-11-25", + capabilities=ServerCapabilities(), + server_info=Implementation(name="stub", version="1"), + )) mock_session_instance.list_tools.side_effect = [ ListToolsResult( tools=[ @@ -4628,3 +4632,74 @@ class TestClientAllowlistOnRestRoutes: assert denied.value.detail["error"] == "Forbidden" assert "'claude-code'" in denied.value.detail["details"] acting.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("revision", ("auto", "2024-11-05", "2025-06-18")) +async def test_preview_client_honors_protocol_metadata(revision: MCPUpstreamProtocol) -> None: + from litellm.experimental_mcp_client.client import MCPClient + + payload: Final = NewMCPServerRequest( + server_name="preview", url="http://127.0.0.1:9/mcp", transport="http", + auth_type=MCPAuth.none, mcp_info={"protocol_version": revision}, + ) + + async def inspect_client(client: MCPClient) -> dict[str, str]: + return {"protocol_version": client.protocol_version} + + result: Final = await rest_endpoints._execute_with_mcp_client(payload, inspect_client) + assert result == {"protocol_version": revision} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("auth_type", (MCPAuth.none, MCPAuth.bearer_token, MCPAuth.oauth2)) +@pytest.mark.parametrize( + ("metadata", "expected"), + ( + (None, "2025-11-25"), + ({}, "2025-11-25"), + ({"description": "edited"}, "2025-11-25"), + ({"protocol_version": "auto"}, "auto"), + ({"protocol_version": "2024-11-05"}, "2024-11-05"), + ({"protocol_version": "2025-06-18"}, "2025-06-18"), + ), +) +async def test_saved_preview_protocol_omission_and_explicit_edits( + monkeypatch: pytest.MonkeyPatch, auth_type: MCPAuth, + metadata: dict[str, str] | None, expected: MCPUpstreamProtocol, +) -> None: + from starlette.datastructures import Headers + + from litellm.experimental_mcp_client.client import MCPClient + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + from litellm.proxy.management_endpoints import mcp_management_endpoints + + saved: Final = MCPServer( + server_id="saved-preview", name="preview", url="https://example.com/mcp", + transport="http", auth_type=auth_type, protocol_version="2025-11-25", + authentication_token="stored-token", + authorization_url="https://example.com/authorize", token_url="https://example.com/token", + ) + manager: Final = MCPServerManager() + manager.registry = {saved.server_id: saved} + monkeypatch.setattr(rest_endpoints, "global_mcp_server_manager", manager) + monkeypatch.setattr(mcp_management_endpoints, "global_mcp_server_manager", manager) + payload: Final = NewMCPServerRequest( + server_id=saved.server_id, server_name=saved.name, url=saved.url, transport="http", + auth_type=auth_type, mcp_info=metadata, + authorization_url=saved.authorization_url, token_url=saved.token_url, + ) + staged: Final = rest_endpoints._stage_server_test( + payload, Headers({"x-litellm-api-key": "sk-admin", "authorization": "Bearer preview-token"}) + ) + + async def inspect_client(client: MCPClient) -> dict[str, str]: + return {"protocol_version": client.protocol_version} + + result: Final = await rest_endpoints._execute_with_mcp_client( + staged.request, inspect_client, + mcp_auth_header=staged.mcp_auth_header, oauth2_headers=staged.oauth2_headers, + ) + assert result == {"protocol_version": expected} + assert saved.protocol_version == "2025-11-25" + assert payload.mcp_info == metadata diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index f8b9043a23f..b1622e0dff0 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -5324,16 +5324,22 @@ def test_build_decode_kwargs_warns_for_unscoped_global_fallback_in_mixed_deploym @pytest.mark.asyncio -async def test_resolve_team_from_header_defers_to_db_membership_only_without_jwt_claims(): +async def test_resolve_team_from_header_accepts_db_teams_provisionally_under_fallback_even_with_jwt_claims(): """With fallback_to_db_teams=True, an x-litellm-team-id header naming an existing - team is accepted provisionally only when the JWT carries no team claims (allowed - set empty). When the JWT does carry team claims, the header must still be validated - against them, and the flag-off behavior must keep rejecting unknown teams.""" + team is accepted provisionally whether or not the JWT carries team claims; the + union of JWT teams and DB memberships is enforced by auth_builder's later + membership check. Unknown values still 403, and the flag-off behavior keeps + rejecting teams outside the JWT's allowed set.""" known_ids = frozenset({"team-from-db"}) deferred, _, _ = await _resolve_header("team-from-db", set(), True, _teams_by_id(known_ids), _team_alias_lookup_404) assert deferred == HeaderTeam(header_value="team-from-db", team_id="team-from-db") + deferred_with_claims, _, _ = await _resolve_header( + "team-from-db", {"team-1"}, True, _teams_by_id(known_ids), _team_alias_lookup_404 + ) + assert deferred_with_claims == HeaderTeam(header_value="team-from-db", team_id="team-from-db") + with pytest.raises(HTTPException) as exc_info: await _resolve_header("team-x", {"team-1", "team-2"}, True, _teams_by_id(known_ids), _team_alias_lookup_404) assert exc_info.value.status_code == 403 @@ -5849,6 +5855,7 @@ async def _run_auth_builder_with_header_team( allowed_team_ids: set, fake_get_team_by_alias=_team_alias_lookup_404, route: str = "/chat/completions", + send_header: bool = True, ): jwt_handler = JWTHandler() jwt_handler.litellm_jwtauth = jwt_auth_config @@ -5909,7 +5916,7 @@ async def _run_auth_builder_with_header_team( user_api_key_cache=None, parent_otel_span=None, proxy_logging_obj=None, - request_headers={"x-litellm-team-id": header_team_id}, + request_headers={"x-litellm-team-id": header_team_id} if send_header else {}, ) @@ -7283,6 +7290,130 @@ async def test_auth_builder_header_alias_under_db_fallback_keeps_the_team_allowe assert allowed["team_id"] == "team_member" +@pytest.mark.asyncio +async def test_auth_builder_header_selects_db_membership_team_when_jwt_also_carries_a_team_claim() -> None: + """Under fallback_to_db_teams, x-litellm-team-id may name a DB-membership + team the JWT does not claim (LIT-8656): the allowed set is the JWT teams + union the user's DB memberships, not the JWT teams alone. The flag-off + path keeps rejecting the same header against the JWT's allowed teams.""" + user_object = LiteLLM_UserTable( + user_id="u_mixed", + user_role=LitellmUserRoles.INTERNAL_USER, + teams=["team_member"], + ) + config = LiteLLM_JWTAuth(fallback_to_db_teams=True, team_id_jwt_field="appid") + token = {"sub": "u_mixed", "scope": "", "appid": "team_claimed"} + fake_get_team = _teams_by_id(frozenset({"team_claimed", "team_member"})) + + by_membership = await _run_auth_builder_with_header_team( + config, token, "team_member", user_object, fake_get_team, {"team_claimed"} + ) + assert by_membership["team_id"] == "team_member" + assert by_membership["team_object"].team_id == "team_member" + + by_claim = await _run_auth_builder_with_header_team( + config, token, "team_claimed", user_object, fake_get_team, {"team_claimed"} + ) + assert by_claim["team_id"] == "team_claimed" + + flag_off = LiteLLM_JWTAuth(fallback_to_db_teams=False, team_id_jwt_field="appid") + with pytest.raises(HTTPException) as exc_info: + await _run_auth_builder_with_header_team( + flag_off, token, "team_member", user_object, fake_get_team, {"team_claimed"} + ) + assert exc_info.value.status_code == 403 + assert "JWT's allowed teams" in exc_info.value.detail + + +@pytest.mark.asyncio +async def test_auth_builder_header_non_member_team_is_denied_when_jwt_also_carries_a_team_claim() -> None: + """A header naming a team the user does not belong to stays a membership + denial even when the JWT carries a team claim, and an existing but + non-member team produces the exact same 403 shape as a nonexistent one so + the response is no oracle for which team ids exist.""" + user_object = LiteLLM_UserTable( + user_id="u_mixed", + user_role=LitellmUserRoles.INTERNAL_USER, + teams=["team_member"], + ) + config = LiteLLM_JWTAuth(fallback_to_db_teams=True, team_id_jwt_field="appid") + token = {"sub": "u_mixed", "scope": "", "appid": "team_claimed"} + fake_get_team = _teams_by_id(frozenset({"team_claimed", "team_member", "team_other"})) + + with pytest.raises(HTTPException) as outsider_exc: + await _run_auth_builder_with_header_team( + config, token, "team_other", user_object, fake_get_team, {"team_claimed"} + ) + with pytest.raises(HTTPException) as missing_exc: + await _run_auth_builder_with_header_team( + config, token, "team_ghost", user_object, fake_get_team, {"team_claimed"} + ) + + assert outsider_exc.value.status_code == 403 + assert missing_exc.value.status_code == 403 + assert outsider_exc.value.detail == ( + "x-litellm-team-id 'team_other' does not resolve to a team id or a unique team alias among your " + "team memberships." + ) + assert missing_exc.value.detail.replace("team_ghost", "") == outsider_exc.value.detail.replace( + "team_other", "" + ) + assert "exist" not in missing_exc.value.detail + + +@pytest.mark.asyncio +async def test_auth_builder_no_header_keeps_the_jwt_team_when_fallback_to_db_teams_is_on() -> None: + """With no x-litellm-team-id header, fallback_to_db_teams must not disturb + the claim path: the JWT's own team claim still binds the request.""" + user_object = LiteLLM_UserTable( + user_id="u_mixed", + user_role=LitellmUserRoles.INTERNAL_USER, + teams=["team_member"], + ) + config = LiteLLM_JWTAuth(fallback_to_db_teams=True, team_id_jwt_field="appid") + token = {"sub": "u_mixed", "scope": "", "appid": "team_claimed"} + + result = await _run_auth_builder_with_header_team( + config, + token, + "team_member", + user_object, + _teams_by_id(frozenset({"team_claimed", "team_member"})), + {"team_claimed"}, + send_header=False, + ) + assert result["team_id"] == "team_claimed" + + +@pytest.mark.asyncio +async def test_auth_builder_team_id_default_does_not_widen_the_header_allowed_set() -> None: + """team_id_default fills in a team for claimless tokens but must not widen + the header's allowed set: a header naming the default team is still held + to DB membership under fallback_to_db_teams.""" + user_object = LiteLLM_UserTable( + user_id="u_default", + user_role=LitellmUserRoles.INTERNAL_USER, + teams=["team_member"], + ) + config = LiteLLM_JWTAuth(fallback_to_db_teams=True, team_id_default="team_default") + token = {"sub": "u_default", "scope": ""} + + with pytest.raises(HTTPException) as exc_info: + await _run_auth_builder_with_header_team( + config, + token, + "team_default", + user_object, + _teams_by_id(frozenset({"team_default", "team_member"})), + set(), + ) + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == ( + "x-litellm-team-id 'team_default' does not resolve to a team id or a unique team alias among your " + "team memberships." + ) + + @pytest.mark.asyncio async def test_sync_user_role_and_teams_singular_claim_only_recognized_under_flag(): """Reading the singular team claim during sync is scoped to fallback_to_db_teams. diff --git a/tests/test_litellm/proxy/common_utils/test_sse_keepalive.py b/tests/test_litellm/proxy/common_utils/test_sse_keepalive.py index 69b92f5e4d7..228fd5bcae6 100644 --- a/tests/test_litellm/proxy/common_utils/test_sse_keepalive.py +++ b/tests/test_litellm/proxy/common_utils/test_sse_keepalive.py @@ -6,10 +6,13 @@ import pytest from fastapi.responses import StreamingResponse from litellm.proxy.common_request_processing import create_response +from litellm.types.utils import ModelResponse from litellm.proxy.common_utils.sse_keepalive import ( ANTHROPIC_PING_SSE_CHUNK, SSE_COMMENT_PING_BYTES, + advance_sse_tail, resolve_ttft_keepalive_interval, + seal_open_sse_frame, split_complete_sse_frames, wrap_passthrough_sse_bytes_with_keepalive_pings, wrap_sse_stream_with_keepalive_pings, @@ -32,6 +35,12 @@ def test_split_complete_sse_frames_holds_bytes_with_no_complete_frame(): assert split_complete_sse_frames(b"data: unterminated") == (b"", b"data: unterminated") +@pytest.mark.parametrize("chunk", [{"content": "hi"}, ModelResponse()]) +def test_advance_sse_tail_ignores_a_chunk_that_is_not_sse_text(chunk: object): + assert advance_sse_tail(b"\n\n", chunk) == b"\n\n" + assert seal_open_sse_frame(advance_sse_tail(b"data: {", chunk)) == "\n" + ANTHROPIC_PING_SSE_CHUNK + + @pytest.mark.asyncio async def test_pings_fill_mid_stream_silence_and_preserve_chunk_order(): async def gappy_stream() -> AsyncGenerator[str, None]: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 1d3d7a452b6..45ad336368c 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -30,7 +30,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockTextContent, ) from litellm.types.utils import CallTypes, ModelResponse -from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe +from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py index bbc8fd539a3..5169d4c9ec6 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py @@ -23,7 +23,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockGuardrailResponse, ) from litellm.types.utils import Choices, Message, ModelResponse -from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe +from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe CONTENT_FILTER_CHECKS = {"contentFilter": {"categories": [{"category": "VIOLENCE"}]}} diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 1230c548281..21c0f565486 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -8029,6 +8029,37 @@ class TestPKCEStateCookieBinding: assert result is not None +@pytest.mark.asyncio +@pytest.mark.parametrize("enable_sso_debug_value", [None, "false", "0"]) +async def test_sso_debug_routes_return_404_unless_explicitly_enabled(enable_sso_debug_value): + """ + /sso/debug/login and /sso/debug/callback must 404 unless ENABLE_SSO_DEBUG is + explicitly set to a truthy value. + """ + from litellm.proxy.management_endpoints.ui_sso import debug_sso_callback, debug_sso_login + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://proxy.example.com/" + mock_request.cookies = {} + mock_request.query_params = {} + + env = {"GENERIC_CLIENT_ID": "test_client_id"} + if enable_sso_debug_value is not None: + env["ENABLE_SSO_DEBUG"] = enable_sso_debug_value + + with patch.dict(os.environ, env, clear=False): + if enable_sso_debug_value is None: + os.environ.pop("ENABLE_SSO_DEBUG", None) + + with pytest.raises(HTTPException) as login_exc: + await debug_sso_login(mock_request) + with pytest.raises(HTTPException) as callback_exc: + await debug_sso_callback(mock_request) + + assert login_exc.value.status_code == 404 + assert callback_exc.value.status_code == 404 + + @pytest.mark.asyncio async def test_debug_sso_callback_renders_full_jwt_claims(): """ @@ -8080,7 +8111,7 @@ async def test_debug_sso_callback_renders_full_jwt_claims(): with ( patch.dict( os.environ, - {"GENERIC_CLIENT_ID": "test_client_id"}, + {"GENERIC_CLIENT_ID": "test_client_id", "ENABLE_SSO_DEBUG": "true"}, clear=False, ), patch( @@ -8165,7 +8196,7 @@ async def test_debug_sso_callback_handles_missing_raw_response(): with ( patch.dict( os.environ, - {"MICROSOFT_CLIENT_ID": "test_microsoft_id"}, + {"MICROSOFT_CLIENT_ID": "test_microsoft_id", "ENABLE_SSO_DEBUG": "true"}, clear=False, ), patch.object( @@ -8213,7 +8244,7 @@ async def _render_debug_page(provider_env, id_jag_registered, force_inert=False) return parsed stack = [ - patch.dict(os.environ, provider_env, clear=False), + patch.dict(os.environ, {**provider_env, "ENABLE_SSO_DEBUG": "true"}, clear=False), patch( # test-quality-ok: endpoint test stubs the upstream generic IdP boundary "litellm.proxy.management_endpoints.ui_sso.get_generic_sso_response", side_effect=fake_generic ), diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 353ffadfa46..227921d6150 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -24,7 +24,7 @@ from starlette.datastructures import FormData import litellm from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing -from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe +from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( BaseOpenAIPassThroughHandler, diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 1604970796a..363bbf41967 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -24,6 +24,7 @@ from starlette.datastructures import UploadFile as StarletteUploadFile import litellm from litellm._logging import verbose_proxy_logger +from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import ProxyException, UserAPIKeyAuth @@ -1168,6 +1169,29 @@ def test_resolve_llm_passthrough_timeout_precedence(): assert resolve_llm_passthrough_timeout() == 6.0 +def test_resolve_llm_passthrough_timeout_honors_explicit_global_request_timeout(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr("litellm.request_timeout", 44.0, raising=False) + monkeypatch.setattr("litellm.request_timeout_explicitly_set", True, raising=False) + + with patch("litellm.proxy.proxy_server.general_settings", {"pass_through_request_timeout": 6}): + assert resolve_llm_passthrough_timeout() == 44.0 + assert resolve_llm_passthrough_timeout(kwargs={"stream": True}) == 44.0 + assert resolve_llm_passthrough_timeout(router_timeout=120) == 120.0 + assert resolve_llm_passthrough_timeout(kwargs={"stream": True}, router_stream_timeout=900) == 900.0 + assert resolve_llm_passthrough_timeout(litellm_params={"timeout": 90}) == 90.0 + assert resolve_llm_passthrough_timeout(kwargs={"timeout": 45}) == 45.0 + + +def test_resolve_llm_passthrough_timeout_skips_unset_global_request_timeout(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr("litellm.request_timeout", float(DEFAULT_REQUEST_TIMEOUT_SECONDS), raising=False) + monkeypatch.setattr("litellm.request_timeout_explicitly_set", False, raising=False) + + with patch("litellm.proxy.proxy_server.general_settings", {"pass_through_request_timeout": 6}): + assert resolve_llm_passthrough_timeout() == 6.0 + with patch("litellm.proxy.proxy_server.general_settings", {}): + assert resolve_llm_passthrough_timeout() == DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS + + def test_resolve_llm_passthrough_timeout_stream_timeout_precedence(): assert ( resolve_llm_passthrough_timeout( 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 7e198bc9131..b2ef327f50e 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -4903,3 +4903,19 @@ async def test_model_refresh_updates_availability_catalog_and_retains_it_on_db_f assert await pc._get_models_from_db(client) == [] assert pc.auto_router_db_catalog == () assert find_many.await_count == 3 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("versions", [None, ["2024-11-05"], [], ["2026-07-28"], ["unknown"]]) +async def test_proxy_config_validates_advertised_mcp_versions_at_load(tmp_path, monkeypatch, versions): + config = tmp_path / "mcp-versions.yaml" + config.write_text(json.dumps({"model_list": [], "general_settings": {"mcp_advertised_versions": versions}})) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + if versions is None or versions == ["2024-11-05"]: + _, _, settings = await ProxyConfig().load_config(router=None, config_file_path=str(config)) + assert settings["mcp_advertised_versions"] == versions + return + with pytest.raises(ValidationError): + await ProxyConfig().load_config(router=None, config_file_path=str(config)) diff --git a/tests/test_litellm/proxy/test__types.py b/tests/test_litellm/proxy/test__types.py index f5abe0561db..b43a75d3323 100644 --- a/tests/test_litellm/proxy/test__types.py +++ b/tests/test_litellm/proxy/test__types.py @@ -377,3 +377,21 @@ def test_change_password_request_passwords_hidden_from_repr(): for rendered in (repr(request), str(request)): assert "hunter2hunter2" not in rendered assert "NewP@ssw0rd-2026" not in rendered +@pytest.mark.parametrize("versions", [[], ["2099-01-01"], ["2026-07-28"]]) +def test_mcp_advertised_versions_reject_unavailable_revisions(versions): + from pydantic import ValidationError + + from litellm.proxy._types import ConfigGeneralSettings + + with pytest.raises(ValidationError): + ConfigGeneralSettings(mcp_advertised_versions=versions) + + +@pytest.mark.parametrize("revision", ["2026-07-28", "unknown", None]) +def test_mcp_metadata_rejects_unavailable_upstream_protocol(revision): + from litellm.proxy._types import NewMCPServerRequest, UpdateMCPServerRequest + + payload = {"server_id": "test", "transport": "http", "url": "https://example.com/mcp", "mcp_info": {"protocol_version": revision}} + for model in (NewMCPServerRequest, UpdateMCPServerRequest): + with pytest.raises(ValidationError): + model.model_validate(payload) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 7b74e69685c..c17f41a8b8f 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -7,6 +7,7 @@ from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, from urllib.parse import unquote_plus from unittest.mock import AsyncMock, MagicMock, patch +import anthropic import httpx import pytest from fastapi import HTTPException, Request, Response, status @@ -14,6 +15,7 @@ from fastapi.responses import JSONResponse, StreamingResponse import litellm from litellm._uuid import uuid +from litellm.anthropic_interface.exceptions import AnthropicErrorSseFrame, anthropic_error_sse_frame from litellm.litellm_core_utils.bug_report import ( DISABLE_ENV_VAR, ISSUE_URL_BASE, @@ -55,6 +57,7 @@ from litellm.proxy.common_request_processing import ( sse_error_payload, ) from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header +from litellm.proxy.common_utils.sse_keepalive import ANTHROPIC_PING_SSE_CHUNK from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import ProxyErrorTypes, ProxyException @@ -2543,6 +2546,63 @@ class TestCommonRequestProcessingHelpers: assert response.headers["x-litellm-call-id"] == "call-8302" assert json.loads(response.body) == {"error": {"code": 403, "message": "forbidden"}} + async def test_a_stream_that_fails_before_its_first_byte_answers_as_an_anthropic_json_error(self): + """A /v1/messages stream whose first chunk is already the error frame has nothing + streamed yet, so the failure answers as JSON with the status the upstream gave, + the shape Anthropic clients raise their status-specific errors on""" + + async def stream(): + yield anthropic_error_sse_frame(status_code=503, raw_message="upstream unavailable") + yield ANTHROPIC_PING_SSE_CHUNK + + generator: Final = stream() + response = await create_response(generator, "text/event-stream", {"x-litellm-call-id": "call-8609"}) + + assert isinstance(response, JSONResponse) + assert response.status_code == 503 + assert response.headers["content-type"] == "application/json" + assert response.headers["x-litellm-call-id"] == "call-8609" + assert json.loads(response.body) == { + "type": "error", + "error": {"type": "api_error", "message": "upstream unavailable"}, + } + assert generator.ag_frame is None + + async def test_a_stream_that_fails_before_its_first_byte_names_the_call_when_opted_in(self): + async def stream(): + yield anthropic_error_sse_frame(status_code=429, raw_message="slow down") + + response = await create_response( + stream(), + "text/event-stream", + {"x-litellm-call-id": "call-8609"}, + general_settings={"include_call_id_in_error_body": True}, + ) + + assert isinstance(response, JSONResponse) + assert response.status_code == 429 + assert json.loads(response.body) == { + "type": "error", + "error": {"type": "rate_limit_error", "message": "slow down", "litellm_call_id": "call-8609"}, + } + + async def test_an_error_event_after_a_keepalive_ping_still_streams(self): + """Once a keepalive ping went out the headers are committed, so the error frame + streams as an event instead of turning into a JSON answer""" + + async def stream(): + yield ANTHROPIC_PING_SSE_CHUNK + yield anthropic_error_sse_frame(status_code=503, raw_message="upstream unavailable") + + response = await create_response(stream(), "text/event-stream", {}) + + assert isinstance(response, StreamingResponse) + assert response.status_code == 200 + assert "".join(await self.consume_stream(response)) == ( + ANTHROPIC_PING_SSE_CHUNK + + 'event: error\ndata: {"type": "error", "error": {"type": "api_error", "message": "upstream unavailable"}}\n\n' + ) + async def test_create_streaming_response_disables_proxy_buffering(self): """Regression for #28384: every StreamingResponse create_response returns must carry the headers that stop nginx/ingress/Envoy from buffering the @@ -9899,3 +9959,270 @@ class TestErrorLogCarriesCallId: record: Final = caplog.records[-1] assert record.litellm_call_id == call_id assert call_id in record.getMessage() + + +class TestAnthropicMessagesStreamErrorFrame: + """A ``/v1/messages`` stream that fails after the headers are out has to say so with an + ``event: error`` frame. Anthropic clients pick events by name, so a bare ``data:`` line is + skipped and the request looks like it ended with nothing in it""" + + @staticmethod + def _sse_generator_failing_with(failure: Exception) -> AsyncGenerator[str, None]: + class FailingUpstream: + def __aiter__(self) -> "FailingUpstream": + return self + + async def __anext__(self) -> object: + raise failure + + ProxyLogging._callback_capabilities_cache.clear() + return ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=FailingUpstream(), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + request_data={"model": "claude-sonnet-4-5"}, + proxy_logging_obj=ProxyLogging(user_api_key_cache=MagicMock()), + ) + + @pytest.mark.parametrize( + "status_code, expected_error_type", + [ + (429, "rate_limit_error"), + (529, "overloaded_error"), + (413, "request_too_large"), + (500, "api_error"), + (502, "api_error"), + (400, "invalid_request_error"), + ], + ) + async def test_mid_stream_failure_arrives_as_an_anthropic_error_event( + self, status_code: int, expected_error_type: str + ) -> None: + class UpstreamFailure(Exception): + def __init__(self) -> None: + super().__init__("upstream stopped sending") + self.status_code: Final = status_code + + frames: Final = [frame async for frame in self._sse_generator_failing_with(UpstreamFailure())] + + assert len(frames) == 1 + event_line, data_line, first_blank, second_blank = frames[0].split("\n") + assert isinstance(frames[0], AnthropicErrorSseFrame) + assert frames[0].status_code == status_code + assert event_line == "event: error" + assert (first_blank, second_blank) == ("", "") + payload: Final = json.loads(data_line.removeprefix("data: ")) + assert payload["type"] == "error" + assert payload["error"]["type"] == expected_error_type + assert "upstream stopped sending" in payload["error"]["message"] + + _CONTENT_DELTA_FRAME: Final = ( + b"event: content_block_delta\n" + b'data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"1\\n2\\n3"}}\n\n' + ) + _TORN_DATA_LINE: Final = ( + b"event: content_block_delta\n" + b'data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"4' + ) + _PING: Final = ANTHROPIC_PING_SSE_CHUNK.encode() + + @staticmethod + def _upstream_failure(status_code: int) -> Exception: + class UpstreamFailure(Exception): + def __init__(self) -> None: + super().__init__("upstream stopped sending") + self.status_code: Final = status_code + + return UpstreamFailure() + + @staticmethod + def _sse_generator_cut_after(relayed: Sequence[bytes], failure: Exception) -> AsyncGenerator[str, None]: + class CutUpstream: + def __init__(self) -> None: + self._remaining: Final = iter(relayed) + + def __aiter__(self) -> "CutUpstream": + return self + + async def __anext__(self) -> object: + chunk: Final = next(self._remaining, None) + if chunk is None: + raise failure + return chunk + + ProxyLogging._callback_capabilities_cache.clear() + return ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=CutUpstream(), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + request_data={"model": "claude-sonnet-4-5"}, + proxy_logging_obj=ProxyLogging(user_api_key_cache=MagicMock()), + ) + + @staticmethod + def _as_bytes(chunk: object) -> bytes: + if isinstance(chunk, bytes): + return chunk + assert isinstance(chunk, str) + return chunk.encode() + + async def _wire_bytes(self, relayed: Sequence[bytes]) -> bytes: + stream: Final = self._sse_generator_cut_after(relayed, self._upstream_failure(500)) + return b"".join([self._as_bytes(chunk) async for chunk in stream]) + + @staticmethod + def _error_frame_after(wire: bytes, relayed: bytes) -> bytes: + assert wire.startswith(relayed), f"the wire did not open with {relayed!r}: {wire!r}" + return wire.removeprefix(relayed) + + @staticmethod + def _assert_error_frame(frame: bytes) -> None: + event_line, data_line, first_blank, second_blank = frame.split(b"\n") + assert event_line == b"event: error" + assert (first_blank, second_blank) == (b"", b"") + payload: Final = json.loads(data_line.removeprefix(b"data: ")) + assert payload["type"] == "error" + assert "upstream stopped sending" in payload["error"]["message"] + + @pytest.mark.parametrize( + "torn, seal", + [ + (_TORN_DATA_LINE, b"\n" + _PING), + (b"event: content_bl", b"\n" + _PING), + (b"event: content_block_delta\n", _PING), + (b'event: content_block_delta\r\ndata: {"type":"content_block_delta"}\r\n', _PING), + ], + ids=["mid_data_line", "mid_event_line", "after_a_complete_line", "after_a_crlf_line"], + ) + async def test_a_frame_the_upstream_tore_is_closed_as_a_ping_before_the_error_event( + self, torn: bytes, seal: bytes + ) -> None: + wire: Final = await self._wire_bytes((self._CONTENT_DELTA_FRAME, torn)) + + self._assert_error_frame(self._error_frame_after(wire, self._CONTENT_DELTA_FRAME + torn + seal)) + + async def test_a_cut_at_a_frame_boundary_gets_the_error_event_alone(self) -> None: + wire: Final = await self._wire_bytes((self._CONTENT_DELTA_FRAME,)) + + self._assert_error_frame(self._error_frame_after(wire, self._CONTENT_DELTA_FRAME)) + + async def test_a_torn_frame_still_raises_the_error_in_the_anthropic_sdk(self) -> None: + wire: Final = await self._wire_bytes((self._CONTENT_DELTA_FRAME, self._TORN_DATA_LINE)) + + def serve(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=wire) + + client: Final = anthropic.Anthropic( + api_key="sk-test", + base_url="http://proxy.test", + http_client=httpx.Client(transport=httpx.MockTransport(serve)), + max_retries=0, + ) + with pytest.raises(anthropic.APIStatusError) as raised: + for _ in client.messages.create( + model="claude-sonnet-4-5", max_tokens=16, messages=[{"role": "user", "content": "count"}], stream=True + ): + pass + body: Final = raised.value.body + assert isinstance(body, dict) + assert body["type"] == "error" + assert "upstream stopped sending" in body["error"]["message"] + + async def test_a_failure_before_the_first_byte_answers_with_its_status_as_json(self) -> None: + response: Final = await create_response( + self._sse_generator_failing_with(self._upstream_failure(502)), "text/event-stream", {} + ) + + assert isinstance(response, JSONResponse) + assert response.status_code == 502 + body: Final = json.loads(response.body) + assert body["type"] == "error" + assert body["error"]["type"] == "api_error" + assert "upstream stopped sending" in body["error"]["message"] + + async def test_a_failure_before_the_first_byte_raises_with_its_status_in_the_anthropic_sdk(self) -> None: + response: Final = await create_response( + self._sse_generator_failing_with(self._upstream_failure(502)), "text/event-stream", {} + ) + assert isinstance(response, JSONResponse) + + def serve(request: httpx.Request) -> httpx.Response: + return httpx.Response(response.status_code, headers=dict(response.headers), content=response.body) + + client: Final = anthropic.Anthropic( + api_key="sk-test", + base_url="http://proxy.test", + http_client=httpx.Client(transport=httpx.MockTransport(serve)), + max_retries=0, + ) + with pytest.raises(anthropic.APIStatusError) as raised: + client.messages.create( + model="claude-sonnet-4-5", max_tokens=16, messages=[{"role": "user", "content": "count"}], stream=True + ) + assert raised.value.status_code == 502 + body: Final = raised.value.body + assert isinstance(body, dict) + assert body["type"] == "error" + assert "upstream stopped sending" in body["error"]["message"] + + +class TestStreamingContainerOwnershipRecordedBeforeDone: + """Regression for LIT-8612: the OpenAI SDK closes the connection at + ``data: [DONE]`` and starlette cancels the body task, so an ownership row + written after the SSE generator is exhausted never lands. The row must be + written before the chunk carrying ``response.completed`` is handed to the + client.""" + + CHUNKS: Final = ( + 'data: {"type":"response.created"}\n\n', + 'data: {"type":"response.output_text.delta"}\n\n', + 'data: {"type":"response.completed"}\n\n', + "data: [DONE]\n\n", + ) + TERMINAL_INDEX: Final = 2 + + @staticmethod + def _completed_event() -> SimpleNamespace: + return SimpleNamespace( + type="response.completed", + response=SimpleNamespace( + id="resp_lit8612", + output=[SimpleNamespace(type="code_interpreter_call", container_id="cntr_lit8612")], + ), + ) + + async def _sse(self, stream: SimpleNamespace, populate_at: int) -> AsyncGenerator[str, None]: + for index, chunk in enumerate(self.CHUNKS): + if index == populate_at: + stream.completed_response = self._completed_event() + yield chunk + if populate_at == len(self.CHUNKS): + stream.completed_response = self._completed_event() + + async def _await_counts_per_chunk(self, populate_at: int) -> tuple[tuple[tuple[str, int], ...], AsyncMock]: + stream: Final = SimpleNamespace(completed_response=None, _hidden_params={"custom_llm_provider": "azure"}) + recorder: Final = AsyncMock(return_value=None) + with patch( + "litellm.proxy.container_endpoints.ownership.record_container_owners_from_responses_response", recorder + ): + wrapped: Final = ProxyBaseLLMRequestProcessing._wrap_responses_stream_for_container_ownership( + original_stream_response=stream, + wrapped_generator=self._sse(stream, populate_at), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test", team_id="team-1"), + ) + observed: Final = tuple([(chunk, recorder.await_count) async for chunk in wrapped]) + return observed, recorder + + async def test_row_is_written_before_the_terminal_chunk_reaches_the_client(self) -> None: + observed, recorder = await self._await_counts_per_chunk(populate_at=self.TERMINAL_INDEX) + + assert tuple(chunk for chunk, _ in observed) == self.CHUNKS + assert tuple(count for _, count in observed) == (0, 0, 1, 1) + recorder.assert_awaited_once() + assert recorder.await_args.kwargs["response"].output[0].container_id == "cntr_lit8612" + assert recorder.await_args.kwargs["user_api_key_dict"].team_id == "team-1" + + async def test_row_is_still_written_when_the_iterator_completes_only_at_exhaustion(self) -> None: + observed, recorder = await self._await_counts_per_chunk(populate_at=len(self.CHUNKS)) + + assert tuple(chunk for chunk, _ in observed) == self.CHUNKS + assert tuple(count for _, count in observed) == (0, 0, 0, 0) + recorder.assert_awaited_once() diff --git a/tests/test_litellm/proxy/test_zerobus_dashboard_config.py b/tests/test_litellm/proxy/test_zerobus_dashboard_config.py new file mode 100644 index 00000000000..d2143767480 --- /dev/null +++ b/tests/test_litellm/proxy/test_zerobus_dashboard_config.py @@ -0,0 +1,52 @@ +from pathlib import Path +from typing import Final + +from pydantic import BaseModel, TypeAdapter + +import litellm +from litellm.integrations.custom_logger import CustomLogger + + +class DashboardField(BaseModel): + type: str + required: bool + + +class DashboardCallbackConfig(BaseModel): + id: str + displayName: str + logo: str + supports_key_team_logging: bool + dynamic_params: dict[str, DashboardField] + + +def _zerobus_config() -> DashboardCallbackConfig: + path: Final = Path(litellm.__file__).parent / "integrations" / "callback_configs.json" + configs: Final = TypeAdapter(tuple[DashboardCallbackConfig, ...]).validate_json(path.read_text()) + return next(config for config in configs if config.id == "zerobus") + + +def test_zerobus_appears_in_the_dashboard_callback_dropdown(): + """The dropdown is served from callback_configs.json, so an entry only in the dashboard source is invisible.""" + entry = _zerobus_config() + + assert entry.displayName == "Databricks Zerobus" + assert entry.supports_key_team_logging is False + assert entry.dynamic_params["ZEROBUS_CLIENT_SECRET"].type == "password" + assert all(field.required is True for field in entry.dynamic_params.values()) + + +def test_the_dropdown_logo_asset_exists(): + """A logo the dashboard cannot resolve degrades silently to a letter tile.""" + logo = _zerobus_config().logo + repo_root = Path(litellm.__file__).parent.parent + asset = repo_root / "ui" / "litellm-dashboard" / "public" / "assets" / "logos" / logo + + assert asset.is_file() + + +def test_the_dropdown_fields_are_the_env_vars_the_logger_reads(): + """Naming the fields as stored means the edit form prefills saved values instead of showing blanks.""" + fields = tuple(_zerobus_config().dynamic_params) + + assert fields == tuple(CustomLogger.get_callback_env_vars("zerobus")) diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index dbf54ec3b9b..9dbbc20591e 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -6,13 +6,14 @@ completion_start_time = end_time.""" import json from datetime import datetime from typing import Final, Optional -from unittest.mock import Mock, patch +from unittest.mock import AsyncMock, Mock, patch import httpx import pytest from pydantic_core import PydanticSerializationError import litellm +from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.streaming_iterator import ( @@ -251,6 +252,171 @@ def test_sync_transport_error_before_completed_event_raises(): pass +_DONE_MARKER: Final = b"data: [DONE]\n\n" +_CREATED_EVENT: Final = _sse_event({"type": "response.created"}) +_IN_PROGRESS_EVENT: Final = _sse_event({"type": "response.in_progress"}) +_PARTIAL_OUTPUT_EVENTS: Final = _COMPLETE_STREAM_EVENTS[:-1] +_PRE_OUTPUT_PREFIXES: Final = [ + pytest.param([], True, id="nothing-yielded"), + pytest.param([_CREATED_EVENT], False, id="created"), + pytest.param([_CREATED_EVENT, _IN_PROGRESS_EVENT], False, id="created-and-in-progress"), +] + + +def _failure_tracking_logging_obj() -> Mock: + logging_obj: Final = _logging_obj_stub() + logging_obj.async_failure_handler = AsyncMock() + return logging_obj + + +def _assert_failure_logged_once(logging_obj: Mock, exception: Exception) -> None: + assert logging_obj.async_failure_handler.await_count == 1 + assert logging_obj.async_failure_handler.await_args.kwargs["exception"] is exception + + +@pytest.mark.asyncio +@pytest.mark.parametrize("prefix, pre_first_chunk", _PRE_OUTPUT_PREFIXES) +@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS, ids=type) +async def test_transport_error_before_any_output_raises_fallback_error(prefix, pre_first_chunk, trailing_error): + """A connection lost while only lifecycle events (response.created / response.in_progress) + have streamed is fallback-eligible, so it must surface as the MidStreamFallbackError the + router re-routes, carrying the raw transport error and no generated content.""" + logging_obj: Final = _failure_tracking_logging_obj() + iterator: Final = _make_iterator(sse_events=prefix, logging_obj=logging_obj, trailing_error=trailing_error) + + with pytest.raises(MidStreamFallbackError) as exc_info: + async for _ in iterator: + pass + + assert exc_info.value.original_exception is trailing_error + assert exc_info.value.is_pre_first_chunk is pre_first_chunk + assert exc_info.value.generated_content == "" + _assert_failure_logged_once(logging_obj, trailing_error) + + +@pytest.mark.asyncio +async def test_transport_error_after_output_started_is_not_fallback_eligible(): + logging_obj: Final = _failure_tracking_logging_obj() + trailing_error: Final = httpx.ReadError("Response payload is not completed") + iterator: Final = _make_iterator( + sse_events=_PARTIAL_OUTPUT_EVENTS, logging_obj=logging_obj, trailing_error=trailing_error + ) + + with pytest.raises(httpx.ReadError) as exc_info: + async for _ in iterator: + pass + + assert exc_info.value is trailing_error + _assert_failure_logged_once(logging_obj, trailing_error) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("trailer", [[], [_DONE_MARKER]], ids=["eof", "done-marker"]) +async def test_stream_ending_after_partial_output_without_terminal_event_raises(trailer): + """A clean EOF or `[DONE]` after output text but with no response.completed / + response.incomplete / response.failed is a truncated answer: the partial events still + reach the caller, then an explicit error follows instead of a normal end of stream.""" + logging_obj: Final = _failure_tracking_logging_obj() + iterator: Final = _make_iterator(sse_events=[*_PARTIAL_OUTPUT_EVENTS, *trailer], logging_obj=logging_obj) + + created: Final = await iterator.__anext__() + delta: Final = await iterator.__anext__() + with pytest.raises(litellm.APIConnectionError) as exc_info: + await iterator.__anext__() + + assert (created.type, delta.type) == ("response.created", "response.output_text.delta") + assert not isinstance(exc_info.value, MidStreamFallbackError) + assert exc_info.value.llm_provider == "openai" + _assert_failure_logged_once(logging_obj, exc_info.value) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("prefix, pre_first_chunk", _PRE_OUTPUT_PREFIXES) +@pytest.mark.parametrize("trailer", [[], [_DONE_MARKER]], ids=["eof", "done-marker"]) +async def test_stream_ending_before_any_output_raises_fallback_error(prefix, pre_first_chunk, trailer): + logging_obj: Final = _failure_tracking_logging_obj() + iterator: Final = _make_iterator(sse_events=[*prefix, *trailer], logging_obj=logging_obj) + + with pytest.raises(MidStreamFallbackError) as exc_info: + async for _ in iterator: + pass + + assert isinstance(exc_info.value.original_exception, litellm.APIConnectionError) + assert exc_info.value.is_pre_first_chunk is pre_first_chunk + assert exc_info.value.generated_content == "" + _assert_failure_logged_once(logging_obj, exc_info.value.original_exception) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("trailer", [[], [_DONE_MARKER]], ids=["eof", "done-marker"]) +async def test_complete_stream_still_ends_normally(trailer): + logging_obj: Final = _failure_tracking_logging_obj() + iterator: Final = _make_iterator(sse_events=[*_COMPLETE_STREAM_EVENTS, *trailer], logging_obj=logging_obj) + + seen: Final = [event.type async for event in iterator] + + assert seen[-1] == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + assert logging_obj.async_failure_handler.await_count == 0 + + +@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS, ids=type) +def test_sync_transport_error_before_any_output_raises_fallback_error(trailing_error): + logging_obj: Final = _failure_tracking_logging_obj() + iterator: Final = _make_sync_iterator( + sse_events=[_CREATED_EVENT, _IN_PROGRESS_EVENT], + logging_obj=logging_obj, + trailing_error=trailing_error, + ) + + with pytest.raises(MidStreamFallbackError) as exc_info: + for _ in iterator: + pass + + assert exc_info.value.original_exception is trailing_error + assert exc_info.value.is_pre_first_chunk is False + assert exc_info.value.generated_content == "" + _assert_failure_logged_once(logging_obj, trailing_error) + + +@pytest.mark.parametrize("trailer", [[], [_DONE_MARKER]], ids=["eof", "done-marker"]) +def test_sync_stream_ending_after_partial_output_without_terminal_event_raises(trailer): + logging_obj: Final = _failure_tracking_logging_obj() + iterator: Final = _make_sync_iterator(sse_events=[*_PARTIAL_OUTPUT_EVENTS, *trailer], logging_obj=logging_obj) + + created: Final = next(iterator) + delta: Final = next(iterator) + with pytest.raises(litellm.APIConnectionError) as exc_info: + next(iterator) + + assert (created.type, delta.type) == ("response.created", "response.output_text.delta") + assert not isinstance(exc_info.value, MidStreamFallbackError) + _assert_failure_logged_once(logging_obj, exc_info.value) + + +def test_sync_stream_ending_before_any_output_raises_fallback_error(): + logging_obj: Final = _failure_tracking_logging_obj() + iterator: Final = _make_sync_iterator(sse_events=[_CREATED_EVENT], logging_obj=logging_obj) + + with pytest.raises(MidStreamFallbackError) as exc_info: + for _ in iterator: + pass + + assert isinstance(exc_info.value.original_exception, litellm.APIConnectionError) + assert exc_info.value.is_pre_first_chunk is False + _assert_failure_logged_once(logging_obj, exc_info.value.original_exception) + + +@pytest.mark.parametrize("trailer", [[], [_DONE_MARKER]], ids=["eof", "done-marker"]) +def test_sync_complete_stream_still_ends_normally(trailer): + logging_obj: Final = _failure_tracking_logging_obj() + iterator: Final = _make_sync_iterator(sse_events=[*_COMPLETE_STREAM_EVENTS, *trailer], logging_obj=logging_obj) + + seen: Final = [event.type for event in iterator] + + assert seen[-1] == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + assert logging_obj.async_failure_handler.await_count == 0 + + def test_stream_cache_write_completes_when_asyncio_run_closes_the_loop(monkeypatch): """ Regression test for LIT-6184 on the /v1/responses streaming surface: the diff --git a/tests/test_litellm/router_strategy/test_budget_limiter.py b/tests/test_litellm/router_strategy/test_budget_limiter.py new file mode 100644 index 00000000000..62de1586fdd --- /dev/null +++ b/tests/test_litellm/router_strategy/test_budget_limiter.py @@ -0,0 +1,137 @@ +""" +Spend tracking in RouterBudgetLimiting.async_log_success_event. + +Only chat completions puts custom_llm_provider into litellm_params. The responses, +anthropic_messages, embedding and rerank surfaces leave it unset, which used to make +the callback raise before any spend was recorded, so those budgets never moved. +""" + +from typing import Final + +import pytest + +from litellm.caching.caching import DualCache +from litellm.router_strategy.budget_limiter import RouterBudgetLimiting + + +@pytest.fixture +def disable_budget_sync(monkeypatch): + async def noop(*args, **kwargs): + return None + + monkeypatch.setattr( + "litellm.router_strategy.budget_limiter.RouterBudgetLimiting.periodic_sync_in_memory_spend_with_redis", + noop, + ) + + +def _success_kwargs( + *, + provider_in_litellm_params: str | None, + provider_in_payload: str | None, + call_type: str = "aresponses", + response_cost: float = 0.25, + model_id: str = "deployment-1", +) -> dict[str, object]: + provider_params: Final[dict[str, str]] = ( + {} if provider_in_litellm_params is None else {"custom_llm_provider": provider_in_litellm_params} + ) + litellm_params: Final[dict[str, str]] = {"model": "openai/gpt-4o", **provider_params} + + return { + "call_type": call_type, + "litellm_params": litellm_params, + "standard_logging_object": { + "response_cost": response_cost, + "model_id": model_id, + "custom_llm_provider": provider_in_payload, + }, + } + + +async def _log_success(limiter: RouterBudgetLimiting, kwargs: dict[str, object]) -> None: + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=None, end_time=None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ["aresponses", "anthropic_messages", "aembedding", "arerank"]) +async def test_provider_spend_tracked_when_litellm_params_omits_provider(disable_budget_sync, call_type): + """Non-chat surfaces carry the provider only on the standard logging payload.""" + limiter = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config={"openai": {"budget_limit": 10.0, "time_period": "1d"}}, + ) + + await _log_success( + limiter, + _success_kwargs( + provider_in_litellm_params=None, + provider_in_payload="openai", + call_type=call_type, + ), + ) + + assert await limiter.dual_cache.async_get_cache("provider_spend:openai:1d") == 0.25 + + +@pytest.mark.asyncio +async def test_chat_completions_spend_still_tracked(disable_budget_sync): + """Chat completions fills in both sources and must keep accumulating.""" + limiter = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config={"openai": {"budget_limit": 10.0, "time_period": "1d"}}, + ) + + await _log_success( + limiter, + _success_kwargs( + provider_in_litellm_params="openai", + provider_in_payload="openai", + call_type="acompletion", + ), + ) + + assert await limiter.dual_cache.async_get_cache("provider_spend:openai:1d") == 0.25 + + +@pytest.mark.asyncio +async def test_budget_of_other_provider_is_untouched(disable_budget_sync): + """A provider without its own budget must not bleed into a configured one.""" + limiter = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config={"openai": {"budget_limit": 10.0, "time_period": "1d"}}, + ) + + await _log_success( + limiter, + _success_kwargs(provider_in_litellm_params=None, provider_in_payload="anthropic"), + ) + + assert await limiter.dual_cache.async_get_cache("provider_spend:openai:1d") in (None, 0.0) + + +@pytest.mark.asyncio +async def test_deployment_budget_tracked_when_provider_is_unresolvable(disable_budget_sync): + """An unresolvable provider must not abort the deployment and tag budgets that follow it.""" + limiter = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config=None, + model_list=[ + { + "model_name": "some-model", + "litellm_params": { + "model": "openai/gpt-4o", + "max_budget": 10.0, + "budget_duration": "1d", + }, + "model_info": {"id": "deployment-1"}, + } + ], + ) + + await _log_success( + limiter, + _success_kwargs(provider_in_litellm_params=None, provider_in_payload=None), + ) + + assert await limiter.dual_cache.async_get_cache("deployment_spend:deployment-1:1d") == 0.25 diff --git a/tests/unit/anthropic_interface/exceptions/test_exception_mapping_utils.py b/tests/unit/anthropic_interface/exceptions/test_exception_mapping_utils.py index ef092b65f28..0d8e7674e7d 100644 --- a/tests/unit/anthropic_interface/exceptions/test_exception_mapping_utils.py +++ b/tests/unit/anthropic_interface/exceptions/test_exception_mapping_utils.py @@ -3,8 +3,15 @@ Tests for AnthropicExceptionMapping class in litellm/anthropic_interface/excepti """ import json +from typing import Final -from litellm.anthropic_interface.exceptions import AnthropicExceptionMapping +import pytest + +from litellm.anthropic_interface.exceptions import ( + AnthropicErrorSseFrame, + AnthropicExceptionMapping, + anthropic_error_sse_frame, +) class TestCreateErrorResponse: @@ -206,3 +213,42 @@ class TestTransformToAnthropicError: ) assert result["type"] == "error" assert result["error"]["message"] == '["error1", "error2"]' + + +class TestAnthropicErrorSseFrame: + @pytest.mark.parametrize( + ("status_code", "expected_error_type"), + [(429, "rate_limit_error"), (503, "api_error"), (400, "invalid_request_error")], + ) + def test_the_frame_is_one_error_event_carrying_the_anthropic_envelope( + self, status_code: int, expected_error_type: str + ) -> None: + frame: Final = anthropic_error_sse_frame(status_code=status_code, raw_message="upstream unavailable") + + event_line, data_line, first_blank, second_blank = frame.split("\n") + assert event_line == "event: error" + assert (first_blank, second_blank) == ("", "") + assert json.loads(data_line.removeprefix("data: ")) == { + "type": "error", + "error": {"type": expected_error_type, "message": "upstream unavailable"}, + } + + def test_the_frame_remembers_the_status_and_body_it_was_built_from(self) -> None: + frame: Final = anthropic_error_sse_frame(status_code=503, raw_message="upstream unavailable") + + assert isinstance(frame, AnthropicErrorSseFrame) + assert frame.status_code == 503 + data_line: Final = frame.split("\n")[1] + assert data_line == f"data: {json.dumps(frame.json_body(call_id=None))}" + + def test_the_json_body_names_the_call_only_when_asked(self) -> None: + frame: Final = anthropic_error_sse_frame(status_code=503, raw_message="upstream unavailable") + + assert frame.json_body(call_id="call-1") == { + "type": "error", + "error": {"type": "api_error", "message": "upstream unavailable", "litellm_call_id": "call-1"}, + } + assert frame.json_body(call_id=None) == { + "type": "error", + "error": {"type": "api_error", "message": "upstream unavailable"}, + } diff --git a/tests/unit/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index dd95addac40..b8b922f72a7 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -464,6 +464,40 @@ def test_total_cost_applies_the_long_context_batch_tier_per_line(): assert result.cost == pytest.approx((300_000 * 2e-6) + (10 * 6e-6) + (100 * 1e-6) + (10 * 4e-6)) +def test_xai_output_lines_bill_reasoning_tokens_as_completion_tokens(): + row = _success_row( + model="grok-4.3", + usage={ + "prompt_tokens": 615, + "completion_tokens": 3, + "total_tokens": 993, + "completion_tokens_details": {"reasoning_tokens": 375}, + }, + ) + + result = bu._aggregate_batch_cost_usage_models( + entries=[row], + custom_llm_provider="xai", + model_info=ModelInfo( + key="xai/grok-4.3", + max_tokens=None, + max_input_tokens=None, + max_output_tokens=None, + input_cost_per_token=1.25e-6, + output_cost_per_token=2.5e-6, + litellm_provider="xai", + mode="chat", + supported_openai_params=None, + input_cost_per_token_batches=1e-6, + output_cost_per_token_batches=2e-6, + ), + ) + + assert result.usage.completion_tokens == 378 + assert result.usage.total_tokens == 993 + assert result.cost == pytest.approx((615 * 1e-6) + (378 * 2e-6)) + + def test_total_usage_empty_is_zero(): result = bu._aggregate_batch_cost_usage_models(entries=[], custom_llm_provider="openai") assert result.cost == 0.0 diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 653b2c9914a..ecea4723bf4 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -1,11 +1,14 @@ import asyncio +import base64 import importlib import os from collections.abc import Coroutine, Iterator +from dataclasses import dataclass, field from pathlib import Path from typing import Final import boto3 +import httpx import pytest from pytest_socket import enable_socket, socket_allow_hosts @@ -15,9 +18,12 @@ import litellm # noqa: E402 # litellm reads LITELLM_LOCAL_MODEL_COST_MAP at im import litellm.router as litellm_router_module # noqa: E402 # same import-time dependency import litellm.utils as litellm_utils_module # noqa: E402 # same import-time dependency from litellm._logging import ALL_LOGGERS # noqa: E402 # same import-time dependency +from litellm.anthropic_beta_headers_manager import reload_beta_headers_config # noqa: E402 # same import-time dependency +from litellm.litellm_core_utils.prompt_templates import factory as prompt_factory_module # noqa: E402 # same import-time dependency from litellm.litellm_core_utils.prompt_templates import ( # noqa: E402 # same import-time dependency image_handling as image_handling_module, ) +from litellm.llms.gemini.chat import transformation as gemini_chat_transformation_module # noqa: E402 # same import-time dependency from litellm.llms.custom_httpx.async_client_cleanup import ( # noqa: E402 # same import-time dependency close_litellm_async_clients, ) @@ -89,6 +95,9 @@ RESTORED_GLOBALS: Final = ( ) MODULE_LEVEL_CLIENTS: Final = ("module_level_client", "module_level_aclient") SESSION_CLIENTS: Final = ("base_llm_aiohttp_handler", "httpx_client", "aclient", "client") +ONE_PIXEL_PNG: Final = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" +) def _allow_loopback_only() -> None: @@ -236,6 +245,47 @@ def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: litellm.get_model_info.cache_clear() +@pytest.fixture +def local_beta_headers_config(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", "True") + reload_beta_headers_config() + yield + monkeypatch.delenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", raising=False) + reload_beta_headers_config() + + +@dataclass(slots=True) +class AsyncOnlyImageFetch: + fetched: list[str] = field(default_factory=list) # mutable-ok: tests assert on the URLs fetched, in order + base64_png: str = base64.b64encode(ONE_PIXEL_PNG).decode() + data_url: str = "data:image/png;base64," + base64.b64encode(ONE_PIXEL_PNG).decode() + + +@pytest.fixture +def async_only_image_fetch(monkeypatch: pytest.MonkeyPatch) -> AsyncOnlyImageFetch: + fetch: Final = AsyncOnlyImageFetch() + + def forbid_sync_fetch(client: object, url: str, **kwargs: object) -> httpx.Response: + raise litellm.ImageFetchError(f"sync image fetch ran on the event loop: {url}") + + async def serve_png(client: object, url: str, **kwargs: object) -> httpx.Response: + fetch.fetched.append(url) + return httpx.Response( + 200, content=ONE_PIXEL_PNG, headers={"content-type": "image/png"}, request=httpx.Request("GET", url) + ) + + def forbid_sync_convert(url: str, *args: object, **kwargs: object) -> str: + if url.startswith(("http://", "https://")): + raise litellm.ImageFetchError(f"sync convert_url_to_base64 ran on the request path: {url}") + return url + + monkeypatch.setattr(image_handling_module, "safe_get", forbid_sync_fetch) + monkeypatch.setattr(image_handling_module, "async_safe_get", serve_png) + for module in (image_handling_module, prompt_factory_module, gemini_chat_transformation_module): + monkeypatch.setattr(module, "convert_url_to_base64", forbid_sync_convert) + return fetch + + @pytest.fixture def no_ambient_azure_credentials(monkeypatch: pytest.MonkeyPatch) -> None: for name in AMBIENT_AZURE_CREDENTIAL_ENV_VARS: diff --git a/tests/test_litellm/integrations/open_telemetry/__init__.py b/tests/unit/expected_fine_tuning_api/__init__.py similarity index 100% rename from tests/test_litellm/integrations/open_telemetry/__init__.py rename to tests/unit/expected_fine_tuning_api/__init__.py diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_cancel_expected_output.json b/tests/unit/expected_fine_tuning_api/azure_cancel_expected_output.json similarity index 100% rename from tests/test_litellm/expected_fine_tuning_api/azure_cancel_expected_output.json rename to tests/unit/expected_fine_tuning_api/azure_cancel_expected_output.json diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_cancel_raw_response.json b/tests/unit/expected_fine_tuning_api/azure_cancel_raw_response.json similarity index 100% rename from tests/test_litellm/expected_fine_tuning_api/azure_cancel_raw_response.json rename to tests/unit/expected_fine_tuning_api/azure_cancel_raw_response.json diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_cancel_request.json b/tests/unit/expected_fine_tuning_api/azure_cancel_request.json similarity index 100% rename from tests/test_litellm/expected_fine_tuning_api/azure_cancel_request.json rename to tests/unit/expected_fine_tuning_api/azure_cancel_request.json diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_create_expected_output.json b/tests/unit/expected_fine_tuning_api/azure_create_expected_output.json similarity index 100% rename from tests/test_litellm/expected_fine_tuning_api/azure_create_expected_output.json rename to tests/unit/expected_fine_tuning_api/azure_create_expected_output.json diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_create_raw_response.json b/tests/unit/expected_fine_tuning_api/azure_create_raw_response.json similarity index 100% rename from tests/test_litellm/expected_fine_tuning_api/azure_create_raw_response.json rename to tests/unit/expected_fine_tuning_api/azure_create_raw_response.json diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_create_request.json b/tests/unit/expected_fine_tuning_api/azure_create_request.json similarity index 100% rename from tests/test_litellm/expected_fine_tuning_api/azure_create_request.json rename to tests/unit/expected_fine_tuning_api/azure_create_request.json diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_list_raw_response.json b/tests/unit/expected_fine_tuning_api/azure_list_raw_response.json similarity index 100% rename from tests/test_litellm/expected_fine_tuning_api/azure_list_raw_response.json rename to tests/unit/expected_fine_tuning_api/azure_list_raw_response.json diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_list_request.json b/tests/unit/expected_fine_tuning_api/azure_list_request.json similarity index 100% rename from tests/test_litellm/expected_fine_tuning_api/azure_list_request.json rename to tests/unit/expected_fine_tuning_api/azure_list_request.json diff --git a/tests/unit/experimental_mcp_client/test_mcp_client.py b/tests/unit/experimental_mcp_client/test_mcp_client.py index 368e34c455d..1a56227b008 100644 --- a/tests/unit/experimental_mcp_client/test_mcp_client.py +++ b/tests/unit/experimental_mcp_client/test_mcp_client.py @@ -58,6 +58,15 @@ from litellm.types.mcp_server.mcp_server_manager import MCPServer _JSONRPC_MESSAGE_ADAPTER: Final = TypeAdapter(JSONRPCMessage) +def _initialized(instructions: str | None = None) -> InitializeResult: + return InitializeResult( + protocol_version=LATEST_HANDSHAKE_VERSION, + capabilities=ServerCapabilities(), + server_info=Implementation(name="test", version="1"), + instructions=instructions, + ) + + class _MockTransportClient(MCPClient): """An MCPClient whose streamable-HTTP transport runs on an httpx2 MockTransport.""" @@ -125,7 +134,7 @@ class TestMCPClient: mock_stdio_client.return_value = mock_stdio_ctx mock_session_instance = AsyncMock() - mock_session_instance.initialize = AsyncMock() + mock_session_instance.initialize = AsyncMock(return_value=_initialized()) mock_session_ctx = AsyncMock() mock_session_ctx.__aenter__.return_value = mock_session_instance mock_session_ctx.__aexit__.return_value = None @@ -168,7 +177,7 @@ class TestMCPClient: # Mock the session with patch("litellm.experimental_mcp_client.client.ClientSession") as mock_session: mock_session_instance = AsyncMock() - mock_session_instance.initialize = AsyncMock() + mock_session_instance.initialize = AsyncMock(return_value=_initialized()) mock_session_ctx = AsyncMock() mock_session_ctx.__aenter__.return_value = mock_session_instance mock_session_ctx.__aexit__.return_value = None @@ -214,7 +223,7 @@ class TestMCPClient: # Mock the session with patch("litellm.experimental_mcp_client.client.ClientSession") as mock_session: mock_session_instance = AsyncMock() - mock_session_instance.initialize = AsyncMock() + mock_session_instance.initialize = AsyncMock(return_value=_initialized()) mock_session_ctx = AsyncMock() mock_session_ctx.__aenter__.return_value = mock_session_instance mock_session_ctx.__aexit__.return_value = None @@ -266,7 +275,7 @@ class TestMCPClient: # Mock the session with patch("litellm.experimental_mcp_client.client.ClientSession") as mock_session: mock_session_instance = AsyncMock() - mock_session_instance.initialize = AsyncMock() + mock_session_instance.initialize = AsyncMock(return_value=_initialized()) mock_session_ctx = AsyncMock() mock_session_ctx.__aenter__.return_value = mock_session_instance mock_session_ctx.__aexit__.return_value = None @@ -413,8 +422,7 @@ class TestMCPClientInstructionsCapture: ) mock_session = AsyncMock() - init_result = MagicMock() - init_result.instructions = " upstream says hello " + init_result = _initialized(" upstream says hello ") mock_session.initialize = AsyncMock(return_value=init_result) session_ctx = MagicMock() @@ -442,8 +450,7 @@ class TestMCPClientInstructionsCapture: ) mock_session = AsyncMock() - init_result = MagicMock() - init_result.instructions = None + init_result = _initialized() mock_session.initialize = AsyncMock(return_value=init_result) session_ctx = MagicMock() @@ -600,8 +607,7 @@ class TestExecuteSessionOperationSurfacesTransportError: @patch("litellm.experimental_mcp_client.client.ClientSession") async def test_cleanup_error_after_success_is_swallowed(self, mock_session_cls): client = MCPClient(server_url="http://example.com/mcp", transport_type="http") - init_result = MagicMock() - init_result.instructions = None + init_result = _initialized() self._make_session(mock_session_cls, AsyncMock(return_value=init_result)) transport_ctx = self._make_transport(_FakeExceptionGroup("late", [httpx2.ConnectError("late cleanup error")])) @@ -634,7 +640,7 @@ class TestExecuteSessionOperationSurfacesTransportError: @pytest.mark.parametrize("original_error", (False, True)) @patch("litellm.experimental_mcp_client.client.ClientSession") async def test_session_exit_cancellation_preserves_original_failure(self, session_class, original_error): - self._make_session(session_class, AsyncMock(return_value=None)) + self._make_session(session_class, AsyncMock(return_value=_initialized())) cancelled: Final = asyncio.CancelledError("cancelled while closing session") session_class.return_value.__aexit__ = AsyncMock(side_effect=cancelled) original: Final = RuntimeError("operation failed") @@ -656,7 +662,7 @@ class TestExecuteSessionOperationSurfacesTransportError: @pytest.mark.parametrize("phase", ("session", "transport")) @patch("litellm.experimental_mcp_client.client.ClientSession") async def test_cleanup_preserves_process_exit(self, session_class, phase, signal_type): - self._make_session(session_class, AsyncMock(return_value=None)) + self._make_session(session_class, AsyncMock(return_value=_initialized())) signal: Final = signal_type("process stopping") if phase == "session": session_class.return_value.__aexit__ = AsyncMock(side_effect=signal) @@ -670,7 +676,7 @@ class TestExecuteSessionOperationSurfacesTransportError: @pytest.mark.asyncio @patch("litellm.experimental_mcp_client.client.ClientSession") async def test_session_and_termination_share_one_cleanup_deadline(self, session_class): - self._make_session(session_class, AsyncMock(return_value=None)) + self._make_session(session_class, AsyncMock(return_value=_initialized())) deleting: Final = asyncio.Event() async def close_session(*args): @@ -1883,16 +1889,17 @@ async def test_sse_read_failure_is_preserved() -> None: @pytest.mark.asyncio +@pytest.mark.parametrize("protocol_version", ["auto", "2025-06-18"]) @pytest.mark.parametrize("transport", [MCPTransport.sse, MCPTransport.stdio]) @pytest.mark.parametrize("mode", ["ok", "closed", "silent"]) -async def test_transport_completion_and_normal_messages(transport: MCPTransport, mode: str) -> None: +async def test_transport_completion_and_normal_messages(transport: MCPTransport, mode: str, protocol_version: str) -> None: from mcp import ClientSession from litellm.proxy._experimental.mcp_server.rest_endpoints import _connection_error_message logging_callback: Final = AsyncMock() read_timeout: Final = 0.2 if mode == "silent" else 30 client: Final = MCPClient( - server_url="https://example.com/sse", transport_type=transport, timeout=read_timeout, logging_callback=logging_callback + server_url="https://example.com/sse", transport_type=transport, timeout=read_timeout, logging_callback=logging_callback, protocol_version=protocol_version ) async def operation(session: ClientSession) -> CallToolResult: @@ -2754,12 +2761,13 @@ async def test_http_close_cancellation_cannot_turn_into_success(original_error: @pytest.mark.asyncio +@pytest.mark.parametrize("protocol_version", ("auto", "2025-06-18")) @pytest.mark.parametrize("cancel_mode", ("scope", "task", "wait_for", "read_timeout")) @pytest.mark.parametrize("concurrency", (1, 5)) @pytest.mark.parametrize("termination", ("ok", "hang", "hang_body")) @pytest.mark.parametrize("raise_on_error", (False, True)) async def test_cancellation_delivers_termination_over_tcp( - cancel_mode: str, concurrency: int, termination: str, raise_on_error: bool + cancel_mode: str, concurrency: int, termination: str, raise_on_error: bool, protocol_version: str ) -> None: started: Final = asyncio.Event() scope_ready: Final[asyncio.Future[anyio.CancelScope]] = asyncio.get_running_loop().create_future() @@ -2813,6 +2821,8 @@ async def test_cancellation_delivers_termination_over_tcp( await stop.wait() return if payload["method"] == "initialize": + if cancel_mode != "read_timeout": + await asyncio.sleep(0.75) response: Final = json.dumps( { "jsonrpc": "2.0", @@ -2839,7 +2849,7 @@ async def test_cancellation_delivers_termination_over_tcp( listener: Final = await asyncio.start_server(handle_connection, "127.0.0.1", 0) port: Final = listener.sockets[0].getsockname()[1] client: Final = MCPClient( - server_url=f"http://127.0.0.1:{port}/mcp", timeout=2 if cancel_mode == "read_timeout" else 0.5 if termination != "ok" else 30 + server_url=f"http://127.0.0.1:{port}/mcp", protocol_version=protocol_version, timeout=2 if cancel_mode == "read_timeout" else 30 ) async def calls(): @@ -2866,7 +2876,7 @@ async def test_cancellation_delivers_termination_over_tcp( try: task: Final = asyncio.create_task(invoke()) - await asyncio.wait_for(started.wait(), 3) + await asyncio.wait_for(started.wait(), 30) if cancel_mode == "scope": (await scope_ready).deadline = anyio.current_time() + 0.2 if cancel_mode == "task": @@ -2901,3 +2911,52 @@ async def test_cancellation_delivers_termination_over_tcp( closed: Final = await asyncio.wait_for(asyncio.gather(*connections, return_exceptions=True), 2) assert all(result is None or isinstance(result, asyncio.CancelledError) for result in closed), closed await asyncio.wait_for(listener.wait_closed(), 2) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("revision", ["2024-11-05", "2025-03-26", "2025-06-18", "2025-11-25", "auto"]) +@pytest.mark.parametrize("accepted", [True, False]) +@pytest.mark.parametrize("callbacks", [False, True]) +async def test_configured_upstream_revision_is_offered_and_checked(revision, accepted, callbacks): + from mcp.types import JSONRPCRequest + from mcp_types.version import LATEST_HANDSHAKE_VERSION + + offered = LATEST_HANDSHAKE_VERSION if revision == "auto" else revision + + def respond(request): + if request.method == "DELETE": + return httpx2.Response(200) + payload = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content) + if not isinstance(payload, JSONRPCRequest): + return httpx2.Response(202) + if payload.method == "initialize": + assert payload.params["protocolVersion"] == offered + assert ("sampling" in payload.params["capabilities"]) == callbacks + assert ("elicitation" in payload.params["capabilities"]) == callbacks + return httpx2.Response(200, json={ + "jsonrpc": "2.0", "id": payload.id, + "result": {"protocolVersion": offered if accepted else "unsupported", + "capabilities": {"tools": {}}, "serverInfo": {"name": "upstream", "version": "1"}}, + }) + assert accepted, "No operation may execute after failed version negotiation" + return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": {"tools": [{"name": "echo", "inputSchema": {"type": "object"}}]}}) + + client = _MockTransportClient( + respond, server_url="https://example.com/mcp", protocol_version=revision, + sampling_callback=AsyncMock() if callbacks else None, + elicitation_callback=AsyncMock() if callbacks else None, + ) + if accepted: + result = await client.list_tools(raise_on_error=True) + assert [tool.name for tool in result] == ["echo"] + else: + with pytest.raises((MCPError, RuntimeError), match="protocol version"): + await client.list_tools(raise_on_error=True) + + +@pytest.mark.parametrize("revision", ["2026-07-28", "unknown", "", None]) +def test_upstream_protocol_configuration_rejects_unavailable_modes(revision): + from pydantic import ValidationError + + with pytest.raises(ValidationError): + MCPClient(protocol_version=revision) diff --git a/tests/test_litellm/llms/anthropic/__init__.py b/tests/unit/integrations/SlackAlerting/__init__.py similarity index 100% rename from tests/test_litellm/llms/anthropic/__init__.py rename to tests/unit/integrations/SlackAlerting/__init__.py diff --git a/tests/test_litellm/integrations/SlackAlerting/test_budget_alert_types.py b/tests/unit/integrations/SlackAlerting/test_budget_alert_types.py similarity index 100% rename from tests/test_litellm/integrations/SlackAlerting/test_budget_alert_types.py rename to tests/unit/integrations/SlackAlerting/test_budget_alert_types.py diff --git a/tests/test_litellm/integrations/SlackAlerting/test_hanging_request_check.py b/tests/unit/integrations/SlackAlerting/test_hanging_request_check.py similarity index 100% rename from tests/test_litellm/integrations/SlackAlerting/test_hanging_request_check.py rename to tests/unit/integrations/SlackAlerting/test_hanging_request_check.py diff --git a/tests/test_litellm/integrations/SlackAlerting/test_model_deprecation_alert.py b/tests/unit/integrations/SlackAlerting/test_model_deprecation_alert.py similarity index 100% rename from tests/test_litellm/integrations/SlackAlerting/test_model_deprecation_alert.py rename to tests/unit/integrations/SlackAlerting/test_model_deprecation_alert.py diff --git a/tests/test_litellm/integrations/SlackAlerting/test_ms_teams.py b/tests/unit/integrations/SlackAlerting/test_ms_teams.py similarity index 100% rename from tests/test_litellm/integrations/SlackAlerting/test_ms_teams.py rename to tests/unit/integrations/SlackAlerting/test_ms_teams.py diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py b/tests/unit/integrations/SlackAlerting/test_slack_alerting.py similarity index 100% rename from tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py rename to tests/unit/integrations/SlackAlerting/test_slack_alerting.py diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting_digest.py b/tests/unit/integrations/SlackAlerting/test_slack_alerting_digest.py similarity index 100% rename from tests/test_litellm/integrations/SlackAlerting/test_slack_alerting_digest.py rename to tests/unit/integrations/SlackAlerting/test_slack_alerting_digest.py diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting_utils.py b/tests/unit/integrations/SlackAlerting/test_slack_alerting_utils.py similarity index 100% rename from tests/test_litellm/integrations/SlackAlerting/test_slack_alerting_utils.py rename to tests/unit/integrations/SlackAlerting/test_slack_alerting_utils.py diff --git a/tests/test_litellm/integrations/SlackAlerting/test_user_spend_alerts.py b/tests/unit/integrations/SlackAlerting/test_user_spend_alerts.py similarity index 100% rename from tests/test_litellm/integrations/SlackAlerting/test_user_spend_alerts.py rename to tests/unit/integrations/SlackAlerting/test_user_spend_alerts.py diff --git a/tests/test_litellm/llms/anthropic/batches/__init__.py b/tests/unit/integrations/arize/__init__.py similarity index 100% rename from tests/test_litellm/llms/anthropic/batches/__init__.py rename to tests/unit/integrations/arize/__init__.py diff --git a/tests/test_litellm/integrations/arize/test_arize.py b/tests/unit/integrations/arize/test_arize.py similarity index 100% rename from tests/test_litellm/integrations/arize/test_arize.py rename to tests/unit/integrations/arize/test_arize.py diff --git a/tests/test_litellm/integrations/arize/test_arize_health_check.py b/tests/unit/integrations/arize/test_arize_health_check.py similarity index 100% rename from tests/test_litellm/integrations/arize/test_arize_health_check.py rename to tests/unit/integrations/arize/test_arize_health_check.py diff --git a/tests/test_litellm/integrations/arize/test_arize_otel_coexistence.py b/tests/unit/integrations/arize/test_arize_otel_coexistence.py similarity index 100% rename from tests/test_litellm/integrations/arize/test_arize_otel_coexistence.py rename to tests/unit/integrations/arize/test_arize_otel_coexistence.py diff --git a/tests/test_litellm/integrations/arize/test_arize_phoenix.py b/tests/unit/integrations/arize/test_arize_phoenix.py similarity index 100% rename from tests/test_litellm/integrations/arize/test_arize_phoenix.py rename to tests/unit/integrations/arize/test_arize_phoenix.py diff --git a/tests/test_litellm/integrations/arize/test_arize_utils.py b/tests/unit/integrations/arize/test_arize_utils.py similarity index 100% rename from tests/test_litellm/integrations/arize/test_arize_utils.py rename to tests/unit/integrations/arize/test_arize_utils.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/__init__.py b/tests/unit/integrations/azure_storage/__init__.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/__init__.py rename to tests/unit/integrations/azure_storage/__init__.py diff --git a/tests/test_litellm/integrations/azure_storage/test_azure_storage.py b/tests/unit/integrations/azure_storage/test_azure_storage.py similarity index 100% rename from tests/test_litellm/integrations/azure_storage/test_azure_storage.py rename to tests/unit/integrations/azure_storage/test_azure_storage.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/__init__.py b/tests/unit/integrations/bitbucket/__init__.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/__init__.py rename to tests/unit/integrations/bitbucket/__init__.py diff --git a/tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py b/tests/unit/integrations/bitbucket/test_bitbucket_integration.py similarity index 100% rename from tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py rename to tests/unit/integrations/bitbucket/test_bitbucket_integration.py diff --git a/tests/test_litellm/integrations/bitbucket/test_bitbucket_prompt_manager.py b/tests/unit/integrations/bitbucket/test_bitbucket_prompt_manager.py similarity index 93% rename from tests/test_litellm/integrations/bitbucket/test_bitbucket_prompt_manager.py rename to tests/unit/integrations/bitbucket/test_bitbucket_prompt_manager.py index d6668bf9ad8..a1a88653ee6 100644 --- a/tests/test_litellm/integrations/bitbucket/test_bitbucket_prompt_manager.py +++ b/tests/unit/integrations/bitbucket/test_bitbucket_prompt_manager.py @@ -307,37 +307,6 @@ def test_bitbucket_prompt_manager_render_template_not_found(): manager.prompt_manager.render_template("nonexistent", {"some": "variable"}) -@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient") -def test_bitbucket_prompt_manager_integration(mock_client_class): - """Test BitBucketPromptManager integration with BitBucketClient.""" - # Mock the BitBucket client - mock_client = MagicMock() - mock_client.get_file_content.return_value = """--- -model: gpt-4 -temperature: 0.7 ---- -Hello {{name}}!""" - mock_client_class.return_value = mock_client - - config = { - "workspace": "test-workspace", - "repository": "test-repo", - "access_token": "test-token", - } - - manager = BitBucketPromptManager(config, prompt_id="test_prompt") - - # Should have loaded the prompt - assert "test_prompt" in manager.prompt_manager.prompts - template = manager.prompt_manager.prompts["test_prompt"] - assert template.model == "gpt-4" - assert template.temperature == 0.7 - - # Test rendering - rendered = manager.prompt_manager.render_template("test_prompt", {"name": "World"}) - assert rendered == "Hello World!" - - def test_bitbucket_prompt_manager_parse_prompt_to_messages(): """Test parsing prompt content into messages.""" config = { diff --git a/tests/test_litellm/llms/anthropic/files/__init__.py b/tests/unit/integrations/cloudzero/__init__.py similarity index 100% rename from tests/test_litellm/llms/anthropic/files/__init__.py rename to tests/unit/integrations/cloudzero/__init__.py diff --git a/tests/test_litellm/integrations/cloudzero/test_cloudzero.py b/tests/unit/integrations/cloudzero/test_cloudzero.py similarity index 100% rename from tests/test_litellm/integrations/cloudzero/test_cloudzero.py rename to tests/unit/integrations/cloudzero/test_cloudzero.py diff --git a/tests/test_litellm/integrations/cloudzero/test_cloudzero_database.py b/tests/unit/integrations/cloudzero/test_cloudzero_database.py similarity index 100% rename from tests/test_litellm/integrations/cloudzero/test_cloudzero_database.py rename to tests/unit/integrations/cloudzero/test_cloudzero_database.py diff --git a/tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py b/tests/unit/integrations/cloudzero/test_cz_stream_api.py similarity index 100% rename from tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py rename to tests/unit/integrations/cloudzero/test_cz_stream_api.py diff --git a/tests/test_litellm/integrations/cloudzero/test_dry_run_endpoint.py b/tests/unit/integrations/cloudzero/test_dry_run_endpoint.py similarity index 100% rename from tests/test_litellm/integrations/cloudzero/test_dry_run_endpoint.py rename to tests/unit/integrations/cloudzero/test_dry_run_endpoint.py diff --git a/tests/test_litellm/integrations/cloudzero/test_transform.py b/tests/unit/integrations/cloudzero/test_transform.py similarity index 100% rename from tests/test_litellm/integrations/cloudzero/test_transform.py rename to tests/unit/integrations/cloudzero/test_transform.py diff --git a/tests/test_litellm/llms/azure/batches/__init__.py b/tests/unit/integrations/code_interpreter_interception/__init__.py similarity index 100% rename from tests/test_litellm/llms/azure/batches/__init__.py rename to tests/unit/integrations/code_interpreter_interception/__init__.py diff --git a/tests/test_litellm/integrations/code_interpreter_interception/test_handler.py b/tests/unit/integrations/code_interpreter_interception/test_handler.py similarity index 100% rename from tests/test_litellm/integrations/code_interpreter_interception/test_handler.py rename to tests/unit/integrations/code_interpreter_interception/test_handler.py diff --git a/tests/test_litellm/integrations/conftest.py b/tests/unit/integrations/conftest.py similarity index 94% rename from tests/test_litellm/integrations/conftest.py rename to tests/unit/integrations/conftest.py index adc8e36e0af..48a01ed8d48 100644 --- a/tests/test_litellm/integrations/conftest.py +++ b/tests/unit/integrations/conftest.py @@ -1,6 +1,7 @@ import functools import http.server import ipaddress +import os import queue import ssl import threading @@ -74,6 +75,14 @@ def write_self_signed_cert(directory: Path, stem: str) -> tuple[Path, Path]: return certificate_path, key_path +@pytest.fixture(autouse=True) +def restore_process_environment() -> Iterator[None]: + original: Final = dict(os.environ) + yield + os.environ.clear() + os.environ.update(original) + + @pytest.fixture def tls_sink(tmp_path: Path) -> Iterator[TlsSink]: certificate_path, key_path = write_self_signed_cert(tmp_path, "sink") diff --git a/tests/test_litellm/llms/azure/vector_stores/__init__.py b/tests/unit/integrations/datadog/__init__.py similarity index 100% rename from tests/test_litellm/llms/azure/vector_stores/__init__.py rename to tests/unit/integrations/datadog/__init__.py diff --git a/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py b/tests/unit/integrations/datadog/test_datadog_cost_management.py similarity index 100% rename from tests/test_litellm/integrations/datadog/test_datadog_cost_management.py rename to tests/unit/integrations/datadog/test_datadog_cost_management.py diff --git a/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py b/tests/unit/integrations/datadog/test_datadog_llm_obs.py similarity index 100% rename from tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py rename to tests/unit/integrations/datadog/test_datadog_llm_obs.py diff --git a/tests/test_litellm/integrations/datadog/test_datadog_llm_obs_agent.py b/tests/unit/integrations/datadog/test_datadog_llm_obs_agent.py similarity index 100% rename from tests/test_litellm/integrations/datadog/test_datadog_llm_obs_agent.py rename to tests/unit/integrations/datadog/test_datadog_llm_obs_agent.py diff --git a/tests/test_litellm/integrations/datadog/test_datadog_logger_batching.py b/tests/unit/integrations/datadog/test_datadog_logger_batching.py similarity index 100% rename from tests/test_litellm/integrations/datadog/test_datadog_logger_batching.py rename to tests/unit/integrations/datadog/test_datadog_logger_batching.py diff --git a/tests/test_litellm/integrations/datadog/test_datadog_metrics.py b/tests/unit/integrations/datadog/test_datadog_metrics.py similarity index 100% rename from tests/test_litellm/integrations/datadog/test_datadog_metrics.py rename to tests/unit/integrations/datadog/test_datadog_metrics.py diff --git a/tests/test_litellm/integrations/datadog/test_datadog_tags_regression.py b/tests/unit/integrations/datadog/test_datadog_tags_regression.py similarity index 100% rename from tests/test_litellm/integrations/datadog/test_datadog_tags_regression.py rename to tests/unit/integrations/datadog/test_datadog_tags_regression.py diff --git a/tests/test_litellm/integrations/datadog/test_datadog_team_handler.py b/tests/unit/integrations/datadog/test_datadog_team_handler.py similarity index 100% rename from tests/test_litellm/integrations/datadog/test_datadog_team_handler.py rename to tests/unit/integrations/datadog/test_datadog_team_handler.py diff --git a/tests/test_litellm/llms/base_llm/__init__.py b/tests/unit/integrations/dotprompt/__init__.py similarity index 100% rename from tests/test_litellm/llms/base_llm/__init__.py rename to tests/unit/integrations/dotprompt/__init__.py diff --git a/tests/test_litellm/integrations/dotprompt/chat_prompt.prompt b/tests/unit/integrations/dotprompt/chat_prompt.prompt similarity index 100% rename from tests/test_litellm/integrations/dotprompt/chat_prompt.prompt rename to tests/unit/integrations/dotprompt/chat_prompt.prompt diff --git a/tests/test_litellm/integrations/dotprompt/chat_prompt.v1.prompt b/tests/unit/integrations/dotprompt/chat_prompt.v1.prompt similarity index 100% rename from tests/test_litellm/integrations/dotprompt/chat_prompt.v1.prompt rename to tests/unit/integrations/dotprompt/chat_prompt.v1.prompt diff --git a/tests/test_litellm/integrations/dotprompt/chat_prompt.v2.prompt b/tests/unit/integrations/dotprompt/chat_prompt.v2.prompt similarity index 100% rename from tests/test_litellm/integrations/dotprompt/chat_prompt.v2.prompt rename to tests/unit/integrations/dotprompt/chat_prompt.v2.prompt diff --git a/tests/test_litellm/integrations/dotprompt/coding_assistant.prompt b/tests/unit/integrations/dotprompt/coding_assistant.prompt similarity index 100% rename from tests/test_litellm/integrations/dotprompt/coding_assistant.prompt rename to tests/unit/integrations/dotprompt/coding_assistant.prompt diff --git a/tests/test_litellm/integrations/dotprompt/sample_prompt.prompt b/tests/unit/integrations/dotprompt/sample_prompt.prompt similarity index 100% rename from tests/test_litellm/integrations/dotprompt/sample_prompt.prompt rename to tests/unit/integrations/dotprompt/sample_prompt.prompt diff --git a/tests/test_litellm/integrations/dotprompt/test_prompt_manager.py b/tests/unit/integrations/dotprompt/test_prompt_manager.py similarity index 100% rename from tests/test_litellm/integrations/dotprompt/test_prompt_manager.py rename to tests/unit/integrations/dotprompt/test_prompt_manager.py diff --git a/tests/test_litellm/llms/base_llm/batches/__init__.py b/tests/unit/integrations/focus/__init__.py similarity index 100% rename from tests/test_litellm/llms/base_llm/batches/__init__.py rename to tests/unit/integrations/focus/__init__.py diff --git a/tests/test_litellm/integrations/focus/test_csv_serializer.py b/tests/unit/integrations/focus/test_csv_serializer.py similarity index 100% rename from tests/test_litellm/integrations/focus/test_csv_serializer.py rename to tests/unit/integrations/focus/test_csv_serializer.py diff --git a/tests/test_litellm/integrations/focus/test_destination_factory.py b/tests/unit/integrations/focus/test_destination_factory.py similarity index 100% rename from tests/test_litellm/integrations/focus/test_destination_factory.py rename to tests/unit/integrations/focus/test_destination_factory.py diff --git a/tests/test_litellm/integrations/focus/test_focus_database.py b/tests/unit/integrations/focus/test_focus_database.py similarity index 100% rename from tests/test_litellm/integrations/focus/test_focus_database.py rename to tests/unit/integrations/focus/test_focus_database.py diff --git a/tests/test_litellm/integrations/focus/test_focus_gcs_destination.py b/tests/unit/integrations/focus/test_focus_gcs_destination.py similarity index 100% rename from tests/test_litellm/integrations/focus/test_focus_gcs_destination.py rename to tests/unit/integrations/focus/test_focus_gcs_destination.py diff --git a/tests/test_litellm/integrations/focus/test_focus_transformer.py b/tests/unit/integrations/focus/test_focus_transformer.py similarity index 100% rename from tests/test_litellm/integrations/focus/test_focus_transformer.py rename to tests/unit/integrations/focus/test_focus_transformer.py diff --git a/tests/test_litellm/integrations/focus/test_mavvrik_destination.py b/tests/unit/integrations/focus/test_mavvrik_destination.py similarity index 100% rename from tests/test_litellm/integrations/focus/test_mavvrik_destination.py rename to tests/unit/integrations/focus/test_mavvrik_destination.py diff --git a/tests/test_litellm/integrations/focus/test_s3_destination.py b/tests/unit/integrations/focus/test_s3_destination.py similarity index 100% rename from tests/test_litellm/integrations/focus/test_s3_destination.py rename to tests/unit/integrations/focus/test_s3_destination.py diff --git a/tests/test_litellm/integrations/focus/test_transformer.py b/tests/unit/integrations/focus/test_transformer.py similarity index 100% rename from tests/test_litellm/integrations/focus/test_transformer.py rename to tests/unit/integrations/focus/test_transformer.py diff --git a/tests/test_litellm/integrations/focus/test_vantage_destination.py b/tests/unit/integrations/focus/test_vantage_destination.py similarity index 100% rename from tests/test_litellm/integrations/focus/test_vantage_destination.py rename to tests/unit/integrations/focus/test_vantage_destination.py diff --git a/tests/test_litellm/llms/base_llm/files/__init__.py b/tests/unit/integrations/gitlab/__init__.py similarity index 100% rename from tests/test_litellm/llms/base_llm/files/__init__.py rename to tests/unit/integrations/gitlab/__init__.py diff --git a/tests/test_litellm/integrations/gitlab/test_gitlab_client.py b/tests/unit/integrations/gitlab/test_gitlab_client.py similarity index 100% rename from tests/test_litellm/integrations/gitlab/test_gitlab_client.py rename to tests/unit/integrations/gitlab/test_gitlab_client.py diff --git a/tests/test_litellm/integrations/gitlab/test_gitlab_integration.py b/tests/unit/integrations/gitlab/test_gitlab_integration.py similarity index 100% rename from tests/test_litellm/integrations/gitlab/test_gitlab_integration.py rename to tests/unit/integrations/gitlab/test_gitlab_integration.py diff --git a/tests/test_litellm/integrations/gitlab/test_gitlab_prompt_manager.py b/tests/unit/integrations/gitlab/test_gitlab_prompt_manager.py similarity index 100% rename from tests/test_litellm/integrations/gitlab/test_gitlab_prompt_manager.py rename to tests/unit/integrations/gitlab/test_gitlab_prompt_manager.py diff --git a/tests/test_litellm/llms/base_llm/realtime/__init__.py b/tests/unit/integrations/langfuse/__init__.py similarity index 100% rename from tests/test_litellm/llms/base_llm/realtime/__init__.py rename to tests/unit/integrations/langfuse/__init__.py diff --git a/tests/test_litellm/integrations/langfuse/test_gemini_cached_tokens.py b/tests/unit/integrations/langfuse/test_gemini_cached_tokens.py similarity index 100% rename from tests/test_litellm/integrations/langfuse/test_gemini_cached_tokens.py rename to tests/unit/integrations/langfuse/test_gemini_cached_tokens.py diff --git a/tests/test_litellm/integrations/langfuse/test_langfuse_prompt_management.py b/tests/unit/integrations/langfuse/test_langfuse_prompt_management.py similarity index 100% rename from tests/test_litellm/integrations/langfuse/test_langfuse_prompt_management.py rename to tests/unit/integrations/langfuse/test_langfuse_prompt_management.py diff --git a/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py b/tests/unit/integrations/langfuse/test_langfuse_sdk.py similarity index 100% rename from tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py rename to tests/unit/integrations/langfuse/test_langfuse_sdk.py diff --git a/tests/test_litellm/llms/bedrock/__init__.py b/tests/unit/integrations/newrelic/__init__.py similarity index 100% rename from tests/test_litellm/llms/bedrock/__init__.py rename to tests/unit/integrations/newrelic/__init__.py diff --git a/tests/test_litellm/integrations/newrelic/test_newrelic.py b/tests/unit/integrations/newrelic/test_newrelic.py similarity index 100% rename from tests/test_litellm/integrations/newrelic/test_newrelic.py rename to tests/unit/integrations/newrelic/test_newrelic.py diff --git a/tests/test_litellm/integrations/newrelic/test_newrelic_metrics.py b/tests/unit/integrations/newrelic/test_newrelic_metrics.py similarity index 100% rename from tests/test_litellm/integrations/newrelic/test_newrelic_metrics.py rename to tests/unit/integrations/newrelic/test_newrelic_metrics.py diff --git a/tests/test_litellm/integrations/newrelic/test_newrelic_team_handler.py b/tests/unit/integrations/newrelic/test_newrelic_team_handler.py similarity index 100% rename from tests/test_litellm/integrations/newrelic/test_newrelic_team_handler.py rename to tests/unit/integrations/newrelic/test_newrelic_team_handler.py diff --git a/tests/test_litellm/llms/bedrock/batches/__init__.py b/tests/unit/integrations/open_telemetry/__init__.py similarity index 100% rename from tests/test_litellm/llms/bedrock/batches/__init__.py rename to tests/unit/integrations/open_telemetry/__init__.py diff --git a/tests/test_litellm/integrations/open_telemetry/_helpers.py b/tests/unit/integrations/open_telemetry/_helpers.py similarity index 100% rename from tests/test_litellm/integrations/open_telemetry/_helpers.py rename to tests/unit/integrations/open_telemetry/_helpers.py diff --git a/tests/test_litellm/integrations/open_telemetry/conftest.py b/tests/unit/integrations/open_telemetry/conftest.py similarity index 100% rename from tests/test_litellm/integrations/open_telemetry/conftest.py rename to tests/unit/integrations/open_telemetry/conftest.py diff --git a/tests/test_litellm/llms/bedrock/chat/agentcore/__init__.py b/tests/unit/integrations/open_telemetry/data/__init__.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/agentcore/__init__.py rename to tests/unit/integrations/open_telemetry/data/__init__.py diff --git a/tests/test_litellm/integrations/open_telemetry/data/captured_kwargs.json b/tests/unit/integrations/open_telemetry/data/captured_kwargs.json similarity index 100% rename from tests/test_litellm/integrations/open_telemetry/data/captured_kwargs.json rename to tests/unit/integrations/open_telemetry/data/captured_kwargs.json diff --git a/tests/test_litellm/integrations/open_telemetry/data/captured_response.json b/tests/unit/integrations/open_telemetry/data/captured_response.json similarity index 100% rename from tests/test_litellm/integrations/open_telemetry/data/captured_response.json rename to tests/unit/integrations/open_telemetry/data/captured_response.json diff --git a/tests/test_litellm/integrations/open_telemetry/test_otel_admin_endpoints.py b/tests/unit/integrations/open_telemetry/test_otel_admin_endpoints.py similarity index 100% rename from tests/test_litellm/integrations/open_telemetry/test_otel_admin_endpoints.py rename to tests/unit/integrations/open_telemetry/test_otel_admin_endpoints.py diff --git a/tests/test_litellm/integrations/open_telemetry/test_otel_exception_handler.py b/tests/unit/integrations/open_telemetry/test_otel_exception_handler.py similarity index 100% rename from tests/test_litellm/integrations/open_telemetry/test_otel_exception_handler.py rename to tests/unit/integrations/open_telemetry/test_otel_exception_handler.py diff --git a/tests/test_litellm/integrations/open_telemetry/test_otel_passthrough_endpoints.py b/tests/unit/integrations/open_telemetry/test_otel_passthrough_endpoints.py similarity index 100% rename from tests/test_litellm/integrations/open_telemetry/test_otel_passthrough_endpoints.py rename to tests/unit/integrations/open_telemetry/test_otel_passthrough_endpoints.py diff --git a/tests/test_litellm/integrations/open_telemetry/test_otel_unified_endpoints.py b/tests/unit/integrations/open_telemetry/test_otel_unified_endpoints.py similarity index 100% rename from tests/test_litellm/integrations/open_telemetry/test_otel_unified_endpoints.py rename to tests/unit/integrations/open_telemetry/test_otel_unified_endpoints.py diff --git a/tests/test_litellm/integrations/open_telemetry/test_passthrough_parent_span.py b/tests/unit/integrations/open_telemetry/test_passthrough_parent_span.py similarity index 100% rename from tests/test_litellm/integrations/open_telemetry/test_passthrough_parent_span.py rename to tests/unit/integrations/open_telemetry/test_passthrough_parent_span.py diff --git a/tests/test_litellm/llms/bedrock/passthrough/guardrail_translation/__init__.py b/tests/unit/integrations/otel/__init__.py similarity index 100% rename from tests/test_litellm/llms/bedrock/passthrough/guardrail_translation/__init__.py rename to tests/unit/integrations/otel/__init__.py diff --git a/tests/test_litellm/integrations/otel/test_db_endpoint.py b/tests/unit/integrations/otel/test_db_endpoint.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_db_endpoint.py rename to tests/unit/integrations/otel/test_db_endpoint.py diff --git a/tests/test_litellm/integrations/otel/test_langfuse_logger.py b/tests/unit/integrations/otel/test_langfuse_logger.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_langfuse_logger.py rename to tests/unit/integrations/otel/test_langfuse_logger.py diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_baggage.py b/tests/unit/integrations/otel/test_otel_v2_baggage.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_otel_v2_baggage.py rename to tests/unit/integrations/otel/test_otel_v2_baggage.py diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_components.py b/tests/unit/integrations/otel/test_otel_v2_components.py similarity index 99% rename from tests/test_litellm/integrations/otel/test_otel_v2_components.py rename to tests/unit/integrations/otel/test_otel_v2_components.py index 07705e17d9a..fd10210c5ba 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_components.py +++ b/tests/unit/integrations/otel/test_otel_v2_components.py @@ -42,7 +42,7 @@ from opentelemetry.trace.propagation.tracecontext import ( # noqa: E402 ) import litellm # noqa: E402 -from conftest import TlsSink # noqa: E402 +from tests.unit.integrations.conftest import TlsSink # noqa: E402 from litellm.integrations.otel.plumbing import context as ctx_mod # noqa: E402 from litellm.integrations.otel.plumbing import providers # noqa: E402 from litellm.integrations.otel.model.config import OpenTelemetryV2Config # noqa: E402 diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_config_baggage_parenting_guardrails.py b/tests/unit/integrations/otel/test_otel_v2_config_baggage_parenting_guardrails.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_otel_v2_config_baggage_parenting_guardrails.py rename to tests/unit/integrations/otel/test_otel_v2_config_baggage_parenting_guardrails.py diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_destinations.py b/tests/unit/integrations/otel/test_otel_v2_destinations.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_otel_v2_destinations.py rename to tests/unit/integrations/otel/test_otel_v2_destinations.py diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_dynamic.py b/tests/unit/integrations/otel/test_otel_v2_dynamic.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_otel_v2_dynamic.py rename to tests/unit/integrations/otel/test_otel_v2_dynamic.py diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py b/tests/unit/integrations/otel/test_otel_v2_emitter.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_otel_v2_emitter.py rename to tests/unit/integrations/otel/test_otel_v2_emitter.py diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py b/tests/unit/integrations/otel/test_otel_v2_logger.py similarity index 99% rename from tests/test_litellm/integrations/otel/test_otel_v2_logger.py rename to tests/unit/integrations/otel/test_otel_v2_logger.py index 00c1343f72e..d478c670e58 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py +++ b/tests/unit/integrations/otel/test_otel_v2_logger.py @@ -2945,14 +2945,6 @@ def test_success_without_pre_call_emits_deferred_span(): assert spans[0].end_time == 101_500_000_000 -def test_no_carrier_and_no_payload_is_noop(): - logger, exporter = _logger() - asyncio.run( - logger.async_log_success_event({"litellm_params": {}}, None, None, None) - ) - assert exporter.get_finished_spans() == () - - def test_second_close_for_same_call_does_not_duplicate_span(): """Success and failure can both fire on one logging object for the same call id. The first close pops the carrier and finishes the boundary span; the diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_metrics.py b/tests/unit/integrations/otel/test_otel_v2_metrics.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_otel_v2_metrics.py rename to tests/unit/integrations/otel/test_otel_v2_metrics.py diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_mount.py b/tests/unit/integrations/otel/test_otel_v2_mount.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_otel_v2_mount.py rename to tests/unit/integrations/otel/test_otel_v2_mount.py diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_multibackend.py b/tests/unit/integrations/otel/test_otel_v2_multibackend.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_otel_v2_multibackend.py rename to tests/unit/integrations/otel/test_otel_v2_multibackend.py diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_presets.py b/tests/unit/integrations/otel/test_otel_v2_presets.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_otel_v2_presets.py rename to tests/unit/integrations/otel/test_otel_v2_presets.py diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py rename to tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py b/tests/unit/integrations/otel/test_otel_v2_vendor_mappers.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py rename to tests/unit/integrations/otel/test_otel_v2_vendor_mappers.py diff --git a/tests/test_litellm/integrations/otel/test_runtime.py b/tests/unit/integrations/otel/test_runtime.py similarity index 100% rename from tests/test_litellm/integrations/otel/test_runtime.py rename to tests/unit/integrations/otel/test_runtime.py diff --git a/tests/test_litellm/integrations/rubrik_test_helpers.py b/tests/unit/integrations/rubrik_test_helpers.py similarity index 100% rename from tests/test_litellm/integrations/rubrik_test_helpers.py rename to tests/unit/integrations/rubrik_test_helpers.py diff --git a/tests/test_litellm/integrations/test_agentops.py b/tests/unit/integrations/test_agentops.py similarity index 100% rename from tests/test_litellm/integrations/test_agentops.py rename to tests/unit/integrations/test_agentops.py diff --git a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py b/tests/unit/integrations/test_anthropic_cache_control_hook.py similarity index 100% rename from tests/test_litellm/integrations/test_anthropic_cache_control_hook.py rename to tests/unit/integrations/test_anthropic_cache_control_hook.py diff --git a/tests/test_litellm/integrations/test_athina.py b/tests/unit/integrations/test_athina.py similarity index 100% rename from tests/test_litellm/integrations/test_athina.py rename to tests/unit/integrations/test_athina.py diff --git a/tests/test_litellm/integrations/test_azure_sentinel.py b/tests/unit/integrations/test_azure_sentinel.py similarity index 100% rename from tests/test_litellm/integrations/test_azure_sentinel.py rename to tests/unit/integrations/test_azure_sentinel.py diff --git a/tests/test_litellm/integrations/test_braintrust_logging.py b/tests/unit/integrations/test_braintrust_logging.py similarity index 100% rename from tests/test_litellm/integrations/test_braintrust_logging.py rename to tests/unit/integrations/test_braintrust_logging.py diff --git a/tests/test_litellm/integrations/test_braintrust_span_name.py b/tests/unit/integrations/test_braintrust_span_name.py similarity index 100% rename from tests/test_litellm/integrations/test_braintrust_span_name.py rename to tests/unit/integrations/test_braintrust_span_name.py diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py similarity index 100% rename from tests/test_litellm/integrations/test_custom_guardrail.py rename to tests/unit/integrations/test_custom_guardrail.py diff --git a/tests/test_litellm/integrations/test_custom_guardrail_recursion.py b/tests/unit/integrations/test_custom_guardrail_recursion.py similarity index 100% rename from tests/test_litellm/integrations/test_custom_guardrail_recursion.py rename to tests/unit/integrations/test_custom_guardrail_recursion.py diff --git a/tests/test_litellm/integrations/test_custom_prompt_management.py b/tests/unit/integrations/test_custom_prompt_management.py similarity index 100% rename from tests/test_litellm/integrations/test_custom_prompt_management.py rename to tests/unit/integrations/test_custom_prompt_management.py diff --git a/tests/test_litellm/integrations/test_deepeval.py b/tests/unit/integrations/test_deepeval.py similarity index 100% rename from tests/test_litellm/integrations/test_deepeval.py rename to tests/unit/integrations/test_deepeval.py diff --git a/tests/test_litellm/integrations/test_galileo.py b/tests/unit/integrations/test_galileo.py similarity index 100% rename from tests/test_litellm/integrations/test_galileo.py rename to tests/unit/integrations/test_galileo.py diff --git a/tests/test_litellm/integrations/test_guardrail_logging_sync.py b/tests/unit/integrations/test_guardrail_logging_sync.py similarity index 100% rename from tests/test_litellm/integrations/test_guardrail_logging_sync.py rename to tests/unit/integrations/test_guardrail_logging_sync.py diff --git a/tests/test_litellm/integrations/test_helicone.py b/tests/unit/integrations/test_helicone.py similarity index 95% rename from tests/test_litellm/integrations/test_helicone.py rename to tests/unit/integrations/test_helicone.py index 64960de050a..99cb1380dd7 100644 --- a/tests/test_litellm/integrations/test_helicone.py +++ b/tests/unit/integrations/test_helicone.py @@ -13,7 +13,7 @@ def _claude_mapping(messages, response_obj): def test_claude_mapping_serializes_custom_tool_calls(monkeypatch): """ Stub the anthropic module unconditionally: the SDK may be absent (it lives in the - proxy-runtime extra), and the tests/test_litellm/llms/anthropic test package can + proxy-runtime extra), and the tests/unit/llms/anthropic test package can shadow it on sys.path, so an import probe proves nothing about the real SDK. """ stub = types.ModuleType("anthropic") diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/unit/integrations/test_langfuse.py similarity index 100% rename from tests/test_litellm/integrations/test_langfuse.py rename to tests/unit/integrations/test_langfuse.py diff --git a/tests/test_litellm/integrations/test_langfuse_otel.py b/tests/unit/integrations/test_langfuse_otel.py similarity index 100% rename from tests/test_litellm/integrations/test_langfuse_otel.py rename to tests/unit/integrations/test_langfuse_otel.py diff --git a/tests/test_litellm/integrations/test_langsmith_init.py b/tests/unit/integrations/test_langsmith_init.py similarity index 100% rename from tests/test_litellm/integrations/test_langsmith_init.py rename to tests/unit/integrations/test_langsmith_init.py diff --git a/tests/test_litellm/integrations/test_lunary.py b/tests/unit/integrations/test_lunary.py similarity index 100% rename from tests/test_litellm/integrations/test_lunary.py rename to tests/unit/integrations/test_lunary.py diff --git a/tests/test_litellm/integrations/test_mlflow.py b/tests/unit/integrations/test_mlflow.py similarity index 100% rename from tests/test_litellm/integrations/test_mlflow.py rename to tests/unit/integrations/test_mlflow.py diff --git a/tests/test_litellm/integrations/test_openmeter.py b/tests/unit/integrations/test_openmeter.py similarity index 100% rename from tests/test_litellm/integrations/test_openmeter.py rename to tests/unit/integrations/test_openmeter.py diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/unit/integrations/test_opentelemetry.py similarity index 99% rename from tests/test_litellm/integrations/test_opentelemetry.py rename to tests/unit/integrations/test_opentelemetry.py index 974961f2eb5..52eeec31e71 100644 --- a/tests/test_litellm/integrations/test_opentelemetry.py +++ b/tests/unit/integrations/test_opentelemetry.py @@ -33,7 +33,7 @@ from parameterized import parameterized import requests -from conftest import TlsSink, write_self_signed_cert +from tests.unit.integrations.conftest import TlsSink, write_self_signed_cert import litellm from litellm.integrations import opentelemetry as otel_module from litellm.integrations.opentelemetry import ( @@ -1244,64 +1244,6 @@ class TestOpenTelemetry(unittest.TestCase): time.sleep(self.POLL_INTERVAL) return [] - @patch("litellm.integrations.opentelemetry.datetime") - def test_create_guardrail_span_with_valid_info(self, mock_datetime): - # Setup - otel = OpenTelemetry() - otel.tracer = MagicMock() - mock_span = MagicMock() - otel.tracer.start_span.return_value = mock_span - - # Create guardrail information - guardrail_info = { - "guardrail_name": "test_guardrail", - "guardrail_mode": "input", - "masked_entity_count": {"CREDIT_CARD": 2}, - "guardrail_response": "filtered_content", - "start_time": 1609459200.0, - "end_time": 1609459201.0, - } - - # Create a kwargs dict with standard_logging_object containing guardrail information - kwargs = { - "standard_logging_object": {"guardrail_information": [guardrail_info]} - } - - # Call the method - otel._create_guardrail_span(kwargs=kwargs, context=None) - - # Assertions - otel.tracer.start_span.assert_called_once() - - # print all calls to mock_span.set_attribute - print("Calls to mock_span.set_attribute:") - for call in mock_span.set_attribute.call_args_list: - print(call) - - # Check that the span has the correct attributes set - mock_span.set_attribute.assert_any_call("guardrail_name", "test_guardrail") - mock_span.set_attribute.assert_any_call("guardrail_mode", "input") - mock_span.set_attribute.assert_any_call( - "guardrail_response", safe_dumps("filtered_content") - ) - mock_span.set_attribute.assert_any_call( - "masked_entity_count", safe_dumps({"CREDIT_CARD": 2}) - ) - - # Verify that the span was ended - mock_span.end.assert_called_once() - - def test_create_guardrail_span_with_no_info(self): - # Setup - otel = OpenTelemetry() - otel.tracer = MagicMock() - - # Test with no guardrail information - kwargs = {"standard_logging_object": {}} - otel._create_guardrail_span(kwargs=kwargs, context=None) - - # Verify that start_span was never called - otel.tracer.start_span.assert_not_called() def test_get_tracer_to_use_for_request_with_dynamic_headers(self): """Test that get_tracer_to_use_for_request returns a dynamic tracer when dynamic headers are present.""" @@ -5461,10 +5403,6 @@ class TestOpenTelemetryPreprocessingDuration(unittest.TestCase): ) assert "litellm.preprocessing.duration_ms" not in self._attr(span, exp) - def test_none_span_is_noop(self): - OpenTelemetry().set_preprocessing_duration_attribute( - None, {"first_api_call_start_time": datetime(2026, 1, 1)} - ) def test_non_dict_container_is_noop(self): otel = OpenTelemetry() diff --git a/tests/test_litellm/integrations/test_opentelemetry_dynamic_imports.py b/tests/unit/integrations/test_opentelemetry_dynamic_imports.py similarity index 100% rename from tests/test_litellm/integrations/test_opentelemetry_dynamic_imports.py rename to tests/unit/integrations/test_opentelemetry_dynamic_imports.py diff --git a/tests/test_litellm/integrations/test_opik_utils.py b/tests/unit/integrations/test_opik_utils.py similarity index 100% rename from tests/test_litellm/integrations/test_opik_utils.py rename to tests/unit/integrations/test_opik_utils.py diff --git a/tests/test_litellm/integrations/test_otel_guardrail_violation_spans.py b/tests/unit/integrations/test_otel_guardrail_violation_spans.py similarity index 100% rename from tests/test_litellm/integrations/test_otel_guardrail_violation_spans.py rename to tests/unit/integrations/test_otel_guardrail_violation_spans.py diff --git a/tests/test_litellm/integrations/test_otel_team_attributes_matrix.py b/tests/unit/integrations/test_otel_team_attributes_matrix.py similarity index 100% rename from tests/test_litellm/integrations/test_otel_team_attributes_matrix.py rename to tests/unit/integrations/test_otel_team_attributes_matrix.py diff --git a/tests/test_litellm/integrations/test_prometheus_api_promql_escape.py b/tests/unit/integrations/test_prometheus_api_promql_escape.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_api_promql_escape.py rename to tests/unit/integrations/test_prometheus_api_promql_escape.py diff --git a/tests/test_litellm/integrations/test_prometheus_budget_metric_guard.py b/tests/unit/integrations/test_prometheus_budget_metric_guard.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_budget_metric_guard.py rename to tests/unit/integrations/test_prometheus_budget_metric_guard.py diff --git a/tests/test_litellm/integrations/test_prometheus_budget_metrics_db_lookups.py b/tests/unit/integrations/test_prometheus_budget_metrics_db_lookups.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_budget_metrics_db_lookups.py rename to tests/unit/integrations/test_prometheus_budget_metrics_db_lookups.py diff --git a/tests/test_litellm/integrations/test_prometheus_budget_metrics_timeout.py b/tests/unit/integrations/test_prometheus_budget_metrics_timeout.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_budget_metrics_timeout.py rename to tests/unit/integrations/test_prometheus_budget_metrics_timeout.py diff --git a/tests/test_litellm/integrations/test_prometheus_cache_metrics.py b/tests/unit/integrations/test_prometheus_cache_metrics.py similarity index 99% rename from tests/test_litellm/integrations/test_prometheus_cache_metrics.py rename to tests/unit/integrations/test_prometheus_cache_metrics.py index aa031bb813b..21f13f0ec5c 100644 --- a/tests/test_litellm/integrations/test_prometheus_cache_metrics.py +++ b/tests/unit/integrations/test_prometheus_cache_metrics.py @@ -1,7 +1,7 @@ """ Unit tests for cache Prometheus metrics. -Run with: uv run pytest tests/test_litellm/integrations/test_prometheus_cache_metrics.py -v +Run with: uv run pytest tests/unit/integrations/test_prometheus_cache_metrics.py -v """ import pytest diff --git a/tests/test_litellm/integrations/test_prometheus_caller_identity.py b/tests/unit/integrations/test_prometheus_caller_identity.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_caller_identity.py rename to tests/unit/integrations/test_prometheus_caller_identity.py diff --git a/tests/test_litellm/integrations/test_prometheus_carried_budget_state.py b/tests/unit/integrations/test_prometheus_carried_budget_state.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_carried_budget_state.py rename to tests/unit/integrations/test_prometheus_carried_budget_state.py diff --git a/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py b/tests/unit/integrations/test_prometheus_client_ip_user_agent.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py rename to tests/unit/integrations/test_prometheus_client_ip_user_agent.py diff --git a/tests/test_litellm/integrations/test_prometheus_custom_metadata_label_counts.py b/tests/unit/integrations/test_prometheus_custom_metadata_label_counts.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_custom_metadata_label_counts.py rename to tests/unit/integrations/test_prometheus_custom_metadata_label_counts.py diff --git a/tests/test_litellm/integrations/test_prometheus_deployment_state_proxy_rejects.py b/tests/unit/integrations/test_prometheus_deployment_state_proxy_rejects.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_deployment_state_proxy_rejects.py rename to tests/unit/integrations/test_prometheus_deployment_state_proxy_rejects.py diff --git a/tests/test_litellm/integrations/test_prometheus_end_user_cardinality.py b/tests/unit/integrations/test_prometheus_end_user_cardinality.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_end_user_cardinality.py rename to tests/unit/integrations/test_prometheus_end_user_cardinality.py diff --git a/tests/test_litellm/integrations/test_prometheus_input_sequence_length_label.py b/tests/unit/integrations/test_prometheus_input_sequence_length_label.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_input_sequence_length_label.py rename to tests/unit/integrations/test_prometheus_input_sequence_length_label.py diff --git a/tests/test_litellm/integrations/test_prometheus_invalid_key_filtering.py b/tests/unit/integrations/test_prometheus_invalid_key_filtering.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_invalid_key_filtering.py rename to tests/unit/integrations/test_prometheus_invalid_key_filtering.py diff --git a/tests/test_litellm/integrations/test_prometheus_labels.py b/tests/unit/integrations/test_prometheus_labels.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_labels.py rename to tests/unit/integrations/test_prometheus_labels.py diff --git a/tests/test_litellm/integrations/test_prometheus_mcp_tool_metrics.py b/tests/unit/integrations/test_prometheus_mcp_tool_metrics.py similarity index 99% rename from tests/test_litellm/integrations/test_prometheus_mcp_tool_metrics.py rename to tests/unit/integrations/test_prometheus_mcp_tool_metrics.py index 22c36f00ca9..da5a0b35e9d 100644 --- a/tests/test_litellm/integrations/test_prometheus_mcp_tool_metrics.py +++ b/tests/unit/integrations/test_prometheus_mcp_tool_metrics.py @@ -5,7 +5,7 @@ These metrics expose ``mcp_tool_call_metadata`` in Prometheus so Grafana dashboards can break down MCP usage by server and tool name. Run with: - uv run pytest tests/test_litellm/integrations/test_prometheus_mcp_tool_metrics.py -v + uv run pytest tests/unit/integrations/test_prometheus_mcp_tool_metrics.py -v """ from typing import get_args diff --git a/tests/test_litellm/integrations/test_prometheus_media_generation_metrics.py b/tests/unit/integrations/test_prometheus_media_generation_metrics.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_media_generation_metrics.py rename to tests/unit/integrations/test_prometheus_media_generation_metrics.py diff --git a/tests/test_litellm/integrations/test_prometheus_metric_name_consistency.py b/tests/unit/integrations/test_prometheus_metric_name_consistency.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_metric_name_consistency.py rename to tests/unit/integrations/test_prometheus_metric_name_consistency.py diff --git a/tests/test_litellm/integrations/test_prometheus_metrics_endpoint.py b/tests/unit/integrations/test_prometheus_metrics_endpoint.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_metrics_endpoint.py rename to tests/unit/integrations/test_prometheus_metrics_endpoint.py diff --git a/tests/test_litellm/integrations/test_prometheus_missing_metrics.py b/tests/unit/integrations/test_prometheus_missing_metrics.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_missing_metrics.py rename to tests/unit/integrations/test_prometheus_missing_metrics.py diff --git a/tests/test_litellm/integrations/test_prometheus_none_metadata.py b/tests/unit/integrations/test_prometheus_none_metadata.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_none_metadata.py rename to tests/unit/integrations/test_prometheus_none_metadata.py diff --git a/tests/test_litellm/integrations/test_prometheus_overhead_with_guardrails.py b/tests/unit/integrations/test_prometheus_overhead_with_guardrails.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_overhead_with_guardrails.py rename to tests/unit/integrations/test_prometheus_overhead_with_guardrails.py diff --git a/tests/test_litellm/integrations/test_prometheus_queue_guardrail_metrics.py b/tests/unit/integrations/test_prometheus_queue_guardrail_metrics.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_queue_guardrail_metrics.py rename to tests/unit/integrations/test_prometheus_queue_guardrail_metrics.py diff --git a/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py b/tests/unit/integrations/test_prometheus_rate_limit_labels.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py rename to tests/unit/integrations/test_prometheus_rate_limit_labels.py diff --git a/tests/test_litellm/integrations/test_prometheus_remaining_tokens_router_fallback.py b/tests/unit/integrations/test_prometheus_remaining_tokens_router_fallback.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_remaining_tokens_router_fallback.py rename to tests/unit/integrations/test_prometheus_remaining_tokens_router_fallback.py diff --git a/tests/test_litellm/integrations/test_prometheus_requested_model_cardinality.py b/tests/unit/integrations/test_prometheus_requested_model_cardinality.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_requested_model_cardinality.py rename to tests/unit/integrations/test_prometheus_requested_model_cardinality.py diff --git a/tests/test_litellm/integrations/test_prometheus_service_tier_label.py b/tests/unit/integrations/test_prometheus_service_tier_label.py similarity index 98% rename from tests/test_litellm/integrations/test_prometheus_service_tier_label.py rename to tests/unit/integrations/test_prometheus_service_tier_label.py index b2212c4ff41..8b8131b5af2 100644 --- a/tests/test_litellm/integrations/test_prometheus_service_tier_label.py +++ b/tests/unit/integrations/test_prometheus_service_tier_label.py @@ -6,7 +6,7 @@ between the tier a provider served and the tier a caller requested, and the end-to-end emit wiring through async_log_success_event. Run with: - uv run pytest tests/test_litellm/integrations/test_prometheus_service_tier_label.py -v + uv run pytest tests/unit/integrations/test_prometheus_service_tier_label.py -v """ import datetime diff --git a/tests/test_litellm/integrations/test_prometheus_services.py b/tests/unit/integrations/test_prometheus_services.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_services.py rename to tests/unit/integrations/test_prometheus_services.py diff --git a/tests/test_litellm/integrations/test_prometheus_spend_capture_rate.py b/tests/unit/integrations/test_prometheus_spend_capture_rate.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_spend_capture_rate.py rename to tests/unit/integrations/test_prometheus_spend_capture_rate.py diff --git a/tests/test_litellm/integrations/test_prometheus_spend_logs_metadata.py b/tests/unit/integrations/test_prometheus_spend_logs_metadata.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_spend_logs_metadata.py rename to tests/unit/integrations/test_prometheus_spend_logs_metadata.py diff --git a/tests/test_litellm/integrations/test_prometheus_stream_label.py b/tests/unit/integrations/test_prometheus_stream_label.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_stream_label.py rename to tests/unit/integrations/test_prometheus_stream_label.py diff --git a/tests/test_litellm/integrations/test_prometheus_token_detail_metrics.py b/tests/unit/integrations/test_prometheus_token_detail_metrics.py similarity index 99% rename from tests/test_litellm/integrations/test_prometheus_token_detail_metrics.py rename to tests/unit/integrations/test_prometheus_token_detail_metrics.py index 5e3846d6fa2..03d67316080 100644 --- a/tests/test_litellm/integrations/test_prometheus_token_detail_metrics.py +++ b/tests/unit/integrations/test_prometheus_token_detail_metrics.py @@ -6,7 +6,7 @@ from the Usage object that providers report. They are sparse — only incremented when the underlying detail is populated and > 0. Run with: - uv run pytest tests/test_litellm/integrations/test_prometheus_token_detail_metrics.py -v + uv run pytest tests/unit/integrations/test_prometheus_token_detail_metrics.py -v """ from typing import get_args diff --git a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py b/tests/unit/integrations/test_prometheus_user_team_metrics.py similarity index 98% rename from tests/test_litellm/integrations/test_prometheus_user_team_metrics.py rename to tests/unit/integrations/test_prometheus_user_team_metrics.py index 0fc91748af2..ab0fb67b52f 100644 --- a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py +++ b/tests/unit/integrations/test_prometheus_user_team_metrics.py @@ -102,27 +102,6 @@ class TestPrometheusUserTeamCountMetrics: f"litellm_teams_count_metric should accept value {value}: {e}" ) - def test_user_count_metric_with_zero(self, prometheus_logger): - """Test that user count metric handles zero users""" - metric = prometheus_logger.litellm_total_users_metric - - # Should handle zero gracefully - try: - metric.set(0) - assert True - except Exception as e: - pytest.fail(f"litellm_total_users_metric should handle zero: {e}") - - def test_team_count_metric_with_zero(self, prometheus_logger): - """Test that team count metric handles zero teams""" - metric = prometheus_logger.litellm_teams_count_metric - - # Should handle zero gracefully - try: - metric.set(0) - assert True - except Exception as e: - pytest.fail(f"litellm_teams_count_metric should handle zero: {e}") def test_metrics_can_be_updated_multiple_times(self, prometheus_logger): """Test that metrics can be updated multiple times (simulating refresh cycle)""" diff --git a/tests/test_litellm/integrations/test_prometheus_zero_cost_metric.py b/tests/unit/integrations/test_prometheus_zero_cost_metric.py similarity index 100% rename from tests/test_litellm/integrations/test_prometheus_zero_cost_metric.py rename to tests/unit/integrations/test_prometheus_zero_cost_metric.py diff --git a/tests/test_litellm/integrations/test_prompt_manager_ssti.py b/tests/unit/integrations/test_prompt_manager_ssti.py similarity index 100% rename from tests/test_litellm/integrations/test_prompt_manager_ssti.py rename to tests/unit/integrations/test_prompt_manager_ssti.py diff --git a/tests/test_litellm/integrations/test_responses_background_cost.py b/tests/unit/integrations/test_responses_background_cost.py similarity index 100% rename from tests/test_litellm/integrations/test_responses_background_cost.py rename to tests/unit/integrations/test_responses_background_cost.py diff --git a/tests/test_litellm/integrations/test_rubrik.py b/tests/unit/integrations/test_rubrik.py similarity index 99% rename from tests/test_litellm/integrations/test_rubrik.py rename to tests/unit/integrations/test_rubrik.py index 4a2ee487c65..f3fea292bde 100644 --- a/tests/test_litellm/integrations/test_rubrik.py +++ b/tests/unit/integrations/test_rubrik.py @@ -19,7 +19,7 @@ from litellm.integrations.rubrik import ( ) from litellm.proxy._types import UserAPIKeyAuth -from tests.test_litellm.integrations.rubrik_test_helpers import ( +from tests.unit.integrations.rubrik_test_helpers import ( make_inputs_with_tools, make_tool_call_dict, ) diff --git a/tests/test_litellm/integrations/test_s3.py b/tests/unit/integrations/test_s3.py similarity index 100% rename from tests/test_litellm/integrations/test_s3.py rename to tests/unit/integrations/test_s3.py diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/unit/integrations/test_s3_v2.py similarity index 100% rename from tests/test_litellm/integrations/test_s3_v2.py rename to tests/unit/integrations/test_s3_v2.py diff --git a/tests/test_litellm/integrations/test_shadow_eval_logger.py b/tests/unit/integrations/test_shadow_eval_logger.py similarity index 100% rename from tests/test_litellm/integrations/test_shadow_eval_logger.py rename to tests/unit/integrations/test_shadow_eval_logger.py diff --git a/tests/test_litellm/integrations/test_weave_otel.py b/tests/unit/integrations/test_weave_otel.py similarity index 100% rename from tests/test_litellm/integrations/test_weave_otel.py rename to tests/unit/integrations/test_weave_otel.py diff --git a/tests/test_litellm/llms/black_forest_labs/__init__.py b/tests/unit/integrations/websearch_interception/__init__.py similarity index 100% rename from tests/test_litellm/llms/black_forest_labs/__init__.py rename to tests/unit/integrations/websearch_interception/__init__.py diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_agentic_loop_cap.py b/tests/unit/integrations/websearch_interception/test_websearch_agentic_loop_cap.py similarity index 100% rename from tests/test_litellm/integrations/websearch_interception/test_websearch_agentic_loop_cap.py rename to tests/unit/integrations/websearch_interception/test_websearch_agentic_loop_cap.py diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_chat_completion.py b/tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py similarity index 100% rename from tests/test_litellm/integrations/websearch_interception/test_websearch_chat_completion.py rename to tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py b/tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py similarity index 100% rename from tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py rename to tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_thinking.py b/tests/unit/integrations/websearch_interception/test_websearch_interception_thinking.py similarity index 100% rename from tests/test_litellm/integrations/websearch_interception/test_websearch_interception_thinking.py rename to tests/unit/integrations/websearch_interception/test_websearch_interception_thinking.py diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py b/tests/unit/integrations/websearch_interception/test_websearch_native_blocks.py similarity index 100% rename from tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py rename to tests/unit/integrations/websearch_interception/test_websearch_native_blocks.py diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_responses.py b/tests/unit/integrations/websearch_interception/test_websearch_responses.py similarity index 100% rename from tests/test_litellm/integrations/websearch_interception/test_websearch_responses.py rename to tests/unit/integrations/websearch_interception/test_websearch_responses.py diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_rich_query_shape.py b/tests/unit/integrations/websearch_interception/test_websearch_rich_query_shape.py similarity index 100% rename from tests/test_litellm/integrations/websearch_interception/test_websearch_rich_query_shape.py rename to tests/unit/integrations/websearch_interception/test_websearch_rich_query_shape.py diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_short_circuit.py b/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py similarity index 100% rename from tests/test_litellm/integrations/websearch_interception/test_websearch_short_circuit.py rename to tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_streaming_wrap.py b/tests/unit/integrations/websearch_interception/test_websearch_streaming_wrap.py similarity index 100% rename from tests/test_litellm/integrations/websearch_interception/test_websearch_streaming_wrap.py rename to tests/unit/integrations/websearch_interception/test_websearch_streaming_wrap.py diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_thinking_constraint.py b/tests/unit/integrations/websearch_interception/test_websearch_thinking_constraint.py similarity index 100% rename from tests/test_litellm/integrations/websearch_interception/test_websearch_thinking_constraint.py rename to tests/unit/integrations/websearch_interception/test_websearch_thinking_constraint.py diff --git a/tests/test_litellm/llms/black_forest_labs/image_edit/__init__.py b/tests/unit/llms/aiml/__init__.py similarity index 100% rename from tests/test_litellm/llms/black_forest_labs/image_edit/__init__.py rename to tests/unit/llms/aiml/__init__.py diff --git a/tests/test_litellm/llms/black_forest_labs/image_generation/__init__.py b/tests/unit/llms/aiml/image_generation/__init__.py similarity index 100% rename from tests/test_litellm/llms/black_forest_labs/image_generation/__init__.py rename to tests/unit/llms/aiml/image_generation/__init__.py diff --git a/tests/test_litellm/llms/aiml/image_generation/test_aiml_image_generation_transformation.py b/tests/unit/llms/aiml/image_generation/test_aiml_image_generation_transformation.py similarity index 100% rename from tests/test_litellm/llms/aiml/image_generation/test_aiml_image_generation_transformation.py rename to tests/unit/llms/aiml/image_generation/test_aiml_image_generation_transformation.py diff --git a/tests/unit/llms/anthropic/batches/test_transformation.py b/tests/unit/llms/anthropic/batches/test_transformation.py index eacd2c9d03b..419fc7740eb 100644 --- a/tests/unit/llms/anthropic/batches/test_transformation.py +++ b/tests/unit/llms/anthropic/batches/test_transformation.py @@ -616,7 +616,7 @@ def test_transform_response_reraises_unexpected_error(config): # automatically. See base_batches_config_test.py. # --------------------------------------------------------------------------- # -from tests.test_litellm.llms.base_llm.batches.base_batches_config_test import ( # noqa: E402 +from tests.unit.llms.base_llm.batches.base_batches_config_test import ( # noqa: E402 BatchesConfigContractTests, ) diff --git a/tests/test_litellm/llms/cerebras/__init__.py b/tests/unit/llms/anthropic/chat/__init__.py similarity index 100% rename from tests/test_litellm/llms/cerebras/__init__.py rename to tests/unit/llms/anthropic/chat/__init__.py diff --git a/tests/test_litellm/llms/anthropic/chat/conftest.py b/tests/unit/llms/anthropic/chat/conftest.py similarity index 100% rename from tests/test_litellm/llms/anthropic/chat/conftest.py rename to tests/unit/llms/anthropic/chat/conftest.py diff --git a/tests/test_litellm/llms/chatgpt/__init__.py b/tests/unit/llms/anthropic/chat/guardrail_translation/__init__.py similarity index 100% rename from tests/test_litellm/llms/chatgpt/__init__.py rename to tests/unit/llms/anthropic/chat/guardrail_translation/__init__.py diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py rename to tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py similarity index 100% rename from tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py rename to tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py rename to tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py diff --git a/tests/test_litellm/llms/anthropic/chat/test_code_interpreter_results_extraction.py b/tests/unit/llms/anthropic/chat/test_code_interpreter_results_extraction.py similarity index 100% rename from tests/test_litellm/llms/anthropic/chat/test_code_interpreter_results_extraction.py rename to tests/unit/llms/anthropic/chat/test_code_interpreter_results_extraction.py diff --git a/tests/test_litellm/llms/chatgpt/chat/__init__.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/__init__.py similarity index 100% rename from tests/test_litellm/llms/chatgpt/chat/__init__.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/__init__.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_prompt_cache_key.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_handler_prompt_cache_key.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_prompt_cache_key.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_handler_prompt_cache_key.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_compaction.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_compaction.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_compaction.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_compaction.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_message_id.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_message_id.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_message_id.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_message_id.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_mid_stream_error.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_mid_stream_error.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_mid_stream_error.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_mid_stream_error.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_stop_reason.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_stop_reason.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_stop_reason.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_stop_reason.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_tool_args.py b/tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_tool_args.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_tool_args.py rename to tests/unit/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_tool_args.py diff --git a/tests/test_litellm/llms/crusoe/__init__.py b/tests/unit/llms/anthropic/experimental_pass_through/context_management/__init__.py similarity index 100% rename from tests/test_litellm/llms/crusoe/__init__.py rename to tests/unit/llms/anthropic/experimental_pass_through/context_management/__init__.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_clear_tool_uses.py b/tests/unit/llms/anthropic/experimental_pass_through/context_management/test_clear_tool_uses.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_clear_tool_uses.py rename to tests/unit/llms/anthropic/experimental_pass_through/context_management/test_clear_tool_uses.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py b/tests/unit/llms/anthropic/experimental_pass_through/context_management/test_compact.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py rename to tests/unit/llms/anthropic/experimental_pass_through/context_management/test_compact.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py b/tests/unit/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py rename to tests/unit/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py diff --git a/tests/test_litellm/llms/databricks/chat/__init__.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/__init__.py similarity index 100% rename from tests/test_litellm/llms/databricks/chat/__init__.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/__init__.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_advisor_integration.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_advisor_integration.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_advisor_integration.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_advisor_integration.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_agentic_streaming_iterator.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_effort.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_encrypted_reasoning.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_encrypted_reasoning.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_encrypted_reasoning.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_encrypted_reasoning.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_per_turn_control.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_per_turn_control.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_per_turn_control.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_per_turn_control.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_speed.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_speed.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_speed.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_speed.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_structured_outputs.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_structured_outputs.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_structured_outputs.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_anthropic_messages_structured_outputs.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_content_after_stop_reason.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_content_after_stop_reason.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_content_after_stop_reason.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_content_after_stop_reason.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mid_conversation_system.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_mid_conversation_system.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mid_conversation_system.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_mid_conversation_system.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_parallel_tool_calls.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_parallel_tool_calls.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_parallel_tool_calls.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_parallel_tool_calls.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_auto_summary_messages.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_reasoning_auto_summary_messages.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_auto_summary_messages.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_reasoning_auto_summary_messages.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_request_optional_param_utils.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_request_optional_param_utils.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_request_optional_param_utils.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_request_optional_param_utils.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_response_cache.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_response_cache.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_sse_wrapper.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_sse_wrapper.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_sse_wrapper.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_sse_wrapper.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/unit/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py rename to tests/unit/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py diff --git a/tests/test_litellm/llms/databricks/responses/__init__.py b/tests/unit/llms/anthropic/experimental_pass_through/responses_adapters/__init__.py similarity index 100% rename from tests/test_litellm/llms/databricks/responses/__init__.py rename to tests/unit/llms/anthropic/experimental_pass_through/responses_adapters/__init__.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py b/tests/unit/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py rename to tests/unit/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py b/tests/unit/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py similarity index 72% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py rename to tests/unit/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py index bfe2d6b7cea..392ecc2bcdd 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py +++ b/tests/unit/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py @@ -4,18 +4,25 @@ Tests for AnthropicResponsesStreamWrapper """ import asyncio +import json import os import sys from types import SimpleNamespace +import pytest + sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../.."))) +import litellm +from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.prompt_templates.common_utils import ( encrypted_reasoning_signature, ) +from litellm.llms.anthropic.experimental_pass_through.messages.utils import INCOMPLETE_STREAM_ERROR_MESSAGE from litellm.llms.anthropic.experimental_pass_through.responses_adapters.streaming_iterator import ( AnthropicResponsesStreamWrapper, ) +from litellm.types.llms.openai import ResponseFailedEvent, ResponsesAPIResponse def _process_all(events: list) -> list: @@ -132,6 +139,7 @@ class TestReasoningItemWithoutSummaryText: {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}}, {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Hello"}, {"type": "response.output_item.done", "item": {"type": "message", "id": "msg_1"}}, + {"type": "response.completed"}, ] def test_reasoning_without_summary_emits_no_thinking_block(self): @@ -144,6 +152,8 @@ class TestReasoningItemWithoutSummaryText: ("content_block_start", 0), ("content_block_delta", 0), ("content_block_stop", 0), + ("message_delta", None), + ("message_stop", None), ] assert chunks[1]["content_block"] == {"type": "text", "text": ""} @@ -166,6 +176,8 @@ class TestReasoningItemWithoutSummaryText: ("content_block_start", 1), ("content_block_delta", 1), ("content_block_stop", 1), + ("message_delta", None), + ("message_stop", None), ] assert chunks[1]["content_block"] == {"type": "thinking", "thinking": "", "signature": ""} assert "".join(c["delta"]["thinking"] for c in chunks[2:4]) == "Weighing options" @@ -215,6 +227,8 @@ class TestEncryptedReasoningIsStreamedForReplay: ("content_block_start", 1), ("content_block_delta", 1), ("content_block_stop", 1), + ("message_delta", None), + ("message_stop", None), ] assert chunks[1]["content_block"] == { "type": "redacted_thinking", @@ -234,9 +248,7 @@ class TestEncryptedReasoningIsStreamedForReplay: ] chunks = _process_all(events) - thinking = "".join( - c["delta"]["thinking"] for c in chunks if c.get("delta", {}).get("type") == "thinking_delta" - ) + thinking = "".join(c["delta"]["thinking"] for c in chunks if c.get("delta", {}).get("type") == "thinking_delta") assert thinking == "First.\n\nSecond." assert [c["type"] for c in chunks].count("content_block_start") == 1 @@ -283,6 +295,7 @@ class TestToolUseBlockClosedExactlyOnce: "type": "response.output_item.done", "item": {"type": "message", "id": "chatcmpl-123", "status": "completed"}, }, + {"type": "response.completed"}, ] def test_one_content_block_stop_per_content_block_start(self): @@ -302,6 +315,8 @@ class TestToolUseBlockClosedExactlyOnce: ("content_block_delta", 0), ("content_block_delta", 0), ("content_block_stop", 0), + ("message_delta", None), + ("message_stop", None), ] assert chunks[1]["content_block"] == { "type": "tool_use", @@ -452,3 +467,158 @@ class TestRefusalStreamEvents: message_delta = next(c for c in chunks if c["type"] == "message_delta") assert message_delta["delta"]["stop_reason"] == "max_tokens" assert "stop_details" not in message_delta["delta"] + + +def _collect(stream) -> list: + async def _run() -> list: + wrapper = AnthropicResponsesStreamWrapper(responses_stream=stream, model="m") + return [chunk async for chunk in wrapper] + + return asyncio.run(_run()) + + +class TestUpstreamFailureEndsStreamWithErrorEvent: + """A provider failure must reach the Anthropic client as an ``error`` event that + ends the stream, never as a fabricated ``end_turn`` or a silent close.""" + + def test_response_failed_event_emits_error_event_and_stops_pulling_upstream(self): + failed = SimpleNamespace( + status="failed", + output=[], + usage=None, + error={"code": "rate_limit_exceeded", "message": "Rate limit reached for gpt-5.5, try again in 20s."}, + ) + + async def _gen(): + yield {"type": "response.created"} + yield {"type": "response.failed", "response": failed} + raise AssertionError("upstream was pulled again after the failure") + + async def _run() -> list: + wrapper = AnthropicResponsesStreamWrapper(responses_stream=_gen(), model="m") + return [frame async for frame in wrapper.async_anthropic_sse_wrapper()] + + frames = asyncio.run(_run()) + assert [frame.split(b"\n", 1)[0] for frame in frames] == [b"event: message_start", b"event: error"] + error_payload = json.loads(frames[1].split(b"data: ", 1)[1]) + assert error_payload["type"] == "error" + assert error_payload["error"] == { + "type": "rate_limit_error", + "message": "Rate limit reached for gpt-5.5, try again in 20s.", + } + + def test_raised_mid_stream_fallback_error_is_unwrapped_to_the_provider_failure(self): + rate_limit = litellm.RateLimitError(message="You have no credits remaining.", llm_provider="openai", model="m") + wrapped = MidStreamFallbackError( + message=str(rate_limit), + model="m", + llm_provider="openai", + original_exception=rate_limit, + is_pre_first_chunk=True, + ) + + async def _gen(): + yield {"type": "response.created"} + raise wrapped + + chunks = _collect(_gen()) + assert [chunk["type"] for chunk in chunks] == ["message_start", "error"] + assert chunks[1]["error"] == {"type": "rate_limit_error", "message": rate_limit.message} + + def test_sync_upstream_transport_error_after_content_becomes_api_error_event(self): + def _events(): + yield {"type": "response.created"} + yield {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}} + yield {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Hi"} + raise ConnectionResetError("Response payload is not completed") + + chunks = _collect(_events()) + assert [chunk["type"] for chunk in chunks] == [ + "message_start", + "content_block_start", + "content_block_delta", + "error", + ] + assert chunks[-1]["error"] == {"type": "api_error", "message": "Response payload is not completed"} + + def test_error_event_message_is_redacted_before_it_reaches_the_client(self): + async def _gen(): + yield {"type": "response.created"} + raise RuntimeError("upstream failed with key sk-proj-abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJ") + + chunks = _collect(_gen()) + assert chunks[-1]["type"] == "error" + assert "sk-proj-" not in chunks[-1]["error"]["message"] + assert chunks[-1]["error"]["message"].startswith("upstream failed with key") + + @pytest.mark.parametrize( + ("raised", "expected_error"), + [ + ( + MidStreamFallbackError(message="boom", model="m", llm_provider="openai"), + {"type": "api_error", "message": "litellm.MidStreamFallbackError: boom"}, + ), + ( + type("StringStatusError", (Exception,), {"status_code": "429"})("throttled"), + {"type": "rate_limit_error", "message": "throttled"}, + ), + ( + type("NonErrorStatusError", (Exception,), {"status_code": 200})("odd status"), + {"type": "api_error", "message": "odd status"}, + ), + ], + ids=["mid-stream-fallback-without-original", "digit-string-status", "status-outside-4xx-5xx"], + ) + def test_raised_failure_status_is_normalized_into_the_error_type(self, raised, expected_error): + async def _gen(): + yield {"type": "response.created"} + raise raised + + chunks = _collect(_gen()) + assert [chunk["type"] for chunk in chunks] == ["message_start", "error"] + assert chunks[1]["error"] == expected_error + + def test_pydantic_response_failed_event_is_mapped_like_a_dict_event(self): + failed = ResponsesAPIResponse( + id="resp_1", + created_at=1, + error={"code": "server_error", "message": "The server had an error while processing your request."}, + status="failed", + output=[], + model="m", + object="response", + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ) + + async def _gen(): + yield {"type": "response.created"} + yield ResponseFailedEvent(type="response.failed", response=failed) + + chunks = _collect(_gen()) + assert [chunk["type"] for chunk in chunks] == ["message_start", "error"] + assert chunks[1]["error"] == { + "type": "api_error", + "message": "The server had an error while processing your request.", + } + + def test_upstream_ending_without_a_terminal_event_is_an_error_not_a_silent_close(self): + async def _gen(): + yield {"type": "response.created"} + yield {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}} + yield {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Hi"} + + chunks = _collect(_gen()) + assert [chunk["type"] for chunk in chunks] == [ + "message_start", + "content_block_start", + "content_block_delta", + "error", + ] + assert chunks[-1]["error"] == {"type": "api_error", "message": INCOMPLETE_STREAM_ERROR_MESSAGE} + + def test_sync_upstream_ending_before_any_event_is_an_error_not_a_silent_close(self): + chunks = _collect(iter(())) + assert [chunk["type"] for chunk in chunks] == ["message_start", "error"] + assert chunks[1]["error"] == {"type": "api_error", "message": INCOMPLETE_STREAM_ERROR_MESSAGE} diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/unit/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py similarity index 100% rename from tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py rename to tests/unit/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/unit/llms/anthropic/test_anthropic_common_utils.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py rename to tests/unit/llms/anthropic/test_anthropic_common_utils.py diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_count_tokens_transformation.py b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_anthropic_count_tokens_transformation.py rename to tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py b/tests/unit/llms/anthropic/test_anthropic_files_and_batches.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py rename to tests/unit/llms/anthropic/test_anthropic_files_and_batches.py diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_output_format_filter.py b/tests/unit/llms/anthropic/test_anthropic_output_format_filter.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_anthropic_output_format_filter.py rename to tests/unit/llms/anthropic/test_anthropic_output_format_filter.py diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_prompt_cache_prediction.py b/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_anthropic_prompt_cache_prediction.py rename to tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py b/tests/unit/llms/anthropic/test_anthropic_reasoning_effort.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py rename to tests/unit/llms/anthropic/test_anthropic_reasoning_effort.py diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_schema_filter.py b/tests/unit/llms/anthropic/test_anthropic_schema_filter.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_anthropic_schema_filter.py rename to tests/unit/llms/anthropic/test_anthropic_schema_filter.py diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_structured_output.py b/tests/unit/llms/anthropic/test_anthropic_structured_output.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_anthropic_structured_output.py rename to tests/unit/llms/anthropic/test_anthropic_structured_output.py diff --git a/tests/test_litellm/llms/anthropic/test_azure_ai_cache_pricing.py b/tests/unit/llms/anthropic/test_azure_ai_cache_pricing.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_azure_ai_cache_pricing.py rename to tests/unit/llms/anthropic/test_azure_ai_cache_pricing.py diff --git a/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py b/tests/unit/llms/anthropic/test_cost_calculation_dict_safety.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py rename to tests/unit/llms/anthropic/test_cost_calculation_dict_safety.py diff --git a/tests/test_litellm/llms/anthropic/test_count_tokens_oauth.py b/tests/unit/llms/anthropic/test_count_tokens_oauth.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_count_tokens_oauth.py rename to tests/unit/llms/anthropic/test_count_tokens_oauth.py diff --git a/tests/test_litellm/llms/anthropic/test_message_sanitization.py b/tests/unit/llms/anthropic/test_message_sanitization.py similarity index 100% rename from tests/test_litellm/llms/anthropic/test_message_sanitization.py rename to tests/unit/llms/anthropic/test_message_sanitization.py diff --git a/tests/test_litellm/llms/deepseek/__init__.py b/tests/unit/llms/azure/batches/__init__.py similarity index 100% rename from tests/test_litellm/llms/deepseek/__init__.py rename to tests/unit/llms/azure/batches/__init__.py diff --git a/tests/test_litellm/llms/azure/batches/test_handler.py b/tests/unit/llms/azure/batches/test_handler.py similarity index 100% rename from tests/test_litellm/llms/azure/batches/test_handler.py rename to tests/unit/llms/azure/batches/test_handler.py diff --git a/tests/test_litellm/llms/deepseek/chat/__init__.py b/tests/unit/llms/azure/chat/__init__.py similarity index 100% rename from tests/test_litellm/llms/deepseek/chat/__init__.py rename to tests/unit/llms/azure/chat/__init__.py diff --git a/tests/test_litellm/llms/azure/chat/test_azure_base_model_routing.py b/tests/unit/llms/azure/chat/test_azure_base_model_routing.py similarity index 100% rename from tests/test_litellm/llms/azure/chat/test_azure_base_model_routing.py rename to tests/unit/llms/azure/chat/test_azure_base_model_routing.py diff --git a/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py b/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py rename to tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py diff --git a/tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py b/tests/unit/llms/azure/chat/test_azure_chat_o_series_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py rename to tests/unit/llms/azure/chat/test_azure_chat_o_series_transformation.py diff --git a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py b/tests/unit/llms/azure/chat/test_azure_gpt5_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py rename to tests/unit/llms/azure/chat/test_azure_gpt5_transformation.py diff --git a/tests/test_litellm/llms/azure/realtime/test_handler.py b/tests/unit/llms/azure/realtime/test_handler.py similarity index 100% rename from tests/test_litellm/llms/azure/realtime/test_handler.py rename to tests/unit/llms/azure/realtime/test_handler.py diff --git a/tests/test_litellm/llms/azure/test_audio_transcriptions.py b/tests/unit/llms/azure/test_audio_transcriptions.py similarity index 100% rename from tests/test_litellm/llms/azure/test_audio_transcriptions.py rename to tests/unit/llms/azure/test_audio_transcriptions.py diff --git a/tests/test_litellm/llms/azure/test_azure.py b/tests/unit/llms/azure/test_azure.py similarity index 100% rename from tests/test_litellm/llms/azure/test_azure.py rename to tests/unit/llms/azure/test_azure.py diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py similarity index 100% rename from tests/test_litellm/llms/azure/test_azure_common_utils.py rename to tests/unit/llms/azure/test_azure_common_utils.py diff --git a/tests/test_litellm/llms/azure/test_azure_cost_calculation.py b/tests/unit/llms/azure/test_azure_cost_calculation.py similarity index 100% rename from tests/test_litellm/llms/azure/test_azure_cost_calculation.py rename to tests/unit/llms/azure/test_azure_cost_calculation.py diff --git a/tests/test_litellm/llms/azure/test_azure_embedding.py b/tests/unit/llms/azure/test_azure_embedding.py similarity index 100% rename from tests/test_litellm/llms/azure/test_azure_embedding.py rename to tests/unit/llms/azure/test_azure_embedding.py diff --git a/tests/test_litellm/llms/azure/test_azure_exception_mapping.py b/tests/unit/llms/azure/test_azure_exception_mapping.py similarity index 100% rename from tests/test_litellm/llms/azure/test_azure_exception_mapping.py rename to tests/unit/llms/azure/test_azure_exception_mapping.py diff --git a/tests/test_litellm/llms/azure/test_azure_fine_tuning_api.py b/tests/unit/llms/azure/test_azure_fine_tuning_api.py similarity index 100% rename from tests/test_litellm/llms/azure/test_azure_fine_tuning_api.py rename to tests/unit/llms/azure/test_azure_fine_tuning_api.py diff --git a/tests/test_litellm/llms/azure/test_azure_speech_audio_transcription.py b/tests/unit/llms/azure/test_azure_speech_audio_transcription.py similarity index 100% rename from tests/test_litellm/llms/azure/test_azure_speech_audio_transcription.py rename to tests/unit/llms/azure/test_azure_speech_audio_transcription.py diff --git a/tests/test_litellm/llms/deepseek/messages/__init__.py b/tests/unit/llms/azure/videos/__init__.py similarity index 100% rename from tests/test_litellm/llms/deepseek/messages/__init__.py rename to tests/unit/llms/azure/videos/__init__.py diff --git a/tests/test_litellm/llms/azure/videos/test_azure_video_transformation.py b/tests/unit/llms/azure/videos/test_azure_video_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure/videos/test_azure_video_transformation.py rename to tests/unit/llms/azure/videos/test_azure_video_transformation.py diff --git a/tests/test_litellm/llms/gemini/__init__.py b/tests/unit/llms/azure_ai/claude/__init__.py similarity index 100% rename from tests/test_litellm/llms/gemini/__init__.py rename to tests/unit/llms/azure_ai/claude/__init__.py diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py rename to tests/unit/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_handler.py b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_handler.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_handler.py rename to tests/unit/llms/azure_ai/claude/test_azure_anthropic_handler.py diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py rename to tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_provider_routing.py b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_provider_routing.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_provider_routing.py rename to tests/unit/llms/azure_ai/claude/test_azure_anthropic_provider_routing.py diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_transformation.py b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_transformation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_transformation.py rename to tests/unit/llms/azure_ai/claude/test_azure_anthropic_transformation.py diff --git a/tests/test_litellm/llms/azure_ai/claude/test_main_azure_anthropic_timeout.py b/tests/unit/llms/azure_ai/claude/test_main_azure_anthropic_timeout.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/claude/test_main_azure_anthropic_timeout.py rename to tests/unit/llms/azure_ai/claude/test_main_azure_anthropic_timeout.py diff --git a/tests/test_litellm/llms/gemini/audio_transcription/__init__.py b/tests/unit/llms/azure_ai/image_generation/__init__.py similarity index 100% rename from tests/test_litellm/llms/gemini/audio_transcription/__init__.py rename to tests/unit/llms/azure_ai/image_generation/__init__.py diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py b/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py rename to tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py b/tests/unit/llms/azure_ai/image_generation/test_mai_image_generation.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py rename to tests/unit/llms/azure_ai/image_generation/test_mai_image_generation.py diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_agents_handler.py b/tests/unit/llms/azure_ai/test_azure_ai_agents_handler.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/test_azure_ai_agents_handler.py rename to tests/unit/llms/azure_ai/test_azure_ai_agents_handler.py diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py b/tests/unit/llms/azure_ai/test_azure_ai_cost_calculator.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py rename to tests/unit/llms/azure_ai/test_azure_ai_cost_calculator.py diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_entra_auth.py b/tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/test_azure_ai_entra_auth.py rename to tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py b/tests/unit/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py rename to tests/unit/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py b/tests/unit/llms/azure_ai/test_azure_ai_fw_models_metadata.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py rename to tests/unit/llms/azure_ai/test_azure_ai_fw_models_metadata.py diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py b/tests/unit/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py similarity index 100% rename from tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py rename to tests/unit/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py diff --git a/tests/test_litellm/llms/base_llm/batches/base_batches_config_test.py b/tests/unit/llms/base_llm/batches/base_batches_config_test.py similarity index 100% rename from tests/test_litellm/llms/base_llm/batches/base_batches_config_test.py rename to tests/unit/llms/base_llm/batches/base_batches_config_test.py diff --git a/tests/test_litellm/llms/gemini/google_genai/__init__.py b/tests/unit/llms/base_llm/files/__init__.py similarity index 100% rename from tests/test_litellm/llms/gemini/google_genai/__init__.py rename to tests/unit/llms/base_llm/files/__init__.py diff --git a/tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py b/tests/unit/llms/base_llm/files/test_azure_blob_storage_backend.py similarity index 100% rename from tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py rename to tests/unit/llms/base_llm/files/test_azure_blob_storage_backend.py diff --git a/tests/test_litellm/llms/base_llm/files/test_litellm_db_storage_backend.py b/tests/unit/llms/base_llm/files/test_litellm_db_storage_backend.py similarity index 100% rename from tests/test_litellm/llms/base_llm/files/test_litellm_db_storage_backend.py rename to tests/unit/llms/base_llm/files/test_litellm_db_storage_backend.py diff --git a/tests/test_litellm/llms/base_llm/files/test_storage_backend_factory.py b/tests/unit/llms/base_llm/files/test_storage_backend_factory.py similarity index 100% rename from tests/test_litellm/llms/base_llm/files/test_storage_backend_factory.py rename to tests/unit/llms/base_llm/files/test_storage_backend_factory.py diff --git a/tests/test_litellm/llms/gemini/google_genai/guardrail_translation/__init__.py b/tests/unit/llms/base_llm/responses/__init__.py similarity index 100% rename from tests/test_litellm/llms/gemini/google_genai/guardrail_translation/__init__.py rename to tests/unit/llms/base_llm/responses/__init__.py diff --git a/tests/test_litellm/llms/base_llm/responses/test_codex_compat.py b/tests/unit/llms/base_llm/responses/test_codex_compat.py similarity index 100% rename from tests/test_litellm/llms/base_llm/responses/test_codex_compat.py rename to tests/unit/llms/base_llm/responses/test_codex_compat.py diff --git a/tests/test_litellm/llms/base_llm/responses/test_transformation.py b/tests/unit/llms/base_llm/responses/test_transformation.py similarity index 100% rename from tests/test_litellm/llms/base_llm/responses/test_transformation.py rename to tests/unit/llms/base_llm/responses/test_transformation.py diff --git a/tests/test_litellm/llms/gemini/image_edit/__init__.py b/tests/unit/llms/base_llm/search/__init__.py similarity index 100% rename from tests/test_litellm/llms/gemini/image_edit/__init__.py rename to tests/unit/llms/base_llm/search/__init__.py diff --git a/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py b/tests/unit/llms/base_llm/search/test_base_search_transformation.py similarity index 100% rename from tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py rename to tests/unit/llms/base_llm/search/test_base_search_transformation.py diff --git a/tests/test_litellm/llms/base_llm/test_base_managed_resource.py b/tests/unit/llms/base_llm/test_base_managed_resource.py similarity index 100% rename from tests/test_litellm/llms/base_llm/test_base_managed_resource.py rename to tests/unit/llms/base_llm/test_base_managed_resource.py diff --git a/tests/test_litellm/llms/base_llm/test_base_model_iterator.py b/tests/unit/llms/base_llm/test_base_model_iterator.py similarity index 100% rename from tests/test_litellm/llms/base_llm/test_base_model_iterator.py rename to tests/unit/llms/base_llm/test_base_model_iterator.py diff --git a/tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py b/tests/unit/llms/base_llm/test_managed_resource_isolation.py similarity index 100% rename from tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py rename to tests/unit/llms/base_llm/test_managed_resource_isolation.py diff --git a/tests/test_litellm/llms/base_llm/test_managed_resources_utils.py b/tests/unit/llms/base_llm/test_managed_resources_utils.py similarity index 100% rename from tests/test_litellm/llms/base_llm/test_managed_resources_utils.py rename to tests/unit/llms/base_llm/test_managed_resources_utils.py diff --git a/tests/test_litellm/llms/gemini/realtime/__init__.py b/tests/unit/llms/bedrock/batches/__init__.py similarity index 100% rename from tests/test_litellm/llms/gemini/realtime/__init__.py rename to tests/unit/llms/bedrock/batches/__init__.py diff --git a/tests/test_litellm/llms/bedrock/batches/test_batch_metadata_sanitization.py b/tests/unit/llms/bedrock/batches/test_batch_metadata_sanitization.py similarity index 100% rename from tests/test_litellm/llms/bedrock/batches/test_batch_metadata_sanitization.py rename to tests/unit/llms/bedrock/batches/test_batch_metadata_sanitization.py diff --git a/tests/test_litellm/llms/bedrock/batches/test_handler.py b/tests/unit/llms/bedrock/batches/test_handler.py similarity index 100% rename from tests/test_litellm/llms/bedrock/batches/test_handler.py rename to tests/unit/llms/bedrock/batches/test_handler.py diff --git a/tests/test_litellm/llms/bedrock/batches/test_transformation.py b/tests/unit/llms/bedrock/batches/test_transformation.py similarity index 99% rename from tests/test_litellm/llms/bedrock/batches/test_transformation.py rename to tests/unit/llms/bedrock/batches/test_transformation.py index 347c459a369..5e987239c54 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_transformation.py +++ b/tests/unit/llms/bedrock/batches/test_transformation.py @@ -878,7 +878,7 @@ def test_validate_environment_passes_headers_through(config): # Shared BaseBatchesConfig contract suite. # --------------------------------------------------------------------------- # -from tests.test_litellm.llms.base_llm.batches.base_batches_config_test import ( # noqa: E402 +from tests.unit.llms.base_llm.batches.base_batches_config_test import ( # noqa: E402 BatchesConfigContractTests, ) diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py similarity index 99% rename from tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py rename to tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py index 67ffe7570a1..08bcac33a35 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -20,7 +20,7 @@ from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.rust_bridge import configuration from litellm.types.utils import ModelResponse -from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe +from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe RESOLVED_CREDENTIALS = Credentials( access_key="AKIARESOLVED", diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py rename to tests/unit/llms/bedrock/chat/test_converse_transformation.py diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation_nova_2.py b/tests/unit/llms/bedrock/chat/test_converse_transformation_nova_2.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/test_converse_transformation_nova_2.py rename to tests/unit/llms/bedrock/chat/test_converse_transformation_nova_2.py diff --git a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py b/tests/unit/llms/bedrock/chat/test_invoke_handler.py similarity index 86% rename from tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py rename to tests/unit/llms/bedrock/chat/test_invoke_handler.py index 466e9b4fda8..ed8b7023977 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/unit/llms/bedrock/chat/test_invoke_handler.py @@ -1,5 +1,6 @@ import base64 import binascii +import itertools import datetime import json import struct @@ -14,10 +15,13 @@ import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.bedrock.chat.invoke_handler import ( + AmazonOpenAICompatibleStreamDecoder, AWSEventStreamDecoder, make_call, make_sync_call, ) +from litellm.exceptions import MidStreamFallbackError +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.utils import ModelResponseStream @@ -799,3 +803,125 @@ async def test_moonshot_invoke_async_stream_yields_openai_shaped_chunks(_aws_tes ) _assert_moonshot_stream_content([chunk async for chunk in stream]) + + +def _truncated_frame() -> bytes: + return _bedrock_event_stream_frame(_openai_stream_chunk({"role": "assistant"}))[:-8] + + +def _event_stream_headers() -> httpx.Headers: + return httpx.Headers({"content-type": "application/vnd.amazon.eventstream", "x-amzn-RequestId": "req-empty-1"}) + + +_UNDECODABLE_STREAM_BODIES: Final = ( + pytest.param(b"", id="empty"), + pytest.param(b"\x00\x00\x00\x05", id="shorter-than-a-prelude"), + pytest.param(_truncated_frame(), id="truncated-first-message"), +) + + +def _assert_no_events_error(error: BedrockError, body: bytes) -> None: + assert error.status_code == 502 + assert "HTTP 200" in error.message + assert "decoded to no events" in error.message + assert f"{len(body)} bytes received" in error.message + assert "application/vnd.amazon.eventstream" in error.message + assert "req-empty-1" in error.message + assert f"first bytes={body[:200]!r}" in error.message + + +@pytest.mark.parametrize("body", _UNDECODABLE_STREAM_BODIES) +def test_iter_bytes_raises_when_a_200_body_decodes_to_no_events(body: bytes) -> None: + decoder: Final = AWSEventStreamDecoder(model="us.moonshotai.kimi-k3") + + with pytest.raises(BedrockError) as exc_info: + list(decoder.iter_bytes(iter([body]), response_headers=_event_stream_headers())) + + _assert_no_events_error(exc_info.value, body) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("body", _UNDECODABLE_STREAM_BODIES) +async def test_aiter_bytes_raises_when_a_200_body_decodes_to_no_events(body: bytes) -> None: + async def _chunks() -> AsyncIterator[bytes]: + yield body + + decoder: Final = AWSEventStreamDecoder(model="us.moonshotai.kimi-k3") + + with pytest.raises(BedrockError) as exc_info: + _ = [chunk async for chunk in decoder.aiter_bytes(_chunks(), response_headers=_event_stream_headers())] + + _assert_no_events_error(exc_info.value, body) + + +def test_iter_bytes_raises_when_the_stream_ends_mid_message() -> None: + decoder: Final = AmazonOpenAICompatibleStreamDecoder(model="moonshot.kimi-k2-thinking", sync_stream=True) + stream: Final = decoder.iter_bytes(iter([_MOONSHOT_RAW_STREAM, _truncated_frame()])) + + chunks: Final = list(itertools.islice(stream, 4)) + with pytest.raises(BedrockError) as exc_info: + next(stream) + + _assert_moonshot_stream_content(chunks) + assert exc_info.value.status_code == 502 + assert f"{len(_truncated_frame())} undecoded bytes after 4 events" in exc_info.value.message + assert "first bytes=" not in exc_info.value.message + + +def test_iter_bytes_yields_a_complete_stream_without_raising() -> None: + decoder: Final = AmazonOpenAICompatibleStreamDecoder(model="moonshot.kimi-k2-thinking", sync_stream=True) + + chunks: Final = list(decoder.iter_bytes(iter([_MOONSHOT_RAW_STREAM[:100], _MOONSHOT_RAW_STREAM[100:]]))) + + _assert_moonshot_stream_content(chunks) + + +def _assert_empty_stream_surfaced_as_bad_gateway(error: MidStreamFallbackError) -> None: + assert error.status_code == 502 + assert error.is_pre_first_chunk is True + assert isinstance(error.original_exception, litellm.BadGatewayError) + assert "decoded to no events" in str(error) + assert "req-empty-1" in str(error) + + +def test_converse_stream_with_an_empty_200_body_raises_instead_of_an_empty_turn(_aws_test_credentials: None) -> None: + response: Final = MagicMock(status_code=200, headers=_event_stream_headers()) + response.iter_bytes = lambda chunk_size=None: iter([b""]) + client: Final = HTTPHandler() + client.post = MagicMock(return_value=response) + + with pytest.raises(MidStreamFallbackError) as exc_info: + list( + litellm.completion( + model="bedrock/us.moonshotai.kimi-k3", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + ) + ) + + _assert_empty_stream_surfaced_as_bad_gateway(exc_info.value) + + +@pytest.mark.asyncio +async def test_async_converse_stream_with_an_empty_200_body_raises_instead_of_an_empty_turn( + _aws_test_credentials: None, +) -> None: + async def _aiter_bytes(chunk_size: int | None = None) -> AsyncIterator[bytes]: + yield b"" + + response: Final = MagicMock(status_code=200, headers=_event_stream_headers()) + response.aiter_bytes = _aiter_bytes + client: Final = AsyncHTTPHandler() + client.post = AsyncMock(return_value=response) + + stream: Final = await litellm.acompletion( + model="bedrock/us.moonshotai.kimi-k3", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + ) + with pytest.raises(MidStreamFallbackError) as exc_info: + _ = [chunk async for chunk in stream] + + _assert_empty_stream_surfaced_as_bad_gateway(exc_info.value) diff --git a/tests/test_litellm/llms/bedrock/chat/test_mistral_config.py b/tests/unit/llms/bedrock/chat/test_mistral_config.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/test_mistral_config.py rename to tests/unit/llms/bedrock/chat/test_mistral_config.py diff --git a/tests/test_litellm/llms/bedrock/chat/test_service_tier.py b/tests/unit/llms/bedrock/chat/test_service_tier.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/test_service_tier.py rename to tests/unit/llms/bedrock/chat/test_service_tier.py diff --git a/tests/test_litellm/llms/bedrock/chat/test_streaming_choice_index.py b/tests/unit/llms/bedrock/chat/test_streaming_choice_index.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/test_streaming_choice_index.py rename to tests/unit/llms/bedrock/chat/test_streaming_choice_index.py diff --git a/tests/test_litellm/llms/bedrock/chat/test_writer_palmyra.py b/tests/unit/llms/bedrock/chat/test_writer_palmyra.py similarity index 100% rename from tests/test_litellm/llms/bedrock/chat/test_writer_palmyra.py rename to tests/unit/llms/bedrock/chat/test_writer_palmyra.py diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py index 3622ce7f212..d67724f261d 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py @@ -7,7 +7,7 @@ from botocore.credentials import RefreshableCredentials from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe +from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe class _ProbedCountTokensHandler(BedrockCountTokensHandler): diff --git a/tests/test_litellm/llms/gigachat/__init__.py b/tests/unit/llms/bedrock/embed/__init__.py similarity index 100% rename from tests/test_litellm/llms/gigachat/__init__.py rename to tests/unit/llms/bedrock/embed/__init__.py diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py b/tests/unit/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py similarity index 99% rename from tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py rename to tests/unit/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py index 18f4b0f6ced..fbcbd0aaea6 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py +++ b/tests/unit/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py @@ -9,7 +9,7 @@ import respx import litellm from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.llms.base import HiddenParams -from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe +from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe # Mock async invoke responses async_invoke_response = { diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py b/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py similarity index 99% rename from tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py rename to tests/unit/llms/bedrock/embed/test_bedrock_embedding.py index e5a460e2f1a..ad21cadaa4b 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py @@ -11,7 +11,7 @@ import litellm from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.bedrock.embed.embedding import BedrockEmbedding -from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe +from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe # Mock responses for different embedding models titan_embedding_response = {"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10} diff --git a/tests/test_litellm/llms/bedrock/embed/test_embedding.py b/tests/unit/llms/bedrock/embed/test_embedding.py similarity index 100% rename from tests/test_litellm/llms/bedrock/embed/test_embedding.py rename to tests/unit/llms/bedrock/embed/test_embedding.py diff --git a/tests/test_litellm/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py b/tests/unit/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py rename to tests/unit/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py diff --git a/tests/test_litellm/llms/bedrock/event_loop_probe.py b/tests/unit/llms/bedrock/event_loop_probe.py similarity index 100% rename from tests/test_litellm/llms/bedrock/event_loop_probe.py rename to tests/unit/llms/bedrock/event_loop_probe.py diff --git a/tests/test_litellm/llms/gigachat/embedding/__init__.py b/tests/unit/llms/bedrock/messages/__init__.py similarity index 100% rename from tests/test_litellm/llms/gigachat/embedding/__init__.py rename to tests/unit/llms/bedrock/messages/__init__.py diff --git a/tests/test_litellm/llms/gigachat/passthrough/__init__.py b/tests/unit/llms/bedrock/messages/invoke_transformations/__init__.py similarity index 100% rename from tests/test_litellm/llms/gigachat/passthrough/__init__.py rename to tests/unit/llms/bedrock/messages/invoke_transformations/__init__.py diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py rename to tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py diff --git a/tests/test_litellm/llms/bedrock/rerank/transformation.py b/tests/unit/llms/bedrock/rerank/transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/rerank/transformation.py rename to tests/unit/llms/bedrock/rerank/transformation.py diff --git a/tests/test_litellm/llms/github_copilot/messages/__init__.py b/tests/unit/llms/bedrock/responses/__init__.py similarity index 100% rename from tests/test_litellm/llms/github_copilot/messages/__init__.py rename to tests/unit/llms/bedrock/responses/__init__.py diff --git a/tests/test_litellm/llms/bedrock/responses/test_bedrock_openai_responses.py b/tests/unit/llms/bedrock/responses/test_bedrock_openai_responses.py similarity index 100% rename from tests/test_litellm/llms/bedrock/responses/test_bedrock_openai_responses.py rename to tests/unit/llms/bedrock/responses/test_bedrock_openai_responses.py diff --git a/tests/test_litellm/llms/gradient_ai/__init__.py b/tests/unit/llms/bedrock/search/__init__.py similarity index 100% rename from tests/test_litellm/llms/gradient_ai/__init__.py rename to tests/unit/llms/bedrock/search/__init__.py diff --git a/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py b/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py rename to tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py diff --git a/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py b/tests/unit/llms/bedrock/test_anthropic_beta_support.py similarity index 100% rename from tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py rename to tests/unit/llms/bedrock/test_anthropic_beta_support.py diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/unit/llms/bedrock/test_base_aws_llm.py similarity index 99% rename from tests/test_litellm/llms/bedrock/test_base_aws_llm.py rename to tests/unit/llms/bedrock/test_base_aws_llm.py index 6b9450afed4..db144ab6d56 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/unit/llms/bedrock/test_base_aws_llm.py @@ -28,7 +28,7 @@ from litellm.llms.bedrock.base_aws_llm import ( run_aws_signing, sign_request_off_loop_if_aws, ) -from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe +from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe # Global variable for the base_aws_llm.py file path diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/unit/llms/bedrock/test_bedrock_common_utils.py similarity index 100% rename from tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py rename to tests/unit/llms/bedrock/test_bedrock_common_utils.py diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_ssl_verify.py b/tests/unit/llms/bedrock/test_bedrock_ssl_verify.py similarity index 100% rename from tests/test_litellm/llms/bedrock/test_bedrock_ssl_verify.py rename to tests/unit/llms/bedrock/test_bedrock_ssl_verify.py diff --git a/tests/test_litellm/llms/bedrock/test_claude_platform_provider.py b/tests/unit/llms/bedrock/test_claude_platform_provider.py similarity index 100% rename from tests/test_litellm/llms/bedrock/test_claude_platform_provider.py rename to tests/unit/llms/bedrock/test_claude_platform_provider.py diff --git a/tests/test_litellm/llms/bedrock/test_converse_context_management.py b/tests/unit/llms/bedrock/test_converse_context_management.py similarity index 100% rename from tests/test_litellm/llms/bedrock/test_converse_context_management.py rename to tests/unit/llms/bedrock/test_converse_context_management.py diff --git a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py b/tests/unit/llms/bedrock/test_cross_region_inference_profile_mapping.py similarity index 100% rename from tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py rename to tests/unit/llms/bedrock/test_cross_region_inference_profile_mapping.py diff --git a/tests/test_litellm/llms/bedrock/test_mantle.py b/tests/unit/llms/bedrock/test_mantle.py similarity index 100% rename from tests/test_litellm/llms/bedrock/test_mantle.py rename to tests/unit/llms/bedrock/test_mantle.py diff --git a/tests/test_litellm/llms/bedrock/test_nova_imported_models.py b/tests/unit/llms/bedrock/test_nova_imported_models.py similarity index 100% rename from tests/test_litellm/llms/bedrock/test_nova_imported_models.py rename to tests/unit/llms/bedrock/test_nova_imported_models.py diff --git a/tests/test_litellm/llms/bedrock/test_request_metadata.py b/tests/unit/llms/bedrock/test_request_metadata.py similarity index 100% rename from tests/test_litellm/llms/bedrock/test_request_metadata.py rename to tests/unit/llms/bedrock/test_request_metadata.py diff --git a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py b/tests/unit/llms/bedrock/test_web_identity_session_policy.py similarity index 100% rename from tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py rename to tests/unit/llms/bedrock/test_web_identity_session_policy.py diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py rename to tests/unit/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py similarity index 100% rename from tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py rename to tests/unit/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py similarity index 99% rename from tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py rename to tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index 0cc3963358f..4bf3dd11fa1 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -19,7 +19,7 @@ import litellm from litellm.llms.bedrock_mantle.chat.transformation import BedrockMantleChatConfig from litellm.llms.bedrock.base_aws_llm import sign_request_off_loop_if_aws from litellm.types.utils import LlmProviders -from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe +from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe @pytest.fixture diff --git a/tests/test_litellm/llms/gradient_ai/chat/__init__.py b/tests/unit/llms/cometapi/__init__.py similarity index 100% rename from tests/test_litellm/llms/gradient_ai/chat/__init__.py rename to tests/unit/llms/cometapi/__init__.py diff --git a/tests/test_litellm/llms/groq/__init__.py b/tests/unit/llms/cometapi/chat/__init__.py similarity index 100% rename from tests/test_litellm/llms/groq/__init__.py rename to tests/unit/llms/cometapi/chat/__init__.py diff --git a/tests/unit/llms/cometapi/chat/test_cometapi_chat_transformation.py b/tests/unit/llms/cometapi/chat/test_cometapi_chat_transformation.py new file mode 100644 index 00000000000..607648dd6c9 --- /dev/null +++ b/tests/unit/llms/cometapi/chat/test_cometapi_chat_transformation.py @@ -0,0 +1,183 @@ +""" +Unit tests for CometAPI Chat Configuration + +Tests the CometAPIChatConfig class methods using mocks +""" + + +import pytest + + +from litellm.llms.cometapi.chat.transformation import ( + CometAPIChatCompletionStreamingHandler, + CometAPIConfig, +) +from litellm.llms.cometapi.common_utils import CometAPIException + + +class TestCometAPIChatCompletionStreamingHandler: + def test_chunk_parser_successful(self): + handler = CometAPIChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ) + + # Test input chunk + chunk = { + "id": "test_id", + "created": 1234567890, + "model": "gpt-3.5-turbo", + "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + "choices": [ + {"delta": {"content": "test content", "reasoning": "test reasoning"}} + ], + } + + # Parse chunk + result = handler.chunk_parser(chunk) + + # Verify response + assert result.id == "test_id" + assert result.object == "chat.completion.chunk" + assert result.created == 1234567890 + assert result.model == "gpt-3.5-turbo" + assert result.usage.prompt_tokens == chunk["usage"]["prompt_tokens"] + assert result.usage.completion_tokens == chunk["usage"]["completion_tokens"] + assert result.usage.total_tokens == chunk["usage"]["total_tokens"] + assert len(result.choices) == 1 + assert result.choices[0]["delta"]["reasoning_content"] == "test reasoning" + + def test_chunk_parser_error_response(self): + handler = CometAPIChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ) + + # Test error chunk + error_chunk = { + "error": { + "message": "test error", + "code": 400, + } + } + + # Verify error handling + with pytest.raises(CometAPIException) as exc_info: + handler.chunk_parser(error_chunk) + + assert "CometAPI Error: test error" in str(exc_info.value) + assert exc_info.value.status_code == 400 + + def test_chunk_parser_key_error(self): + handler = CometAPIChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ) + + # Test invalid chunk missing required fields + invalid_chunk = {"incomplete": "data"} + + # Verify KeyError handling + with pytest.raises(CometAPIException) as exc_info: + handler.chunk_parser(invalid_chunk) + + assert "KeyError" in str(exc_info.value) + assert exc_info.value.status_code == 400 + + +class TestCometAPIConfig: + def test_transform_request_basic(self): + """Test basic request transformation""" + config = CometAPIConfig() + + transformed_request = config.transform_request( + model="cometapi/gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello, world!"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert transformed_request["model"] == "cometapi/gpt-3.5-turbo" + assert transformed_request["messages"] == [ + {"role": "user", "content": "Hello, world!"} + ] + + def test_transform_request_with_extra_body(self): + """Test request transformation with extra_body parameters""" + config = CometAPIConfig() + + transformed_request = config.transform_request( + model="cometapi/gpt-4", + messages=[{"role": "user", "content": "Hello, world!"}], + optional_params={"extra_body": {"custom_param": "custom_value"}}, + litellm_params={}, + headers={}, + ) + + # Validate that extra_body parameters are merged into the request + assert transformed_request["custom_param"] == "custom_value" + assert transformed_request["messages"] == [ + {"role": "user", "content": "Hello, world!"} + ] + + def test_cache_control_flag_removal(self): + """Test cache control flag removal from messages""" + config = CometAPIConfig() + + transformed_request = config.transform_request( + model="cometapi/gpt-3.5-turbo", + messages=[ + { + "role": "user", + "content": "Hello, world!", + "cache_control": {"type": "ephemeral"}, + } + ], + optional_params={}, + litellm_params={}, + headers={}, + ) + + # CometAPI should remove cache_control flags by default + assert transformed_request["messages"][0].get("cache_control") is None + + def test_map_openai_params(self): + """Test OpenAI parameter mapping""" + config = CometAPIConfig() + + non_default_params = { + "temperature": 0.7, + "max_tokens": 100, + "top_p": 0.9, + } + + mapped_params = config.map_openai_params( + non_default_params=non_default_params, + optional_params={}, + model="cometapi/gpt-3.5-turbo", + drop_params=False, + ) + + assert mapped_params["temperature"] == 0.7 + assert mapped_params["max_tokens"] == 100 + assert mapped_params["top_p"] == 0.9 + + def test_get_error_class(self): + """Test error class creation""" + config = CometAPIConfig() + + error = config.get_error_class( + error_message="Test error", + status_code=400, + headers={"Content-Type": "application/json"}, + ) + + assert isinstance(error, CometAPIException) + assert error.message == "Test error" + assert error.status_code == 400 + + +# Integration test example (requires real API key) + + +if __name__ == "__main__": + # Quick test runner + pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/llms/groq/chat/__init__.py b/tests/unit/llms/compactifai/__init__.py similarity index 100% rename from tests/test_litellm/llms/groq/chat/__init__.py rename to tests/unit/llms/compactifai/__init__.py diff --git a/tests/test_litellm/llms/compactifai/test_compactifai.py b/tests/unit/llms/compactifai/test_compactifai.py similarity index 84% rename from tests/test_litellm/llms/compactifai/test_compactifai.py rename to tests/unit/llms/compactifai/test_compactifai.py index fd31049731a..1367c703fda 100644 --- a/tests/test_litellm/llms/compactifai/test_compactifai.py +++ b/tests/unit/llms/compactifai/test_compactifai.py @@ -104,56 +104,6 @@ def test_compactifai_completion_streaming(respx_mock): assert chunks[0].choices[0].delta.content == "Hello" -@pytest.mark.respx() -def test_compactifai_models_endpoint(respx_mock): - """Test CompactifAI models listing""" - litellm.disable_aiohttp_transport = True - - mock_response = { - "object": "list", - "data": [ - { - "id": "cai-llama-3-1-8b-slim", - "object": "model", - "created": 1677610602, - "owned_by": "compactifai", - }, - { - "id": "mistral-7b-compressed", - "object": "model", - "created": 1677610602, - "owned_by": "compactifai", - }, - ], - } - - respx_mock.post("https://api.compactif.ai/v1/chat/completions").respond( - json={ - "id": "chatcmpl-123", - "object": "chat.completion", - "created": 1677652288, - "model": "cai-llama-3-1-8b-slim", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": "Test response"}, - "finish_reason": "stop", - } - ], - "usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15}, - }, - status_code=200, - ) - - # This would be tested if litellm had a models() function - # For now, we'll test that the provider is properly configured - response = litellm.completion( - model="compactifai/cai-llama-3-1-8b-slim", - messages=[{"role": "user", "content": "test"}], - api_key="test-key", - ) - - @pytest.mark.respx() def test_compactifai_authentication_error(respx_mock): """Test CompactifAI authentication error handling""" diff --git a/tests/test_litellm/llms/huggingface/__init__.py b/tests/unit/llms/custom_httpx/__init__.py similarity index 100% rename from tests/test_litellm/llms/huggingface/__init__.py rename to tests/unit/llms/custom_httpx/__init__.py diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py b/tests/unit/llms/custom_httpx/test_aiohttp_cleanup_closed.py similarity index 100% rename from tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py rename to tests/unit/llms/custom_httpx/test_aiohttp_cleanup_closed.py diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_handler.py b/tests/unit/llms/custom_httpx/test_aiohttp_handler.py similarity index 100% rename from tests/test_litellm/llms/custom_httpx/test_aiohttp_handler.py rename to tests/unit/llms/custom_httpx/test_aiohttp_handler.py diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py b/tests/unit/llms/custom_httpx/test_aiohttp_so_keepalive.py similarity index 100% rename from tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py rename to tests/unit/llms/custom_httpx/test_aiohttp_so_keepalive.py diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py b/tests/unit/llms/custom_httpx/test_aiohttp_transport.py similarity index 100% rename from tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py rename to tests/unit/llms/custom_httpx/test_aiohttp_transport.py diff --git a/tests/test_litellm/llms/custom_httpx/test_asgi_handler.py b/tests/unit/llms/custom_httpx/test_asgi_handler.py similarity index 100% rename from tests/test_litellm/llms/custom_httpx/test_asgi_handler.py rename to tests/unit/llms/custom_httpx/test_asgi_handler.py diff --git a/tests/test_litellm/llms/custom_httpx/test_async_client_cleanup.py b/tests/unit/llms/custom_httpx/test_async_client_cleanup.py similarity index 100% rename from tests/test_litellm/llms/custom_httpx/test_async_client_cleanup.py rename to tests/unit/llms/custom_httpx/test_async_client_cleanup.py diff --git a/tests/test_litellm/llms/custom_httpx/test_container_handler.py b/tests/unit/llms/custom_httpx/test_container_handler.py similarity index 100% rename from tests/test_litellm/llms/custom_httpx/test_container_handler.py rename to tests/unit/llms/custom_httpx/test_container_handler.py diff --git a/tests/test_litellm/llms/custom_httpx/test_credential_leak_prevention.py b/tests/unit/llms/custom_httpx/test_credential_leak_prevention.py similarity index 100% rename from tests/test_litellm/llms/custom_httpx/test_credential_leak_prevention.py rename to tests/unit/llms/custom_httpx/test_credential_leak_prevention.py diff --git a/tests/test_litellm/llms/custom_httpx/test_gemini_session_leak.py b/tests/unit/llms/custom_httpx/test_gemini_session_leak.py similarity index 100% rename from tests/test_litellm/llms/custom_httpx/test_gemini_session_leak.py rename to tests/unit/llms/custom_httpx/test_gemini_session_leak.py diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py similarity index 100% rename from tests/test_litellm/llms/custom_httpx/test_http_handler.py rename to tests/unit/llms/custom_httpx/test_http_handler.py diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py similarity index 99% rename from tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py rename to tests/unit/llms/custom_httpx/test_llm_http_handler.py index 0350ca74904..399e4dbf206 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -45,7 +45,7 @@ from litellm.llms.tinyfish.search.transformation import TinyfishSearchConfig from litellm.types.llms.openai import HttpxBinaryResponseContent, ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, TranscriptionResponse -from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe +from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe _ACTIVE_KEY = "_code_interpreter_interception_active" _SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" diff --git a/tests/test_litellm/llms/custom_httpx/test_mock_transport.py b/tests/unit/llms/custom_httpx/test_mock_transport.py similarity index 100% rename from tests/test_litellm/llms/custom_httpx/test_mock_transport.py rename to tests/unit/llms/custom_httpx/test_mock_transport.py diff --git a/tests/test_litellm/llms/inception/__init__.py b/tests/unit/llms/dashscope/__init__.py similarity index 100% rename from tests/test_litellm/llms/inception/__init__.py rename to tests/unit/llms/dashscope/__init__.py diff --git a/tests/test_litellm/llms/dashscope/test_dashscope_chat_transformation.py b/tests/unit/llms/dashscope/test_dashscope_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/dashscope/test_dashscope_chat_transformation.py rename to tests/unit/llms/dashscope/test_dashscope_chat_transformation.py diff --git a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py b/tests/unit/llms/dashscope/test_dashscope_cost_calculator.py similarity index 100% rename from tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py rename to tests/unit/llms/dashscope/test_dashscope_cost_calculator.py diff --git a/tests/test_litellm/llms/dashscope/test_dashscope_embedding_transformation.py b/tests/unit/llms/dashscope/test_dashscope_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/dashscope/test_dashscope_embedding_transformation.py rename to tests/unit/llms/dashscope/test_dashscope_embedding_transformation.py diff --git a/tests/test_litellm/llms/dashscope/test_dashscope_rerank_transformation.py b/tests/unit/llms/dashscope/test_dashscope_rerank_transformation.py similarity index 100% rename from tests/test_litellm/llms/dashscope/test_dashscope_rerank_transformation.py rename to tests/unit/llms/dashscope/test_dashscope_rerank_transformation.py diff --git a/tests/test_litellm/llms/dashscope/test_qwen_brand_aliases.py b/tests/unit/llms/dashscope/test_qwen_brand_aliases.py similarity index 100% rename from tests/test_litellm/llms/dashscope/test_qwen_brand_aliases.py rename to tests/unit/llms/dashscope/test_qwen_brand_aliases.py diff --git a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py index 52bb89fed5a..9cd17bd3580 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -16,6 +16,9 @@ from litellm.llms.databricks.chat.transformation import ( DatabricksConfig, _sanitize_empty_content, ) +from typing import Final +import httpx +import respx @pytest.fixture() @@ -808,3 +811,75 @@ def test_chunk_parser_surfaces_top_level_reasoning_delta(reasoning_key: str) -> assert parsed.choices[0].delta.reasoning_content == "We need answer" assert parsed.choices[0].delta.content is None + + +def test_completion_merges_leading_system_and_developer_messages_for_chat_template_models( + respx_mock: respx.MockRouter, +): + upstream: Final = respx_mock.post("https://example.databricks.test/serving-endpoints/chat/completions").mock( + return_value=httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "my-custom-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "Answer"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 9, "completion_tokens": 1, "total_tokens": 10}, + }, + ) + ) + + response: Final = litellm.completion( + model="databricks/my-custom-model", + messages=[ + {"role": "system", "content": "You are terse."}, + {"role": "developer", "content": "Skills: none."}, + {"role": "user", "content": "Hello"}, + ], + api_base="https://example.databricks.test/serving-endpoints", + api_key="fake-databricks-api-key", + num_retries=0, + ) + + assert upstream.call_count == 1 + request_body: Final = json.loads(upstream.calls[0].request.read()) + assert request_body["messages"] == [ + {"role": "system", "content": "You are terse.\n\nSkills: none."}, + {"role": "user", "content": "Hello"}, + ] + assert response.choices[0].message.content == "Answer" + + +def test_completion_merges_system_messages_when_one_has_empty_content(respx_mock: respx.MockRouter): + upstream: Final = respx_mock.post("https://example.databricks.test/serving-endpoints/chat/completions").mock( + return_value=httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "my-custom-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "Answer"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 9, "completion_tokens": 1, "total_tokens": 10}, + }, + ) + ) + + litellm.completion( + model="databricks/my-custom-model", + messages=[ + {"role": "system", "content": "You are terse."}, + {"role": "system", "content": ""}, + {"role": "user", "content": "Hello"}, + ], + api_base="https://example.databricks.test/serving-endpoints", + api_key="fake-databricks-api-key", + num_retries=0, + ) + + request_body: Final = json.loads(upstream.calls[0].request.read()) + assert request_body["messages"] == [ + {"role": "system", "content": "You are terse."}, + {"role": "user", "content": "Hello"}, + ] diff --git a/tests/test_litellm/llms/databricks/test_databricks_common_utils.py b/tests/unit/llms/databricks/test_databricks_common_utils.py similarity index 100% rename from tests/test_litellm/llms/databricks/test_databricks_common_utils.py rename to tests/unit/llms/databricks/test_databricks_common_utils.py diff --git a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py b/tests/unit/llms/databricks/test_databricks_cost_calculator.py similarity index 100% rename from tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py rename to tests/unit/llms/databricks/test_databricks_cost_calculator.py diff --git a/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py b/tests/unit/llms/databricks/test_databricks_partner_integration.py similarity index 100% rename from tests/test_litellm/llms/databricks/test_databricks_partner_integration.py rename to tests/unit/llms/databricks/test_databricks_partner_integration.py diff --git a/tests/test_litellm/llms/databricks/test_databricks_streaming_utils.py b/tests/unit/llms/databricks/test_databricks_streaming_utils.py similarity index 100% rename from tests/test_litellm/llms/databricks/test_databricks_streaming_utils.py rename to tests/unit/llms/databricks/test_databricks_streaming_utils.py diff --git a/tests/test_litellm/llms/mistral/batches/__init__.py b/tests/unit/llms/deepgram/__init__.py similarity index 100% rename from tests/test_litellm/llms/mistral/batches/__init__.py rename to tests/unit/llms/deepgram/__init__.py diff --git a/tests/test_litellm/llms/mistral/files/__init__.py b/tests/unit/llms/deepgram/audio_transcription/__init__.py similarity index 100% rename from tests/test_litellm/llms/mistral/files/__init__.py rename to tests/unit/llms/deepgram/audio_transcription/__init__.py diff --git a/tests/test_litellm/llms/deepgram/audio_transcription/test_deepgram_audio_transcription_transformation.py b/tests/unit/llms/deepgram/audio_transcription/test_deepgram_audio_transcription_transformation.py similarity index 100% rename from tests/test_litellm/llms/deepgram/audio_transcription/test_deepgram_audio_transcription_transformation.py rename to tests/unit/llms/deepgram/audio_transcription/test_deepgram_audio_transcription_transformation.py diff --git a/tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py b/tests/unit/llms/deepgram/test_deepgram_common_utils.py similarity index 100% rename from tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py rename to tests/unit/llms/deepgram/test_deepgram_common_utils.py diff --git a/tests/test_litellm/llms/deepgram/test_deepgram_mock_transcription.py b/tests/unit/llms/deepgram/test_deepgram_mock_transcription.py similarity index 100% rename from tests/test_litellm/llms/deepgram/test_deepgram_mock_transcription.py rename to tests/unit/llms/deepgram/test_deepgram_mock_transcription.py diff --git a/tests/test_litellm/llms/nvidia_riva/__init__.py b/tests/unit/llms/deepinfra/__init__.py similarity index 100% rename from tests/test_litellm/llms/nvidia_riva/__init__.py rename to tests/unit/llms/deepinfra/__init__.py diff --git a/tests/test_litellm/llms/deepinfra/test_deepinfra_chat_transformation.py b/tests/unit/llms/deepinfra/test_deepinfra_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/deepinfra/test_deepinfra_chat_transformation.py rename to tests/unit/llms/deepinfra/test_deepinfra_chat_transformation.py diff --git a/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank.py b/tests/unit/llms/deepinfra/test_deepinfra_rerank.py similarity index 100% rename from tests/test_litellm/llms/deepinfra/test_deepinfra_rerank.py rename to tests/unit/llms/deepinfra/test_deepinfra_rerank.py diff --git a/tests/unit/llms/deepinfra/test_deepinfra_rerank_integration.py b/tests/unit/llms/deepinfra/test_deepinfra_rerank_integration.py new file mode 100644 index 00000000000..8a2a1d09cb6 --- /dev/null +++ b/tests/unit/llms/deepinfra/test_deepinfra_rerank_integration.py @@ -0,0 +1,159 @@ +""" +Integration tests for DeepInfra rerank functionality. +Tests the full rerank flow following the repository patterns. +""" + +import asyncio +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import litellm + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post") +@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") +def test_deepinfra_rerank_with_queries_param( + mock_sync_post, mock_async_post, sync_mode +): + """Test DeepInfra rerank with multiple queries parameter.""" + mock_response_data = { + "scores": [0.8, 0.6, 0.2], + "input_tokens": 35, + "request_id": "deepinfra-multi-query-123", + "inference_status": {"status": "success", "runtime_ms": 200}, + } + + def return_val(): + return mock_response_data + + if sync_mode: + mock_response = MagicMock() + mock_response.json = return_val + mock_response.status_code = 200 + mock_response.headers = {"content-type": "application/json"} + mock_response.text = json.dumps(mock_response_data) + mock_sync_post.return_value = mock_response + + response = litellm.rerank( + model="deepinfra/Qwen/Qwen3-Reranker-4B", + query="hello", + documents=["hello", "world", "test"], + queries=["hello", "hi there"], # DeepInfra specific param + custom_llm_provider="deepinfra", + api_key="test_key", + api_base="https://api.deepinfra.com", + ) + + mock_sync_post.assert_called_once() + # Verify that queries parameter was passed in request + call_data = json.loads(mock_sync_post.call_args.kwargs["data"]) + assert "queries" in call_data + assert call_data["queries"] == ["hello", "hi there"] + else: + mock_response = AsyncMock() + mock_response.json = return_val + mock_response.status_code = 200 + mock_response.headers = {"content-type": "application/json"} + mock_response.text = json.dumps(mock_response_data) + mock_async_post.return_value = mock_response + + response = asyncio.run( + litellm.arerank( + model="deepinfra/Qwen/Qwen3-Reranker-4B", + query="hello", + documents=["hello", "world", "test"], + queries=["hello", "hi there"], + custom_llm_provider="deepinfra", + api_key="test_key", + api_base="https://api.deepinfra.com", + ) + ) + + mock_async_post.assert_called_once() + call_data = json.loads(mock_async_post.call_args.kwargs["data"]) + assert "queries" in call_data + assert call_data["queries"] == ["hello", "hi there"] + + assert response.results is not None + assert len(response.results) == 3 + + +@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") +def test_deepinfra_rerank_with_env_vars(mock_post, monkeypatch): + """Test DeepInfra rerank with environment variable configuration.""" + monkeypatch.setenv("DEEPINFRA_API_KEY", "env_test_key") + monkeypatch.setenv("DEEPINFRA_API_BASE", "https://custom-deepinfra.com") + + mock_response_data = { + "scores": [0.88, 0.22], + "input_tokens": 28, + "request_id": "env-test-123", + } + + def return_val(): + return mock_response_data + + mock_response = MagicMock() + mock_response.json = return_val + mock_response.status_code = 200 + mock_response.headers = {"content-type": "application/json"} + mock_response.text = json.dumps(mock_response_data) + mock_post.return_value = mock_response + + response = litellm.rerank( + model="deepinfra/Qwen/Qwen3-Reranker-0.6B", + query="hello", + documents=["hello", "world"], + custom_llm_provider="deepinfra", + ) + + mock_post.assert_called_once() + + # Verify headers contain env API key + headers = mock_post.call_args.kwargs.get("headers", {}) + assert "Bearer env_test_key" in headers.get("Authorization", "") + + assert response.results is not None + + +@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") +def test_deepinfra_rerank_defaults_api_base_when_missing(mock_post, monkeypatch): + """With no api_base anywhere, the call still goes out against DeepInfra's own base.""" + monkeypatch.delenv("DEEPINFRA_API_BASE", raising=False) + + mock_response = MagicMock() + mock_response.json = lambda: {"scores": [0.9, 0.1], "input_tokens": 20} + mock_response.status_code = 200 + mock_response.headers = {"content-type": "application/json"} + mock_post.return_value = mock_response + + response = litellm.rerank( + model="deepinfra/Qwen/Qwen3-Reranker-0.6B", + query="hello", + documents=["hello", "world"], + custom_llm_provider="deepinfra", + api_key="test_key", + # api_base is intentionally missing + ) + + assert "api.deepinfra.com" in mock_post.call_args.kwargs["url"] + assert [result["relevance_score"] for result in response.results] == [0.9, 0.1] + + +def test_deepinfra_rerank_models(): + """Test that DeepInfra Qwen rerank models are recognized.""" + # These should not raise errors during model validation + models = [ + "deepinfra/Qwen/Qwen3-Reranker-0.6B", + "deepinfra/Qwen/Qwen3-Reranker-4B", + "deepinfra/Qwen/Qwen3-Reranker-8B", + ] + + for model in models: + resolved_model, provider, _, api_base = litellm.get_llm_provider(model=model) + assert provider == "deepinfra" + assert resolved_model == model.removeprefix("deepinfra/") + assert api_base == "https://api.deepinfra.com/v1/openai" diff --git a/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py b/tests/unit/llms/deepinfra/test_deepinfra_rerank_transformation.py similarity index 100% rename from tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py rename to tests/unit/llms/deepinfra/test_deepinfra_rerank_transformation.py diff --git a/tests/test_litellm/llms/oci/rerank/__init__.py b/tests/unit/llms/edenai/__init__.py similarity index 100% rename from tests/test_litellm/llms/oci/rerank/__init__.py rename to tests/unit/llms/edenai/__init__.py diff --git a/tests/test_litellm/llms/ocr/__init__.py b/tests/unit/llms/edenai/audio_transcription/__init__.py similarity index 100% rename from tests/test_litellm/llms/ocr/__init__.py rename to tests/unit/llms/edenai/audio_transcription/__init__.py diff --git a/tests/test_litellm/llms/edenai/audio_transcription/test_edenai_audio_transcription_transformation.py b/tests/unit/llms/edenai/audio_transcription/test_edenai_audio_transcription_transformation.py similarity index 100% rename from tests/test_litellm/llms/edenai/audio_transcription/test_edenai_audio_transcription_transformation.py rename to tests/unit/llms/edenai/audio_transcription/test_edenai_audio_transcription_transformation.py diff --git a/tests/test_litellm/llms/openai_like/responses/__init__.py b/tests/unit/llms/edenai/chat/__init__.py similarity index 100% rename from tests/test_litellm/llms/openai_like/responses/__init__.py rename to tests/unit/llms/edenai/chat/__init__.py diff --git a/tests/test_litellm/llms/edenai/chat/test_edenai_chat_transformation.py b/tests/unit/llms/edenai/chat/test_edenai_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/edenai/chat/test_edenai_chat_transformation.py rename to tests/unit/llms/edenai/chat/test_edenai_chat_transformation.py diff --git a/tests/test_litellm/llms/edenai/conftest.py b/tests/unit/llms/edenai/conftest.py similarity index 100% rename from tests/test_litellm/llms/edenai/conftest.py rename to tests/unit/llms/edenai/conftest.py diff --git a/tests/test_litellm/llms/parallel_ai/__init__.py b/tests/unit/llms/edenai/embedding/__init__.py similarity index 100% rename from tests/test_litellm/llms/parallel_ai/__init__.py rename to tests/unit/llms/edenai/embedding/__init__.py diff --git a/tests/test_litellm/llms/edenai/embedding/test_edenai_embedding_transformation.py b/tests/unit/llms/edenai/embedding/test_edenai_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/edenai/embedding/test_edenai_embedding_transformation.py rename to tests/unit/llms/edenai/embedding/test_edenai_embedding_transformation.py diff --git a/tests/test_litellm/llms/pass_through/__init__.py b/tests/unit/llms/edenai/image_generation/__init__.py similarity index 100% rename from tests/test_litellm/llms/pass_through/__init__.py rename to tests/unit/llms/edenai/image_generation/__init__.py diff --git a/tests/test_litellm/llms/edenai/image_generation/test_edenai_image_generation_transformation.py b/tests/unit/llms/edenai/image_generation/test_edenai_image_generation_transformation.py similarity index 100% rename from tests/test_litellm/llms/edenai/image_generation/test_edenai_image_generation_transformation.py rename to tests/unit/llms/edenai/image_generation/test_edenai_image_generation_transformation.py diff --git a/tests/test_litellm/llms/pass_through/guardrail_translation/__init__.py b/tests/unit/llms/edenai/messages/__init__.py similarity index 100% rename from tests/test_litellm/llms/pass_through/guardrail_translation/__init__.py rename to tests/unit/llms/edenai/messages/__init__.py diff --git a/tests/test_litellm/llms/edenai/messages/test_edenai_anthropic_messages_transformation.py b/tests/unit/llms/edenai/messages/test_edenai_anthropic_messages_transformation.py similarity index 100% rename from tests/test_litellm/llms/edenai/messages/test_edenai_anthropic_messages_transformation.py rename to tests/unit/llms/edenai/messages/test_edenai_anthropic_messages_transformation.py diff --git a/tests/test_litellm/llms/perplexity/__init__.py b/tests/unit/llms/edenai/responses/__init__.py similarity index 100% rename from tests/test_litellm/llms/perplexity/__init__.py rename to tests/unit/llms/edenai/responses/__init__.py diff --git a/tests/test_litellm/llms/edenai/responses/test_edenai_responses_transformation.py b/tests/unit/llms/edenai/responses/test_edenai_responses_transformation.py similarity index 100% rename from tests/test_litellm/llms/edenai/responses/test_edenai_responses_transformation.py rename to tests/unit/llms/edenai/responses/test_edenai_responses_transformation.py diff --git a/tests/test_litellm/llms/edenai/test_edenai_common_utils.py b/tests/unit/llms/edenai/test_edenai_common_utils.py similarity index 100% rename from tests/test_litellm/llms/edenai/test_edenai_common_utils.py rename to tests/unit/llms/edenai/test_edenai_common_utils.py diff --git a/tests/test_litellm/llms/perplexity/embedding/__init__.py b/tests/unit/llms/edenai/text_to_speech/__init__.py similarity index 100% rename from tests/test_litellm/llms/perplexity/embedding/__init__.py rename to tests/unit/llms/edenai/text_to_speech/__init__.py diff --git a/tests/test_litellm/llms/edenai/text_to_speech/test_edenai_text_to_speech_transformation.py b/tests/unit/llms/edenai/text_to_speech/test_edenai_text_to_speech_transformation.py similarity index 100% rename from tests/test_litellm/llms/edenai/text_to_speech/test_edenai_text_to_speech_transformation.py rename to tests/unit/llms/edenai/text_to_speech/test_edenai_text_to_speech_transformation.py diff --git a/tests/test_litellm/llms/stability/__init__.py b/tests/unit/llms/edenai/videos/__init__.py similarity index 100% rename from tests/test_litellm/llms/stability/__init__.py rename to tests/unit/llms/edenai/videos/__init__.py diff --git a/tests/test_litellm/llms/edenai/videos/test_edenai_video_transformation.py b/tests/unit/llms/edenai/videos/test_edenai_video_transformation.py similarity index 100% rename from tests/test_litellm/llms/edenai/videos/test_edenai_video_transformation.py rename to tests/unit/llms/edenai/videos/test_edenai_video_transformation.py diff --git a/tests/test_litellm/llms/stability/image_generation/__init__.py b/tests/unit/llms/fal_ai/__init__.py similarity index 100% rename from tests/test_litellm/llms/stability/image_generation/__init__.py rename to tests/unit/llms/fal_ai/__init__.py diff --git a/tests/test_litellm/llms/tencent/__init__.py b/tests/unit/llms/fal_ai/chat/__init__.py similarity index 100% rename from tests/test_litellm/llms/tencent/__init__.py rename to tests/unit/llms/fal_ai/chat/__init__.py diff --git a/tests/test_litellm/llms/fal_ai/chat/test_fal_ai_chat_transformation.py b/tests/unit/llms/fal_ai/chat/test_fal_ai_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/fal_ai/chat/test_fal_ai_chat_transformation.py rename to tests/unit/llms/fal_ai/chat/test_fal_ai_chat_transformation.py diff --git a/tests/test_litellm/llms/tencent/chat/__init__.py b/tests/unit/llms/fal_ai/image_edit/__init__.py similarity index 100% rename from tests/test_litellm/llms/tencent/chat/__init__.py rename to tests/unit/llms/fal_ai/image_edit/__init__.py diff --git a/tests/test_litellm/llms/fal_ai/image_edit/test_fal_ai_flux_lora_depth_transformation.py b/tests/unit/llms/fal_ai/image_edit/test_fal_ai_flux_lora_depth_transformation.py similarity index 100% rename from tests/test_litellm/llms/fal_ai/image_edit/test_fal_ai_flux_lora_depth_transformation.py rename to tests/unit/llms/fal_ai/image_edit/test_fal_ai_flux_lora_depth_transformation.py diff --git a/tests/test_litellm/llms/fal_ai/image_edit/test_fal_ai_image_edit_transformation.py b/tests/unit/llms/fal_ai/image_edit/test_fal_ai_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/fal_ai/image_edit/test_fal_ai_image_edit_transformation.py rename to tests/unit/llms/fal_ai/image_edit/test_fal_ai_image_edit_transformation.py diff --git a/tests/test_litellm/llms/tencent/messages/__init__.py b/tests/unit/llms/fal_ai/image_generation/__init__.py similarity index 100% rename from tests/test_litellm/llms/tencent/messages/__init__.py rename to tests/unit/llms/fal_ai/image_generation/__init__.py diff --git a/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_flux_dev_transformation.py b/tests/unit/llms/fal_ai/image_generation/test_fal_ai_flux_dev_transformation.py similarity index 100% rename from tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_flux_dev_transformation.py rename to tests/unit/llms/fal_ai/image_generation/test_fal_ai_flux_dev_transformation.py diff --git a/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_gpt_image_2_transformation.py b/tests/unit/llms/fal_ai/image_generation/test_fal_ai_gpt_image_2_transformation.py similarity index 100% rename from tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_gpt_image_2_transformation.py rename to tests/unit/llms/fal_ai/image_generation/test_fal_ai_gpt_image_2_transformation.py diff --git a/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_nano_banana_transformation.py b/tests/unit/llms/fal_ai/image_generation/test_fal_ai_nano_banana_transformation.py similarity index 100% rename from tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_nano_banana_transformation.py rename to tests/unit/llms/fal_ai/image_generation/test_fal_ai_nano_banana_transformation.py diff --git a/tests/test_litellm/llms/fal_ai/test_cost_calculator.py b/tests/unit/llms/fal_ai/test_cost_calculator.py similarity index 100% rename from tests/test_litellm/llms/fal_ai/test_cost_calculator.py rename to tests/unit/llms/fal_ai/test_cost_calculator.py diff --git a/tests/test_litellm/llms/vercel_ai_gateway/embedding/__init__.py b/tests/unit/llms/fal_ai/videos/__init__.py similarity index 100% rename from tests/test_litellm/llms/vercel_ai_gateway/embedding/__init__.py rename to tests/unit/llms/fal_ai/videos/__init__.py diff --git a/tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py b/tests/unit/llms/fal_ai/videos/test_fal_ai_video_transformation.py similarity index 100% rename from tests/test_litellm/llms/fal_ai/videos/test_fal_ai_video_transformation.py rename to tests/unit/llms/fal_ai/videos/test_fal_ai_video_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/agent_engine/__init__.py b/tests/unit/llms/featherless_ai/__init__.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/agent_engine/__init__.py rename to tests/unit/llms/featherless_ai/__init__.py diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/__init__.py b/tests/unit/llms/featherless_ai/chat/__init__.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/audio_transcription/__init__.py rename to tests/unit/llms/featherless_ai/chat/__init__.py diff --git a/tests/test_litellm/llms/featherless_ai/chat/test_featherless_chat_transformation.py b/tests/unit/llms/featherless_ai/chat/test_featherless_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/featherless_ai/chat/test_featherless_chat_transformation.py rename to tests/unit/llms/featherless_ai/chat/test_featherless_chat_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/batches/__init__.py b/tests/unit/llms/fireworks_ai/completion/__init__.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/batches/__init__.py rename to tests/unit/llms/fireworks_ai/completion/__init__.py diff --git a/tests/test_litellm/llms/fireworks_ai/completion/test_fireworks_ai_completion_transformation.py b/tests/unit/llms/fireworks_ai/completion/test_fireworks_ai_completion_transformation.py similarity index 100% rename from tests/test_litellm/llms/fireworks_ai/completion/test_fireworks_ai_completion_transformation.py rename to tests/unit/llms/fireworks_ai/completion/test_fireworks_ai_completion_transformation.py diff --git a/tests/test_litellm/llms/fireworks_ai/completion/test_fireworks_ai_text_completion_transformation.py b/tests/unit/llms/fireworks_ai/completion/test_fireworks_ai_text_completion_transformation.py similarity index 100% rename from tests/test_litellm/llms/fireworks_ai/completion/test_fireworks_ai_text_completion_transformation.py rename to tests/unit/llms/fireworks_ai/completion/test_fireworks_ai_text_completion_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/files/__init__.py b/tests/unit/llms/gdc/__init__.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/files/__init__.py rename to tests/unit/llms/gdc/__init__.py diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/__init__.py b/tests/unit/llms/gdc/chat/__init__.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/gemini_embeddings/__init__.py rename to tests/unit/llms/gdc/chat/__init__.py diff --git a/tests/test_litellm/llms/gdc/chat/test_gdc_chat_transformation.py b/tests/unit/llms/gdc/chat/test_gdc_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/gdc/chat/test_gdc_chat_transformation.py rename to tests/unit/llms/gdc/chat/test_gdc_chat_transformation.py diff --git a/tests/test_litellm/llms/gemini/test_cost_calculator.py b/tests/unit/llms/gemini/test_cost_calculator.py similarity index 100% rename from tests/test_litellm/llms/gemini/test_cost_calculator.py rename to tests/unit/llms/gemini/test_cost_calculator.py diff --git a/tests/test_litellm/llms/gemini/test_gemini_client_setup.py b/tests/unit/llms/gemini/test_gemini_client_setup.py similarity index 100% rename from tests/test_litellm/llms/gemini/test_gemini_client_setup.py rename to tests/unit/llms/gemini/test_gemini_client_setup.py diff --git a/tests/test_litellm/llms/gemini/test_gemini_common_utils.py b/tests/unit/llms/gemini/test_gemini_common_utils.py similarity index 100% rename from tests/test_litellm/llms/gemini/test_gemini_common_utils.py rename to tests/unit/llms/gemini/test_gemini_common_utils.py diff --git a/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py b/tests/unit/llms/gemini/test_gemini_image_generation_transformation.py similarity index 100% rename from tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py rename to tests/unit/llms/gemini/test_gemini_image_generation_transformation.py diff --git a/tests/test_litellm/llms/gemini/test_gemini_tts.py b/tests/unit/llms/gemini/test_gemini_tts.py similarity index 100% rename from tests/test_litellm/llms/gemini/test_gemini_tts.py rename to tests/unit/llms/gemini/test_gemini_tts.py diff --git a/tests/test_litellm/llms/github_copilot/test_github_copilot_authenticator.py b/tests/unit/llms/github_copilot/test_github_copilot_authenticator.py similarity index 100% rename from tests/test_litellm/llms/github_copilot/test_github_copilot_authenticator.py rename to tests/unit/llms/github_copilot/test_github_copilot_authenticator.py diff --git a/tests/test_litellm/llms/github_copilot/test_github_copilot_transformation.py b/tests/unit/llms/github_copilot/test_github_copilot_transformation.py similarity index 100% rename from tests/test_litellm/llms/github_copilot/test_github_copilot_transformation.py rename to tests/unit/llms/github_copilot/test_github_copilot_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/text_to_speech/__init__.py b/tests/unit/llms/heroku/__init__.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/text_to_speech/__init__.py rename to tests/unit/llms/heroku/__init__.py diff --git a/tests/test_litellm/llms/heroku/test_heroku_chat_transformation.py b/tests/unit/llms/heroku/test_heroku_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/heroku/test_heroku_chat_transformation.py rename to tests/unit/llms/heroku/test_heroku_chat_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/__init__.py b/tests/unit/llms/huggingface/embedding/__init__.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/__init__.py rename to tests/unit/llms/huggingface/embedding/__init__.py diff --git a/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py b/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py similarity index 100% rename from tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py rename to tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py diff --git a/tests/test_litellm/llms/langflow/test_langflow_a2a.py b/tests/unit/llms/langflow/test_langflow_a2a.py similarity index 100% rename from tests/test_litellm/llms/langflow/test_langflow_a2a.py rename to tests/unit/llms/langflow/test_langflow_a2a.py diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/__init__.py b/tests/unit/llms/lemonade/__init__.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/__init__.py rename to tests/unit/llms/lemonade/__init__.py diff --git a/tests/test_litellm/llms/lemonade/test_lemonade.py b/tests/unit/llms/lemonade/test_lemonade.py similarity index 100% rename from tests/test_litellm/llms/lemonade/test_lemonade.py rename to tests/unit/llms/lemonade/test_lemonade.py diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/__init__.py b/tests/unit/llms/lm_studio/__init__.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/__init__.py rename to tests/unit/llms/lm_studio/__init__.py diff --git a/tests/test_litellm/llms/lm_studio/test_lm_studio_chat_transformation.py b/tests/unit/llms/lm_studio/test_lm_studio_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/lm_studio/test_lm_studio_chat_transformation.py rename to tests/unit/llms/lm_studio/test_lm_studio_chat_transformation.py diff --git a/tests/test_litellm/llms/voyage/rerank/__init__.py b/tests/unit/llms/mistral/audio_transcription/__init__.py similarity index 100% rename from tests/test_litellm/llms/voyage/rerank/__init__.py rename to tests/unit/llms/mistral/audio_transcription/__init__.py diff --git a/tests/unit/llms/mistral/audio_transcription/test_mistral_audio_transcription_transformation.py b/tests/unit/llms/mistral/audio_transcription/test_mistral_audio_transcription_transformation.py new file mode 100644 index 00000000000..68875ff6d32 --- /dev/null +++ b/tests/unit/llms/mistral/audio_transcription/test_mistral_audio_transcription_transformation.py @@ -0,0 +1,195 @@ +import os +from unittest.mock import MagicMock + +import httpx +import litellm + +from litellm.llms.base_llm.audio_transcription.transformation import ( + BaseAudioTranscriptionConfig, +) +from litellm.llms.mistral.audio_transcription.transformation import ( + MistralAudioTranscriptionConfig, +) +from litellm.types.utils import TranscriptionResponse +from litellm.utils import ProviderConfigManager + + +def test_mistral_audio_transcription_config_installed(): + """Ensure Mistral audio transcription config is registered with ProviderConfigManager.""" + config = ProviderConfigManager.get_provider_audio_transcription_config( + model="mistral/voxtral-mini-latest", + provider=litellm.LlmProviders.MISTRAL, + ) + assert config is not None + assert isinstance(config, BaseAudioTranscriptionConfig) + assert isinstance(config, MistralAudioTranscriptionConfig) + + +def test_mistral_audio_transcription_get_complete_url(): + config = MistralAudioTranscriptionConfig() + url = config.get_complete_url( + api_base=None, + api_key="fake-key", + model="voxtral-mini-latest", + optional_params={}, + litellm_params={}, + ) + assert url == "https://api.mistral.ai/v1/audio/transcriptions" + + +def test_mistral_audio_transcription_get_complete_url_custom_base(): + config = MistralAudioTranscriptionConfig() + url = config.get_complete_url( + api_base="https://custom.api.example.com/v1/", + api_key="fake-key", + model="voxtral-mini-latest", + optional_params={}, + litellm_params={}, + ) + assert url == "https://custom.api.example.com/v1/audio/transcriptions" + + +def test_mistral_audio_transcription_validate_environment(): + config = MistralAudioTranscriptionConfig() + headers = config.validate_environment( + headers={}, + model="voxtral-mini-latest", + messages=[], + optional_params={}, + litellm_params={}, + api_key="test-key-123", + ) + assert headers["Authorization"] == "Bearer test-key-123" + assert headers["accept"] == "application/json" + + +def test_mistral_audio_transcription_supported_params(): + config = MistralAudioTranscriptionConfig() + params = config.get_supported_openai_params("voxtral-mini-latest") + assert "language" in params + assert "temperature" in params + assert "response_format" in params + assert "timestamp_granularities" in params + + +def test_mistral_audio_transcription_request_transform(): + config = MistralAudioTranscriptionConfig() + + wav_path = os.path.join( + os.path.dirname(__file__), + "../../../../..", + "tests", + "llm_translation", + "gettysburg.wav", + ) + audio_file = open(wav_path, "rb") + + result = config.transform_audio_transcription_request( + model="voxtral-mini-latest", + audio_file=audio_file, + optional_params={"language": "en", "temperature": 0.0}, + litellm_params={}, + ) + + audio_file.close() + + assert isinstance(result.data, dict) + assert result.data["model"] == "voxtral-mini-latest" + assert result.data["language"] == "en" + assert result.data["temperature"] == 0.0 + assert result.files is not None + assert "file" in result.files + + +def test_mistral_audio_transcription_request_with_diarize(): + """Test that Mistral-specific params like diarize are passed through.""" + config = MistralAudioTranscriptionConfig() + + wav_path = os.path.join( + os.path.dirname(__file__), + "../../../../..", + "tests", + "llm_translation", + "gettysburg.wav", + ) + audio_file = open(wav_path, "rb") + + result = config.transform_audio_transcription_request( + model="voxtral-mini-latest", + audio_file=audio_file, + optional_params={"diarize": True}, + litellm_params={}, + ) + + audio_file.close() + + assert isinstance(result.data, dict) + assert result.data["diarize"] == "true" + + +def test_mistral_audio_transcription_response_transform(): + config = MistralAudioTranscriptionConfig() + + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.return_value = {"text": "Four score and seven years ago..."} + + response = config.transform_audio_transcription_response(mock_response) + + assert isinstance(response, TranscriptionResponse) + assert response.text == "Four score and seven years ago..." + + +def test_mistral_audio_transcription_response_transform_diarized(): + """Test that diarized responses preserve segments and language.""" + config = MistralAudioTranscriptionConfig() + + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.return_value = { + "model": "voxtral-mini-latest", + "text": "Hello, how are you? I am fine.", + "language": None, + "segments": [ + { + "text": "Hello, how are you?", + "start": 0.3, + "end": 2.1, + "speaker_id": "speaker_1", + "type": "transcription_segment", + }, + { + "text": "I am fine.", + "start": 2.5, + "end": 3.8, + "speaker_id": "speaker_2", + "type": "transcription_segment", + }, + ], + "usage": { + "prompt_audio_seconds": 4, + "prompt_tokens": 5, + "total_tokens": 50, + "completion_tokens": 20, + }, + } + + response = config.transform_audio_transcription_response(mock_response) + + assert isinstance(response, TranscriptionResponse) + assert response.text == "Hello, how are you? I am fine." + assert response["segments"] is not None + assert len(response["segments"]) == 2 + assert response["segments"][0]["speaker_id"] == "speaker_1" + assert response["segments"][1]["speaker_id"] == "speaker_2" + assert response["language"] is None + + +def test_mistral_audio_transcription_response_transform_empty(): + config = MistralAudioTranscriptionConfig() + + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.return_value = {} + + response = config.transform_audio_transcription_response(mock_response) + + assert isinstance(response, TranscriptionResponse) + assert response.text == "" diff --git a/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py b/tests/unit/llms/mistral/test_mistral_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py rename to tests/unit/llms/mistral/test_mistral_chat_transformation.py diff --git a/tests/test_litellm/llms/mistral/test_mistral_completion.py b/tests/unit/llms/mistral/test_mistral_completion.py similarity index 100% rename from tests/test_litellm/llms/mistral/test_mistral_completion.py rename to tests/unit/llms/mistral/test_mistral_completion.py diff --git a/tests/test_litellm/llms/watsonx/__init__.py b/tests/unit/llms/modelscope/chat/__init__.py similarity index 100% rename from tests/test_litellm/llms/watsonx/__init__.py rename to tests/unit/llms/modelscope/chat/__init__.py diff --git a/tests/test_litellm/llms/modelscope/chat/test_modelscope_chat_transformation.py b/tests/unit/llms/modelscope/chat/test_modelscope_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/modelscope/chat/test_modelscope_chat_transformation.py rename to tests/unit/llms/modelscope/chat/test_modelscope_chat_transformation.py diff --git a/tests/test_litellm/llms/watsonx/audio_transcription/__init__.py b/tests/unit/llms/nadir/__init__.py similarity index 100% rename from tests/test_litellm/llms/watsonx/audio_transcription/__init__.py rename to tests/unit/llms/nadir/__init__.py diff --git a/tests/test_litellm/llms/nadir/test_nadir.py b/tests/unit/llms/nadir/test_nadir.py similarity index 100% rename from tests/test_litellm/llms/nadir/test_nadir.py rename to tests/unit/llms/nadir/test_nadir.py diff --git a/tests/test_litellm/llms/watsonx/rerank/__init__.py b/tests/unit/llms/nebius/__init__.py similarity index 100% rename from tests/test_litellm/llms/watsonx/rerank/__init__.py rename to tests/unit/llms/nebius/__init__.py diff --git a/tests/test_litellm/llms/nebius/test_nebius_chat_transformation.py b/tests/unit/llms/nebius/test_nebius_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/nebius/test_nebius_chat_transformation.py rename to tests/unit/llms/nebius/test_nebius_chat_transformation.py diff --git a/tests/test_litellm/llms/nebius/test_nebius_embedding_transformation.py b/tests/unit/llms/nebius/test_nebius_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/nebius/test_nebius_embedding_transformation.py rename to tests/unit/llms/nebius/test_nebius_embedding_transformation.py diff --git a/tests/test_litellm/llms/you_com/__init__.py b/tests/unit/llms/oci/rerank/__init__.py similarity index 100% rename from tests/test_litellm/llms/you_com/__init__.py rename to tests/unit/llms/oci/rerank/__init__.py diff --git a/tests/test_litellm/llms/oci/test_oci_common_utils.py b/tests/unit/llms/oci/test_oci_common_utils.py similarity index 100% rename from tests/test_litellm/llms/oci/test_oci_common_utils.py rename to tests/unit/llms/oci/test_oci_common_utils.py diff --git a/tests/test_litellm/llms/oci/test_oci_coverage_boost.py b/tests/unit/llms/oci/test_oci_coverage_boost.py similarity index 100% rename from tests/test_litellm/llms/oci/test_oci_coverage_boost.py rename to tests/unit/llms/oci/test_oci_coverage_boost.py diff --git a/tests/unit/llms/ollama/__init__.py b/tests/unit/llms/ollama/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py b/tests/unit/llms/ollama/test_ollama_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py rename to tests/unit/llms/ollama/test_ollama_chat_transformation.py diff --git a/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py b/tests/unit/llms/ollama/test_ollama_completion_transformation.py similarity index 100% rename from tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py rename to tests/unit/llms/ollama/test_ollama_completion_transformation.py diff --git a/tests/test_litellm/llms/ollama/test_ollama_embedding.py b/tests/unit/llms/ollama/test_ollama_embedding.py similarity index 100% rename from tests/test_litellm/llms/ollama/test_ollama_embedding.py rename to tests/unit/llms/ollama/test_ollama_embedding.py diff --git a/tests/test_litellm/llms/ollama/test_ollama_model_info.py b/tests/unit/llms/ollama/test_ollama_model_info.py similarity index 100% rename from tests/test_litellm/llms/ollama/test_ollama_model_info.py rename to tests/unit/llms/ollama/test_ollama_model_info.py diff --git a/tests/test_litellm/llms/openai/realtime/README.md b/tests/unit/llms/openai/realtime/README.md similarity index 100% rename from tests/test_litellm/llms/openai/realtime/README.md rename to tests/unit/llms/openai/realtime/README.md diff --git a/tests/unit/llms/openai/realtime/__init__.py b/tests/unit/llms/openai/realtime/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/openai/realtime/test_openai_realtime_handler.py b/tests/unit/llms/openai/realtime/test_openai_realtime_handler.py similarity index 100% rename from tests/test_litellm/llms/openai/realtime/test_openai_realtime_handler.py rename to tests/unit/llms/openai/realtime/test_openai_realtime_handler.py diff --git a/tests/test_litellm/llms/openai/realtime/test_transcription_sessions.py b/tests/unit/llms/openai/realtime/test_transcription_sessions.py similarity index 100% rename from tests/test_litellm/llms/openai/realtime/test_transcription_sessions.py rename to tests/unit/llms/openai/realtime/test_transcription_sessions.py diff --git a/tests/unit/llms/openai/responses/__init__.py b/tests/unit/llms/openai/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/openai/responses/test_openai_count_tokens_transformation.py b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/responses/test_openai_count_tokens_transformation.py rename to tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_data_residency.py b/tests/unit/llms/openai/responses/test_openai_responses_data_residency.py similarity index 100% rename from tests/test_litellm/llms/openai/responses/test_openai_responses_data_residency.py rename to tests/unit/llms/openai/responses/test_openai_responses_data_residency.py diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py rename to tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py similarity index 100% rename from tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py rename to tests/unit/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/unit/llms/openai/responses/test_openai_responses_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py rename to tests/unit/llms/openai/responses/test_openai_responses_transformation.py diff --git a/tests/test_litellm/llms/openai/test_cost_calculation.py b/tests/unit/llms/openai/test_cost_calculation.py similarity index 100% rename from tests/test_litellm/llms/openai/test_cost_calculation.py rename to tests/unit/llms/openai/test_cost_calculation.py diff --git a/tests/test_litellm/llms/openai/test_data_residency.py b/tests/unit/llms/openai/test_data_residency.py similarity index 100% rename from tests/test_litellm/llms/openai/test_data_residency.py rename to tests/unit/llms/openai/test_data_residency.py diff --git a/tests/test_litellm/llms/openai/test_gpt5_transformation.py b/tests/unit/llms/openai/test_gpt5_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/test_gpt5_transformation.py rename to tests/unit/llms/openai/test_gpt5_transformation.py diff --git a/tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py b/tests/unit/llms/openai/test_is_model_gpt_5_model.py similarity index 100% rename from tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py rename to tests/unit/llms/openai/test_is_model_gpt_5_model.py diff --git a/tests/test_litellm/llms/openai/test_o_series_transformation.py b/tests/unit/llms/openai/test_o_series_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/test_o_series_transformation.py rename to tests/unit/llms/openai/test_o_series_transformation.py diff --git a/tests/test_litellm/llms/openai/test_openai.py b/tests/unit/llms/openai/test_openai.py similarity index 100% rename from tests/test_litellm/llms/openai/test_openai.py rename to tests/unit/llms/openai/test_openai.py diff --git a/tests/test_litellm/llms/openai/test_openai_common_utils.py b/tests/unit/llms/openai/test_openai_common_utils.py similarity index 100% rename from tests/test_litellm/llms/openai/test_openai_common_utils.py rename to tests/unit/llms/openai/test_openai_common_utils.py diff --git a/tests/test_litellm/llms/openai/test_openai_empty_response.py b/tests/unit/llms/openai/test_openai_empty_response.py similarity index 100% rename from tests/test_litellm/llms/openai/test_openai_empty_response.py rename to tests/unit/llms/openai/test_openai_empty_response.py diff --git a/tests/test_litellm/llms/openai/test_openai_file_content_streaming.py b/tests/unit/llms/openai/test_openai_file_content_streaming.py similarity index 100% rename from tests/test_litellm/llms/openai/test_openai_file_content_streaming.py rename to tests/unit/llms/openai/test_openai_file_content_streaming.py diff --git a/tests/test_litellm/llms/openai/test_openai_image_edit_transformation.py b/tests/unit/llms/openai/test_openai_image_edit_transformation.py similarity index 100% rename from tests/test_litellm/llms/openai/test_openai_image_edit_transformation.py rename to tests/unit/llms/openai/test_openai_image_edit_transformation.py diff --git a/tests/test_litellm/llms/openai/test_openai_workload_identity.py b/tests/unit/llms/openai/test_openai_workload_identity.py similarity index 100% rename from tests/test_litellm/llms/openai/test_openai_workload_identity.py rename to tests/unit/llms/openai/test_openai_workload_identity.py diff --git a/tests/test_litellm/llms/openai/test_organization_costs.py b/tests/unit/llms/openai/test_organization_costs.py similarity index 100% rename from tests/test_litellm/llms/openai/test_organization_costs.py rename to tests/unit/llms/openai/test_organization_costs.py diff --git a/tests/test_litellm/llms/openai/test_use_chat_completions_api_no_leak.py b/tests/unit/llms/openai/test_use_chat_completions_api_no_leak.py similarity index 100% rename from tests/test_litellm/llms/openai/test_use_chat_completions_api_no_leak.py rename to tests/unit/llms/openai/test_use_chat_completions_api_no_leak.py diff --git a/tests/test_litellm/llms/openai/transcriptions/test_openai_transcriptions_handler.py b/tests/unit/llms/openai/transcriptions/test_openai_transcriptions_handler.py similarity index 100% rename from tests/test_litellm/llms/openai/transcriptions/test_openai_transcriptions_handler.py rename to tests/unit/llms/openai/transcriptions/test_openai_transcriptions_handler.py diff --git a/tests/unit/llms/openai_like/responses/__init__.py b/tests/unit/llms/openai_like/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/openai_like/responses/test_openai_like_responses.py b/tests/unit/llms/openai_like/responses/test_openai_like_responses.py similarity index 100% rename from tests/test_litellm/llms/openai_like/responses/test_openai_like_responses.py rename to tests/unit/llms/openai_like/responses/test_openai_like_responses.py diff --git a/tests/test_litellm/llms/openai_like/test_abliteration_provider.py b/tests/unit/llms/openai_like/test_abliteration_provider.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_abliteration_provider.py rename to tests/unit/llms/openai_like/test_abliteration_provider.py diff --git a/tests/test_litellm/llms/openai_like/test_assemblyai_provider.py b/tests/unit/llms/openai_like/test_assemblyai_provider.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_assemblyai_provider.py rename to tests/unit/llms/openai_like/test_assemblyai_provider.py diff --git a/tests/test_litellm/llms/openai_like/test_charity_engine.py b/tests/unit/llms/openai_like/test_charity_engine.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_charity_engine.py rename to tests/unit/llms/openai_like/test_charity_engine.py diff --git a/tests/test_litellm/llms/openai_like/test_cognition_provider.py b/tests/unit/llms/openai_like/test_cognition_provider.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_cognition_provider.py rename to tests/unit/llms/openai_like/test_cognition_provider.py diff --git a/tests/test_litellm/llms/openai_like/test_dynamic_config.py b/tests/unit/llms/openai_like/test_dynamic_config.py similarity index 96% rename from tests/test_litellm/llms/openai_like/test_dynamic_config.py rename to tests/unit/llms/openai_like/test_dynamic_config.py index 55e1a1679de..de70f98c3f1 100644 --- a/tests/test_litellm/llms/openai_like/test_dynamic_config.py +++ b/tests/unit/llms/openai_like/test_dynamic_config.py @@ -20,9 +20,6 @@ def _isolate_generated_class_cache(): class TestClassCaching: - def test_same_slug_returns_the_identical_class_object(self): - provider = _provider("cache_same_slug") - assert create_responses_config_class(provider) is create_responses_config_class(provider) def test_cache_is_keyed_on_slug_not_on_the_provider_instance(self): first = create_responses_config_class(_provider("cache_by_slug")) diff --git a/tests/test_litellm/llms/openai_like/test_empiriolabs_provider.py b/tests/unit/llms/openai_like/test_empiriolabs_provider.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_empiriolabs_provider.py rename to tests/unit/llms/openai_like/test_empiriolabs_provider.py diff --git a/tests/unit/llms/openai_like/test_json_providers.py b/tests/unit/llms/openai_like/test_json_providers.py new file mode 100644 index 00000000000..a56108ca9ac --- /dev/null +++ b/tests/unit/llms/openai_like/test_json_providers.py @@ -0,0 +1,317 @@ +""" +Tests for JSON-based provider configuration system. +""" + +import os +import sys +from unittest.mock import patch + +try: + import pytest +except ImportError: + # pytest not available, will run as standalone script + pytest = None + +# Add workspace to path +workspace_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) +sys.path.insert(0, workspace_path) + + + +class TestJSONProviderLoader: + """Test JSON provider loading and configuration""" + + def test_load_json_providers(self): + """Test that JSON providers load correctly""" + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + # Verify publicai is loaded + assert JSONProviderRegistry.exists("publicai") + + # Get publicai config + publicai = JSONProviderRegistry.get("publicai") + assert publicai is not None + assert publicai.base_url == "https://api.publicai.co/v1" + assert publicai.api_key_env == "PUBLICAI_API_KEY" + assert publicai.api_base_env == "PUBLICAI_API_BASE" + assert publicai.param_mappings.get("max_completion_tokens") == "max_tokens" + + def test_dynamic_config_generation(self): + """Test dynamic config class creation""" + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("publicai") + config_class = create_config_class(provider) + config = config_class() + + # Test API info resolution + api_base, api_key = config._get_openai_compatible_provider_info(None, None) + assert api_base == "https://api.publicai.co/v1" + + # Test with custom base + api_base, api_key = config._get_openai_compatible_provider_info( + "https://custom.api.com", "test-key" + ) + assert api_base == "https://custom.api.com" + assert api_key == "test-key" + + def test_parameter_mapping(self): + """Test parameter mapping works""" + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("publicai") + config_class = create_config_class(provider) + config = config_class() + + # Test parameter mapping + optional_params = {} + non_default_params = {"max_completion_tokens": 100, "temperature": 0.7} + result = config.map_openai_params( + non_default_params, optional_params, "gpt-4", False + ) + + # max_completion_tokens should be mapped to max_tokens + assert "max_tokens" in result + assert result["max_tokens"] == 100 + assert "max_completion_tokens" not in result + + # temperature should be passed through + assert result["temperature"] == 0.7 + + def test_supported_params(self): + """Test that config returns supported params""" + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("publicai") + config_class = create_config_class(provider) + config = config_class() + + # Get supported params + supported = config.get_supported_openai_params("gpt-4") + + # Should have standard OpenAI params + assert isinstance(supported, list) + assert len(supported) > 0 + + def test_tool_params_excluded_when_function_calling_not_supported(self): + """Test that tool-related params are excluded for models that don't support + function calling. Regression test for https://github.com/BerriAI/litellm/issues/21125 + """ + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("publicai") + config_class = create_config_class(provider) + config = config_class() + + # Mock supports_function_calling to return False + with patch("litellm.utils.supports_function_calling", return_value=False): + supported = config.get_supported_openai_params("some-model-without-fc") + + tool_params = [ + "tools", + "tool_choice", + "function_call", + "functions", + "parallel_tool_calls", + ] + for param in tool_params: + assert ( + param not in supported + ), f"'{param}' should not be in supported params when function calling is not supported" + + # Non-tool params should still be present + assert "temperature" in supported + assert "max_tokens" in supported + assert "stop" in supported + + def test_tool_params_included_when_function_calling_supported(self): + """Test that tool-related params are included for models that support function calling.""" + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("publicai") + config_class = create_config_class(provider) + config = config_class() + + # Mock supports_function_calling to return True + with patch("litellm.utils.supports_function_calling", return_value=True): + supported = config.get_supported_openai_params("some-model-with-fc") + + assert "tools" in supported + assert "tool_choice" in supported + + def test_provider_resolution(self): + """Test that provider resolution finds JSON providers""" + from litellm.litellm_core_utils.get_llm_provider_logic import ( + get_llm_provider, + ) + + model, provider, api_key, api_base = get_llm_provider( + model="publicai/gpt-4", + custom_llm_provider=None, + api_base=None, + api_key=None, + ) + + assert model == "gpt-4" + assert provider == "publicai" + assert api_base == "https://api.publicai.co/v1" + + def test_provider_config_manager(self): + """Test that ProviderConfigManager returns JSON-based configs""" + from litellm import LlmProviders + from litellm.utils import ProviderConfigManager + + config = ProviderConfigManager.get_provider_chat_config( + model="gpt-4", provider=LlmProviders.PUBLICAI + ) + + assert config is not None + assert config.custom_llm_provider == "publicai" + + +class TestPinstripes: + """Tests for Pinstripes JSON-configured provider""" + + def test_pinstripes_json_config_exists(self): + """Test that pinstripes is configured in providers.json""" + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + assert JSONProviderRegistry.exists("pinstripes") + + pinstripes = JSONProviderRegistry.get("pinstripes") + assert pinstripes is not None + assert pinstripes.base_url == "https://pinstripes.io/v1" + assert pinstripes.api_key_env == "PINSTRIPES_API_KEY" + assert pinstripes.param_mappings.get("max_completion_tokens") == "max_tokens" + + def test_pinstripes_provider_resolution(self): + """Test that provider resolution finds pinstripes and returns the default base URL""" + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + model, provider, api_key, api_base = get_llm_provider( + model="pinstripes/ps/glm-4.5-air", + custom_llm_provider=None, + api_base=None, + api_key=None, + ) + + assert model == "ps/glm-4.5-air" + assert provider == "pinstripes" + assert api_base == "https://pinstripes.io/v1" + + def test_pinstripes_dynamic_config(self): + """Test dynamic config class creation for pinstripes""" + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("pinstripes") + config_class = create_config_class(provider) + config = config_class() + + api_base, api_key = config._get_openai_compatible_provider_info(None, None) + assert api_base == "https://pinstripes.io/v1" + + api_base, api_key = config._get_openai_compatible_provider_info( + "https://custom.pinstripes.io/v1", "test-key" + ) + assert api_base == "https://custom.pinstripes.io/v1" + assert api_key == "test-key" + + def test_pinstripes_parameter_mapping(self): + """Test that max_completion_tokens is mapped to max_tokens for pinstripes""" + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("pinstripes") + config_class = create_config_class(provider) + config = config_class() + + optional_params = {} + non_default_params = {"max_completion_tokens": 100, "temperature": 0.7} + result = config.map_openai_params( + non_default_params, optional_params, "ps/glm-4.5-air", False + ) + + assert "max_tokens" in result + assert result["max_tokens"] == 100 + assert "max_completion_tokens" not in result + assert result["temperature"] == 0.7 + + +class TestDarkbloom: + def test_darkbloom_json_config_exists(self): + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + darkbloom = JSONProviderRegistry.get("darkbloom") + assert darkbloom is not None + assert darkbloom.base_url == "https://api.darkbloom.dev/v1" + assert darkbloom.api_key_env == "DARKBLOOM_API_KEY" + assert darkbloom.api_base_env == "DARKBLOOM_API_BASE" + assert darkbloom.param_mappings.get("max_completion_tokens") == "max_tokens" + + def test_darkbloom_provider_resolution(self): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + model, provider, api_key, api_base = get_llm_provider( + model="darkbloom/gemma-4-26b", + custom_llm_provider=None, + api_base=None, + api_key=None, + ) + + assert model == "gemma-4-26b" + assert provider == "darkbloom" + assert api_key is None + assert api_base == "https://api.darkbloom.dev/v1" + + def test_darkbloom_dynamic_config(self): + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("darkbloom") + config_class = create_config_class(provider) + config = config_class() + + api_base, api_key = config._get_openai_compatible_provider_info(None, None) + assert api_base == "https://api.darkbloom.dev/v1" + + api_base, api_key = config._get_openai_compatible_provider_info( + "https://custom.darkbloom.dev/v1", "test-key" + ) + assert api_base == "https://custom.darkbloom.dev/v1" + assert api_key == "test-key" + + def test_darkbloom_complete_url_appends_endpoint(self): + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("darkbloom") + config_class = create_config_class(provider) + config = config_class() + + url = config.get_complete_url( + api_base="https://api.darkbloom.dev/v1", + api_key="test-key", + model="darkbloom/gemma-4-26b", + optional_params={}, + litellm_params={}, + stream=True, + ) + + assert url == "https://api.darkbloom.dev/v1/chat/completions" + + def test_darkbloom_provider_config_manager(self): + from litellm import LlmProviders + from litellm.utils import ProviderConfigManager + + config = ProviderConfigManager.get_provider_chat_config( + model="gemma-4-26b", provider=LlmProviders.DARKBLOOM + ) + + assert config is not None + assert config.custom_llm_provider == "darkbloom" diff --git a/tests/test_litellm/llms/openai_like/test_libertai_provider.py b/tests/unit/llms/openai_like/test_libertai_provider.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_libertai_provider.py rename to tests/unit/llms/openai_like/test_libertai_provider.py diff --git a/tests/test_litellm/llms/openai_like/test_meta_provider.py b/tests/unit/llms/openai_like/test_meta_provider.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_meta_provider.py rename to tests/unit/llms/openai_like/test_meta_provider.py diff --git a/tests/test_litellm/llms/openai_like/test_model_info.py b/tests/unit/llms/openai_like/test_model_info.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_model_info.py rename to tests/unit/llms/openai_like/test_model_info.py diff --git a/tests/test_litellm/llms/openai_like/test_pinstripes_provider.py b/tests/unit/llms/openai_like/test_pinstripes_provider.py similarity index 68% rename from tests/test_litellm/llms/openai_like/test_pinstripes_provider.py rename to tests/unit/llms/openai_like/test_pinstripes_provider.py index 70bb786b2e6..e7a2dfb92dc 100644 --- a/tests/test_litellm/llms/openai_like/test_pinstripes_provider.py +++ b/tests/unit/llms/openai_like/test_pinstripes_provider.py @@ -16,17 +16,6 @@ class TestPinstripeProviderConfig: assert LlmProviders.PINSTRIPES.value == "pinstripes" assert "pinstripes" in litellm.provider_list - def test_pinstripes_json_config_exists(self): - """Test that pinstripes is configured in providers.json""" - from litellm.llms.openai_like.json_loader import JSONProviderRegistry - - assert JSONProviderRegistry.exists("pinstripes") - - pinstripes = JSONProviderRegistry.get("pinstripes") - assert pinstripes is not None - assert pinstripes.base_url == "https://pinstripes.io/v1" - assert pinstripes.api_key_env == "PINSTRIPES_API_KEY" - assert pinstripes.param_mappings.get("max_completion_tokens") == "max_tokens" def test_pinstripes_in_openai_compatible_providers(self): """Test that pinstripes is in the openai_compatible_providers list""" @@ -34,20 +23,6 @@ class TestPinstripeProviderConfig: assert "pinstripes" in openai_compatible_providers - def test_pinstripes_provider_resolution(self): - """Test that provider resolution finds pinstripes and returns the default base URL""" - from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider - - model, provider, api_key, api_base = get_llm_provider( - model="pinstripes/ps/glm-4.5-air", - custom_llm_provider=None, - api_base=None, - api_key=None, - ) - - assert model == "ps/glm-4.5-air" - assert provider == "pinstripes" - assert api_base == "https://pinstripes.io/v1" def test_pinstripes_api_base_override(self): """Test that an explicit api_base / api_key overrides the default""" diff --git a/tests/test_litellm/llms/openai_like/test_provider_affinity_forwarding.py b/tests/unit/llms/openai_like/test_provider_affinity_forwarding.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_provider_affinity_forwarding.py rename to tests/unit/llms/openai_like/test_provider_affinity_forwarding.py diff --git a/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py b/tests/unit/llms/openai_like/test_scx_ai_provider.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_scx_ai_provider.py rename to tests/unit/llms/openai_like/test_scx_ai_provider.py diff --git a/tests/test_litellm/llms/openai_like/test_tensormesh_provider.py b/tests/unit/llms/openai_like/test_tensormesh_provider.py similarity index 100% rename from tests/test_litellm/llms/openai_like/test_tensormesh_provider.py rename to tests/unit/llms/openai_like/test_tensormesh_provider.py diff --git a/tests/unit/llms/openai_like/test_xiaomi_mimo.py b/tests/unit/llms/openai_like/test_xiaomi_mimo.py new file mode 100644 index 00000000000..a642cc91f90 --- /dev/null +++ b/tests/unit/llms/openai_like/test_xiaomi_mimo.py @@ -0,0 +1,84 @@ +""" +Tests for Xiaomi MiMo provider configuration and integration. +Related to issue #18794 +""" + +import os +import sys +from unittest.mock import MagicMock, patch + +try: + import pytest +except ImportError: + pytest = None + +# Add workspace to path +workspace_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) +sys.path.insert(0, workspace_path) + +import litellm + + +class TestXiaomiMiMoProviderConfig: + """Test Xiaomi MiMo provider configuration""" + + def test_xiaomi_mimo_in_provider_list(self): + """Test that xiaomi_mimo is in the provider list (fixes #18794)""" + from litellm import LlmProviders + + # Verify xiaomi_mimo is in the enum + assert hasattr(LlmProviders, "XIAOMI_MIMO") + assert LlmProviders.XIAOMI_MIMO.value == "xiaomi_mimo" + + # Verify it's in the provider list + assert "xiaomi_mimo" in litellm.provider_list + + def test_xiaomi_mimo_json_config_exists(self): + """Test that xiaomi_mimo is configured in providers.json""" + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + # Verify xiaomi_mimo is loaded + assert JSONProviderRegistry.exists("xiaomi_mimo") + + # Get xiaomi_mimo config + xiaomi_mimo = JSONProviderRegistry.get("xiaomi_mimo") + assert xiaomi_mimo is not None + assert xiaomi_mimo.base_url == "https://api.xiaomimimo.com/v1" + assert xiaomi_mimo.api_key_env == "XIAOMI_MIMO_API_KEY" + assert xiaomi_mimo.param_mappings.get("max_completion_tokens") == "max_tokens" + + def test_xiaomi_mimo_provider_resolution(self): + """Test that provider resolution finds xiaomi_mimo""" + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + model, provider, api_key, api_base = get_llm_provider( + model="xiaomi_mimo/mimo-v2-flash", + custom_llm_provider=None, + api_base=None, + api_key=None, + ) + + assert model == "mimo-v2-flash" + assert provider == "xiaomi_mimo" + assert api_base == "https://api.xiaomimimo.com/v1" + + def test_xiaomi_mimo_router_config(self): + """Test that xiaomi_mimo can be used in Router configuration (fixes #18794)""" + from litellm import Router + + # This should not raise "Unsupported provider - xiaomi_mimo" + router = Router( + model_list=[ + { + "model_name": "mimo-v2-flash", + "litellm_params": { + "model": "xiaomi_mimo/mimo-v2-flash", + "api_key": "test-key", + }, + } + ] + ) + + # Verify the deployment was created successfully + assert len(router.model_list) == 1 + assert router.model_list[0]["model_name"] == "mimo-v2-flash" diff --git a/tests/unit/llms/ovhcloud/__init__.py b/tests/unit/llms/ovhcloud/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/ovhcloud/test_ovhcloud_audio_transcription_transformation.py b/tests/unit/llms/ovhcloud/test_ovhcloud_audio_transcription_transformation.py new file mode 100644 index 00000000000..87e54dfba9b --- /dev/null +++ b/tests/unit/llms/ovhcloud/test_ovhcloud_audio_transcription_transformation.py @@ -0,0 +1,58 @@ + + + + +class TestOVHCloudDurationFieldMigration: + """Tests for OVHCloud duration -> seconds field migration.""" + + def test_seconds_field_mapped_to_duration(self): + """New `seconds` field should be normalized to `duration`.""" + from litellm.llms.ovhcloud.audio_transcription.transformation import ( + OVHCloudAudioTranscriptionConfig, + ) + from unittest.mock import MagicMock + + config = OVHCloudAudioTranscriptionConfig() + mock_response = MagicMock() + mock_response.json.return_value = { + "text": "Hello world", + "seconds": 3.14, + } + + result = config.transform_audio_transcription_response(mock_response) + + assert result.text == "Hello world" + assert result._hidden_params["duration"] == 3.14 + + def test_legacy_duration_field_still_works(self): + """Legacy `duration` field should still be accepted.""" + from litellm.llms.ovhcloud.audio_transcription.transformation import ( + OVHCloudAudioTranscriptionConfig, + ) + from unittest.mock import MagicMock + + config = OVHCloudAudioTranscriptionConfig() + mock_response = MagicMock() + mock_response.json.return_value = { + "text": "Hello world", + "duration": 2.71, + } + + result = config.transform_audio_transcription_response(mock_response) + + assert result.text == "Hello world" + assert result._hidden_params["duration"] == 2.71 + + + def test_seconds_zero_mapped_to_duration(self): + """seconds=0.0 must not be treated as falsy and lost.""" + from litellm.llms.ovhcloud.audio_transcription.transformation import ( + OVHCloudAudioTranscriptionConfig, + ) + from unittest.mock import MagicMock + + config = OVHCloudAudioTranscriptionConfig() + mock_response = MagicMock() + mock_response.json.return_value = {"text": "silence", "seconds": 0.0} + result = config.transform_audio_transcription_response(mock_response) + assert result._hidden_params["duration"] == 0.0 diff --git a/tests/unit/llms/ovhcloud/test_ovhcloud_chat_transformation.py b/tests/unit/llms/ovhcloud/test_ovhcloud_chat_transformation.py new file mode 100644 index 00000000000..c2bc4ee4a4c --- /dev/null +++ b/tests/unit/llms/ovhcloud/test_ovhcloud_chat_transformation.py @@ -0,0 +1,250 @@ +""" +Unit tests for OVHCloud AI Endpoints chat integration. +""" + + +import pytest + +from litellm.llms.ovhcloud.utils import OVHCloudException +from litellm.utils import get_optional_params + + +from litellm.llms.ovhcloud.chat.transformation import ( + OVHCloudChatCompletionStreamingHandler, + OVHCloudChatConfig, +) + +config = OVHCloudChatConfig() +model = "ovhcloud/Mistral-7B-Instruct-v0.3" + + +class TestOvhCloudChatCompletionStreamingHandler: + def test_chunk_parser_successful(self): + handler = OVHCloudChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ) + + chunk = { + "id": "test_id", + "created": 1234567890, + "model": "gpt-oss-20b", + "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + "choices": [ + {"delta": {"content": "test content", "reasoning": "test reasoning"}} + ], + } + + result = handler.chunk_parser(chunk) + + assert result.id == "test_id" + assert result.object == "chat.completion.chunk" + assert result.created == 1234567890 + assert result.model == "gpt-oss-20b" + assert result.usage.prompt_tokens == chunk["usage"]["prompt_tokens"] + assert result.usage.completion_tokens == chunk["usage"]["completion_tokens"] + assert result.usage.total_tokens == chunk["usage"]["total_tokens"] + assert len(result.choices) == 1 + assert result.choices[0]["delta"]["reasoning_content"] == "test reasoning" + + def test_chunk_parser_error_response(self): + handler = OVHCloudChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ) + + error_chunk = { + "error": { + "message": "test error", + "code": 400, + } + } + + with pytest.raises(OVHCloudException) as exc_info: + handler.chunk_parser(error_chunk) + + assert "OVHCloud Error: test error" in str(exc_info.value) + assert exc_info.value.status_code == 400 + + def test_chunk_parser_key_error(self): + handler = OVHCloudChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ) + + invalid_chunk = {"incomplete": "data"} + + with pytest.raises(OVHCloudException) as exc_info: + handler.chunk_parser(invalid_chunk) + + assert "KeyError" in str(exc_info.value) + assert exc_info.value.status_code == 400 + + +class TestOVHCloudConfig: + def test_transform_request_basic(self): + """Test basic request transformation""" + transformed_request = config.transform_request( + model, + messages=[{"role": "user", "content": "Hello, world!"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert transformed_request["model"] == model + assert transformed_request["messages"] == [ + {"role": "user", "content": "Hello, world!"} + ] + + def test_transform_request_with_extra_body(self): + """Test request transformation with extra_body parameters""" + transformed_request = config.transform_request( + model, + messages=[{"role": "user", "content": "Hello, world!"}], + optional_params={"extra_body": {"custom_param": "custom_value"}}, + litellm_params={}, + headers={}, + ) + + assert transformed_request["custom_param"] == "custom_value" + assert transformed_request["messages"] == [ + {"role": "user", "content": "Hello, world!"} + ] + + def test_map_openai_params(self): + """Test OpenAI parameter mapping""" + non_default_params = { + "temperature": 0.7, + "max_tokens": 100, + "top_p": 0.9, + } + + mapped_params = config.map_openai_params( + non_default_params=non_default_params, + optional_params={}, + model=model, + drop_params=False, + ) + + assert mapped_params["temperature"] == 0.7 + assert mapped_params["max_tokens"] == 100 + assert mapped_params["top_p"] == 0.9 + + def test_get_error_class(self): + """Test error class creation""" + error = config.get_error_class( + error_message="Test error", + status_code=400, + headers={"Content-Type": "application/json"}, + ) + + assert isinstance(error, OVHCloudException) + assert error.message == "Test error" + assert error.status_code == 400 + + @pytest.mark.parametrize( + "model", + [ + "Meta-Llama-3_3-70B-Instruct", + "Meta-Llama-3_1-70B-Instruct", + "Mixtral-8x7B-Instruct-v0.1", + "gpt-oss-120b", + "some-model-not-in-the-cost-map", + ], + ) + def test_tools_not_filtered_by_static_model_map(self, model): + """ + OVHCloud AI Endpoints are OpenAI-compatible; tools/tool_choice must pass + through for any model. The server is responsible for rejecting unsupported + tool calls — LiteLLM must not strip them based on a stale static catalog. + """ + + params = get_optional_params( + model=model, + custom_llm_provider="ovhcloud", + tools=[ + { + "type": "function", + "function": {"name": "x", "parameters": {}}, + } + ], + tool_choice="auto", + ) + + assert "tools" in params + assert "tool_choice" in params + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) + + +class TestOVHCloudReasoningFieldMigration: + """Tests for OVHCloud reasoning_content -> reasoning field migration.""" + + def test_streaming_new_reasoning_field(self): + """New `reasoning` field should be mapped to `reasoning_content`.""" + handler = OVHCloudChatCompletionStreamingHandler( + streaming_response=iter([]), + sync_stream=True, + ) + chunk = { + "id": "test-id", + "created": 1234567890, + "model": "test-model", + "choices": [ + { + "delta": { + "role": "assistant", + "reasoning": "Let me think...", + }, + "index": 0, + } + ], + } + result = handler.chunk_parser(chunk) + assert result.choices[0]["delta"]["reasoning_content"] == "Let me think..." + + def test_streaming_legacy_reasoning_content_unchanged(self): + """Legacy `reasoning_content` field should pass through untouched.""" + handler = OVHCloudChatCompletionStreamingHandler( + streaming_response=iter([]), + sync_stream=True, + ) + chunk = { + "id": "test-id", + "created": 1234567890, + "model": "test-model", + "choices": [ + { + "delta": { + "role": "assistant", + "reasoning_content": "Already correct field.", + }, + "index": 0, + } + ], + } + result = handler.chunk_parser(chunk) + assert result.choices[0]["delta"]["reasoning_content"] == "Already correct field." + + def test_streaming_both_fields_legacy_wins(self): + """When both fields present, existing `reasoning_content` is not overwritten.""" + handler = OVHCloudChatCompletionStreamingHandler( + streaming_response=iter([]), + sync_stream=True, + ) + chunk = { + "id": "test-id", + "created": 1234567890, + "model": "test-model", + "choices": [ + { + "delta": { + "reasoning": "new field", + "reasoning_content": "legacy field", + }, + "index": 0, + } + ], + } + result = handler.chunk_parser(chunk) + assert result.choices[0]["delta"]["reasoning_content"] == "legacy field" diff --git a/tests/test_litellm/llms/ovhcloud/test_ovhcloud_embeddings_transformation.py b/tests/unit/llms/ovhcloud/test_ovhcloud_embeddings_transformation.py similarity index 100% rename from tests/test_litellm/llms/ovhcloud/test_ovhcloud_embeddings_transformation.py rename to tests/unit/llms/ovhcloud/test_ovhcloud_embeddings_transformation.py diff --git a/tests/unit/llms/pass_through/__init__.py b/tests/unit/llms/pass_through/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/pass_through/guardrail_translation/__init__.py b/tests/unit/llms/pass_through/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/perplexity/test_perplexity.py b/tests/unit/llms/perplexity/test_perplexity.py similarity index 100% rename from tests/test_litellm/llms/perplexity/test_perplexity.py rename to tests/unit/llms/perplexity/test_perplexity.py diff --git a/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py b/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py similarity index 100% rename from tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py rename to tests/unit/llms/perplexity/test_perplexity_cost_calculator.py diff --git a/tests/test_litellm/llms/perplexity/test_perplexity_integration.py b/tests/unit/llms/perplexity/test_perplexity_integration.py similarity index 100% rename from tests/test_litellm/llms/perplexity/test_perplexity_integration.py rename to tests/unit/llms/perplexity/test_perplexity_integration.py diff --git a/tests/unit/llms/pg_vector/__init__.py b/tests/unit/llms/pg_vector/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/pg_vector/vector_stores/__init__.py b/tests/unit/llms/pg_vector/vector_stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/pg_vector/vector_stores/test_pg_vector_transformation.py b/tests/unit/llms/pg_vector/vector_stores/test_pg_vector_transformation.py similarity index 100% rename from tests/test_litellm/llms/pg_vector/vector_stores/test_pg_vector_transformation.py rename to tests/unit/llms/pg_vector/vector_stores/test_pg_vector_transformation.py diff --git a/tests/test_litellm/llms/azure/realtime/__init__.py b/tests/unit/llms/reducto/__init__.py similarity index 100% rename from tests/test_litellm/llms/azure/realtime/__init__.py rename to tests/unit/llms/reducto/__init__.py diff --git a/tests/test_litellm/llms/reducto/conftest.py b/tests/unit/llms/reducto/conftest.py similarity index 100% rename from tests/test_litellm/llms/reducto/conftest.py rename to tests/unit/llms/reducto/conftest.py diff --git a/tests/test_litellm/llms/reducto/test_cost.py b/tests/unit/llms/reducto/test_cost.py similarity index 100% rename from tests/test_litellm/llms/reducto/test_cost.py rename to tests/unit/llms/reducto/test_cost.py diff --git a/tests/test_litellm/llms/reducto/test_model_info.py b/tests/unit/llms/reducto/test_model_info.py similarity index 100% rename from tests/test_litellm/llms/reducto/test_model_info.py rename to tests/unit/llms/reducto/test_model_info.py diff --git a/tests/test_litellm/llms/reducto/test_parse_legacy.py b/tests/unit/llms/reducto/test_parse_legacy.py similarity index 100% rename from tests/test_litellm/llms/reducto/test_parse_legacy.py rename to tests/unit/llms/reducto/test_parse_legacy.py diff --git a/tests/test_litellm/llms/reducto/test_parse_v3.py b/tests/unit/llms/reducto/test_parse_v3.py similarity index 100% rename from tests/test_litellm/llms/reducto/test_parse_v3.py rename to tests/unit/llms/reducto/test_parse_v3.py diff --git a/tests/test_litellm/llms/reducto/test_upload.py b/tests/unit/llms/reducto/test_upload.py similarity index 100% rename from tests/test_litellm/llms/reducto/test_upload.py rename to tests/unit/llms/reducto/test_upload.py diff --git a/tests/unit/llms/sagemaker/__init__.py b/tests/unit/llms/sagemaker/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_handler.py b/tests/unit/llms/sagemaker/test_sagemaker_chat_handler.py similarity index 100% rename from tests/test_litellm/llms/sagemaker/test_sagemaker_chat_handler.py rename to tests/unit/llms/sagemaker/test_sagemaker_chat_handler.py diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py b/tests/unit/llms/sagemaker/test_sagemaker_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py rename to tests/unit/llms/sagemaker/test_sagemaker_chat_transformation.py diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_common_utils.py b/tests/unit/llms/sagemaker/test_sagemaker_common_utils.py similarity index 100% rename from tests/test_litellm/llms/sagemaker/test_sagemaker_common_utils.py rename to tests/unit/llms/sagemaker/test_sagemaker_common_utils.py diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py b/tests/unit/llms/sagemaker/test_sagemaker_completion_handler.py similarity index 100% rename from tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py rename to tests/unit/llms/sagemaker/test_sagemaker_completion_handler.py diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_role_assumption.py b/tests/unit/llms/sagemaker/test_sagemaker_embedding_role_assumption.py similarity index 100% rename from tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_role_assumption.py rename to tests/unit/llms/sagemaker/test_sagemaker_embedding_role_assumption.py diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_voyage.py b/tests/unit/llms/sagemaker/test_sagemaker_embedding_voyage.py similarity index 100% rename from tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_voyage.py rename to tests/unit/llms/sagemaker/test_sagemaker_embedding_voyage.py diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_nova_transformation.py b/tests/unit/llms/sagemaker/test_sagemaker_nova_transformation.py similarity index 100% rename from tests/test_litellm/llms/sagemaker/test_sagemaker_nova_transformation.py rename to tests/unit/llms/sagemaker/test_sagemaker_nova_transformation.py diff --git a/tests/unit/llms/sambanova/__init__.py b/tests/unit/llms/sambanova/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/sambanova/tests_sambanova_embedding_transformation.py b/tests/unit/llms/sambanova/tests_sambanova_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/sambanova/tests_sambanova_embedding_transformation.py rename to tests/unit/llms/sambanova/tests_sambanova_embedding_transformation.py diff --git a/tests/unit/llms/sap/chat/__init__.py b/tests/unit/llms/sap/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/sap/chat/test_sap_chat_calls.py b/tests/unit/llms/sap/chat/test_sap_chat_calls.py similarity index 100% rename from tests/test_litellm/llms/sap/chat/test_sap_chat_calls.py rename to tests/unit/llms/sap/chat/test_sap_chat_calls.py diff --git a/tests/test_litellm/llms/sap/chat/test_sap_langchain_strict_param.py b/tests/unit/llms/sap/chat/test_sap_langchain_strict_param.py similarity index 100% rename from tests/test_litellm/llms/sap/chat/test_sap_langchain_strict_param.py rename to tests/unit/llms/sap/chat/test_sap_langchain_strict_param.py diff --git a/tests/test_litellm/llms/sap/chat/test_sap_response_format.py b/tests/unit/llms/sap/chat/test_sap_response_format.py similarity index 100% rename from tests/test_litellm/llms/sap/chat/test_sap_response_format.py rename to tests/unit/llms/sap/chat/test_sap_response_format.py diff --git a/tests/test_litellm/llms/sap/chat/test_sap_tool_parameters.py b/tests/unit/llms/sap/chat/test_sap_tool_parameters.py similarity index 100% rename from tests/test_litellm/llms/sap/chat/test_sap_tool_parameters.py rename to tests/unit/llms/sap/chat/test_sap_tool_parameters.py diff --git a/tests/test_litellm/llms/sap/chat/test_sap_transformation.py b/tests/unit/llms/sap/chat/test_sap_transformation.py similarity index 100% rename from tests/test_litellm/llms/sap/chat/test_sap_transformation.py rename to tests/unit/llms/sap/chat/test_sap_transformation.py diff --git a/tests/unit/llms/sap/embed/__init__.py b/tests/unit/llms/sap/embed/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/sap/embed/test_sap_embed_transformation.py b/tests/unit/llms/sap/embed/test_sap_embed_transformation.py similarity index 100% rename from tests/test_litellm/llms/sap/embed/test_sap_embed_transformation.py rename to tests/unit/llms/sap/embed/test_sap_embed_transformation.py diff --git a/tests/test_litellm/llms/sap/embed/test_sap_embedding.py b/tests/unit/llms/sap/embed/test_sap_embedding.py similarity index 100% rename from tests/test_litellm/llms/sap/embed/test_sap_embedding.py rename to tests/unit/llms/sap/embed/test_sap_embedding.py diff --git a/tests/unit/llms/snowflake/chat/__init__.py b/tests/unit/llms/snowflake/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py b/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py rename to tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py diff --git a/tests/unit/llms/snowflake/embedding/__init__.py b/tests/unit/llms/snowflake/embedding/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/snowflake/embedding/test_snowflake_embedding.py b/tests/unit/llms/snowflake/embedding/test_snowflake_embedding.py similarity index 100% rename from tests/test_litellm/llms/snowflake/embedding/test_snowflake_embedding.py rename to tests/unit/llms/snowflake/embedding/test_snowflake_embedding.py diff --git a/tests/unit/llms/snowflake/test_snowflake_native_endpoints.py b/tests/unit/llms/snowflake/test_snowflake_native_endpoints.py index 7970f7771fc..344b8e5573d 100644 --- a/tests/unit/llms/snowflake/test_snowflake_native_endpoints.py +++ b/tests/unit/llms/snowflake/test_snowflake_native_endpoints.py @@ -7,7 +7,7 @@ Covers: - Claude models → /messages (Anthropic format) Run: - pytest tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py -v + pytest tests/unit/llms/snowflake/test_snowflake_native_endpoints.py -v """ import json diff --git a/tests/test_litellm/llms/soniox/audio_transcription/__init__.py b/tests/unit/llms/soniox/audio_transcription/__init__.py similarity index 100% rename from tests/test_litellm/llms/soniox/audio_transcription/__init__.py rename to tests/unit/llms/soniox/audio_transcription/__init__.py diff --git a/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py b/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py similarity index 100% rename from tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py rename to tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py diff --git a/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py b/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py similarity index 100% rename from tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py rename to tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py diff --git a/tests/test_litellm/llms/test_cache_control_and_reasoning.py b/tests/unit/llms/test_cache_control_and_reasoning.py similarity index 100% rename from tests/test_litellm/llms/test_cache_control_and_reasoning.py rename to tests/unit/llms/test_cache_control_and_reasoning.py diff --git a/tests/test_litellm/llms/test_file_content_block.py b/tests/unit/llms/test_file_content_block.py similarity index 100% rename from tests/test_litellm/llms/test_file_content_block.py rename to tests/unit/llms/test_file_content_block.py diff --git a/tests/test_litellm/llms/test_file_search_responses.py b/tests/unit/llms/test_file_search_responses.py similarity index 100% rename from tests/test_litellm/llms/test_file_search_responses.py rename to tests/unit/llms/test_file_search_responses.py diff --git a/tests/test_litellm/llms/test_lifecycle_fix.py b/tests/unit/llms/test_lifecycle_fix.py similarity index 100% rename from tests/test_litellm/llms/test_lifecycle_fix.py rename to tests/unit/llms/test_lifecycle_fix.py diff --git a/tests/test_litellm/llms/test_polling_url_origin_match.py b/tests/unit/llms/test_polling_url_origin_match.py similarity index 100% rename from tests/test_litellm/llms/test_polling_url_origin_match.py rename to tests/unit/llms/test_polling_url_origin_match.py diff --git a/tests/test_litellm/llms/test_predibase_transformation.py b/tests/unit/llms/test_predibase_transformation.py similarity index 100% rename from tests/test_litellm/llms/test_predibase_transformation.py rename to tests/unit/llms/test_predibase_transformation.py diff --git a/tests/unit/llms/tinyfish/__init__.py b/tests/unit/llms/tinyfish/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py b/tests/unit/llms/tinyfish/test_tinyfish_search.py similarity index 100% rename from tests/test_litellm/llms/tinyfish/test_tinyfish_search.py rename to tests/unit/llms/tinyfish/test_tinyfish_search.py diff --git a/tests/test_litellm/llms/vercel_ai_gateway/test_vercel_ai_gateway.py b/tests/unit/llms/vercel_ai_gateway/test_vercel_ai_gateway.py similarity index 100% rename from tests/test_litellm/llms/vercel_ai_gateway/test_vercel_ai_gateway.py rename to tests/unit/llms/vercel_ai_gateway/test_vercel_ai_gateway.py diff --git a/tests/unit/llms/vertex_ai/audio_transcription/__init__.py b/tests/unit/llms/vertex_ai/audio_transcription/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_audio_transcription_transformation.py b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_audio_transcription_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_audio_transcription_transformation.py rename to tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_audio_transcription_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py rename to tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py rename to tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py rename to tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py diff --git a/tests/unit/llms/vertex_ai/batches/__init__.py b/tests/unit/llms/vertex_ai/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/unit/llms/vertex_ai/batches/test_handler.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/batches/test_handler.py rename to tests/unit/llms/vertex_ai/batches/test_handler.py diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/unit/llms/vertex_ai/batches/test_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/batches/test_transformation.py rename to tests/unit/llms/vertex_ai/batches/test_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/files/test_transformation.py b/tests/unit/llms/vertex_ai/files/test_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/files/test_transformation.py rename to tests/unit/llms/vertex_ai/files/test_transformation.py diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_handler.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_handler.py index e0f0b7e5c0b..9b7cd127b83 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_handler.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_handler.py @@ -208,6 +208,7 @@ class TestVertexAIFilesHandler: assert service_account == "/model/sa.json" def test_resolve_read_gcs_config_falls_back_to_env(self, monkeypatch): + monkeypatch.delenv("GCS_BATCH_BUCKET_NAME", raising=False) monkeypatch.setenv("GCS_BUCKET_NAME", "env-default-bucket") monkeypatch.setenv("GCS_PATH_SERVICE_ACCOUNT", "/env/sa.json") @@ -216,6 +217,40 @@ class TestVertexAIFilesHandler: assert bucket == "env-default-bucket" assert service_account == "/env/sa.json" + def test_resolve_read_gcs_config_prefers_batch_env_over_logging_env(self, monkeypatch): + monkeypatch.setenv("GCS_BATCH_BUCKET_NAME", "batch-bucket") + monkeypatch.setenv("GCS_BUCKET_NAME", "logging-bucket") + + bucket, _ = self.handler._resolve_read_gcs_config(litellm_params={}, vertex_credentials=None) + + assert bucket == "batch-bucket" + + def test_resolve_read_gcs_config_prefers_per_model_bucket_over_batch_env(self, monkeypatch): + monkeypatch.setenv("GCS_BATCH_BUCKET_NAME", "batch-bucket") + + bucket, _ = self.handler._resolve_read_gcs_config( + litellm_params={"gcs_bucket_name": "my-model-bucket"}, + vertex_credentials=None, + ) + + assert bucket == "my-model-bucket" + + def test_resolve_read_gcs_config_prefers_gcs_bucket_name_over_legacy(self): + bucket, _ = self.handler._resolve_read_gcs_config( + litellm_params={"gcs_bucket_name": "my-model-bucket", "bucket_name": "legacy-bucket"}, + vertex_credentials=None, + ) + + assert bucket == "my-model-bucket" + + def test_resolve_read_gcs_config_accepts_legacy_bucket_name_alone(self): + bucket, _ = self.handler._resolve_read_gcs_config( + litellm_params={"bucket_name": "legacy-bucket"}, + vertex_credentials=None, + ) + + assert bucket == "legacy-bucket" + def test_resolve_read_gcs_config_serializes_dict_credentials(self, monkeypatch): monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 7434eae72a4..6f18a391f7b 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -1186,10 +1186,21 @@ class TestConfiguredBucketNameResolution: assert config._get_configured_bucket_name({"gcs_bucket_name": "new", "bucket_name": "legacy"}) == "new" def test_should_fall_back_to_env(self, config, monkeypatch): + monkeypatch.delenv("GCS_BATCH_BUCKET_NAME", raising=False) monkeypatch.setenv("GCS_BUCKET_NAME", "env-bucket") assert config._get_configured_bucket_name({}) == "env-bucket" + def test_should_prefer_batch_env_over_logging_env(self, config, monkeypatch): + monkeypatch.setenv("GCS_BATCH_BUCKET_NAME", "batch-bucket") + monkeypatch.setenv("GCS_BUCKET_NAME", "logging-bucket") + assert config._get_configured_bucket_name({}) == "batch-bucket" + + def test_should_prefer_litellm_params_over_batch_env(self, config, monkeypatch): + monkeypatch.setenv("GCS_BATCH_BUCKET_NAME", "batch-bucket") + assert config._get_configured_bucket_name({"gcs_bucket_name": "per-model-bucket"}) == "per-model-bucket" + def test_should_raise_when_no_bucket_anywhere(self, config, monkeypatch): + monkeypatch.delenv("GCS_BATCH_BUCKET_NAME", raising=False) monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) with pytest.raises(ValueError, match="GCS bucket_name is required"): config._get_configured_bucket_name({}) diff --git a/tests/unit/llms/vertex_ai/gemini/__init__.py b/tests/unit/llms/vertex_ai/gemini/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_context_circulation.py b/tests/unit/llms/vertex_ai/gemini/test_context_circulation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/gemini/test_context_circulation.py rename to tests/unit/llms/vertex_ai/gemini/test_context_circulation.py diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_function_call_args_serialization.py b/tests/unit/llms/vertex_ai/gemini/test_function_call_args_serialization.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/gemini/test_function_call_args_serialization.py rename to tests/unit/llms/vertex_ai/gemini/test_function_call_args_serialization.py diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_gemini_image_url_missing_field.py b/tests/unit/llms/vertex_ai/gemini/test_gemini_image_url_missing_field.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/gemini/test_gemini_image_url_missing_field.py rename to tests/unit/llms/vertex_ai/gemini/test_gemini_image_url_missing_field.py diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_gemini_streaming_tool_call_finish_reason.py b/tests/unit/llms/vertex_ai/gemini/test_gemini_streaming_tool_call_finish_reason.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/gemini/test_gemini_streaming_tool_call_finish_reason.py rename to tests/unit/llms/vertex_ai/gemini/test_gemini_streaming_tool_call_finish_reason.py diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_grounding_requests.py b/tests/unit/llms/vertex_ai/gemini/test_grounding_requests.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/gemini/test_grounding_requests.py rename to tests/unit/llms/vertex_ai/gemini/test_grounding_requests.py diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py b/tests/unit/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py rename to tests/unit/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_tool_call_followed_by_text_assistant.py b/tests/unit/llms/vertex_ai/gemini/test_tool_call_followed_by_text_assistant.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/gemini/test_tool_call_followed_by_text_assistant.py rename to tests/unit/llms/vertex_ai/gemini/test_tool_call_followed_by_text_assistant.py diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py b/tests/unit/llms/vertex_ai/gemini/test_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py rename to tests/unit/llms/vertex_ai/gemini/test_transformation.py diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py new file mode 100644 index 00000000000..4f23ac1773a --- /dev/null +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -0,0 +1,2729 @@ +import base64 + +import pytest + +from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_result, +) +from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + _transform_request_body, + check_if_part_exists_in_parts, + _get_highest_media_resolution, + _extract_max_media_resolution_from_messages, +) +from litellm.types.llms.vertex_ai import BlobType +from litellm.types.utils import Message + + +def test_check_if_part_exists_in_parts(): + parts = [ + {"text": "Hello", "thought": True}, + {"text": "World", "thought": False}, + ] + part = {"text": "Hello", "thought": True} + new_part = {"text": "Hello World", "thought": True} + assert check_if_part_exists_in_parts(parts, part) + assert not check_if_part_exists_in_parts(parts, new_part, ["thought"]) + assert check_if_part_exists_in_parts(parts, new_part, ["text"]) + + +def test_check_if_part_exists_in_parts_camel_case_snake_case(): + """Test that function handles both camelCase and snake_case key variations""" + # Test snake_case to camelCase matching + parts_with_snake_case = [ + { + "function_call": { + "name": "get_current_weather", + "args": {"location": "San Francisco, CA"}, + } + }, + {"text": "Some other content"}, + ] + + part_with_camel_case = { + "functionCall": { + "name": "get_current_weather", + "args": {"location": "San Francisco, CA"}, + } + } + + # Should find match between function_call and functionCall + assert check_if_part_exists_in_parts(parts_with_snake_case, part_with_camel_case) + + # Test camelCase to snake_case matching + parts_with_camel_case = [ + {"functionCall": {"name": "calculate_sum", "args": {"a": 1, "b": 2}}} + ] + + part_with_snake_case = { + "function_call": {"name": "calculate_sum", "args": {"a": 1, "b": 2}} + } + + # Should find match between functionCall and function_call + assert check_if_part_exists_in_parts(parts_with_camel_case, part_with_snake_case) + + # Test no match when values differ + part_with_different_values = { + "function_call": {"name": "different_function", "args": {"x": 5}} + } + + assert not check_if_part_exists_in_parts( + parts_with_snake_case, part_with_different_values + ) + + # Test multiple keys with mixed casing + parts_mixed = [ + { + "function_call": {"name": "test"}, + "thoughtSignature": "reasoning", + "text": "content", + } + ] + + part_mixed_casing = { + "functionCall": {"name": "test"}, + "thought_signature": "reasoning", + "text": "content", + } + + assert check_if_part_exists_in_parts(parts_mixed, part_mixed_casing) + + +def test_cached_content_respects_modify_params_for_cache_incompatible_fields(): + """Regression: cachedContent drops system/tools/toolConfig only when modify_params=True.""" + import litellm + + cache_name = "projects/p/locations/us-central1/cachedContents/abc123" + messages = [ + {"role": "system", "content": "You are helpful"}, + {"role": "user", "content": "hi"}, + ] + optional_params = { + "tools": [ + { + "functionDeclarations": [ + {"name": "get_weather", "description": "Get weather"}, + ] + } + ], + "tool_choice": {"functionCallingConfig": {"mode": "AUTO"}}, + } + + original_modify_params = litellm.modify_params + try: + # With modify_params=False (default), keep fields even with cachedContent. + litellm.modify_params = False + result = _transform_request_body( + messages=list(messages), + model="gemini-2.5-pro", + optional_params=dict(optional_params), + custom_llm_provider="vertex_ai", + litellm_params={}, + cached_content=cache_name, + ) + assert result.get("cachedContent") == cache_name + assert "system_instruction" in result + assert "tools" in result + assert "toolConfig" in result + assert "contents" in result + + # With modify_params=True, drop cache-incompatible fields. + litellm.modify_params = True + result_modify_true = _transform_request_body( + messages=list(messages), + model="gemini-2.5-pro", + optional_params=dict(optional_params), + custom_llm_provider="vertex_ai", + litellm_params={}, + cached_content=cache_name, + ) + assert result_modify_true.get("cachedContent") == cache_name + assert "system_instruction" not in result_modify_true + assert "tools" not in result_modify_true + assert "toolConfig" not in result_modify_true + assert "contents" in result_modify_true + + # Without cache, fields are always included. + result_no_cache = _transform_request_body( + messages=list(messages), + model="gemini-2.5-pro", + optional_params=dict(optional_params), + custom_llm_provider="vertex_ai", + litellm_params={}, + cached_content=None, + ) + assert "system_instruction" in result_no_cache + assert "tools" in result_no_cache + assert "toolConfig" in result_no_cache + finally: + litellm.modify_params = original_modify_params + + +# Tests for issue #14556: Labels field provider-aware filtering +def test_google_genai_excludes_labels(): + """Test that Google GenAI/AI Studio endpoints exclude labels when custom_llm_provider='gemini'""" + messages = [{"role": "user", "content": "test"}] + optional_params = {"labels": {"project": "test", "team": "ai"}} + litellm_params = {} + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params, + custom_llm_provider="gemini", + litellm_params=litellm_params, + cached_content=None, + ) + + # Google GenAI/AI Studio should NOT include labels + assert "labels" not in result + assert "contents" in result + + +def test_vertex_ai_includes_labels(): + """Test that Vertex AI endpoints include labels when custom_llm_provider='vertex_ai'""" + messages = [{"role": "user", "content": "test"}] + optional_params = {"labels": {"project": "test", "team": "ai"}} + litellm_params = {} + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params, + custom_llm_provider="vertex_ai", + litellm_params=litellm_params, + cached_content=None, + ) + + # Vertex AI SHOULD include labels + assert "labels" in result + assert result["labels"] == {"project": "test", "team": "ai"} + + +def test_service_tier_forwarded_to_vertex_ai(): + """Test that service_tier in optional_params is mapped to serviceTier in request body.""" + messages = [{"role": "user", "content": "test"}] + optional_params = {"service_tier": "flex"} + litellm_params = {} + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params, + custom_llm_provider="vertex_ai", + litellm_params=litellm_params, + cached_content=None, + ) + + assert "serviceTier" in result + assert result["serviceTier"] == "flex" + + +def test_extra_body_cache_not_forwarded_to_vertex_ai(): + """ + 'cache' inside extra_body is a LiteLLM-internal proxy caching control. + It must NOT be forwarded to the Vertex AI request body. + + Regression test for: "Invalid JSON payload received. Unknown name \"cache\": Cannot find field." + Vertex AI enforces a strict JSON schema and rejects any unknown field. + """ + messages = [{"role": "user", "content": "test"}] + optional_params = { + "extra_body": { + "cache": {"use-cache": True, "ttl": 86400}, # LiteLLM-internal + "some_vertex_param": "value", # legitimate provider extra + }, + } + litellm_params = {} + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params, + custom_llm_provider="vertex_ai", + litellm_params=litellm_params, + cached_content=None, + ) + + # 'cache' must be stripped — Vertex AI has no such field + assert "cache" not in result, ( + "extra_body.cache must not be forwarded to Vertex AI. " + 'Vertex AI rejects it with 400: Unknown name "cache": Cannot find field.' + ) + + # Other legitimate extra_body keys should still pass through + assert "some_vertex_param" in result + assert result["some_vertex_param"] == "value" + + # Core request fields must be present + assert "contents" in result + + +def test_extra_body_tags_not_forwarded_to_vertex_ai(): + """ + 'tags' inside extra_body is a LiteLLM-internal param for logging/tracking. + It must NOT be forwarded to the Vertex AI request body. + Documented in litellm_proxy.md: "Send tags by including them in the extra_body parameter" + """ + messages = [{"role": "user", "content": "test"}] + optional_params = { + "extra_body": { + "tags": ["user:alice", "env:prod"], + "custom_param": "allowed", + }, + } + litellm_params = {} + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params, + custom_llm_provider="vertex_ai", + litellm_params=litellm_params, + cached_content=None, + ) + + assert "tags" not in result + assert "custom_param" in result + assert result["custom_param"] == "allowed" + + +def test_extra_body_google_maps_rewrites_json_response_format(): + messages = [{"role": "user", "content": "test"}] + optional_params = { + "response_mime_type": "application/json", + "response_schema": { + "type": "object", + "properties": {"answer": {"type": "string"}}, + }, + "extra_body": { + "tools": [{"googleMaps": {}}], + }, + } + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params, + custom_llm_provider="vertex_ai", + litellm_params={}, + cached_content=None, + ) + + generation_config = result["generationConfig"] + assert "response_mime_type" not in generation_config + assert generation_config["responseFormat"] == { + "text": { + "mimeType": "APPLICATION_JSON", + "schema": { + "type": "object", + "properties": {"answer": {"type": "string"}}, + }, + } + } + + +def test_extra_body_generation_config_cannot_restore_google_maps_json_mime_type(): + messages = [{"role": "user", "content": "test"}] + optional_params = { + "tools": [{"googleMaps": {}}], + "response_mime_type": "application/json", + "extra_body": { + "generationConfig": { + "response_mime_type": "application/json", + "response_json_schema": { + "type": "object", + "properties": {"answer": {"type": "string"}}, + }, + }, + }, + } + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params, + custom_llm_provider="vertex_ai", + litellm_params={}, + cached_content=None, + ) + + generation_config = result["generationConfig"] + assert "response_mime_type" not in generation_config + assert "response_json_schema" not in generation_config + assert generation_config["responseFormat"] == { + "text": { + "mimeType": "APPLICATION_JSON", + "schema": { + "type": "object", + "properties": {"answer": {"type": "string"}}, + }, + } + } + + +def test_metadata_to_labels_vertex_only(): + """Test that metadata->labels conversion only happens for Vertex AI""" + messages = [{"role": "user", "content": "test"}] + optional_params = {} + litellm_params = { + "metadata": { + "requester_metadata": {"user": "john_doe", "project": "test-project"} + } + } + + # Google GenAI/AI Studio should not include labels from metadata + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params.copy(), + custom_llm_provider="gemini", + litellm_params=litellm_params.copy(), + cached_content=None, + ) + assert "labels" not in result + + # Vertex AI should include labels from metadata + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params.copy(), + custom_llm_provider="vertex_ai", + litellm_params=litellm_params.copy(), + cached_content=None, + ) + assert "labels" in result + assert result["labels"] == {"user": "john_doe", "project": "test-project"} + + +def test_empty_content_handling(): + """Test that empty content strings are properly handled in Gemini message transformation""" + # Test with empty content in user message + messages = [{"content": "", "role": "user"}] + + contents = _gemini_convert_messages_with_history(messages=messages) + + # Verify that the content was properly transformed + assert len(contents) == 1 + assert contents[0]["role"] == "user" + assert len(contents[0]["parts"]) == 1 + assert "text" in contents[0]["parts"][0] + assert contents[0]["parts"][0]["text"] == "" + + +def test_thought_signature_extraction_from_response(): + """Test that thought signatures are extracted from Gemini response parts and stored in provider_specific_fields""" + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + from litellm.types.llms.vertex_ai import HttpxPartType + + # Test case: Single function call with thought signature + test_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" + + parts_with_signature = [ + HttpxPartType( + functionCall={ + "name": "get_current_temperature", + "args": {"location": "Paris"}, + }, + thoughtSignature=test_signature, + ) + ] + + function, tools, _ = VertexGeminiConfig._transform_parts( + parts=parts_with_signature, + cumulative_tool_call_idx=0, + is_function_call=False, + ) + + # Verify thought signature is stored in provider_specific_fields + assert tools is not None + assert len(tools) == 1 + assert "provider_specific_fields" in tools[0] + assert tools[0]["provider_specific_fields"]["thought_signature"] == test_signature + + +def test_thought_signature_parallel_function_calls(): + """Test that only the first function call in parallel calls has thought signature""" + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + from litellm.types.llms.vertex_ai import HttpxPartType + + test_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" + + # Parallel function calls - only first has signature + parts_parallel = [ + HttpxPartType( + functionCall={ + "name": "get_current_temperature", + "args": {"location": "Paris"}, + }, + thoughtSignature=test_signature, # First FC has signature + ), + HttpxPartType( + functionCall={ + "name": "get_current_temperature", + "args": {"location": "London"}, + }, + # Second FC has no signature (parallel call) + ), + ] + + function, tools, _ = VertexGeminiConfig._transform_parts( + parts=parts_parallel, + cumulative_tool_call_idx=0, + is_function_call=False, + ) + + # Verify only first tool call has thought signature + assert tools is not None + assert len(tools) == 2 + assert "provider_specific_fields" in tools[0] + assert tools[0]["provider_specific_fields"]["thought_signature"] == test_signature + # Second tool call should not have thought signature + assert "provider_specific_fields" not in tools[ + 1 + ] or "thought_signature" not in tools[1].get("provider_specific_fields", {}) + + +def test_thought_signature_preservation_in_conversion(): + """Test that thought signatures are preserved when converting assistant messages back to Gemini format""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + test_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" + + # Assistant message with tool calls containing thought signatures + assistant_message = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc123", + "type": "function", + "function": { + "name": "get_current_temperature", + "arguments": '{"location": "Paris"}', + }, + "index": 0, + "provider_specific_fields": { + "thought_signature": test_signature, + }, + }, + { + "id": "call_def456", + "type": "function", + "function": { + "name": "get_current_temperature", + "arguments": '{"location": "London"}', + }, + "index": 1, + # No thought signature for parallel call + }, + ], + } + + gemini_parts = convert_to_gemini_tool_call_invoke(assistant_message) + + # Verify thought signature is preserved in first function call part + assert len(gemini_parts) == 2 + assert "function_call" in gemini_parts[0] + assert "thoughtSignature" in gemini_parts[0] + assert gemini_parts[0]["thoughtSignature"] == test_signature + + # Verify second function call part does not have thought signature + assert "function_call" in gemini_parts[1] + assert "thoughtSignature" not in gemini_parts[1] + + +def test_thought_signature_sequential_function_calls(): + """Test that each sequential function call preserves its own thought signature""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + signature_1 = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" + signature_2 = "DifferentSignatureForSecondCall1234567890ABCDEFGHIJKLMNOPQRSTUVWXYZ" + + # Sequential function calls - each has its own signature + # This simulates a multi-step conversation where each step has a signature + assistant_message_step1 = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_step1", + "type": "function", + "function": { + "name": "check_flight", + "arguments": '{"flight": "AA100"}', + }, + "index": 0, + "provider_specific_fields": { + "thought_signature": signature_1, + }, + }, + ], + } + + assistant_message_step2 = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_step2", + "type": "function", + "function": { + "name": "book_taxi", + "arguments": '{"destination": "airport"}', + }, + "index": 0, + "provider_specific_fields": { + "thought_signature": signature_2, + }, + }, + ], + } + + gemini_parts_step1 = convert_to_gemini_tool_call_invoke(assistant_message_step1) + gemini_parts_step2 = convert_to_gemini_tool_call_invoke(assistant_message_step2) + + # Verify each step preserves its own signature + assert len(gemini_parts_step1) == 1 + assert gemini_parts_step1[0]["thoughtSignature"] == signature_1 + + assert len(gemini_parts_step2) == 1 + assert gemini_parts_step2[0]["thoughtSignature"] == signature_2 + + +def test_thought_signature_with_function_call_mode(): + """Test thought signature extraction in function_call mode (is_function_call=True)""" + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + from litellm.types.llms.vertex_ai import HttpxPartType + + test_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" + + parts_with_signature = [ + HttpxPartType( + functionCall={ + "name": "get_current_weather", + "args": {"location": "Tokyo"}, + }, + thoughtSignature=test_signature, + ) + ] + + function, tools, _ = VertexGeminiConfig._transform_parts( + parts=parts_with_signature, + cumulative_tool_call_idx=0, + is_function_call=True, + ) + + # Verify thought signature is stored in function's provider_specific_fields + assert function is not None + # Function should be dict-like (TypedDict or dict) + assert hasattr(function, "__getitem__") or isinstance(function, dict) + assert "provider_specific_fields" in function + assert function["provider_specific_fields"]["thought_signature"] == test_signature + assert tools is None + + +def test_dummy_signature_added_for_gemini_3_conversation_history(): + """Test that dummy signatures are added when transferring conversation history from older models (like gemini-2.5-flash) to gemini-3.""" + import base64 + + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + # Simulate conversation history from gemini-2.5-flash (no thought signature) + assistant_message_from_older_model = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc123", + "type": "function", + "function": { + "name": "get_current_temperature", + "arguments": '{"location": "Paris"}', + }, + "index": 0, + # No provider_specific_fields - older model doesn't provide signatures + }, + ], + } + + # Convert to Gemini format for gemini-3-pro-preview (should add dummy signature) + gemini_parts = convert_to_gemini_tool_call_invoke( + assistant_message_from_older_model, model="gemini-3-pro-preview" + ) + + # Verify dummy signature is added + assert len(gemini_parts) == 1 + assert "function_call" in gemini_parts[0] + assert "thoughtSignature" in gemini_parts[0] + + # Verify it's the expected dummy signature (base64 encoded "skip_thought_signature_validator") + expected_dummy = base64.b64encode(b"skip_thought_signature_validator").decode( + "utf-8" + ) + assert gemini_parts[0]["thoughtSignature"] == expected_dummy + + +def test_dummy_signature_not_added_for_gemini_2_5(): + """Test that dummy signatures are NOT added when target model is not gemini-3.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + # Simulate conversation history from gemini-2.5-flash (no thought signature) + assistant_message = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc123", + "type": "function", + "function": { + "name": "get_current_temperature", + "arguments": '{"location": "Paris"}', + }, + "index": 0, + # No provider_specific_fields + }, + ], + } + + # Convert to Gemini format for gemini-2.5-flash (should NOT add dummy signature) + gemini_parts = convert_to_gemini_tool_call_invoke( + assistant_message, model="gemini-2.5-flash" + ) + + # Verify no dummy signature is added for non-gemini-3 models + assert len(gemini_parts) == 1 + assert "function_call" in gemini_parts[0] + assert "thoughtSignature" not in gemini_parts[0] + + +def test_dummy_signature_not_added_when_signature_exists(): + """Test that dummy signatures are NOT added when a real signature already exists.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + real_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" + + # Assistant message with existing thought signature + assistant_message_with_signature = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc123", + "type": "function", + "function": { + "name": "get_current_temperature", + "arguments": '{"location": "Paris"}', + "provider_specific_fields": { + "thought_signature": real_signature, + }, + }, + "index": 0, + }, + ], + } + + # Convert to Gemini format for gemini-3-pro-preview + gemini_parts = convert_to_gemini_tool_call_invoke( + assistant_message_with_signature, model="gemini-3-pro-preview" + ) + + # Verify real signature is preserved, not replaced with dummy + assert len(gemini_parts) == 1 + assert "function_call" in gemini_parts[0] + assert "thoughtSignature" in gemini_parts[0] + assert gemini_parts[0]["thoughtSignature"] == real_signature + + +def test_dummy_signature_with_function_call_mode(): + """Test that dummy signatures are added for function_call mode when converting to gemini-3.""" + import base64 + + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + # Assistant message with function_call (not tool_calls) and no signature + assistant_message_function_call = { + "role": "assistant", + "content": None, + "function_call": { + "name": "get_current_temperature", + "arguments": '{"location": "Paris"}', + # No provider_specific_fields + }, + } + + # Convert to Gemini format for gemini-3-pro-preview + gemini_parts = convert_to_gemini_tool_call_invoke( + assistant_message_function_call, model="gemini-3-pro-preview" + ) + + # Verify dummy signature is added + assert len(gemini_parts) == 1 + assert "function_call" in gemini_parts[0] + assert "thoughtSignature" in gemini_parts[0] + + # Verify it's the expected dummy signature + expected_dummy = base64.b64encode(b"skip_thought_signature_validator").decode( + "utf-8" + ) + assert gemini_parts[0]["thoughtSignature"] == expected_dummy + + +def _parallel_tool_calls(*signatures): + return [ + { + "id": f"call_{idx}", + "type": "function", + "function": { + "name": f"tool_{idx}", + "arguments": '{"location": "Paris"}', + **( + {"provider_specific_fields": {"thought_signature": signature}} + if signature is not None + else {} + ), + }, + "index": idx, + } + for idx, signature in enumerate(signatures) + ] + + +def _parallel_tool_calls_signed_via_id(*signatures): + """Parallel tool calls in the shape LiteLLM actually hands back to clients. + + The signature rides in the tool call id behind __thought__, which is what an + OpenAI-format client echoes back on the next turn. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + _encode_tool_call_id_with_signature, + ) + + return [ + { + "id": _encode_tool_call_id_with_signature(f"call_{idx}", signature), + "type": "function", + "function": {"name": f"tool_{idx}", "arguments": '{"location": "Paris"}'}, + "index": idx, + } + for idx, signature in enumerate(signatures) + ] + + +REAL_THOUGHT_SIGNATURE = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n" +PLACEHOLDER_SIGNATURE = base64.b64encode(b"skip_thought_signature_validator").decode( + "utf-8" +) + + +def test_dummy_signature_only_on_first_parallel_tool_call(): + """Google documents the placeholder as a last resort that degrades quality, so an unsigned + parallel turn replayed to gemini-3 gets a budget of exactly one.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(None, None, None), + }, + model="gemini-3-pro-preview", + ) + + assert len(gemini_parts) == 3 + assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + assert "thoughtSignature" not in gemini_parts[2] + + +def test_real_signature_on_first_parallel_tool_call_leaves_siblings_empty(): + """Gemini signs only the first of N parallel function calls, so a faithful replay has + nothing to attach to the siblings.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(REAL_THOUGHT_SIGNATURE, None, None), + }, + model="gemini-3-pro-preview", + ) + + assert len(gemini_parts) == 3 + assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + assert "thoughtSignature" not in gemini_parts[2] + + +def test_real_signature_on_later_parallel_tool_call_is_preserved(): + """Clients may reorder or drop calls, so a signature that lands on a non-first call is + still the model's own and must survive the round trip.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(None, REAL_THOUGHT_SIGNATURE), + }, + model="gemini-3-pro-preview", + ) + + assert len(gemini_parts) == 2 + assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE + assert gemini_parts[1]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + + +def test_no_signatures_on_parallel_tool_calls_for_gemini_2_5(): + """Non-gemini-3 models never get a placeholder signature, on any call.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(None, None), + }, + model="gemini-2.5-flash", + ) + + assert len(gemini_parts) == 2 + assert all("thoughtSignature" not in part for part in gemini_parts) + + +def test_signature_embedded_in_tool_call_id_only_on_first_parallel_call(): + """The production shape: the signature arrives inside the first call's id, siblings have bare ids.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls_signed_via_id( + REAL_THOUGHT_SIGNATURE, None, None + ), + }, + model="gemini-3-pro-preview", + ) + + assert len(gemini_parts) == 3 + assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + assert "thoughtSignature" not in gemini_parts[2] + + +def test_tool_level_provider_specific_fields_signature_leaves_siblings_empty(): + """A signature on the tool call itself, rather than on its function, behaves the same way.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + tool_calls = _parallel_tool_calls(None, None) + tool_calls[0]["provider_specific_fields"] = { + "thought_signature": REAL_THOUGHT_SIGNATURE + } + + gemini_parts = convert_to_gemini_tool_call_invoke( + {"role": "assistant", "content": None, "tool_calls": tool_calls}, + model="gemini-3-pro-preview", + ) + + assert len(gemini_parts) == 2 + assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + + +def test_placeholder_lands_on_first_emitted_part_not_first_tool_call_entry(): + """A non-function entry (e.g. an OpenAI custom tool call) emits no part, so it must not + consume the one placeholder slot and leave the real first function call bare.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + tool_calls = [ + {"id": "call_custom", "type": "custom", "custom": {"name": "noop", "input": ""}} + ] + _parallel_tool_calls(None, None) + + gemini_parts = convert_to_gemini_tool_call_invoke( + {"role": "assistant", "content": None, "tool_calls": tool_calls}, + model="gemini-3-pro-preview", + ) + + assert len(gemini_parts) == 2 + assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + + +def test_no_placeholder_when_model_is_unknown(): + """Without a model there is nothing to prove the target needs a placeholder, so none is added.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(None, None), + }, + ) + + assert len(gemini_parts) == 2 + assert all("thoughtSignature" not in part for part in gemini_parts) + + +def test_real_signature_forwarded_to_gemini_2_5_without_placeholder_siblings(): + """Older models still receive a real signature that a client replays, and still get no placeholder.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(REAL_THOUGHT_SIGNATURE, None), + }, + model="gemini-2.5-flash", + ) + + assert len(gemini_parts) == 2 + assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + + +def test_parallel_tool_call_history_replayed_through_full_message_conversion(): + """End to end through the message-history converter, the path a real /chat/completions replay takes.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + messages = [ + {"role": "user", "content": "Weather in Paris, London and Tokyo?"}, + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls_signed_via_id( + REAL_THOUGHT_SIGNATURE, None, None + ), + }, + ] + + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-3-pro-preview" + ) + + model_parts = contents[1]["parts"] + assert len(model_parts) == 3 + assert model_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in model_parts[1] + assert "thoughtSignature" not in model_parts[2] + + +@pytest.mark.parametrize( + "model", + ["gemini-3.5-flash", "vertex_ai/gemini-3.5-flash", "gemini/gemini-3.5-flash"], +) +def test_natively_signed_parallel_turn_never_carries_a_placeholder(model): + """A native gemini-3.5 parallel turn replays with zero skip_thought_signature_validator parts. + + Fabricating the placeholder alongside a real signature is what produced empty text responses + on gemini-3.5 parallel function calling, so the whole payload has to stay placeholder-free. + """ + import json + + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + messages = [ + {"role": "user", "content": "Weather in Paris, London and Tokyo?"}, + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls_signed_via_id( + REAL_THOUGHT_SIGNATURE, None, None + ), + }, + ] + + contents = _gemini_convert_messages_with_history(messages=messages, model=model) + + model_parts = contents[1]["parts"] + assert len(model_parts) == 3 + assert model_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in model_parts[1] + assert "thoughtSignature" not in model_parts[2] + assert PLACEHOLDER_SIGNATURE not in json.dumps(contents) + + +@pytest.mark.parametrize( + "model", + [ + "gemini-3-pro-preview", + "gemini-3-flash-preview", + "gemini-3.1-pro-preview", + "gemini-3.5-flash", + "gemini-3.6-flash", + "gemini-3.7-flash", + "gemini-3.8-flash", + "vertex_ai/gemini-3.5-flash", + "vertex_ai/gemini-3.7-flash", + "vertex_ai/gemini-3.8-flash", + "gemini/gemini-3.5-flash", + "gemini/gemini-3.7-flash", + "gemini/gemini-3.8-flash", + ], +) +def test_placeholder_scoped_to_first_call_across_gemini_3_variants(model): + """The gemini-3 gate is a substring match, so every family member and prefix form has to + land on the same one-placeholder budget rather than only the versions we happened to try.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(None, None, None), + }, + model=model, + ) + + assert len(gemini_parts) == 3 + assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + assert "thoughtSignature" not in gemini_parts[2] + + +def test_signed_text_part_survives_alongside_unsigned_parallel_tool_calls(): + """Text-part and function-call signatures are collected by separate code paths, so scoping the + placeholder must not disturb a real signature that arrived on the text part.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": "Checking all three cities.", + "provider_specific_fields": {"thought_signatures": ["real_25_signature"]}, + "tool_calls": _parallel_tool_calls(None, None, None), + } + + parts = _gemini_convert_messages_with_history( + messages=[msg], model="gemini-3-pro-preview" + )[0]["parts"] + + assert parts[0]["text"] == "Checking all three cities." + assert parts[0]["thoughtSignature"] == "real_25_signature" + assert parts[1]["thoughtSignature"] == PLACEHOLDER_SIGNATURE + assert "thoughtSignature" not in parts[2] + assert "thoughtSignature" not in parts[3] + + +# Tests for media_resolution (detail parameter) handling - Issue #17084 +class TestMediaResolution: + """Tests for media_resolution handling in Gemini 2.x models""" + + def test_get_highest_media_resolution_high_wins(self): + """Test that 'high' resolution takes precedence over 'low'""" + assert _get_highest_media_resolution("low", "high") == "high" + assert _get_highest_media_resolution("high", "low") == "high" + assert _get_highest_media_resolution(None, "high") == "high" + assert _get_highest_media_resolution("high", None) == "high" + + def test_get_highest_media_resolution_low_over_none(self): + """Test that 'low' resolution takes precedence over None""" + assert _get_highest_media_resolution(None, "low") == "low" + assert _get_highest_media_resolution("low", None) == "low" + + def test_get_highest_media_resolution_same_values(self): + """Test handling of same resolution values""" + assert _get_highest_media_resolution("high", "high") == "high" + assert _get_highest_media_resolution("low", "low") == "low" + assert _get_highest_media_resolution(None, None) is None + + def test_get_highest_media_resolution_medium(self): + """Test that 'medium' resolution is correctly ranked between 'low' and 'high'""" + assert _get_highest_media_resolution("low", "medium") == "medium" + assert _get_highest_media_resolution("medium", "low") == "medium" + assert _get_highest_media_resolution("medium", "high") == "high" + assert _get_highest_media_resolution("high", "medium") == "high" + assert _get_highest_media_resolution(None, "medium") == "medium" + assert _get_highest_media_resolution("medium", None) == "medium" + + def test_get_highest_media_resolution_ultra_high(self): + """Test that 'ultra_high' resolution takes precedence over all others""" + assert _get_highest_media_resolution("high", "ultra_high") == "ultra_high" + assert _get_highest_media_resolution("ultra_high", "high") == "ultra_high" + assert _get_highest_media_resolution("medium", "ultra_high") == "ultra_high" + assert _get_highest_media_resolution("low", "ultra_high") == "ultra_high" + assert _get_highest_media_resolution(None, "ultra_high") == "ultra_high" + assert _get_highest_media_resolution("ultra_high", None) == "ultra_high" + + def test_extract_max_media_resolution_single_image_high(self): + """Test extraction of media resolution from single image with detail=high""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,abc123", + "detail": "high", + }, + }, + ], + } + ] + assert _extract_max_media_resolution_from_messages(messages) == "high" + + def test_extract_max_media_resolution_single_image_low(self): + """Test extraction of media resolution from single image with detail=low""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,abc123", + "detail": "low", + }, + }, + ], + } + ] + assert _extract_max_media_resolution_from_messages(messages) == "low" + + def test_extract_max_media_resolution_no_detail(self): + """Test extraction when no detail parameter is provided""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,abc123"}, + }, + ], + } + ] + assert _extract_max_media_resolution_from_messages(messages) is None + + def test_extract_max_media_resolution_multiple_images_mixed(self): + """Test that highest resolution is returned when multiple images have different details""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Compare these images"}, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,abc123", + "detail": "low", + }, + }, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,def456", + "detail": "high", + }, + }, + ], + } + ] + assert _extract_max_media_resolution_from_messages(messages) == "high" + + def test_extract_max_media_resolution_text_only(self): + """Test extraction from messages with no images""" + messages = [ + {"role": "user", "content": "Hello, how are you?"}, + {"role": "assistant", "content": "I'm doing well!"}, + ] + assert _extract_max_media_resolution_from_messages(messages) is None + + def test_transform_request_body_gemini_2x_adds_media_resolution(self): + """Test that media_resolution is added to generationConfig for Gemini 2.x models""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,iVBORw0KGgo=", + "detail": "high", + }, + }, + ], + } + ] + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-flash", + optional_params={}, + custom_llm_provider="gemini", + litellm_params={}, + cached_content=None, + ) + + assert "generationConfig" in result + assert "mediaResolution" in result["generationConfig"] + assert result["generationConfig"]["mediaResolution"] == "MEDIA_RESOLUTION_HIGH" + + def test_transform_request_body_gemini_2x_low_resolution(self): + """Test that low media_resolution is correctly added for Gemini 2.x""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,iVBORw0KGgo=", + "detail": "low", + }, + }, + ], + } + ] + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-flash", + optional_params={}, + custom_llm_provider="gemini", + litellm_params={}, + cached_content=None, + ) + + assert "generationConfig" in result + assert "mediaResolution" in result["generationConfig"] + assert result["generationConfig"]["mediaResolution"] == "MEDIA_RESOLUTION_LOW" + + def test_transform_request_body_gemini_3_no_global_media_resolution(self): + """Test that Gemini 3 models don't add media_resolution to generationConfig (they use per-part)""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,iVBORw0KGgo=", + "detail": "high", + }, + }, + ], + } + ] + + result = _transform_request_body( + messages=messages, + model="gemini-3-pro-preview", + optional_params={}, + custom_llm_provider="gemini", + litellm_params={}, + cached_content=None, + ) + + # Gemini 3 should NOT have mediaResolution in generationConfig + # (it's handled per-part in the content transformation) + if "generationConfig" in result: + assert "mediaResolution" not in result["generationConfig"] + + def test_transform_request_body_no_detail_no_media_resolution(self): + """Test that no mediaResolution is added when detail is not specified""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}, + }, + ], + } + ] + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-flash", + optional_params={}, + custom_llm_provider="gemini", + litellm_params={}, + cached_content=None, + ) + + # When no detail is specified, mediaResolution should not be in generationConfig + if "generationConfig" in result: + assert "mediaResolution" not in result["generationConfig"] + + def test_extract_max_media_resolution_file_type_with_detail(self): + """Test that detail is extracted from file content type, not just image_url""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is in this file?"}, + { + "type": "file", + "file": { + "url": "data:image/png;base64,abc123", + "detail": "high", + }, + }, + ], + } + ] + assert _extract_max_media_resolution_from_messages(messages) == "high" + + def test_extract_max_media_resolution_mixed_image_and_file(self): + """Test that highest detail is returned across both image_url and file types""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Compare these"}, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,abc123", + "detail": "low", + }, + }, + { + "type": "file", + "file": { + "url": "data:image/png;base64,def456", + "detail": "high", + }, + }, + ], + } + ] + assert _extract_max_media_resolution_from_messages(messages) == "high" + + def test_transform_request_body_gemini_1x_no_media_resolution(self): + """Test that Gemini 1.x models don't get mediaResolution in generationConfig""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,iVBORw0KGgo=", + "detail": "high", + }, + }, + ], + } + ] + + result = _transform_request_body( + messages=messages, + model="gemini-1.5-pro", + optional_params={}, + custom_llm_provider="gemini", + litellm_params={}, + cached_content=None, + ) + + # Gemini 1.x should NOT have mediaResolution (not supported) + if "generationConfig" in result: + assert "mediaResolution" not in result["generationConfig"] + + +# Tests for VideoMetadata support across all Gemini models (Issue #25474) +class TestVideoMetadataAllGeminiModels: + """Tests that video_metadata (fps, start_offset, end_offset) works for all Gemini models""" + + def _make_video_messages(self, video_metadata: dict) -> list: + return [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Analyze this video"}, + { + "type": "file", + "file": { + "file_id": "gs://bucket/video.mp4", + "format": "video/mp4", + "video_metadata": video_metadata, + }, + }, + ], + } + ] + + def _get_file_part(self, contents: list) -> dict: + for part in contents[0]["parts"]: + if "file_data" in part: + return part + raise AssertionError("No file part found in contents") + + def test_video_metadata_fps_gemini_2_5_flash(self): + """Gemini 2.5 Flash: fps in video_metadata should be forwarded (Issue #25474)""" + messages = self._make_video_messages({"fps": 5}) + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-2.5-flash" + ) + file_part = self._get_file_part(contents) + assert "video_metadata" in file_part + assert file_part["video_metadata"]["fps"] == 5 + + def test_video_metadata_fps_gemini_2_5_pro(self): + """Gemini 2.5 Pro: fps in video_metadata should be forwarded (Issue #25474)""" + messages = self._make_video_messages({"fps": 10}) + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-2.5-pro" + ) + file_part = self._get_file_part(contents) + assert "video_metadata" in file_part + assert file_part["video_metadata"]["fps"] == 10 + + def test_video_metadata_offsets_gemini_2_5_flash(self): + """Gemini 2.5 Flash: start_offset/end_offset converted to camelCase (Issue #25474)""" + messages = self._make_video_messages( + {"start_offset": "5s", "end_offset": "30s"} + ) + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-2.5-flash" + ) + file_part = self._get_file_part(contents) + assert "video_metadata" in file_part + vm = file_part["video_metadata"] + assert vm["startOffset"] == "5s" + assert vm["endOffset"] == "30s" + + def test_video_metadata_all_fields_gemini_2_5_flash(self): + """Gemini 2.5 Flash: all video_metadata fields forwarded correctly (Issue #25474)""" + messages = self._make_video_messages( + {"fps": 5, "start_offset": "10s", "end_offset": "60s"} + ) + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-2.5-flash" + ) + file_part = self._get_file_part(contents) + assert "video_metadata" in file_part + vm = file_part["video_metadata"] + assert vm["fps"] == 5 + assert vm["startOffset"] == "10s" + assert vm["endOffset"] == "60s" + + def test_video_metadata_gemini_1_5_pro(self): + """Gemini 1.5 Pro: video_metadata should also be forwarded (Issue #25474)""" + messages = self._make_video_messages({"fps": 2}) + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-1.5-pro" + ) + file_part = self._get_file_part(contents) + assert "video_metadata" in file_part + assert file_part["video_metadata"]["fps"] == 2 + + +def test_convert_tool_response_with_base64_image(): + """Test tool response with base64 data URI image.""" + # Create a small test image (1x1 red pixel PNG) + test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + image_data_uri = f"data:image/png;base64,{test_image_base64}" + + # Create tool message with image + tool_message = { + "role": "tool", + "tool_call_id": "call_test123", + "content": [ + { + "type": "text", + "text": '{"url": "https://example.com", "status": "success"}', + }, + {"type": "input_image", "image_url": image_data_uri}, + ], + } + + # Mock last message with tool calls + last_message_with_tool_calls = { + "tool_calls": [ + { + "id": "call_test123", + "function": {"name": "click_at", "arguments": '{"x": 100, "y": 200}'}, + } + ] + } + + # Convert tool response with nested multimodal functionResponse.parts. + result = convert_to_gemini_tool_call_result( + tool_message, last_message_with_tool_calls + ) + + assert isinstance(result, list), "Should return a parts list when media is present" + assert len(result) == 1, "Should return one function_response part" + result_part = result[0] + assert "function_response" in result_part + assert "inline_data" not in result_part + function_response = result_part["function_response"] + assert function_response["name"] == "click_at" + assert "response" in function_response + # Verify JSON response is parsed correctly + assert "url" in function_response["response"] + assert function_response["response"]["url"] == "https://example.com" + + # Check inline_data is nested under functionResponse.parts. + assert "parts" in function_response + assert len(function_response["parts"]) == 1 + inline_data: BlobType = function_response["parts"][0]["inline_data"] + assert "data" in inline_data + assert "mime_type" in inline_data + assert inline_data["mime_type"] == "image/png" + assert inline_data["data"] == test_image_base64 + + +def test_gemini_history_nests_multimodal_tool_response_parts(): + """Full history conversion should not emit sibling inline_data tool result parts.""" + test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + messages = [ + {"role": "user", "content": "Get me an image"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_get_image", + "type": "function", + "function": {"name": "get_image", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_get_image", + "content": [ + {"type": "text", "text": '{"image_ref": "inline"}'}, + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": test_image_base64, + }, + }, + ], + }, + ] + + contents = _gemini_convert_messages_with_history(messages=messages) + + tool_response_parts = contents[-1]["parts"] + assert len(tool_response_parts) == 1 + assert "inline_data" not in tool_response_parts[0] + function_response = tool_response_parts[0]["function_response"] + assert function_response["parts"] == [ + { + "inline_data": { + "data": test_image_base64, + "mime_type": "image/png", + } + } + ] + + +def test_convert_tool_response_text_only(): + """Test tool response with only text (no image).""" + tool_message = { + "role": "tool", + "tool_call_id": "call_test789", + "content": [ + {"type": "text", "text": '{"status": "completed", "result": "success"}'} + ], + } + + last_message_with_tool_calls = { + "tool_calls": [ + { + "id": "call_test789", + "function": {"name": "wait_5_seconds", "arguments": "{}"}, + } + ] + } + + result = convert_to_gemini_tool_call_result( + tool_message, last_message_with_tool_calls + ) + + # Should be a single part (no list) when no image + assert not isinstance(result, list), "Should return single part when no image" + + # Check function_response exists + assert "function_response" in result + function_response = result["function_response"] + assert function_response["name"] == "wait_5_seconds" + # Verify JSON response is parsed correctly + assert "status" in function_response["response"] + assert function_response["response"]["status"] == "completed" + + # Check inline_data does NOT exist (no image provided) + assert "inline_data" not in result + + +def test_file_data_field_order(): + """ + Test that file_data fields are in the correct order (mime_type before file_uri). + + The Gemini API is sensitive to field order in the file_data object. + This test verifies that mime_type comes before file_uri in both: + 1. Dictionary key order + 2. JSON serialization + + Related issue: Gemini API returns 400 INVALID_ARGUMENT when fields are in wrong order. + """ + import json + + from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media + + # Test with HTTPS URL and explicit format (audio file) + file_url = "https://generativelanguage.googleapis.com/v1beta/files/test123" + format = "audio/mpeg" + + result = _process_gemini_media(image_url=file_url, format=format) + + # Verify the result has file_data + assert "file_data" in result + file_data = result["file_data"] + + # Verify both fields are present + assert "mime_type" in file_data + assert "file_uri" in file_data + assert file_data["mime_type"] == "audio/mpeg" + assert file_data["file_uri"] == file_url + + # Verify field order by checking dictionary keys + # In Python 3.7+, dict maintains insertion order + file_data_keys = list(file_data.keys()) + assert file_data_keys.index("mime_type") < file_data_keys.index( + "file_uri" + ), "mime_type must come before file_uri in the file_data dict" + + # Also verify by serializing to JSON string + json_str = json.dumps(file_data) + mime_type_pos = json_str.find('"mime_type"') + file_uri_pos = json_str.find('"file_uri"') + assert ( + mime_type_pos < file_uri_pos + ), "mime_type must appear before file_uri in JSON serialization" + + +def test_file_data_field_order_gcs_urls(): + """Test that GCS URLs also maintain correct field order.""" + import json + + from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media + + # Test with GCS URL + gcs_url = "gs://bucket/audio.mp3" + + result = _process_gemini_media(image_url=gcs_url) + + # Verify the result has file_data + assert "file_data" in result + file_data = result["file_data"] + + # Verify both fields are present + assert "mime_type" in file_data + assert "file_uri" in file_data + + # Verify field order + file_data_keys = list(file_data.keys()) + assert file_data_keys.index("mime_type") < file_data_keys.index( + "file_uri" + ), "mime_type must come before file_uri in the file_data dict" + + +def test_gemini_files_api_uri_without_format(): + """ + Test that Gemini Files API URIs work WITHOUT an explicit format/mime_type. + + When a user uploads a file via the Gemini Files API and then references it + by URI (https://generativelanguage.googleapis.com/v1beta/files/...), + the file is already on Google's servers. These URLs return 403 when + fetched directly, so _process_gemini_media must NOT try to resolve the + MIME type via HTTP. Instead it should pass the URI through as file_data + and let the Gemini API resolve the type from its stored metadata. + + Related issue: https://github.com/BerriAI/litellm/issues/24907 + """ + from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media + + file_url = "https://generativelanguage.googleapis.com/v1beta/files/37eh7rsw1vfe" + + # Should NOT raise — previously this hit the generic https:// handler + # which called _get_image_mime_type_from_url() and got a 403. + result = _process_gemini_media(image_url=file_url) + + assert "file_data" in result + file_data = result["file_data"] + assert file_data["file_uri"] == file_url + # When no format is provided, mime_type should be absent so the + # Gemini API infers it from the stored file metadata. + assert "mime_type" not in file_data + + +def test_gemini_files_api_uri_with_format(): + """ + Test that Gemini Files API URIs correctly forward an explicit format. + + Related issue: https://github.com/BerriAI/litellm/issues/24907 + """ + from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media + + file_url = "https://generativelanguage.googleapis.com/v1beta/files/n1vhxa28lyaw" + + result = _process_gemini_media(image_url=file_url, format="text/plain") + + assert "file_data" in result + file_data = result["file_data"] + assert file_data["file_uri"] == file_url + assert file_data["mime_type"] == "text/plain" + + +def test_extract_file_data_with_path_object(): + """ + Test that filename is correctly extracted from Path objects for MIME type detection. + + When uploading files using Path objects (e.g., Path("speech.mp3")), the filename + must be extracted to enable proper MIME type detection. Without this, files get + uploaded with 'application/octet-stream' instead of the correct MIME type. + + Related issue: Files uploaded with wrong MIME type cause Gemini API to reject + requests where the specified format doesn't match the uploaded file's MIME type. + """ + import os + import tempfile + from pathlib import Path + + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + extract_file_data, + ) + + # Create a temporary MP3 file + with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as tmp: + tmp.write(b"fake mp3 content") + tmp_path = tmp.name + + try: + # Test with Path object + path_obj = Path(tmp_path) + extracted = extract_file_data(path_obj) + + # Verify filename was extracted + assert extracted["filename"] is not None + assert extracted["filename"].endswith(".mp3") + + # Verify MIME type was correctly detected + assert ( + extracted["content_type"] == "audio/mpeg" + ), f"Expected 'audio/mpeg' but got '{extracted['content_type']}'" + + # Verify content was read + assert extracted["content"] == b"fake mp3 content" + + finally: + # Clean up temporary file + os.unlink(tmp_path) + + +def test_extract_file_data_with_pathlib_path(): + """Test that filename is correctly extracted from pathlib.Path inputs. + Bare str paths are rejected — when this runs in a proxy request handler + the value is attacker-controlled and opening it as a path is an LFI.""" + import os + import tempfile + from pathlib import Path + + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + extract_file_data, + ) + + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: + tmp.write(b"fake wav content") + tmp_path = Path(tmp.name) + + try: + extracted = extract_file_data(tmp_path) + + assert extracted["filename"] is not None + assert extracted["filename"].endswith(".wav") + assert extracted["content_type"] in [ + "audio/wav", + "audio/x-wav", + ], f"Expected 'audio/wav' or 'audio/x-wav' but got '{extracted['content_type']}'" + assert extracted["content"] == b"fake wav content" + finally: + os.unlink(str(tmp_path)) + + +def test_extract_file_data_with_tuple_format(): + """Test that tuple format (with explicit content_type) still works correctly.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + extract_file_data, + ) + + # Test with tuple format: (filename, content, content_type) + filename = "test_audio.mp3" + content = b"test audio content" + content_type = "audio/mpeg" + + extracted = extract_file_data((filename, content, content_type)) + + # Verify all fields are correct + assert extracted["filename"] == filename + assert extracted["content"] == content + assert extracted["content_type"] == content_type + + +def test_extract_file_data_fallback_to_octet_stream(): + """Unknown file types fall back to application/octet-stream.""" + import os + import tempfile + from pathlib import Path + + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + extract_file_data, + ) + + with tempfile.NamedTemporaryFile(suffix=".xyz123", delete=False) as tmp: + tmp.write(b"unknown content") + tmp_path = Path(tmp.name) + + try: + extracted = extract_file_data(tmp_path) + + assert extracted["filename"] is not None + assert extracted["filename"].endswith(".xyz123") + assert ( + extracted["content_type"] == "application/octet-stream" + ), f"Expected 'application/octet-stream' for unknown type, got '{extracted['content_type']}'" + finally: + os.unlink(str(tmp_path)) + + +def test_convert_tool_response_with_pdf_file(): + """Test tool response with PDF file content using file_data field.""" + # Create a minimal test PDF (base64 encoded) + test_pdf_base64 = "JVBERi0xLjQKJeLjz9MKMSAwIG9iago8PC9UeXBlL0NhdGFsb2cvUGFnZXMgMiAwIFI+PgplbmRvYmoKdHJhaWxlcgo8PC9TaXplIDQvUm9vdCAxIDAgUj4+CnN0YXJ0eHJlZgoyMTYKJSVFT0Y=" + file_data_uri = f"data:application/pdf;base64,{test_pdf_base64}" + + # Create tool message with file + tool_message = { + "role": "tool", + "tool_call_id": "call_pdf_test", + "content": [ + {"type": "text", "text": '{"status": "success", "pages": 1}'}, + {"type": "file", "file_data": file_data_uri}, + ], + } + + # Mock last message with tool calls + last_message_with_tool_calls = { + "tool_calls": [ + { + "id": "call_pdf_test", + "function": { + "name": "analyze_document", + "arguments": '{"path": "/tmp/doc.pdf"}', + }, + } + ] + } + + # Convert tool response with nested multimodal functionResponse.parts. + result = convert_to_gemini_tool_call_result( + tool_message, last_message_with_tool_calls + ) + + assert isinstance(result, list), "Should return a parts list when media is present" + assert len(result) == 1, "Should return one function_response part" + result_part = result[0] + assert "function_response" in result_part + assert "inline_data" not in result_part + function_response = result_part["function_response"] + assert function_response["name"] == "analyze_document" + assert "response" in function_response + # Verify JSON response is parsed correctly + assert "status" in function_response["response"] + assert function_response["response"]["status"] == "success" + + # Check inline_data is nested under functionResponse.parts. + assert "parts" in function_response + assert len(function_response["parts"]) == 1 + inline_data: BlobType = function_response["parts"][0]["inline_data"] + assert "data" in inline_data + assert "mime_type" in inline_data + assert inline_data["mime_type"] == "application/pdf" + assert inline_data["data"] == test_pdf_base64 + + +def test_convert_tool_response_with_input_file_type(): + """Test tool response with input_file content type (Responses API format).""" + # Create a minimal test PDF (base64 encoded) + test_pdf_base64 = "JVBERi0xLjQKJeLjz9MKMSAwIG9iago8PC9UeXBlL0NhdGFsb2cvUGFnZXMgMiAwIFI+PgplbmRvYmoKdHJhaWxlcgo8PC9TaXplIDQvUm9vdCAxIDAgUj4+CnN0YXJ0eHJlZgoyMTYKJSVFT0Y=" + file_data_uri = f"data:application/pdf;base64,{test_pdf_base64}" + + # Create tool message with input_file type + tool_message = { + "role": "tool", + "tool_call_id": "call_input_file_test", + "content": [{"type": "input_file", "file_data": file_data_uri}], + } + + # Mock last message with tool calls + last_message_with_tool_calls = { + "tool_calls": [ + { + "id": "call_input_file_test", + "function": {"name": "read_file", "arguments": "{}"}, + } + ] + } + + # Convert tool response + result = convert_to_gemini_tool_call_result( + tool_message, last_message_with_tool_calls + ) + + # Check inline_data is nested under functionResponse.parts. + assert isinstance(result, list), "Should return a parts list when media is present" + assert len(result) == 1, "Should return one function_response part" + function_response = result[0]["function_response"] + assert ( + function_response["parts"][0]["inline_data"]["mime_type"] == "application/pdf" + ) + + +def test_convert_tool_response_with_nested_file_object(): + """Test tool response with file content using nested file object format.""" + # Create a minimal test PDF (base64 encoded) + test_pdf_base64 = "JVBERi0xLjQKJeLjz9MKMSAwIG9iago8PC9UeXBlL0NhdGFsb2cvUGFnZXMgMiAwIFI+PgplbmRvYmoKdHJhaWxlcgo8PC9TaXplIDQvUm9vdCAxIDAgUj4+CnN0YXJ0eHJlZgoyMTYKJSVFT0Y=" + file_data_uri = f"data:application/pdf;base64,{test_pdf_base64}" + + # Create tool message with nested file object (OpenAI Agents SDK format) + tool_message = { + "role": "tool", + "tool_call_id": "call_nested_test", + "content": [{"type": "file", "file": {"file_data": file_data_uri}}], + } + + # Mock last message with tool calls + last_message_with_tool_calls = { + "tool_calls": [ + { + "id": "call_nested_test", + "function": {"name": "process_document", "arguments": "{}"}, + } + ] + } + + # Convert tool response + result = convert_to_gemini_tool_call_result( + tool_message, last_message_with_tool_calls + ) + + # Check inline_data is nested under functionResponse.parts. + assert isinstance(result, list), "Should return a parts list when media is present" + assert len(result) == 1, "Should return one function_response part" + function_response = result[0]["function_response"] + inline_data: BlobType = function_response["parts"][0]["inline_data"] + assert "data" in inline_data + assert "mime_type" in inline_data + assert inline_data["mime_type"] == "application/pdf" + assert inline_data["data"] == test_pdf_base64 + + +def test_assistant_message_with_images_field(): + """ + Test that assistant messages with images field are properly converted to Gemini format. + + This handles the case where an assistant message contains generated images in the + `images` field (e.g., from image generation models like gemini-2.5-flash-image). + The images should be converted to inline_data parts in the Gemini format. + """ + # Create a small test image (1x1 red pixel PNG) + test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + image_data_uri = f"data:image/png;base64,{test_image_base64}" + + # Create messages with assistant message containing images field + messages = [ + { + "role": "user", + "content": "Generate an image of a banana wearing a costume that says LiteLLM", + }, + { + "role": "assistant", + "content": "Here's your banana in a LiteLLM costume!", + "images": [ + { + "image_url": {"url": image_data_uri, "detail": "auto"}, + "index": 0, + "type": "image_url", + } + ], + }, + ] + + # Convert messages to Gemini format + contents = _gemini_convert_messages_with_history(messages=messages) + + # Verify structure + assert len(contents) == 2, f"Expected 2 content blocks, got {len(contents)}" + + # Verify user message + assert contents[0]["role"] == "user" + assert len(contents[0]["parts"]) == 1 + assert ( + contents[0]["parts"][0]["text"] + == "Generate an image of a banana wearing a costume that says LiteLLM" + ) + + # Verify assistant message + assert contents[1]["role"] == "model" + assert ( + len(contents[1]["parts"]) == 2 + ), f"Expected 2 parts (text + image), got {len(contents[1]['parts'])}" + + # Find text part and inline_data part + text_part = None + inline_data_part = None + for part in contents[1]["parts"]: + if "text" in part: + text_part = part + elif "inline_data" in part: + inline_data_part = part + + # Verify text part + assert text_part is not None, "Missing text part in assistant message" + assert text_part["text"] == "Here's your banana in a LiteLLM costume!" + + # Verify inline_data part (image) + assert inline_data_part is not None, "Missing inline_data part in assistant message" + inline_data: BlobType = inline_data_part["inline_data"] + assert "data" in inline_data + assert "mime_type" in inline_data + assert inline_data["mime_type"] == "image/png" + assert inline_data["data"] == test_image_base64 + + +def test_assistant_message_with_multiple_images(): + """Test that assistant messages with multiple images are properly converted.""" + # Create two test images + test_image1_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + test_image2_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg==" + image1_data_uri = f"data:image/png;base64,{test_image1_base64}" + image2_data_uri = f"data:image/jpeg;base64,{test_image2_base64}" + + messages = [ + {"role": "user", "content": "Generate two images"}, + { + "role": "assistant", + "content": "Here are your images:", + "images": [ + { + "image_url": {"url": image1_data_uri, "detail": "auto"}, + "index": 0, + "type": "image_url", + }, + { + "image_url": {"url": image2_data_uri, "detail": "high"}, + "index": 1, + "type": "image_url", + }, + ], + }, + ] + + # Convert messages to Gemini format + contents = _gemini_convert_messages_with_history(messages=messages) + + # Verify assistant message has 3 parts (1 text + 2 images) + assert contents[1]["role"] == "model" + assert ( + len(contents[1]["parts"]) == 3 + ), f"Expected 3 parts (text + 2 images), got {len(contents[1]['parts'])}" + + # Count inline_data parts + inline_data_parts = [part for part in contents[1]["parts"] if "inline_data" in part] + assert ( + len(inline_data_parts) == 2 + ), f"Expected 2 inline_data parts, got {len(inline_data_parts)}" + + # Verify first image + assert inline_data_parts[0]["inline_data"]["mime_type"] == "image/png" + assert inline_data_parts[0]["inline_data"]["data"] == test_image1_base64 + + # Verify second image + assert inline_data_parts[1]["inline_data"]["mime_type"] == "image/jpeg" + assert inline_data_parts[1]["inline_data"]["data"] == test_image2_base64 + + +def test_assistant_message_with_images_using_message_object(): + """Test that Message objects with images field are properly converted.""" + # Create a small test image + test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + image_data_uri = f"data:image/png;base64,{test_image_base64}" + + # Create messages using Message object (as returned by LiteLLM) + user_message = {"role": "user", "content": "Generate an image"} + + assistant_message = Message( + content="Here's your image!", + role="assistant", + tool_calls=None, + function_call=None, + images=[ + { + "image_url": {"url": image_data_uri, "detail": "auto"}, + "index": 0, + "type": "image_url", + } + ], + ) + + messages = [user_message, assistant_message] + + # Convert messages to Gemini format + contents = _gemini_convert_messages_with_history(messages=messages) + + # Verify assistant message has both text and image + assert contents[1]["role"] == "model" + assert len(contents[1]["parts"]) == 2 + + # Verify image was converted + inline_data_parts = [part for part in contents[1]["parts"] if "inline_data" in part] + assert len(inline_data_parts) == 1 + assert inline_data_parts[0]["inline_data"]["mime_type"] == "image/png" + assert inline_data_parts[0]["inline_data"]["data"] == test_image_base64 + + +def test_assistant_message_with_images_in_conversation_history(): + """ + Test multi-turn conversation where assistant message with images is in history. + + This simulates the real use case where: + 1. User asks for image generation + 2. Assistant generates image (with images field) + 3. User asks follow-up question about the image + """ + test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + image_data_uri = f"data:image/png;base64,{test_image_base64}" + + messages = [ + {"role": "user", "content": "Generate an image of a cat"}, + { + "role": "assistant", + "content": "Here's a cat image:", + "images": [ + { + "image_url": {"url": image_data_uri, "detail": "auto"}, + "index": 0, + "type": "image_url", + } + ], + }, + {"role": "user", "content": "Can you make it more colorful?"}, + ] + + # Convert messages to Gemini format + contents = _gemini_convert_messages_with_history(messages=messages) + + # Verify structure: user -> model (with image) -> user + assert len(contents) == 3 + assert contents[0]["role"] == "user" + assert contents[1]["role"] == "model" + assert contents[2]["role"] == "user" + + # Verify assistant message has image in history + inline_data_parts = [part for part in contents[1]["parts"] if "inline_data" in part] + assert len(inline_data_parts) == 1 + assert inline_data_parts[0]["inline_data"]["mime_type"] == "image/png" + + +def test_function_response_has_user_role(): + """ + Test that function response ContentType blocks include role="user". + + Gemini API only accepts two roles: "user" and "model". Function responses + must be sent with role="user". Previously, LiteLLM omitted the role field + entirely, causing 400 errors from the Gemini API. + + Fixes: https://github.com/BerriAI/litellm/issues/22003 + Fixes: https://github.com/BerriAI/litellm/issues/20690 + """ + messages = [ + {"role": "user", "content": "What is the weather in Berlin?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc123", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "Berlin"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_abc123", + "content": '{"temperature": "15°C", "condition": "Cloudy"}', + }, + ] + + contents = _gemini_convert_messages_with_history(messages=messages) + + # Expect: user -> model (functionCall) -> user (functionResponse) + assert len(contents) == 3 + + assert contents[0]["role"] == "user" + assert contents[1]["role"] == "model" + assert "function_call" in contents[1]["parts"][0] + + # The critical assertion: function response must have role="user" + assert contents[2]["role"] == "user" + assert "function_response" in contents[2]["parts"][0] + + +def test_multi_turn_function_calling_roles(): + """ + Test a full multi-turn function calling conversation produces correct roles. + + Simulates: user asks → model calls tool → tool responds → model answers → user asks again. + Every content block must have an explicit role of "user" or "model". + + Fixes: https://github.com/BerriAI/litellm/issues/22003 + """ + messages = [ + {"role": "user", "content": "What is the weather in Berlin?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_001", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "Berlin"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_001", + "content": '{"temperature": "15°C"}', + }, + { + "role": "assistant", + "content": "The weather in Berlin is 15°C.", + }, + {"role": "user", "content": "And in Paris?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_002", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "Paris"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_002", + "content": '{"temperature": "18°C"}', + }, + ] + + contents = _gemini_convert_messages_with_history(messages=messages) + + # Every content block must have a valid role + for i, content in enumerate(contents): + assert "role" in content, f"Content block {i} missing 'role' field" + assert content["role"] in ( + "user", + "model", + ), f"Content block {i} has invalid role: {content.get('role')}" + + # Verify the function response blocks specifically have role="user" + for i, content in enumerate(contents): + for part in content["parts"]: + if "function_response" in part: + assert ( + content["role"] == "user" + ), f"Content block {i} with function_response has role='{content['role']}', expected 'user'" + + +def test_gemini_thought_signature_preservation_real_response(): + """Test that thought signatures are preserved on the text part if originally there, without dropping or duplicating (real response case).""" + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + real_candidate = { + "content": { + "parts": [ + { + "text": "I will explain and then list files.", + "thoughtSignature": "mock_signature_from_text_part", + }, + { + "functionCall": { + "name": "list_files", + "args": {}, + } + }, + ] + } + } + + parts = real_candidate["content"]["parts"] + + content, reasoning_content = ( + VertexGeminiConfig().get_assistant_content_message(parts=parts) + ) + thought_signatures = ( + VertexGeminiConfig()._extract_thought_signatures_from_parts( + parts=parts + ) + ) + functions, tools, _ = VertexGeminiConfig._transform_parts( + parts=parts, + cumulative_tool_call_idx=0, + is_function_call=False, + ) + + msg: dict = {"role": "assistant"} + if content is not None: + msg["content"] = content + if tools: + msg["tool_calls"] = tools + if functions is not None: + msg["function_call"] = functions + if thought_signatures is not None: + msg["provider_specific_fields"] = { + "thought_signatures": thought_signatures + } + + converted_real = _gemini_convert_messages_with_history( + messages=[msg], + model="gemini-2.5-pro", + ) + + assert len(converted_real) == 1 + assert "parts" in converted_real[0] + parts_out = converted_real[0]["parts"] + assert len(parts_out) == 2 + assert "text" in parts_out[0] + assert ( + parts_out[0]["thoughtSignature"] == "mock_signature_from_text_part" + ) + assert "function_call" in parts_out[1] + assert "thoughtSignature" not in parts_out[1] + + +def test_gemini_thought_signature_deduplication_assumed_response(): + """Test that thought signatures are deduplicated and not attached to the text part if already present in the tool call (assumed response case).""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + pr_assumed_msg = { + "role": "assistant", + "content": "I will list the directory.", + "provider_specific_fields": { + "thought_signatures": ["mock_signature_63k"] + }, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "list_files", "arguments": "{}"}, + "provider_specific_fields": { + "thought_signature": "mock_signature_63k" + }, + } + ], + } + + converted_pr = _gemini_convert_messages_with_history( + messages=[pr_assumed_msg], + model="gemini-2.5-pro", + ) + + assert len(converted_pr) == 1 + assert "parts" in converted_pr[0] + parts_out = converted_pr[0]["parts"] + assert len(parts_out) == 2 + assert "text" in parts_out[0] + assert "thoughtSignature" not in parts_out[0] + assert "function_call" in parts_out[1] + assert parts_out[1]["thoughtSignature"] == "mock_signature_63k" + + +def test_gemini_thought_signature_pure_text(): + """Test that thought signatures are preserved on the text part for responses with no tool calls.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": "Hello, I am a model.", + "provider_specific_fields": { + "thought_signatures": ["pure_text_signature"] + }, + } + + converted = _gemini_convert_messages_with_history( + messages=[msg], + model="gemini-2.5-pro", + ) + + assert len(converted) == 1 + assert "parts" in converted[0] + parts_out = converted[0]["parts"] + assert len(parts_out) == 1 + assert "text" in parts_out[0] + assert parts_out[0]["thoughtSignature"] == "pure_text_signature" + + +def test_gemini_thought_signature_pure_tool_call(): + """Test that thought signatures are preserved on the tool call for responses with no intermediate text.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": None, + "provider_specific_fields": { + "thought_signatures": ["pure_tool_signature"] + }, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "list_files", "arguments": "{}"}, + "provider_specific_fields": { + "thought_signature": "pure_tool_signature" + }, + } + ], + } + + converted = _gemini_convert_messages_with_history( + messages=[msg], + model="gemini-2.5-pro", + ) + + assert len(converted) == 1 + assert "parts" in converted[0] + parts_out = converted[0]["parts"] + assert len(parts_out) == 1 + assert "function_call" in parts_out[0] + assert parts_out[0]["thoughtSignature"] == "pure_tool_signature" + + +def test_gemini_distinct_text_and_tool_signatures_are_both_preserved(): + """A text-part signature that differs from the tool-call signature must stay on the text part.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": "Some analysis.", + "provider_specific_fields": { + "thought_signatures": ["text_signature", "tool_signature"] + }, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "list_files", "arguments": "{}"}, + "provider_specific_fields": {"thought_signature": "tool_signature"}, + } + ], + } + + parts = _gemini_convert_messages_with_history( + messages=[msg], model="gemini-2.5-pro" + )[0]["parts"] + + assert parts[0]["text"] == "Some analysis." + assert parts[0]["thoughtSignature"] == "text_signature" + assert "function_call" in parts[1] + assert parts[1]["thoughtSignature"] == "tool_signature" + + +def test_gemini_25_text_signature_survives_replay_to_gemini_3(): + """gemini-2.5 history (signed text, unsigned tool call) replayed to gemini-3 keeps the real + text signature; the dummy signature synthesized for the unsigned tool call must not suppress it.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + _get_dummy_thought_signature, + ) + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": "I will list the directory.", + "provider_specific_fields": {"thought_signatures": ["real_25_signature"]}, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "list_files", "arguments": "{}"}, + } + ], + } + + parts = _gemini_convert_messages_with_history(messages=[msg], model="gemini-3-pro")[ + 0 + ]["parts"] + + assert parts[0]["text"] == "I will list the directory." + assert parts[0]["thoughtSignature"] == "real_25_signature" + assert "function_call" in parts[1] + assert parts[1]["thoughtSignature"] == _get_dummy_thought_signature() + + +def test_gemini_function_call_signature_round_trip_no_duplicate(): + """End to end: a gemini-3-style response (unsigned text + signed functionCall) parsed and + re-serialized sends the signature exactly once, on the function-call part.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + response_parts = [ + {"text": "I will calculate the result for you."}, + { + "functionCall": {"name": "add_numbers", "args": {"a": 17, "b": 25}}, + "thoughtSignature": "signature_from_function_call", + }, + ] + + config = VertexGeminiConfig() + content, _ = config.get_assistant_content_message(parts=response_parts) + thought_signatures = config._extract_thought_signatures_from_parts( + parts=response_parts + ) + _, tools, _ = VertexGeminiConfig._transform_parts( + parts=response_parts, cumulative_tool_call_idx=0, is_function_call=False + ) + + msg = { + "role": "assistant", + "content": content, + "tool_calls": tools, + "provider_specific_fields": {"thought_signatures": thought_signatures}, + } + + parts = _gemini_convert_messages_with_history(messages=[msg], model="gemini-3-pro")[ + 0 + ]["parts"] + + signatures = [p["thoughtSignature"] for p in parts if "thoughtSignature" in p] + assert signatures == ["signature_from_function_call"] + assert "thoughtSignature" not in parts[0] + assert "function_call" in parts[1] + + +def test_gemini_server_side_tool_signature_not_duplicated_on_text(): + """A signature already re-injected on a server-side toolCall part is not attached to the text part again.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": "The weather in Buenos Aires is sunny.", + "provider_specific_fields": { + "thought_signatures": ["server_side_signature"], + "server_side_tool_invocations": [ + { + "tool_type": "GOOGLE_SEARCH_WEB", + "id": "abc123", + "args": {"queries": ["weather Buenos Aires"]}, + "response": {"weather": "Sunny"}, + "thought_signature": "server_side_signature", + } + ], + }, + } + + parts = _gemini_convert_messages_with_history( + messages=[msg], model="gemini-2.5-pro" + )[0]["parts"] + + text_part = next(p for p in parts if "text" in p) + assert "thoughtSignature" not in text_part + tool_call_part = next(p for p in parts if "toolCall" in p) + assert tool_call_part["thoughtSignature"] == "server_side_signature" diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py similarity index 99% rename from tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py rename to tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 88ba7fc37d9..739744336a1 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -1894,42 +1894,6 @@ def test_vertex_ai_tool_call_id_format(): ), f"All 10 IDs should be unique, got {len(ids_generated)} unique IDs" -def test_vertex_ai_code_line_length(): - """ - Test that the specific code line generating tool call IDs is within character limit. - - This is a meta-test to ensure the code change meets the 40-character requirement. - """ - import inspect - - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexGeminiConfig, - ) - - # Get the source code of the _transform_parts method - source_lines = inspect.getsource(VertexGeminiConfig._transform_parts).split("\n") - - # Find the line that generates the ID - id_line = None - for line in source_lines: - if '"id": f"call_' in line and "uuid.uuid4().hex[:28]" in line: - id_line = line.strip() # Remove indentation for length check - break - - assert id_line is not None, "Could not find the ID generation line in source code" - - # Check that the line is 40 characters or less (excluding indentation) - line_length = len(id_line) - assert ( - line_length <= 40 - ), f"ID generation line is {line_length} characters, should be ≤40: {id_line}" - - # Verify it contains the expected UUID format - assert ( - "uuid.uuid4().hex[:28]" in id_line - ), f"Line should contain shortened UUID format: {id_line}" - - def test_vertex_ai_map_google_maps_tool_simple(): """ Test googleMaps tool transformation without location data. @@ -2530,8 +2494,6 @@ def test_fine_tuned_endpoint_and_gemma_get_no_gemini_3_default_temperature(model assert "temperature" not in mapped - - def _tool_call_messages(tool_call_id: str): return [ {"role": "user", "content": "hi"}, diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_gemini_unbound_local_error.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_gemini_unbound_local_error.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/gemini/test_vertex_gemini_unbound_local_error.py rename to tests/unit/llms/vertex_ai/gemini/test_vertex_gemini_unbound_local_error.py diff --git a/tests/unit/llms/vertex_ai/image_generation/__init__.py b/tests/unit/llms/vertex_ai/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py rename to tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py diff --git a/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py new file mode 100644 index 00000000000..a72a570c2a2 --- /dev/null +++ b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py @@ -0,0 +1,637 @@ +from unittest.mock import MagicMock, patch + +import httpx + + +from litellm.llms.vertex_ai.image_generation import ( + get_vertex_ai_image_generation_config, +) +from litellm.llms.vertex_ai.image_generation.vertex_gemini_transformation import ( + VertexAIGeminiImageGenerationConfig, +) +from litellm.llms.vertex_ai.image_generation.vertex_imagen_transformation import ( + VertexAIImagenImageGenerationConfig, +) + + +class TestVertexAIGeminiImageGenerationConfig: + def setup_method(self): + """Set up test fixtures""" + self.config = VertexAIGeminiImageGenerationConfig() + + def test_get_supported_openai_params(self): + """Test get_supported_openai_params returns correct params""" + supported = self.config.get_supported_openai_params("gemini-2.5-flash-image") + assert "n" in supported + assert "size" in supported + + def test_map_openai_params_n(self): + """Test mapping n parameter to candidate_count""" + non_default_params = {"n": 3} + optional_params = {} + result = self.config.map_openai_params(non_default_params, optional_params, "gemini-2.5-flash-image", False) + assert result.get("candidate_count") == 3 + + def test_map_openai_params_size(self): + """Test mapping size parameter to aspectRatio""" + non_default_params = {"size": "1024x1024"} + optional_params = {} + result = self.config.map_openai_params(non_default_params, optional_params, "gemini-2.5-flash-image", False) + assert result.get("aspectRatio") == "1:1" + + def test_map_openai_params_size_16_9(self): + """Test mapping 16:9 size""" + non_default_params = {"size": "1792x1024"} + optional_params = {} + result = self.config.map_openai_params(non_default_params, optional_params, "gemini-2.5-flash-image", False) + assert result.get("aspectRatio") == "16:9" + + def test_map_size_to_aspect_ratio(self): + """Test size to aspect ratio mapping""" + assert self.config._map_size_to_aspect_ratio("1024x1024") == "1:1" + assert self.config._map_size_to_aspect_ratio("1792x1024") == "16:9" + assert self.config._map_size_to_aspect_ratio("1024x1792") == "9:16" + assert self.config._map_size_to_aspect_ratio("1280x896") == "4:3" + assert self.config._map_size_to_aspect_ratio("896x1280") == "3:4" + assert self.config._map_size_to_aspect_ratio("unknown") == "1:1" # default + + def test_get_supported_openai_params_includes_native_gemini_params(self): + """Test that native Gemini imageConfig params are supported""" + supported = self.config.get_supported_openai_params("gemini-3-pro-image-preview") + assert "aspectRatio" in supported + assert "aspect_ratio" in supported + assert "imageSize" in supported + assert "image_size" in supported + assert "imageConfig" in supported + + def test_map_openai_params_aspect_ratio_camel_case(self): + """Test mapping native aspectRatio parameter""" + result = self.config.map_openai_params({"aspectRatio": "9:16"}, {}, "gemini-3-pro-image-preview", False) + assert result["aspectRatio"] == "9:16" + + def test_map_openai_params_aspect_ratio_snake_case(self): + """Test mapping native aspect_ratio parameter""" + result = self.config.map_openai_params({"aspect_ratio": "16:9"}, {}, "gemini-3-pro-image-preview", False) + assert result["aspectRatio"] == "16:9" + + def test_map_openai_params_image_size_camel_case(self): + """Test mapping native imageSize parameter""" + result = self.config.map_openai_params({"imageSize": "4K"}, {}, "gemini-3-pro-image-preview", False) + assert result["imageSize"] == "4K" + + def test_map_openai_params_image_size_snake_case(self): + """Test mapping native image_size parameter""" + result = self.config.map_openai_params({"image_size": "2K"}, {}, "gemini-3-pro-image-preview", False) + assert result["imageSize"] == "2K" + + def test_map_openai_params_image_config_dict_stored_whole(self): + """imageConfig dict is stored as-is so all fields survive""" + result = self.config.map_openai_params( + {"imageConfig": {"aspectRatio": "16:9", "imageSize": "2K"}}, + {}, + "gemini-3.1-flash-image", + False, + ) + assert result["imageConfig"] == {"aspectRatio": "16:9", "imageSize": "2K"} + + def test_map_openai_params_image_config_all_fields(self): + """All ImageConfig fields (personGeneration, imageOutputOptions) pass through""" + payload = { + "imageConfig": { + "aspectRatio": "9:16", + "imageSize": "4K", + "personGeneration": "DONT_ALLOW", + "imageOutputOptions": { + "mimeType": "image/jpeg", + "compressionQuality": 80, + }, + } + } + result = self.config.map_openai_params(payload, {}, "gemini-3.1-flash-image", False) + assert result["imageConfig"] == payload["imageConfig"] + + def test_map_openai_params_image_config_non_dict_warns_and_drops(self): + """Non-dict imageConfig is dropped with a warning, not silently discarded""" + with patch("litellm.llms.vertex_ai.image_generation.vertex_gemini_transformation.verbose_logger") as mock_log: + result = self.config.map_openai_params( + {"imageConfig": "bad-string-value"}, {}, "gemini-3.1-flash-image", False + ) + assert "imageConfig" not in result + mock_log.warning.assert_called_once() + + def test_transform_image_generation_request_from_image_config(self): + """Full imageConfig dict is forwarded verbatim into generationConfig""" + full_config = { + "aspectRatio": "16:9", + "imageSize": "2K", + "personGeneration": "DONT_ALLOW", + "imageOutputOptions": {"mimeType": "image/jpeg", "compressionQuality": 85}, + } + mapped = self.config.map_openai_params( + {"imageConfig": full_config}, + {}, + "gemini-3.1-flash-image", + False, + ) + request = self.config.transform_image_generation_request( + model="gemini-3.1-flash-image", + prompt="A nano banana on a desk", + optional_params=mapped, + litellm_params={}, + headers={}, + ) + assert request["generationConfig"]["imageConfig"] == full_config + + def test_transform_image_generation_flat_params_override_image_config(self): + """Explicit flat params win over the same key inside imageConfig""" + request = self.config.transform_image_generation_request( + model="gemini-3.1-flash-image", + prompt="A nano banana", + optional_params={ + "imageConfig": {"aspectRatio": "1:1", "personGeneration": "DONT_ALLOW"}, + "aspectRatio": "16:9", # should win + }, + litellm_params={}, + headers={}, + ) + assert request["generationConfig"]["imageConfig"]["aspectRatio"] == "16:9" + assert request["generationConfig"]["imageConfig"]["personGeneration"] == "DONT_ALLOW" + + def test_transform_image_generation_request_basic(self): + """Test basic request transformation""" + request = self.config.transform_image_generation_request( + model="gemini-2.5-flash-image", + prompt="A nano banana", + optional_params={}, + litellm_params={}, + headers={}, + ) + assert "contents" in request + assert "generationConfig" in request + assert request["generationConfig"]["responseModalities"] == ["IMAGE"] + assert request["contents"][0]["parts"][0]["text"] == "A nano banana" + + def test_transform_image_generation_request_with_aspect_ratio(self): + """Test request transformation with aspectRatio""" + request = self.config.transform_image_generation_request( + model="gemini-2.5-flash-image", + prompt="A nano banana", + optional_params={"aspectRatio": "16:9"}, + litellm_params={}, + headers={}, + ) + assert request["generationConfig"]["imageConfig"]["aspectRatio"] == "16:9" + + def test_transform_image_generation_request_with_image_size(self): + """Test request transformation with imageSize (Gemini 3 Pro)""" + request = self.config.transform_image_generation_request( + model="gemini-3-pro-image-preview", + prompt="A nano banana", + optional_params={"imageSize": "4K"}, + litellm_params={}, + headers={}, + ) + assert request["generationConfig"]["imageConfig"]["imageSize"] == "4K" + + def test_map_openai_params_web_search_options(self): + """Test web_search_options maps to googleSearch tool""" + result = self.config.map_openai_params({"web_search_options": {}}, {}, "gemini-3.1-flash-image-preview", False) + assert result["tools"] == [{"googleSearch": {}}] + + def test_transform_image_generation_request_with_web_search_tools(self): + """Test request transformation includes googleSearch tools""" + request = self.config.transform_image_generation_request( + model="gemini-3.1-flash-image-preview", + prompt="Generate an image of the latest iPhone", + optional_params={"tools": [{"googleSearch": {}}]}, + litellm_params={}, + headers={}, + ) + assert request["tools"] == [{"googleSearch": {}}] + + def test_transform_image_generation_request_forwards_tool_config(self): + """Test request transformation forwards toolConfig side-effects from tool mapping""" + mapped = self.config.map_openai_params( + {"tools": [{"googleMaps": {"latitude": 37.7, "longitude": -122.4}}]}, + {}, + "gemini-3.1-flash-image-preview", + False, + ) + request = self.config.transform_image_generation_request( + model="gemini-3.1-flash-image-preview", + prompt="Generate an image of a coffee shop nearby", + optional_params=mapped, + litellm_params={}, + headers={}, + ) + assert request["tools"] == [{"googleMaps": {}}] + assert request["toolConfig"] == {"retrievalConfig": {"latLng": {"latitude": 37.7, "longitude": -122.4}}} + + def test_transform_image_generation_request_with_candidate_count(self): + """Test request transformation with candidate_count""" + request = self.config.transform_image_generation_request( + model="gemini-2.5-flash-image", + prompt="A nano banana", + optional_params={"candidate_count": 2}, + litellm_params={}, + headers={}, + ) + assert request["generationConfig"]["candidateCount"] == 2 + + def test_transform_image_generation_request_with_n(self): + """Test request transformation with n parameter""" + request = self.config.transform_image_generation_request( + model="gemini-2.5-flash-image", + prompt="A nano banana", + optional_params={"n": 2}, + litellm_params={}, + headers={}, + ) + assert request["generationConfig"]["candidateCount"] == 2 + + def test_transform_image_generation_response(self): + """Test response transformation""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = { + "candidates": [ + { + "content": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": "base64_encoded_image_data", + } + } + ] + } + } + ], + "usageMetadata": { + "promptTokenCount": 93, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 54, + }, + { + "modality": "IMAGE", + "tokenCount": 39, + }, + ], + "candidatesTokenCount": 17, + "totalTokenCount": 110, + }, + } + mock_response.headers = {} + + from litellm.types.utils import ImageResponse + + model_response = ImageResponse() + result = self.config.transform_image_generation_response( + model="gemini-2.5-flash-image", + raw_response=mock_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert len(result.data) == 1 + assert result.data[0].b64_json == "base64_encoded_image_data" + assert result.data[0].url is None + assert result.usage.input_tokens == 93 + assert result.usage.input_tokens_details.text_tokens == 54 + assert result.usage.input_tokens_details.image_tokens == 39 + assert result.usage.output_tokens == 17 + assert result.usage.total_tokens == 110 + + def test_transform_image_generation_response_multiple_images(self): + """Test response transformation with multiple images""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = { + "candidates": [ + { + "content": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": "image1", + } + }, + { + "inlineData": { + "mimeType": "image/png", + "data": "image2", + } + }, + ] + } + } + ] + } + mock_response.headers = {} + + from litellm.types.utils import ImageResponse + + model_response = ImageResponse() + result = self.config.transform_image_generation_response( + model="gemini-2.5-flash-image", + raw_response=mock_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert len(result.data) == 2 + assert result.data[0].b64_json == "image1" + assert result.data[1].b64_json == "image2" + + def test_transform_image_generation_response_signature(self): + """Test response transformation includes thoughtSignature for Gemini 3 Pro""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = { + "candidates": [ + { + "content": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": "base64_encoded_image_data", + }, + "thoughtSignature": "test_signature_abc123", + } + ] + } + } + ] + } + mock_response.headers = {} + + from litellm.types.utils import ImageResponse + + model_response = ImageResponse() + result = self.config.transform_image_generation_response( + model="gemini-3-pro-image-preview", + raw_response=mock_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert len(result.data) == 1 + assert result.data[0].b64_json == "base64_encoded_image_data" + assert result.data[0].provider_specific_fields["thought_signature"] == "test_signature_abc123" + + def test_transform_image_generation_response_tracks_web_search_requests(self): + """Grounding queries are carried onto usage so search spend can be billed""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = { + "candidates": [ + { + "content": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": "base64_encoded_image_data", + } + } + ] + }, + "groundingMetadata": {"webSearchQueries": ["eiffel tower", "paris skyline"]}, + } + ], + "usageMetadata": { + "promptTokenCount": 93, + "candidatesTokenCount": 17, + "totalTokenCount": 110, + }, + } + mock_response.headers = {} + + from litellm.types.utils import ImageResponse + + result = self.config.transform_image_generation_response( + model="gemini-2.5-flash-image", + raw_response=mock_response, + model_response=ImageResponse(), + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert result.usage.web_search_requests == 2 + + +class TestVertexAIImagenImageGenerationConfig: + def setup_method(self): + """Set up test fixtures""" + self.config = VertexAIImagenImageGenerationConfig() + + def test_get_supported_openai_params(self): + """Test get_supported_openai_params returns correct params""" + supported = self.config.get_supported_openai_params("imagegeneration@006") + assert "n" in supported + assert "size" in supported + + def test_map_openai_params_n(self): + """Test mapping n parameter to sampleCount""" + non_default_params = {"n": 3} + optional_params = {} + result = self.config.map_openai_params(non_default_params, optional_params, "imagegeneration@006", False) + assert result.get("sampleCount") == 3 + + def test_map_openai_params_size(self): + """Test mapping size parameter to aspectRatio""" + non_default_params = {"size": "1024x1024"} + optional_params = {} + result = self.config.map_openai_params(non_default_params, optional_params, "imagegeneration@006", False) + assert result.get("aspectRatio") == "1:1" + + def test_map_size_to_aspect_ratio(self): + """Test size to aspect ratio mapping""" + assert self.config._map_size_to_aspect_ratio("1024x1024") == "1:1" + assert self.config._map_size_to_aspect_ratio("1792x1024") == "16:9" + assert self.config._map_size_to_aspect_ratio("unknown") == "1:1" # default + + def test_transform_image_generation_request_basic(self): + """Test basic request transformation""" + request = self.config.transform_image_generation_request( + model="imagegeneration@006", + prompt="A cat", + optional_params={}, + litellm_params={}, + headers={}, + ) + assert "instances" in request + assert "parameters" in request + assert request["instances"][0]["prompt"] == "A cat" + assert request["parameters"]["sampleCount"] == 1 + + def test_transform_image_generation_request_with_params(self): + """Test request transformation with parameters""" + request = self.config.transform_image_generation_request( + model="imagegeneration@006", + prompt="A cat", + optional_params={"sampleCount": 2, "aspectRatio": "16:9"}, + litellm_params={}, + headers={}, + ) + assert request["parameters"]["sampleCount"] == 2 + assert request["parameters"]["aspectRatio"] == "16:9" + + def test_transform_image_generation_request_labels_from_metadata(self): + """Billing labels from litellm_params.metadata.requester_metadata on predict body.""" + request = self.config.transform_image_generation_request( + model="imagegeneration@006", + prompt="A cat", + optional_params={}, + litellm_params={"metadata": {"requester_metadata": {"team": "platform", "env": "prod"}}}, + headers={}, + ) + assert request["labels"] == {"team": "platform", "env": "prod"} + assert "labels" not in request["parameters"] + + def test_transform_image_generation_response(self): + """Test response transformation""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = {"predictions": [{"bytesBase64Encoded": "base64_encoded_image_data"}]} + mock_response.headers = {} + + from litellm.types.utils import ImageResponse + + model_response = ImageResponse() + result = self.config.transform_image_generation_response( + model="imagegeneration@006", + raw_response=mock_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert len(result.data) == 1 + assert result.data[0].b64_json == "base64_encoded_image_data" + assert result.data[0].url is None + + def test_transform_image_generation_response_multiple_images(self): + """Test response transformation with multiple images""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = { + "predictions": [ + {"bytesBase64Encoded": "image1"}, + {"bytesBase64Encoded": "image2"}, + ] + } + mock_response.headers = {} + + from litellm.types.utils import ImageResponse + + model_response = ImageResponse() + result = self.config.transform_image_generation_response( + model="imagegeneration@006", + raw_response=mock_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert len(result.data) == 2 + assert result.data[0].b64_json == "image1" + assert result.data[1].b64_json == "image2" + + +class TestGetVertexAIImageGenerationConfig: + """Test the router function that selects the correct config""" + + def test_get_gemini_model_config(self): + """Test that Gemini models return Gemini config""" + config = get_vertex_ai_image_generation_config("gemini-2.5-flash-image") + assert isinstance(config, VertexAIGeminiImageGenerationConfig) + + config = get_vertex_ai_image_generation_config("gemini-3-pro-image-preview") + assert isinstance(config, VertexAIGeminiImageGenerationConfig) + + config = get_vertex_ai_image_generation_config("vertex_ai/gemini-2.5-flash-image") + assert isinstance(config, VertexAIGeminiImageGenerationConfig) + + def test_get_imagen_model_config(self): + """Test that Imagen models return Imagen config""" + config = get_vertex_ai_image_generation_config("imagegeneration@006") + assert isinstance(config, VertexAIImagenImageGenerationConfig) + + config = get_vertex_ai_image_generation_config("imagen-4.0-generate-001") + assert isinstance(config, VertexAIImagenImageGenerationConfig) + + config = get_vertex_ai_image_generation_config("vertex_ai/imagegeneration@006") + assert isinstance(config, VertexAIImagenImageGenerationConfig) + + def test_get_non_gemini_model_config(self): + """Test that non-Gemini models default to Imagen config""" + config = get_vertex_ai_image_generation_config("some-other-model") + assert isinstance(config, VertexAIImagenImageGenerationConfig) + + +class TestVertexAIImageGenerationIntegration: + """Integration tests for Vertex AI image generation""" + + + def test_gemini_get_complete_url(self): + """Test Gemini config URL generation""" + config = VertexAIGeminiImageGenerationConfig() + url = config.get_complete_url( + api_base=None, + api_key=None, + model="gemini-2.5-flash-image", + optional_params={}, + litellm_params={ + "vertex_project": "test-project", + "vertex_location": "us-central1", + }, + ) + assert "test-project" in url + assert "us-central1" in url + assert "gemini-2.5-flash-image" in url + assert "generateContent" in url + + def test_imagen_get_complete_url(self): + """Test Imagen config URL generation""" + config = VertexAIImagenImageGenerationConfig() + url = config.get_complete_url( + api_base=None, + api_key=None, + model="imagegeneration@006", + optional_params={}, + litellm_params={ + "vertex_project": "test-project", + "vertex_location": "us-central1", + }, + ) + assert "test-project" in url + assert "us-central1" in url + assert "imagegeneration@006" in url + assert "predict" in url diff --git a/tests/unit/llms/vertex_ai/rag_engine/__init__.py b/tests/unit/llms/vertex_ai/rag_engine/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/vertex_ai/rag_engine/test_ingestion.py b/tests/unit/llms/vertex_ai/rag_engine/test_ingestion.py new file mode 100644 index 00000000000..3acabc4d14e --- /dev/null +++ b/tests/unit/llms/vertex_ai/rag_engine/test_ingestion.py @@ -0,0 +1,76 @@ +import asyncio +import sys +from types import ModuleType, SimpleNamespace + +import litellm +from litellm.litellm_core_utils.get_litellm_params import get_litellm_params +from litellm.llms.vertex_ai.files.transformation import VertexAIFilesConfig +from litellm.llms.vertex_ai.rag_engine.ingestion import VertexAIRAGIngestion + + +def _ingestion_for_bucket(bucket: str) -> VertexAIRAGIngestion: + return VertexAIRAGIngestion( + { + "vector_store": { + "custom_llm_provider": "vertex_ai", + "vector_store_id": "corpus-123", + "vertex_project": "test-project", + "vertex_location": "us-central1", + "gcs_bucket": bucket, + } + } + ) + + +def test_upload_lands_in_the_corpus_bucket_when_batch_bucket_env_is_set(monkeypatch): + monkeypatch.setenv("GCS_BATCH_BUCKET_NAME", "batch-bucket") + monkeypatch.setenv("GCS_BUCKET_NAME", "logging-bucket") + resolver = VertexAIFilesConfig() + + async def acreate_file_through_real_bucket_resolver(**kwargs): + bucket = resolver._get_configured_bucket_name(get_litellm_params(**kwargs)) + return SimpleNamespace(id=f"gs://{bucket}/{kwargs['file'][0]}") + + monkeypatch.setattr(litellm, "acreate_file", acreate_file_through_real_bucket_resolver) + + uri = asyncio.run(_ingestion_for_bucket("rag-bucket")._upload_file_to_gcs(b"doc", "doc.txt", "text/plain")) + + assert uri == "gs://rag-bucket/doc.txt" + + +def _vertexai_sdk_stub(import_calls: list[dict[str, object]]) -> ModuleType: + rag = ModuleType("vertexai.rag") + rag.TransformationConfig = lambda chunking_config: chunking_config + rag.ChunkingConfig = lambda chunk_size, chunk_overlap: (chunk_size, chunk_overlap) + + def import_files(**kwargs): + import_calls.append(kwargs) + return SimpleNamespace(imported_rag_files_count=1) + + rag.import_files = import_files + vertexai = ModuleType("vertexai") + vertexai.init = lambda project, location: None + vertexai.rag = rag + return vertexai + + +def test_ingest_runs_end_to_end_through_the_base_pipeline(monkeypatch): + monkeypatch.setenv("GCS_BATCH_BUCKET_NAME", "batch-bucket") + resolver = VertexAIFilesConfig() + import_calls: list[dict[str, object]] = [] + stub = _vertexai_sdk_stub(import_calls) + monkeypatch.setitem(sys.modules, "vertexai", stub) + monkeypatch.setitem(sys.modules, "vertexai.rag", stub.rag) + + async def acreate_file_through_real_bucket_resolver(**kwargs): + bucket = resolver._get_configured_bucket_name(get_litellm_params(**kwargs)) + return SimpleNamespace(id=f"gs://{bucket}/{kwargs['file'][0]}") + + monkeypatch.setattr(litellm, "acreate_file", acreate_file_through_real_bucket_resolver) + + result = asyncio.run(_ingestion_for_bucket("rag-bucket").ingest(file_data=("doc.txt", b"doc", "text/plain"))) + + assert (result["status"], result["vector_store_id"], result["file_id"]) == ("completed", "corpus-123", "gs://rag-bucket/doc.txt") + assert [(c["corpus_name"], c["paths"]) for c in import_calls] == [ + ("projects/test-project/locations/us-central1/ragCorpora/corpus-123", ["gs://rag-bucket/doc.txt"]) + ] diff --git a/tests/unit/llms/vertex_ai/rerank/__init__.py b/tests/unit/llms/vertex_ai/rerank/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py b/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py rename to tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py diff --git a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py b/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py rename to tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_userlabels_e2e.py b/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_userlabels_e2e.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_userlabels_e2e.py rename to tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_userlabels_e2e.py diff --git a/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py b/tests/unit/llms/vertex_ai/test_bge_embedding.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_bge_embedding.py rename to tests/unit/llms/vertex_ai/test_bge_embedding.py diff --git a/tests/test_litellm/llms/vertex_ai/test_bge_response_transformation.py b/tests/unit/llms/vertex_ai/test_bge_response_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_bge_response_transformation.py rename to tests/unit/llms/vertex_ai/test_bge_response_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/test_gemini_batch_embeddings.py b/tests/unit/llms/vertex_ai/test_gemini_batch_embeddings.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_gemini_batch_embeddings.py rename to tests/unit/llms/vertex_ai/test_gemini_batch_embeddings.py diff --git a/tests/test_litellm/llms/vertex_ai/test_gemini_empty_properties.py b/tests/unit/llms/vertex_ai/test_gemini_empty_properties.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_gemini_empty_properties.py rename to tests/unit/llms/vertex_ai/test_gemini_empty_properties.py diff --git a/tests/test_litellm/llms/vertex_ai/test_gemini_header_forwarding.py b/tests/unit/llms/vertex_ai/test_gemini_header_forwarding.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_gemini_header_forwarding.py rename to tests/unit/llms/vertex_ai/test_gemini_header_forwarding.py diff --git a/tests/test_litellm/llms/vertex_ai/test_http_status_201.py b/tests/unit/llms/vertex_ai/test_http_status_201.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_http_status_201.py rename to tests/unit/llms/vertex_ai/test_http_status_201.py diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex.py b/tests/unit/llms/vertex_ai/test_vertex.py similarity index 97% rename from tests/test_litellm/llms/vertex_ai/test_vertex.py rename to tests/unit/llms/vertex_ai/test_vertex.py index e3007bac7f3..ab8bf123ab2 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex.py +++ b/tests/unit/llms/vertex_ai/test_vertex.py @@ -1193,7 +1193,6 @@ def test_logprobs(): def test_process_gemini_media(): """Test the _process_gemini_media function for different image sources""" - from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media from litellm.types.llms.vertex_ai import FileDataType # Test GCS URI @@ -1271,7 +1270,6 @@ def test_process_gemini_media(): assert base64_result["inline_data"]["data"] == "/9j/4AAQSkZJRg..." - def test_get_image_mime_type_from_url(): """Test the _get_image_mime_type_from_url function for different image URLs""" from litellm.llms.vertex_ai.gemini.transformation import ( @@ -1372,46 +1370,6 @@ def encoded_images(): return [encode_image_to_base64(path) for path in image_paths] -@pytest.fixture -def mock_convert_url_to_base64(): - with patch( - "litellm.litellm_core_utils.prompt_templates.factory.convert_url_to_base64", - ) as mock: - # Setup the mock to return a valid image object - mock.return_value = "data:image/jpeg;base64,/9j/4AAQSkZJRg..." - yield mock - - -@pytest.fixture -def mock_blob(): - return Mock(spec=BlobType) - - -@pytest.mark.parametrize( - "http_url", - [ - "http://img1.etsystatic.com/260/0/7813604/il_fullxfull.4226713999_q86e.jpg", - "http://example.com/image.jpg", - "http://subdomain.domain.com/path/to/image.png", - ], -) -def test_process_gemini_media_http_url( - http_url: str, mock_convert_url_to_base64: Mock, mock_blob: Mock -) -> None: - """ - Test that _process_gemini_media correctly handles HTTP URLs. - - Args: - http_url: Test HTTP URL - mock_convert_to_anthropic: Mocked convert_to_anthropic_image_obj function - mock_blob: Mocked BlobType instance - - Vertex AI supports image urls. Ensure no network requests are made. - """ - expected_image_data = "data:image/jpeg;base64,/9j/4AAQSkZJRg..." - mock_convert_url_to_base64.return_value = expected_image_data - # Act - result = _process_gemini_media(http_url) # assert result["file_data"]["file_uri"] == http_url diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py b/tests/unit/llms/vertex_ai/test_vertex_ai_batch_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py rename to tests/unit/llms/vertex_ai/test_vertex_ai_batch_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py rename to tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py b/tests/unit/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py rename to tests/unit/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py b/tests/unit/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py rename to tests/unit/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_gemini_gcs_uri_mime.py b/tests/unit/llms/vertex_ai/test_vertex_gemini_gcs_uri_mime.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_vertex_gemini_gcs_uri_mime.py rename to tests/unit/llms/vertex_ai/test_vertex_gemini_gcs_uri_mime.py diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_global_url_support.py b/tests/unit/llms/vertex_ai/test_vertex_global_url_support.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_vertex_global_url_support.py rename to tests/unit/llms/vertex_ai/test_vertex_global_url_support.py diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_image_generation.py b/tests/unit/llms/vertex_ai/test_vertex_image_generation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_vertex_image_generation.py rename to tests/unit/llms/vertex_ai/test_vertex_image_generation.py diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py b/tests/unit/llms/vertex_ai/test_vertex_llm_base.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py rename to tests/unit/llms/vertex_ai/test_vertex_llm_base.py diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_model_garden_openapi.py b/tests/unit/llms/vertex_ai/test_vertex_model_garden_openapi.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_vertex_model_garden_openapi.py rename to tests/unit/llms/vertex_ai/test_vertex_model_garden_openapi.py diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py b/tests/unit/llms/vertex_ai/test_vertex_passthrough_logging_handler.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py rename to tests/unit/llms/vertex_ai/test_vertex_passthrough_logging_handler.py diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/__init__.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py 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/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/gemma/__init__.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/gemma/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/__init__.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/qwen/__init__.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/qwen/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/test_partner_models_credential_reuse.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/test_partner_models_credential_reuse.py similarity index 100% rename from tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/test_partner_models_credential_reuse.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/test_partner_models_credential_reuse.py diff --git a/tests/test_litellm/llms/volcengine/embedding/__init__.py b/tests/unit/llms/volcengine/embedding/__init__.py similarity index 100% rename from tests/test_litellm/llms/volcengine/embedding/__init__.py rename to tests/unit/llms/volcengine/embedding/__init__.py diff --git a/tests/test_litellm/llms/volcengine/test_volcengine.py b/tests/unit/llms/volcengine/test_volcengine.py similarity index 100% rename from tests/test_litellm/llms/volcengine/test_volcengine.py rename to tests/unit/llms/volcengine/test_volcengine.py diff --git a/tests/unit/llms/wandb/__init__.py b/tests/unit/llms/wandb/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/wandb/test_wandb_chat_transformation.py b/tests/unit/llms/wandb/test_wandb_chat_transformation.py similarity index 100% rename from tests/test_litellm/llms/wandb/test_wandb_chat_transformation.py rename to tests/unit/llms/wandb/test_wandb_chat_transformation.py diff --git a/tests/unit/llms/xai/batches/__init__.py b/tests/unit/llms/xai/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/xai/batches/test_xai_batches_handler.py b/tests/unit/llms/xai/batches/test_xai_batches_handler.py new file mode 100644 index 00000000000..6dcdf06e7ab --- /dev/null +++ b/tests/unit/llms/xai/batches/test_xai_batches_handler.py @@ -0,0 +1,344 @@ +import json +from typing import Final + +import httpx +import pytest +import respx + +import litellm +from litellm.llms.xai.batches.transformation import XAIBatchesError +from litellm.types.utils import LiteLLMBatch + +API_BASE: Final = "https://api.x.ai" +KEY: Final = "xai-test-key" + + +@pytest.fixture(autouse=True) +def _httpx_transport_so_respx_can_intercept(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + + +_XAI_BATCH: Final = { + "batch_id": "batch_1", + "name": "litellm-batch", + "create_time": "2026-09-23", + "expire_time": "2026-10-23", + "cancel_time": None, + "cancel_by_xai_message": None, + "state": {"num_requests": 2, "num_pending": 0, "num_success": 2, "num_error": 0, "num_cancelled": 0}, + "input_file_id": "file_1", +} + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@respx.mock +async def test_create_batch_posts_input_file_id_with_bearer_auth(sync_mode: bool) -> None: + route: Final = respx.post(f"{API_BASE}/v1/batches").respond(200, json=_XAI_BATCH) + + kwargs: Final = { + "completion_window": "24h", + "endpoint": "/v1/embeddings", + "input_file_id": "file_1", + "custom_llm_provider": "xai", + "api_key": KEY, + "api_base": API_BASE, + } + batch: Final = litellm.create_batch(**kwargs) if sync_mode else await litellm.acreate_batch(**kwargs) + + assert isinstance(batch, LiteLLMBatch) + request: Final = route.calls.last.request + assert request.headers["authorization"] == f"Bearer {KEY}" + assert json.loads(request.content) == {"name": "litellm-batch", "input_file_id": "file_1"} + assert (batch.id, batch.endpoint, batch.status, batch.output_file_id) == ( + "batch_1", + "/v1/embeddings", + "completed", + "batch_1", + ) + + +@pytest.mark.parametrize( + "endpoint", + [ + "/v1/chat/completions", + "/v1/embeddings", + "/v1/completions", + "/v1/responses", + "/v1/ocr", + "/v1/images/generations", + "/v1/images/edits", + "/v1/videos/generations", + "/v1/videos", + "/v1/videos/edits", + "/v1/videos/extensions", + ], +) +@respx.mock +async def test_create_batch_keeps_image_and_video_endpoints_on_the_batch(endpoint: str) -> None: + respx.post(f"{API_BASE}/v1/batches").respond(200, json=_XAI_BATCH) + + batch: Final = await litellm.acreate_batch( + completion_window="24h", + endpoint=endpoint, + input_file_id="file_1", + custom_llm_provider="xai", + api_key=KEY, + api_base=API_BASE, + ) + + assert isinstance(batch, LiteLLMBatch) + assert batch.endpoint == endpoint + assert json.loads(respx.calls.last.request.content) == {"name": "litellm-batch", "input_file_id": "file_1"} + + +@respx.mock +async def test_retrieve_after_a_non_chat_create_reports_chat() -> None: + respx.post(f"{API_BASE}/v1/batches").respond(200, json=_XAI_BATCH) + respx.get(f"{API_BASE}/v1/batches/batch_1").respond(200, json=_XAI_BATCH) + + created: Final = await litellm.acreate_batch( + completion_window="24h", + endpoint="/v1/embeddings", + input_file_id="file_1", + custom_llm_provider="xai", + api_key=KEY, + api_base=API_BASE, + ) + retrieved: Final = await litellm.aretrieve_batch( + batch_id="batch_1", custom_llm_provider="xai", api_key=KEY, api_base=API_BASE + ) + + assert isinstance(created, LiteLLMBatch) and isinstance(retrieved, LiteLLMBatch) + assert (created.endpoint, retrieved.endpoint) == ("/v1/embeddings", "/v1/chat/completions") + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@respx.mock +async def test_retrieve_batch_reads_native_batch_route(sync_mode: bool) -> None: + respx.get(f"{API_BASE}/v1/batches/batch_1").respond( + 200, json={**_XAI_BATCH, "state": {"num_requests": 2, "num_pending": 2}} + ) + + kwargs: Final = {"batch_id": "batch_1", "custom_llm_provider": "xai", "api_key": KEY, "api_base": API_BASE} + batch: Final = litellm.retrieve_batch(**kwargs) if sync_mode else await litellm.aretrieve_batch(**kwargs) + + assert isinstance(batch, LiteLLMBatch) + assert (batch.status, batch.output_file_id, batch.input_file_id, batch.endpoint) == ( + "in_progress", + None, + "file_1", + "/v1/chat/completions", + ) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@respx.mock +async def test_cancel_batch_uses_colon_cancel_route(sync_mode: bool) -> None: + route: Final = respx.post(f"{API_BASE}/v1/batches/batch_1:cancel").respond( + 200, json={**_XAI_BATCH, "cancel_time": "2026-09-23", "state": {}} + ) + + kwargs: Final = {"batch_id": "batch_1", "custom_llm_provider": "xai", "api_key": KEY, "api_base": API_BASE} + batch: Final = litellm.cancel_batch(**kwargs) if sync_mode else await litellm.acancel_batch(**kwargs) + + assert route.called + assert isinstance(batch, LiteLLMBatch) + assert (batch.status, batch.endpoint) == ("cancelled", "/v1/chat/completions") + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@respx.mock +async def test_list_batches_forwards_cursor_and_returns_openai_list(sync_mode: bool) -> None: + route: Final = respx.get(f"{API_BASE}/v1/batches").respond( + 200, json={"batches": [_XAI_BATCH], "pagination_token": "next"} + ) + + kwargs: Final = {"custom_llm_provider": "xai", "api_key": KEY, "api_base": API_BASE, "after": "cur", "limit": 5} + listed: Final = litellm.list_batches(**kwargs) if sync_mode else await litellm.alist_batches(**kwargs) + + assert dict(route.calls.last.request.url.params) == {"limit": "5", "pagination_token": "cur"} + assert listed.object == "list" + assert [(b.id, b.endpoint) for b in listed.data] == [("batch_1", "/v1/chat/completions")] + assert (listed.has_more, listed.next_page_token) == (True, "next") + + +@respx.mock +async def test_list_batches_treats_empty_pagination_token_as_last_page() -> None: + respx.get(f"{API_BASE}/v1/batches").respond(200, json={"batches": [_XAI_BATCH], "pagination_token": ""}) + + listed: Final = await litellm.alist_batches(custom_llm_provider="xai", api_key=KEY, api_base=API_BASE) + + assert (listed.has_more, listed.next_page_token) == (False, None) + assert [batch.endpoint for batch in listed.data] == ["/v1/chat/completions"] + + +@respx.mock +async def test_file_content_stops_paging_on_empty_pagination_token() -> None: + route: Final = respx.get(f"{API_BASE}/v1/batches/batch_1/results").respond( + 200, + json={ + "results": [{"batch_request_id": "r1", "batch_result": {"error": {"code": 3, "message": "boom"}}}], + "pagination_token": "", + }, + ) + + content: Final = await litellm.afile_content( + file_id="batch_1", custom_llm_provider="xai", api_key=KEY, api_base=API_BASE + ) + + assert route.call_count == 1 + assert len(content.content.decode().splitlines()) == 1 + + +@pytest.mark.parametrize("operation", ["create", "retrieve", "cancel", "list", "file_content"]) +@respx.mock +async def test_batch_calls_fall_back_to_litellm_xai_key(operation: str, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("XAI_API_KEY", raising=False) + monkeypatch.setattr(litellm, "xai_key", "configured-xai-key") + monkeypatch.setattr(litellm, "api_key", "generic-key-must-not-be-used") + routes: Final = { + "create": respx.post(f"{API_BASE}/v1/batches").respond(200, json=_XAI_BATCH), + "retrieve": respx.get(f"{API_BASE}/v1/batches/batch_1").respond(200, json=_XAI_BATCH), + "cancel": respx.post(f"{API_BASE}/v1/batches/batch_1:cancel").respond(200, json=_XAI_BATCH), + "list": respx.get(f"{API_BASE}/v1/batches").respond( + 200, json={"batches": [_XAI_BATCH], "pagination_token": None} + ), + "file_content": respx.get(f"{API_BASE}/v1/batches/batch_1/results").respond( + 200, json={"results": [], "pagination_token": None} + ), + } + + if operation == "create": + await litellm.acreate_batch( + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file_1", + custom_llm_provider="xai", + api_base=API_BASE, + ) + elif operation == "retrieve": + await litellm.aretrieve_batch(batch_id="batch_1", custom_llm_provider="xai", api_base=API_BASE) + elif operation == "cancel": + await litellm.acancel_batch(batch_id="batch_1", custom_llm_provider="xai", api_base=API_BASE) + elif operation == "list": + await litellm.alist_batches(custom_llm_provider="xai", api_base=API_BASE) + else: + await litellm.afile_content(file_id="batch_1", custom_llm_provider="xai", api_base=API_BASE) + + assert routes[operation].calls.last.request.headers["authorization"] == "Bearer configured-xai-key" + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@respx.mock +async def test_file_content_of_a_batch_id_walks_every_results_page(sync_mode: bool) -> None: + def _page(request: httpx.Request) -> httpx.Response: + token: Final = request.url.params.get("pagination_token") + if token is None: + return httpx.Response( + 200, + json={ + "results": [ + { + "batch_request_id": "r1", + "batch_result": {"response": {"chat_get_completion": {"id": "c1", "choices": []}}}, + } + ], + "pagination_token": "r1", + }, + ) + assert token == "r1" + return httpx.Response( + 200, + json={ + "results": [ + {"batch_request_id": "r2", "batch_result": {"error": {"code": 3, "message": "boom"}}}, + ], + "pagination_token": None, + }, + ) + + route: Final = respx.get(f"{API_BASE}/v1/batches/batch_1/results").mock(side_effect=_page) + + kwargs: Final = {"file_id": "batch_1", "custom_llm_provider": "xai", "api_key": KEY, "api_base": API_BASE} + content: Final = litellm.file_content(**kwargs) if sync_mode else await litellm.afile_content(**kwargs) + + assert route.call_count == 2 + assert [dict(c.request.url.params) for c in route.calls] == [ + {"limit": "1000"}, + {"limit": "1000", "pagination_token": "r1"}, + ] + assert [json.loads(line) for line in content.content.decode().splitlines()] == [ + { + "id": "batch_req_r1", + "custom_id": "r1", + "response": {"status_code": 200, "request_id": "c1", "body": {"id": "c1", "choices": []}}, + "error": None, + }, + {"id": "batch_req_r2", "custom_id": "r2", "response": None, "error": {"code": "3", "message": "boom"}}, + ] + + +@respx.mock +async def test_file_content_unwraps_image_and_video_result_bodies() -> None: + respx.get(f"{API_BASE}/v1/batches/batch_1/results").respond( + 200, + json={ + "results": [ + { + "batch_request_id": "img", + "batch_result": { + "response": {"image_generation": {"data": [{"url": "https://cdn.example/img.png"}]}} + }, + }, + { + "batch_request_id": "vid", + "batch_result": { + "response": {"video_generation": {"id": "vid_1", "url": "https://cdn.example/clip.mp4"}} + }, + }, + ], + "pagination_token": None, + }, + ) + + content: Final = await litellm.afile_content( + file_id="batch_1", custom_llm_provider="xai", api_key=KEY, api_base=API_BASE + ) + + assert [json.loads(line)["response"]["body"] for line in content.content.decode().splitlines()] == [ + {"data": [{"url": "https://cdn.example/img.png"}]}, + {"id": "vid_1", "url": "https://cdn.example/clip.mp4"}, + ] + + +@respx.mock +async def test_missing_xai_key_is_a_401_before_any_request(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("XAI_API_KEY", raising=False) + monkeypatch.setattr(litellm, "xai_key", None) + monkeypatch.setattr(litellm, "api_key", "generic-key-must-not-be-used") + route: Final = respx.post(f"{API_BASE}/v1/batches").respond(200, json=_XAI_BATCH) + + with pytest.raises(XAIBatchesError) as exc: + await litellm.acreate_batch( + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file_1", + custom_llm_provider="xai", + api_base=API_BASE, + ) + + assert exc.value.status_code == 401 + assert route.called is False + + +@respx.mock +async def test_upstream_error_surfaces_status_code_and_body() -> None: + respx.get(f"{API_BASE}/v1/batches/batch_missing").respond(404, json={"code": "404", "error": "not found"}) + + with pytest.raises(XAIBatchesError) as exc: + await litellm.aretrieve_batch( + batch_id="batch_missing", custom_llm_provider="xai", api_key=KEY, api_base=API_BASE + ) + + assert exc.value.status_code == 404 + assert "not found" in exc.value.message diff --git a/tests/unit/llms/xai/batches/test_xai_batches_transformation.py b/tests/unit/llms/xai/batches/test_xai_batches_transformation.py new file mode 100644 index 00000000000..5f2bb6a33ce --- /dev/null +++ b/tests/unit/llms/xai/batches/test_xai_batches_transformation.py @@ -0,0 +1,224 @@ +import json +from typing import Final + +import pytest + +from litellm.llms.xai.batches.transformation import ( + XAIBatch, + XAIBatchesError, + XAIBatchList, + XAIBatchResult, + XAIBatchResultsPage, + get_xai_api_base, + results_to_openai_jsonl, + to_create_batch_body, + to_litellm_batch, + to_openai_batch_list, + xai_batches_url, +) +from litellm.types.llms.openai import CreateBatchRequest + +SEPT_23_2026_UTC: Final = 1790121600 + + +def _xai_batch(**overrides: object) -> XAIBatch: + return XAIBatch.model_validate( + { + "batch_id": "batch_9bdf", + "name": "nightly", + "create_time": "2026-09-23", + "expire_time": "2026-10-23", + "cancel_time": None, + "cancel_by_xai_message": None, + "state": {"num_requests": 2, "num_pending": 0, "num_success": 2, "num_error": 0, "num_cancelled": 0}, + "input_file_id": "file_07", + **overrides, + } + ) + + +def test_completed_batch_exposes_batch_id_as_output_file_and_maps_counts() -> None: + batch: Final = to_litellm_batch(_xai_batch()) + + assert batch.model_dump(exclude_none=True) == { + "id": "batch_9bdf", + "object": "batch", + "endpoint": "/v1/chat/completions", + "input_file_id": "file_07", + "completion_window": "24h", + "status": "completed", + "created_at": SEPT_23_2026_UTC, + "expires_at": SEPT_23_2026_UTC + 30 * 86400, + "output_file_id": "batch_9bdf", + "request_counts": {"total": 2, "completed": 2, "failed": 0}, + "metadata": {"name": "nightly"}, + } + + +def test_pending_requests_mean_in_progress_and_no_output_file() -> None: + batch: Final = to_litellm_batch( + _xai_batch(state={"num_requests": 3, "num_pending": 1, "num_success": 1, "num_error": 1, "num_cancelled": 0}) + ) + + assert (batch.status, batch.output_file_id) == ("in_progress", None) + assert batch.request_counts is not None + assert batch.request_counts.model_dump() == {"total": 3, "completed": 1, "failed": 1} + + +def test_empty_batch_is_still_validating() -> None: + assert to_litellm_batch(_xai_batch(state={})).status == "validating" + + +def test_batch_cancelled_by_xai_validation_is_failed_with_the_message() -> None: + batch: Final = to_litellm_batch( + _xai_batch( + state={}, + cancel_time="2026-09-23T10:00:00Z", + cancel_by_xai_message="JSONL file validation failed: Model grok-nope is not supported", + ) + ) + + assert batch.status == "failed" + assert batch.failed_at == SEPT_23_2026_UTC + 10 * 3600 + assert batch.cancelled_at is None + assert batch.errors is not None and batch.errors.data is not None + assert [e.message for e in batch.errors.data] == ["JSONL file validation failed: Model grok-nope is not supported"] + + +def test_batch_cancelled_by_caller_is_cancelled() -> None: + batch: Final = to_litellm_batch(_xai_batch(cancel_time="2026-09-23")) + + assert (batch.status, batch.cancelled_at, batch.errors) == ("cancelled", SEPT_23_2026_UTC, None) + + +@pytest.mark.parametrize( + "endpoint", + [ + "/v1/images/generations", + "/v1/images/edits", + "/v1/videos/generations", + "/v1/videos/edits", + "/v1/videos/extensions", + ], +) +def test_create_body_accepts_image_and_video_endpoints(endpoint: str) -> None: + body: Final = to_create_batch_body( + CreateBatchRequest(completion_window="24h", endpoint=endpoint, input_file_id="file_07") + ) + + assert dict(body) == {"name": "litellm-batch", "input_file_id": "file_07"} + + +def test_create_body_uses_input_file_id_and_metadata_name() -> None: + body: Final = to_create_batch_body( + CreateBatchRequest( + completion_window="24h", endpoint="/v1/chat/completions", input_file_id="file_07", metadata={"name": "n1"} + ) + ) + + assert dict(body) == {"name": "n1", "input_file_id": "file_07"} + + +def test_create_body_without_input_file_id_is_a_400() -> None: + with pytest.raises(XAIBatchesError) as exc: + to_create_batch_body(CreateBatchRequest(completion_window="24h", endpoint="/v1/chat/completions")) + + assert exc.value.status_code == 400 + + +def test_results_render_as_openai_output_jsonl_with_errors_per_line() -> None: + page: Final = XAIBatchResultsPage.model_validate( + { + "results": [ + { + "batch_request_id": "r1", + "batch_result": { + "response": { + "chat_get_completion": {"id": "c1", "object": "chat.completion", "choices": [], "usage": {}} + } + }, + }, + {"batch_request_id": "r2", "batch_result": {"error": {"code": 3, "message": "bad model"}}}, + {"batch_request_id": "r3", "batch_result": {}}, + ], + "pagination_token": None, + } + ) + + lines: Final = [json.loads(line) for line in results_to_openai_jsonl(page.results).decode().splitlines()] + + assert lines == [ + { + "id": "batch_req_r1", + "custom_id": "r1", + "response": { + "status_code": 200, + "request_id": "c1", + "body": {"id": "c1", "object": "chat.completion", "choices": [], "usage": {}}, + }, + "error": None, + }, + {"id": "batch_req_r2", "custom_id": "r2", "response": None, "error": {"code": "3", "message": "bad model"}}, + { + "id": "batch_req_r3", + "custom_id": "r3", + "response": None, + "error": {"code": "request_failed", "message": "xAI returned no response for this request"}, + }, + ] + + +@pytest.mark.parametrize( + ("response_key", "body"), + [ + ("responses", {"id": "resp_1", "output": []}), + ("image_generation", {"created": 1, "data": [{"url": "https://cdn.example/img.png"}]}), + ("video_generation", {"id": "vid_1", "url": "https://cdn.example/clip.mp4"}), + ], +) +def test_result_unwraps_the_single_response_key_into_the_openai_body( + response_key: str, body: dict[str, object] +) -> None: + result: Final = XAIBatchResult.model_validate( + {"batch_request_id": "r", "batch_result": {"response": {response_key: body}}} + ) + + line: Final = json.loads(results_to_openai_jsonl((result,)).decode()) + assert line["response"]["body"] == body + assert line["response"]["request_id"] == body.get("id") + assert response_key not in line["response"]["body"] + + +def test_retrieve_and_list_report_chat_because_xai_has_no_batch_endpoint() -> None: + retrieved: Final = to_litellm_batch(_xai_batch()) + listed: Final = to_openai_batch_list(XAIBatchList.model_validate({"batches": [_xai_batch().model_dump()]})) + + assert retrieved.endpoint == "/v1/chat/completions" + assert [batch.endpoint for batch in listed.data] == ["/v1/chat/completions"] + assert retrieved.metadata == {"name": "nightly"} + + +def test_list_page_maps_to_openai_list_with_cursor_flags() -> None: + page: Final = XAIBatchList.model_validate( + {"batches": [_xai_batch().model_dump(), _xai_batch(batch_id="batch_2").model_dump()], "pagination_token": "t"} + ) + + listed: Final = to_openai_batch_list(page) + + assert (listed.object, listed.first_id, listed.last_id, listed.has_more, listed.next_page_token) == ( + "list", + "batch_9bdf", + "batch_2", + True, + "t", + ) + assert [b.id for b in listed.data] == ["batch_9bdf", "batch_2"] + + +@pytest.mark.parametrize( + "api_base", ["https://api.x.ai", "https://api.x.ai/", "https://api.x.ai/v1", "https://api.x.ai/v1/"] +) +def test_api_base_never_doubles_the_v1_segment(api_base: str) -> None: + assert get_xai_api_base(api_base) == "https://api.x.ai" + assert xai_batches_url(api_base, "batch_1", ":cancel") == "https://api.x.ai/v1/batches/batch_1:cancel" + assert xai_batches_url(api_base) == "https://api.x.ai/v1/batches" diff --git a/tests/unit/llms/xai/files/__init__.py b/tests/unit/llms/xai/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/xai/files/test_xai_files_transformation.py b/tests/unit/llms/xai/files/test_xai_files_transformation.py new file mode 100644 index 00000000000..5a7d86bdfb7 --- /dev/null +++ b/tests/unit/llms/xai/files/test_xai_files_transformation.py @@ -0,0 +1,144 @@ +from typing import Final + +import httpx +import pytest +import respx +from pydantic import TypeAdapter + +import litellm +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.types.llms.openai import OpenAIFileObject + +API_BASE: Final = "https://api.x.ai" +KEY: Final = "xai-test-key" + + +@pytest.fixture(autouse=True) +def _httpx_transport_so_respx_can_intercept(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + + +_XAI_FILE: Final = { + "bytes": 337, + "created_at": 1790197740, + "expires_at": None, + "filename": "batch.jsonl", + "id": "file_07", + "object": "file", + "purpose": "", +} + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@respx.mock +async def test_create_file_uploads_multipart_to_xai_and_reports_batch_purpose(sync_mode: bool) -> None: + route: Final = respx.post(f"{API_BASE}/v1/files").respond(200, json=_XAI_FILE) + + kwargs: Final = { + "file": ("batch.jsonl", b'{"custom_id":"r1"}\n', "application/jsonl"), + "purpose": "batch", + "custom_llm_provider": "xai", + "api_key": KEY, + "api_base": API_BASE, + } + created: Final = litellm.create_file(**kwargs) if sync_mode else await litellm.acreate_file(**kwargs) + + request: Final = route.calls.last.request + assert request.headers["authorization"] == f"Bearer {KEY}" + assert request.headers["content-type"].startswith("multipart/form-data") + assert b'filename="batch.jsonl"' in request.content + assert b'{"custom_id":"r1"}' in request.content + assert created.model_dump(exclude_none=True) == { + "id": "file_07", + "bytes": 337, + "created_at": 1790197740, + "filename": "batch.jsonl", + "object": "file", + "purpose": "batch", + "status": "uploaded", + } + + +@respx.mock +async def test_file_content_of_an_uploaded_file_downloads_original_bytes() -> None: + respx.get(f"{API_BASE}/v1/files/file_07/content").respond(200, content=b'{"custom_id":"r1"}\n') + + content: Final = await litellm.afile_content( + file_id="file_07", custom_llm_provider="xai", api_key=KEY, api_base=API_BASE + ) + + assert content.content == b'{"custom_id":"r1"}\n' + + +@respx.mock +async def test_delete_file_maps_xai_deleted_object() -> None: + respx.delete(f"{API_BASE}/v1/files/file_07").respond(200, json={"id": "file_07", "deleted": True, "object": "file"}) + + deleted: Final = await litellm.afile_delete( + file_id="file_07", custom_llm_provider="xai", api_key=KEY, api_base=API_BASE + ) + + assert deleted.model_dump() == {"id": "file_07", "deleted": True, "object": "file"} + + +@respx.mock +async def test_create_file_falls_back_to_litellm_xai_key(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("XAI_API_KEY", raising=False) + monkeypatch.setattr(litellm, "xai_key", "configured-xai-key") + monkeypatch.setattr(litellm, "api_key", "generic-key-must-not-be-used") + route: Final = respx.post(f"{API_BASE}/v1/files").respond(200, json=_XAI_FILE) + + await litellm.acreate_file( + file=("batch.jsonl", b'{"custom_id":"r1"}\n', "application/jsonl"), + purpose="batch", + custom_llm_provider="xai", + api_base=API_BASE, + ) + + assert route.calls.last.request.headers["authorization"] == "Bearer configured-xai-key" + + +@respx.mock +async def test_list_files_reads_data_array() -> None: + respx.get(f"{API_BASE}/v1/files").respond(200, json={"data": [_XAI_FILE], "pagination_token": None}) + + listed: Final = await litellm.afile_list(custom_llm_provider="xai", api_key=KEY, api_base=API_BASE) + + files: Final = TypeAdapter(tuple[OpenAIFileObject, ...]).validate_python(listed) + assert [f.id for f in files] == ["file_07"] + + +@respx.mock +async def test_list_files_walks_every_page_by_pagination_token() -> None: + route: Final = respx.get(f"{API_BASE}/v1/files").mock( + side_effect=[ + httpx.Response(200, json={"data": [_XAI_FILE], "pagination_token": "file_07"}), + httpx.Response(200, json={"data": [{**_XAI_FILE, "id": "file_08"}], "pagination_token": "file_08"}), + httpx.Response(200, json={"data": [], "pagination_token": "file_08"}), + ] + ) + + listed: Final = await litellm.afile_list(custom_llm_provider="xai", api_key=KEY, api_base=API_BASE) + + files: Final = TypeAdapter(tuple[OpenAIFileObject, ...]).validate_python(listed) + assert [f.id for f in files] == ["file_07", "file_08"] + assert [call.request.url.params.get("pagination_token") for call in route.calls] == [None, "file_07", "file_08"] + + +async def _retrieve_file(sync_mode: bool, file_id: str) -> None: + if sync_mode: + litellm.file_retrieve(file_id=file_id, custom_llm_provider="xai", api_key=KEY, api_base=API_BASE) + return + await litellm.afile_retrieve(file_id=file_id, custom_llm_provider="xai", api_key=KEY, api_base=API_BASE) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@respx.mock +async def test_retrieve_file_maps_xai_not_found_to_a_404_error(sync_mode: bool) -> None: + respx.get(f"{API_BASE}/v1/files/file_gone").respond(404, json={"code": "not-found", "error": "File not found"}) + + with pytest.raises(BaseLLMException) as raised: + await _retrieve_file(sync_mode, "file_gone") + + assert raised.value.status_code == 404 + assert "File not found" in str(raised.value) diff --git a/tests/test_litellm/llms/xai/test_xai_audio_transcription_transformation.py b/tests/unit/llms/xai/test_xai_audio_transcription_transformation.py similarity index 100% rename from tests/test_litellm/llms/xai/test_xai_audio_transcription_transformation.py rename to tests/unit/llms/xai/test_xai_audio_transcription_transformation.py diff --git a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py b/tests/unit/llms/xai/test_xai_chat_transformation.py similarity index 96% rename from tests/test_litellm/llms/xai/test_xai_chat_transformation.py rename to tests/unit/llms/xai/test_xai_chat_transformation.py index 3fd666e4f50..704d0061103 100644 --- a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py +++ b/tests/unit/llms/xai/test_xai_chat_transformation.py @@ -16,7 +16,7 @@ from litellm.types.utils import ( class TestXAIReasoningTokenFolding: - """``_fold_reasoning_tokens_into_completion`` re-aligns xAI Usage to the OpenAI invariant.""" + """``fold_reasoning_tokens_into_completion`` re-aligns xAI Usage to the OpenAI invariant.""" @staticmethod def _make_response( @@ -45,7 +45,7 @@ class TestXAIReasoningTokenFolding: reasoning_tokens=312, ) - XAIChatConfig._fold_reasoning_tokens_into_completion(response) + XAIChatConfig.fold_reasoning_tokens_into_completion(response) usage = response.usage assert usage.completion_tokens == 322 @@ -59,7 +59,7 @@ class TestXAIReasoningTokenFolding: reasoning_tokens=312, ) - XAIChatConfig._fold_reasoning_tokens_into_completion(response) + XAIChatConfig.fold_reasoning_tokens_into_completion(response) assert response.usage.completion_tokens == 322 @@ -71,7 +71,7 @@ class TestXAIReasoningTokenFolding: reasoning_tokens=0, ) - XAIChatConfig._fold_reasoning_tokens_into_completion(response) + XAIChatConfig.fold_reasoning_tokens_into_completion(response) assert response.usage.completion_tokens == 10 @@ -84,7 +84,7 @@ class TestXAIReasoningTokenFolding: reasoning_tokens=312, ) - XAIChatConfig._fold_reasoning_tokens_into_completion(response) + XAIChatConfig.fold_reasoning_tokens_into_completion(response) assert response.usage.completion_tokens == 10 assert response.usage.total_tokens == 999 diff --git a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py b/tests/unit/llms/xai/test_xai_cost_calculator.py similarity index 100% rename from tests/test_litellm/llms/xai/test_xai_cost_calculator.py rename to tests/unit/llms/xai/test_xai_cost_calculator.py diff --git a/tests/test_litellm/llms/xai/test_xai_key_fallback.py b/tests/unit/llms/xai/test_xai_key_fallback.py similarity index 100% rename from tests/test_litellm/llms/xai/test_xai_key_fallback.py rename to tests/unit/llms/xai/test_xai_key_fallback.py diff --git a/tests/test_litellm/llms/xai/test_xai_model_registry.py b/tests/unit/llms/xai/test_xai_model_registry.py similarity index 100% rename from tests/test_litellm/llms/xai/test_xai_model_registry.py rename to tests/unit/llms/xai/test_xai_model_registry.py diff --git a/tests/test_litellm/llms/xai/test_xai_oauth.py b/tests/unit/llms/xai/test_xai_oauth.py similarity index 100% rename from tests/test_litellm/llms/xai/test_xai_oauth.py rename to tests/unit/llms/xai/test_xai_oauth.py diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_client_unit.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_client_unit.py index 6438525706a..1a592aa1c9a 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_client_unit.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_client_unit.py @@ -12,7 +12,7 @@ import litellm.experimental_mcp_client.client as mcp_client_module from litellm.experimental_mcp_client.client import MCPClient from litellm.types.mcp import MCPAuth, MCPTransport from mcp.types import CallToolResult as MCPCallToolResult -from mcp.types import ListToolsResult, PaginatedRequestParams +from mcp.types import Implementation, InitializeResult, ListToolsResult, PaginatedRequestParams, ServerCapabilities from mcp.types import Tool as MCPTool @@ -128,6 +128,11 @@ class TestMCPClientUnitTests: mock_session_ctx = AsyncMock() mock_session_class.return_value = mock_session_ctx mock_session_instance = AsyncMock() + mock_session_instance.initialize.return_value = InitializeResult( + protocol_version="2025-11-25", + capabilities=ServerCapabilities(), + server_info=Implementation(name="test-peer", version="1"), + ) mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance) client = MCPClient( @@ -163,6 +168,11 @@ class TestMCPClientUnitTests: mock_session_ctx = AsyncMock() mock_session_class.return_value = mock_session_ctx mock_session_instance = AsyncMock() + mock_session_instance.initialize.return_value = InitializeResult( + protocol_version="2025-11-25", + capabilities=ServerCapabilities(), + server_info=Implementation(name="test-peer", version="1"), + ) mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance) mock_tools = [ @@ -204,6 +214,11 @@ class TestMCPClientUnitTests: mock_session_ctx = AsyncMock() mock_session_class.return_value = mock_session_ctx mock_session_instance = AsyncMock() + mock_session_instance.initialize.return_value = InitializeResult( + protocol_version="2025-11-25", + capabilities=ServerCapabilities(), + server_info=Implementation(name="test-peer", version="1"), + ) mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance) first_page_tools = [ @@ -245,6 +260,11 @@ class TestMCPClientUnitTests: mock_session_ctx = AsyncMock() mock_session_class.return_value = mock_session_ctx mock_session_instance = AsyncMock() + mock_session_instance.initialize.return_value = InitializeResult( + protocol_version="2025-11-25", + capabilities=ServerCapabilities(), + server_info=Implementation(name="test-peer", version="1"), + ) mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance) mock_session_instance.list_tools.side_effect = [ @@ -277,6 +297,11 @@ class TestMCPClientUnitTests: mock_session_ctx = AsyncMock() mock_session_class.return_value = mock_session_ctx mock_session_instance = AsyncMock() + mock_session_instance.initialize.return_value = InitializeResult( + protocol_version="2025-11-25", + capabilities=ServerCapabilities(), + server_info=Implementation(name="test-peer", version="1"), + ) mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance) mock_result = MCPCallToolResult(content=[]) @@ -289,7 +314,7 @@ class TestMCPClientUnitTests: assert result == mock_result mock_session_instance.initialize.assert_called_once() mock_session_instance.call_tool.assert_called_once_with( - name="test_tool", arguments={"arg1": "value1"}, progress_callback=ANY + name="test_tool", arguments={"arg1": "value1"}, progress_callback=ANY, allow_input_required=False ) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py index dfd250338ff..26674373b4e 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py @@ -489,7 +489,7 @@ async def test_sse_mcp_handler_mock(): ) with ( - patch("litellm.proxy._experimental.mcp_server.server.server.run", run), + patch("litellm.proxy._experimental.mcp_server.server.serve_loop", run), patch( "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True, @@ -511,7 +511,7 @@ async def test_sse_mcp_handler_mock(): # Call the handler await handle_sse_mcp(mock_scope, mock_receive, mock_send) - assert run.await_args.args[:2] == (read_stream, write_stream) + assert run.await_args.args[1:3] == (read_stream, write_stream) assert mock_sse.connect_sse.call_args.args[0]["path"] == "/mcp/sse" diff --git a/tests/unit/secret_managers/__init__.py b/tests/unit/secret_managers/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/secret_managers/hashicorp_vault_parity.json b/tests/unit/secret_managers/hashicorp_vault_parity.json similarity index 100% rename from tests/test_litellm/secret_managers/hashicorp_vault_parity.json rename to tests/unit/secret_managers/hashicorp_vault_parity.json diff --git a/tests/test_litellm/secret_managers/test_aws_secret_manager_replication.py b/tests/unit/secret_managers/test_aws_secret_manager_replication.py similarity index 100% rename from tests/test_litellm/secret_managers/test_aws_secret_manager_replication.py rename to tests/unit/secret_managers/test_aws_secret_manager_replication.py diff --git a/tests/test_litellm/secret_managers/test_aws_secret_manager_rotation.py b/tests/unit/secret_managers/test_aws_secret_manager_rotation.py similarity index 100% rename from tests/test_litellm/secret_managers/test_aws_secret_manager_rotation.py rename to tests/unit/secret_managers/test_aws_secret_manager_rotation.py diff --git a/tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py b/tests/unit/secret_managers/test_aws_secret_manager_v2.py similarity index 100% rename from tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py rename to tests/unit/secret_managers/test_aws_secret_manager_v2.py diff --git a/tests/test_litellm/secret_managers/test_base_secret_manager.py b/tests/unit/secret_managers/test_base_secret_manager.py similarity index 100% rename from tests/test_litellm/secret_managers/test_base_secret_manager.py rename to tests/unit/secret_managers/test_base_secret_manager.py diff --git a/tests/test_litellm/secret_managers/test_custom_secret_manager.py b/tests/unit/secret_managers/test_custom_secret_manager.py similarity index 100% rename from tests/test_litellm/secret_managers/test_custom_secret_manager.py rename to tests/unit/secret_managers/test_custom_secret_manager.py diff --git a/tests/test_litellm/secret_managers/test_cyberark_secret_manager.py b/tests/unit/secret_managers/test_cyberark_secret_manager.py similarity index 100% rename from tests/test_litellm/secret_managers/test_cyberark_secret_manager.py rename to tests/unit/secret_managers/test_cyberark_secret_manager.py diff --git a/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py b/tests/unit/secret_managers/test_get_azure_ad_token_provider.py similarity index 100% rename from tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py rename to tests/unit/secret_managers/test_get_azure_ad_token_provider.py diff --git a/tests/test_litellm/secret_managers/test_hashicorp_secret_manager.py b/tests/unit/secret_managers/test_hashicorp_secret_manager.py similarity index 100% rename from tests/test_litellm/secret_managers/test_hashicorp_secret_manager.py rename to tests/unit/secret_managers/test_hashicorp_secret_manager.py diff --git a/tests/test_litellm/secret_managers/test_secret_manager_handler.py b/tests/unit/secret_managers/test_secret_manager_handler.py similarity index 100% rename from tests/test_litellm/secret_managers/test_secret_manager_handler.py rename to tests/unit/secret_managers/test_secret_manager_handler.py diff --git a/tests/test_litellm/secret_managers/test_secret_managers_main.py b/tests/unit/secret_managers/test_secret_managers_main.py similarity index 100% rename from tests/test_litellm/secret_managers/test_secret_managers_main.py rename to tests/unit/secret_managers/test_secret_managers_main.py diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 99dea6366f9..62ef9f11c2e 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -4581,6 +4581,55 @@ def test_every_openai_entry_with_a_long_context_rate_and_a_batch_rate_declares_t assert undeclared == [] +@pytest.mark.parametrize("prefix", _BATCH_RATE_PREFIXES) +def test_every_xai_entry_with_a_long_context_rate_and_a_batch_rate_declares_the_batch_tier( + _local_model_cost_map: None, prefix: str +) -> None: + undeclared: Final = [ + name + for name, entry in litellm.model_cost.items() + if isinstance(entry, dict) + and entry.get("litellm_provider") == "xai" + and entry.get(f"{prefix}_above_200k_tokens") is not None + and entry.get(f"{prefix}_batches") is not None + and entry.get(f"{prefix}_above_200k_tokens_batches") is None + ] + + assert undeclared == [] + + +_XAI_TIERED_BATCH_MODEL: Final = "xai/grok-4.3" + + +def test_xai_batch_tier_discounts_the_long_context_rate_like_the_flat_batch_rate(_local_model_cost_map: None) -> None: + info: Final = litellm.get_model_info(_XAI_TIERED_BATCH_MODEL, custom_llm_provider="xai") + flat_discount: Final = info["input_cost_per_token_batches"] / info["input_cost_per_token"] + + for prefix in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost"): + tier_discount = info[f"{prefix}_above_200k_tokens_batches"] / info[f"{prefix}_above_200k_tokens"] + assert tier_discount == pytest.approx(flat_discount) + assert info[f"{prefix}_above_200k_tokens_batches"] < info[f"{prefix}_above_200k_tokens"] + + +@pytest.mark.parametrize( + ("prompt_tokens", "tier"), [(200_000, "_above_200k_tokens_batches"), (199_999, "_batches")] +) +def test_xai_batch_cost_calculator_bills_the_200k_batch_tier_inclusively( + _local_model_cost_map: None, prompt_tokens: int, tier: str +) -> None: + from litellm.cost_calculator import batch_cost_calculator + + info: Final = litellm.get_model_info(_XAI_TIERED_BATCH_MODEL, custom_llm_provider="xai") + usage: Final = Usage(prompt_tokens=prompt_tokens, completion_tokens=64, total_tokens=prompt_tokens + 64) + + prompt_cost, completion_cost_value = batch_cost_calculator( + usage=usage, model=_XAI_TIERED_BATCH_MODEL, custom_llm_provider="xai" + ) + + assert prompt_cost == pytest.approx(prompt_tokens * info[f"input_cost_per_token{tier}"]) + assert completion_cost_value == pytest.approx(64 * info[f"output_cost_per_token{tier}"]) + + def test_batch_cost_calculator_ignores_malformed_batch_tier_keys(): from litellm.cost_calculator import batch_cost_calculator diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 80131534183..3393c2f0d3c 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -4486,6 +4486,107 @@ async def test_aresponses_streaming_iterator_pre_first_chunk_skips_continuation( assert fbk["input"] == "Hello" # original input, no continuation messages +def _make_native_responses_iterator(*, sse_payloads: tuple[dict[str, str], ...], trailing_error: Exception | None): + """A real ResponsesAPIStreamingIterator over canned SSE bytes, so the router test covers the + iterator's own transport-error classification instead of a hand-built MidStreamFallbackError.""" + from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig + from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator + + async def aiter_bytes(): + for payload in sse_payloads: + yield f"data: {json.dumps(payload)}\n\n".encode() + if trailing_error is not None: + raise trailing_error + + def transform(model, parsed_chunk, logging_obj): + return MagicMock(type=parsed_chunk["type"]) + + response: Final = MagicMock() + response.headers = {} + response.aiter_bytes = aiter_bytes + config: Final = MagicMock(spec=BaseResponsesAPIConfig) + config.transform_streaming_response.side_effect = transform + logging_obj: Final = MagicMock(spec=LiteLLMLogging) + logging_obj.completion_start_time = None + logging_obj.model_call_details = {"litellm_params": {}} + return ResponsesAPIStreamingIterator( + response=response, + model="gpt-4", + responses_api_provider_config=config, + logging_obj=logging_obj, + litellm_metadata={}, + custom_llm_provider="openai", + ) + + +_RESPONSES_LIFECYCLE_PAYLOADS: Final = ({"type": "response.created"}, {"type": "response.in_progress"}) + + +@pytest.mark.asyncio +async def test_aresponses_streaming_iterator_falls_back_on_transport_drop_before_output(): + """A connection lost after response.created but before any output item is re-routed to the + fallback with the original input, the same as a provider error event would be.""" + router: Final = _make_router_with_fallback() + src: Final = _make_native_responses_iterator( + sse_payloads=_RESPONSES_LIFECYCLE_PAYLOADS, + trailing_error=httpx.ReadError("Response payload is not completed"), + ) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + return_value=_AsyncList([MagicMock(type="response.completed")]), + ) as mock_fallback_utils: + wrapped: Final = await router._aresponses_streaming_iterator( + response=src, + initial_kwargs={ + "model": "gpt-4", + "stream": True, + "input": "Hello", + "original_generic_function": litellm.aresponses, + }, + ) + seen: Final = [chunk.type async for chunk in wrapped] + + assert seen == ["response.created", "response.in_progress", "response.completed"] + assert isinstance(mock_fallback_utils.call_args.kwargs["e"], MidStreamFallbackError) + assert mock_fallback_utils.call_args.kwargs["kwargs"]["input"] == "Hello" + + +@pytest.mark.asyncio +async def test_aresponses_streaming_iterator_surfaces_transport_drop_when_no_fallback_lands(): + transport_error: Final = httpx.ReadError("Response payload is not completed") + router: Final = _make_router_with_fallback() + src: Final = _make_native_responses_iterator( + sse_payloads=_RESPONSES_LIFECYCLE_PAYLOADS, trailing_error=transport_error + ) + + async def reraise_trigger(**kwargs): + raise kwargs["e"] + + with patch.object( + router, "async_function_with_fallbacks_common_utils", new=AsyncMock(side_effect=reraise_trigger) + ) as mock_fallback_utils: + wrapped: Final = await router._aresponses_streaming_iterator( + response=src, + initial_kwargs={ + "model": "gpt-4", + "stream": True, + "input": "Hello", + "original_generic_function": litellm.aresponses, + }, + ) + with pytest.raises(httpx.ReadError) as exc_info: + async for _ in wrapped: + pass + + assert exc_info.value is transport_error + assert mock_fallback_utils.await_count == 1 + trigger: Final = mock_fallback_utils.await_args.kwargs["e"] + assert isinstance(trigger, MidStreamFallbackError) + assert trigger.original_exception is transport_error + + @pytest.mark.asyncio async def test_aresponses_streaming_iterator_partial_content_injects_continuation(): """Mid-stream error: input is rewritten to include user prompt + @@ -6090,6 +6191,32 @@ def test_update_kwargs_with_deployment_passthrough_router_stream_timeout_sources assert _passthrough_timeout(default_router, default_router.model_list[0], stream=False) == 120.0 +def test_update_kwargs_with_deployment_passthrough_honors_global_request_timeout(monkeypatch: pytest.MonkeyPatch): + """litellm_settings.request_timeout must bound the native responses route when neither the + deployment nor the router carries a timeout, while a deployment timeout keeps winning.""" + monkeypatch.setattr("litellm.request_timeout", 44.0, raising=False) + monkeypatch.setattr("litellm.request_timeout_explicitly_set", True, raising=False) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "responses-global-timeout", + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "fake-key"}, + }, + { + "model_name": "responses-deployment-timeout", + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "fake-key", "timeout": 3}, + }, + ], + ) + global_only, per_deployment = router.model_list + + with patch("litellm.proxy.proxy_server.general_settings", {"pass_through_request_timeout": 6}): + assert _passthrough_timeout(router, global_only, stream=True) == 44.0 + assert _passthrough_timeout(router, global_only, stream=False) == 44.0 + assert _passthrough_timeout(router, per_deployment, stream=True) == 3.0 + assert _passthrough_timeout(router, per_deployment, stream=False) == 3.0 + + @pytest.mark.asyncio async def test_router_acompletion_with_unknown_model_and_default_fallback(): """ @@ -6343,6 +6470,51 @@ def test_get_deployment_credentials_with_provider_includes_bucket_name(): assert credentials["custom_llm_provider"] == "vertex_ai" +def test_get_deployment_credentials_with_provider_keeps_legacy_bucket_name(): + router = litellm.Router( + model_list=[ + { + "model_name": "vertex-gemini", + "litellm_params": { + "model": "vertex_ai/gemini-3.5-flash", + "vertex_project": "my-project", + "vertex_location": "global", + "bucket_name": "my-legacy-bucket", + }, + } + ], + ) + + credentials = router.get_deployment_credentials_with_provider(model_id="vertex-gemini") + + assert credentials is not None + assert credentials["bucket_name"] == "my-legacy-bucket" + assert "gcs_bucket_name" not in credentials + + +def test_get_deployment_credentials_with_provider_keeps_both_bucket_keys(): + router = litellm.Router( + model_list=[ + { + "model_name": "vertex-gemini", + "litellm_params": { + "model": "vertex_ai/gemini-3.5-flash", + "vertex_project": "my-project", + "vertex_location": "global", + "gcs_bucket_name": "new-bucket", + "bucket_name": "legacy-bucket", + }, + } + ], + ) + + credentials = router.get_deployment_credentials_with_provider(model_id="vertex-gemini") + + assert credentials is not None + assert credentials["gcs_bucket_name"] == "new-bucket" + assert credentials["bucket_name"] == "legacy-bucket" + + def test_get_deployment_credentials_with_provider_resolves_credential_name(): """ Test that get_deployment_credentials_with_provider correctly resolves diff --git a/tests/unit/test_unit_shard_missing_paths.py b/tests/unit/test_unit_shard_missing_paths.py index 4fa9c5bd3c1..0360a227142 100644 --- a/tests/unit/test_unit_shard_missing_paths.py +++ b/tests/unit/test_unit_shard_missing_paths.py @@ -38,6 +38,7 @@ def _run_shard(tmp_path: Path, test_path: str, workers: str) -> subprocess.Compl "PATH": f"{shim_dir}{os.pathsep}{os.environ['PATH']}", "GITHUB_OUTPUT": str(tmp_path / "github_output"), "TEST_PATH": test_path, + "UNIT_FLAG": "", "WORKERS": workers, }, capture_output=True, diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 768d8955b8e..0cdc52c9a93 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -774,6 +774,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "cache_read_input_token_cost_above_32k_tokens": {"type": "number"}, "cache_read_input_token_cost_above_128k_tokens": {"type": "number"}, "cache_read_input_token_cost_above_200k_tokens": {"type": "number"}, + "cache_read_input_token_cost_above_200k_tokens_batches": {"type": "number"}, "cache_read_input_token_cost_above_256k_tokens": {"type": "number"}, "cache_read_input_token_cost_above_272k_tokens": {"type": "number"}, "cache_read_input_token_cost_above_272k_tokens_flex": {"type": "number"}, @@ -797,6 +798,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "input_cost_per_video_token": {"type": "number"}, "input_cost_per_token_above_32k_tokens": {"type": "number"}, "input_cost_per_token_above_200k_tokens": {"type": "number"}, + "input_cost_per_token_above_200k_tokens_batches": {"type": "number"}, "input_cost_per_token_above_256k_tokens": {"type": "number"}, "input_cost_per_token_above_272k_tokens": {"type": "number"}, "input_cost_per_token_above_512k_tokens": {"type": "number"}, @@ -897,6 +899,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "output_cost_per_token_above_32k_tokens": {"type": "number"}, "output_cost_per_token_above_128k_tokens": {"type": "number"}, "output_cost_per_token_above_200k_tokens": {"type": "number"}, + "output_cost_per_token_above_200k_tokens_batches": {"type": "number"}, "output_cost_per_token_above_256k_tokens": {"type": "number"}, "output_cost_per_token_above_272k_tokens": {"type": "number"}, "output_cost_per_token_above_512k_tokens": {"type": "number"}, diff --git a/ui/litellm-dashboard/src/components/callback_info_helpers.tsx b/ui/litellm-dashboard/src/components/callback_info_helpers.tsx index f5138d55b5d..bc9889da724 100644 --- a/ui/litellm-dashboard/src/components/callback_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/callback_info_helpers.tsx @@ -10,6 +10,7 @@ import newrelicLogo from "../../public/assets/logos/newrelic.png"; import openmeterLogo from "../../public/assets/logos/openmeter.png"; import otelLogo from "../../public/assets/logos/otel.png"; import pointfiveLogo from "../../public/assets/logos/pointfive.png"; +import databricksLogo from "../../public/assets/logos/databricks.svg"; interface CallbackConfig { id: string; @@ -181,6 +182,20 @@ export const CALLBACK_CONFIGS: CallbackConfig[] = [ }, description: "PointFive Logging Integration", }, + { + id: "zerobus", + displayName: "Databricks Zerobus", + logo: databricksLogo.src, + supports_key_team_logging: false, + dynamic_params: { + ZEROBUS_WORKSPACE_URL: "text", + ZEROBUS_SERVER_ENDPOINT: "text", + ZEROBUS_CLIENT_ID: "text", + ZEROBUS_CLIENT_SECRET: "password", + ZEROBUS_TABLE_NAME: "text", + }, + description: "Databricks Zerobus Ingest Logging Integration", + }, { id: "s3", displayName: "S3", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index bba67bcf6c2..1e0aed46923 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -28350,6 +28350,11 @@ export interface components { * @description Maximum retention period for spend logs (e.g., '7d' for 7 days). Logs older than this will be deleted. */ maximum_spend_logs_retention_period?: string | null; + /** + * Mcp Advertised Versions + * @description MCP revisions enabled by the gateway. Defaults to all completed legacy revisions. Modern protocol serving and Apps/Tasks remain disabled. + */ + mcp_advertised_versions?: ("2024-11-05" | "2025-03-26" | "2025-06-18" | "2025-11-25")[] | null; /** * Mcp Allowed Clients * @description MCP client applications admitted by the gateway, each an {alias, value} pair where alias is the name shown in the dashboard and logs and value is the identity that must match exactly. When set, every MCP request must carry a client identity equal to one of the values: a JWT caller is identified by the claim named in litellm_jwtauth.mcp_client_id_jwt_field, any other caller by the header named in mcp_client_id_header. A request with no resolvable identity, or an unlisted one, is rejected with 403. Unset means every client is admitted. @@ -32940,6 +32945,8 @@ export interface components { azure_username?: string | null; /** Bedrock Tags */ bedrock_tags?: unknown[] | null; + /** Bucket Name */ + bucket_name?: string | null; /** Budget Duration */ budget_duration?: string | null; /** Cache Creation Input Audio Token Cost */ @@ -32974,6 +32981,8 @@ export interface components { cache_read_input_token_cost?: number | null; /** Cache Read Input Token Cost Above 200K Tokens */ cache_read_input_token_cost_above_200k_tokens?: number | null; + /** Cache Read Input Token Cost Above 200K Tokens Batches */ + cache_read_input_token_cost_above_200k_tokens_batches?: number | null; /** Cache Read Input Token Cost Above 200K Tokens Priority */ cache_read_input_token_cost_above_200k_tokens_priority?: number | null; /** Cache Read Input Token Cost Above 272K Tokens */ @@ -33052,6 +33061,8 @@ export interface components { input_cost_per_token_above_128k_tokens?: number | null; /** Input Cost Per Token Above 200K Tokens */ input_cost_per_token_above_200k_tokens?: number | null; + /** Input Cost Per Token Above 200K Tokens Batches */ + input_cost_per_token_above_200k_tokens_batches?: number | null; /** Input Cost Per Token Above 200K Tokens Priority */ input_cost_per_token_above_200k_tokens_priority?: number | null; /** Input Cost Per Token Above 272K Tokens */ @@ -33175,6 +33186,8 @@ export interface components { output_cost_per_token_above_128k_tokens?: number | null; /** Output Cost Per Token Above 200K Tokens */ output_cost_per_token_above_200k_tokens?: number | null; + /** Output Cost Per Token Above 200K Tokens Batches */ + output_cost_per_token_above_200k_tokens_batches?: number | null; /** Output Cost Per Token Above 200K Tokens Priority */ output_cost_per_token_above_200k_tokens_priority?: number | null; /** Output Cost Per Token Above 272K Tokens */ @@ -46725,6 +46738,8 @@ export interface components { azure_username?: string | null; /** Bedrock Tags */ bedrock_tags?: unknown[] | null; + /** Bucket Name */ + bucket_name?: string | null; /** Budget Duration */ budget_duration?: string | null; /** Cache Creation Input Audio Token Cost */ @@ -46759,6 +46774,8 @@ export interface components { cache_read_input_token_cost?: number | null; /** Cache Read Input Token Cost Above 200K Tokens */ cache_read_input_token_cost_above_200k_tokens?: number | null; + /** Cache Read Input Token Cost Above 200K Tokens Batches */ + cache_read_input_token_cost_above_200k_tokens_batches?: number | null; /** Cache Read Input Token Cost Above 200K Tokens Priority */ cache_read_input_token_cost_above_200k_tokens_priority?: number | null; /** Cache Read Input Token Cost Above 272K Tokens */ @@ -46837,6 +46854,8 @@ export interface components { input_cost_per_token_above_128k_tokens?: number | null; /** Input Cost Per Token Above 200K Tokens */ input_cost_per_token_above_200k_tokens?: number | null; + /** Input Cost Per Token Above 200K Tokens Batches */ + input_cost_per_token_above_200k_tokens_batches?: number | null; /** Input Cost Per Token Above 200K Tokens Priority */ input_cost_per_token_above_200k_tokens_priority?: number | null; /** Input Cost Per Token Above 272K Tokens */ @@ -46960,6 +46979,8 @@ export interface components { output_cost_per_token_above_128k_tokens?: number | null; /** Output Cost Per Token Above 200K Tokens */ output_cost_per_token_above_200k_tokens?: number | null; + /** Output Cost Per Token Above 200K Tokens Batches */ + output_cost_per_token_above_200k_tokens_batches?: number | null; /** Output Cost Per Token Above 200K Tokens Priority */ output_cost_per_token_above_200k_tokens_priority?: number | null; /** Output Cost Per Token Above 272K Tokens */ diff --git a/uv.lock b/uv.lock index c235171ecb2..527f53bd372 100644 --- a/uv.lock +++ b/uv.lock @@ -4743,6 +4743,7 @@ proxy-dev = [ { name = "opentelemetry-sdk" }, { name = "prisma" }, { name = "prometheus-client" }, + { name = "sentry-sdk" }, ] [package.metadata] @@ -4956,6 +4957,7 @@ proxy-dev = [ { name = "opentelemetry-sdk", specifier = "==1.33.1" }, { name = "prisma", specifier = "==0.11.0" }, { name = "prometheus-client", specifier = "==0.20.0" }, + { name = "sentry-sdk", specifier = "==2.21.0" }, ] [[package]]