diff --git a/.github/actions/setup-uv-with-retries/action.yml b/.github/actions/setup-uv-with-retries/action.yml index 1627038dc3d..98ff91f0283 100644 --- a/.github/actions/setup-uv-with-retries/action.yml +++ b/.github/actions/setup-uv-with-retries/action.yml @@ -1,11 +1,7 @@ name: "Set up uv with retries" description: >- - Install uv via astral-sh/setup-uv, retrying on transient failures. Even with - an exact pinned version, the action resolves the artifact URL by fetching - https://raw.githubusercontent.com/astral-sh/versions/main/v1/uv.ndjson in a - single request with no retry, timeout, or fallback, so one connection-level - network error ("fetch failed") fails the whole job before any test runs. - Retrying the full step covers the manifest fetch and the binary download. + Install uv via astral-sh/setup-uv, retrying the full setup step so manifest + resolution and binary downloads get fresh attempts after transient failures. inputs: version: @@ -18,7 +14,7 @@ runs: - name: Set up uv (attempt 1) id: attempt-1 continue-on-error: true - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 with: version: ${{ inputs.version }} @@ -31,7 +27,7 @@ runs: id: attempt-2 if: steps.attempt-1.outcome == 'failure' continue-on-error: true - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 with: version: ${{ inputs.version }} @@ -42,6 +38,6 @@ runs: - name: Set up uv (attempt 3) if: steps.attempt-2.outcome == 'failure' - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 with: version: ${{ inputs.version }} diff --git a/.github/ci-coverage-allowlist.yml b/.github/ci-coverage-allowlist.yml index 32232de381c..0da07038152 100644 --- a/.github/ci-coverage-allowlist.yml +++ b/.github/ci-coverage-allowlist.yml @@ -79,6 +79,11 @@ test_paths: - tests/load_tests/test_otel_load_test.py - tests/load_tests/test_vertex_embeddings_load_test.py - tests/load_tests/test_vertex_load_tests.py + - reason: >- + Env-gated saturation benchmark requires a live proxy and provider credentials, so it is run + locally rather than in pull-request jobs + paths: + - tests/load_tests/test_granian_admission_saturation.py - reason: >- A local-only agent rig: test_a2a_completion_bridge.py needs a LangGraph server on localhost:2024 and test_a2a.py drives a live A2A endpoint, so neither can run in a diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 6bc44995804..33245ec5b5f 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -96,7 +96,6 @@ jobs: - shard: misc artifact-name: misc test-path: >- - tests/sdk_function_trace tests/test_litellm/batches tests/test_litellm/secret_managers tests/test_litellm/a2a_protocol diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 9a9b1138a7d..cb3756575f7 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -3,7 +3,7 @@ "limit": 14074 }, "reportArgumentType": { - "limit": 2214 + "limit": 2206 }, "reportAssignmentType": { "limit": 319 @@ -57,7 +57,7 @@ "limit": 5601 }, "reportMissingTypeArgument": { - "limit": 15287 + "limit": 15285 }, "reportMissingTypeStubs": { "limit": 40 @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44362 + "limit": 44360 }, "reportUnknownLambdaType": { "limit": 109 }, "reportUnknownMemberType": { - "limit": 38323 + "limit": 38311 }, "reportUnknownParameterType": { "limit": 19624 }, "reportUnknownVariableType": { - "limit": 29861 + "limit": 29847 }, "reportUnnecessaryCast": { "limit": 111 diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index b3dac5ca935..803df633c27 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1435,14 +1435,17 @@ dependencies = [ "aws-sigv4", "aws-smithy-runtime-api", "aws-types", + "base64", "rand 0.8.7", "reqwest", + "rstest", "serde", "serde_json", "sha2 0.10.9", "thiserror 2.0.19", "tokio", "tracing", + "tracing-subscriber", ] [[package]] @@ -1461,7 +1464,6 @@ dependencies = [ "tokio", "tokio-tungstenite", "tracing", - "tracing-subscriber", ] [[package]] @@ -1511,6 +1513,16 @@ version = "0.3.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" +[[package]] +name = "mime_guess" +version = "2.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f7c44f8e672c00fe5308fa235f821cb4198414e1c77935c1ab6948d3fd78550e" +dependencies = [ + "mime", + "unicase", +] + [[package]] name = "mio" version = "1.2.2" @@ -1969,6 +1981,7 @@ dependencies = [ "hyper-util", "js-sys", "log", + "mime_guess", "percent-encoding", "pin-project-lite", "quinn", @@ -2736,6 +2749,12 @@ version = "1.20.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" +[[package]] +name = "unicase" +version = "2.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dbc4bc3a9f746d862c45cb89d705aa10f187bb96c76001afab07a0d35ce60142" + [[package]] name = "unicode-ident" version = "1.0.24" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 643ad985251..62f62872dd7 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -24,7 +24,7 @@ pyo3 = "0.29.2" pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } pythonize = "0.29.0" rand = "0.8" -reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] } +reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] } rstest = "0.26.1" serde = { version = "1.0", features = ["derive"] } serde_json = { version = "1.0", features = ["float_roundtrip"] } diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml index e3dbdf24ce6..eef2bf55a07 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -32,6 +32,7 @@ subtle = { workspace = true, optional = true } # SHA-256 hash_token) so the plaintext credential never enters a log payload. sha2 = { workspace = true, optional = true } pyo3 = { workspace = true, features = ["auto-initialize"], optional = true } +tower = { version = "0.5.3", features = ["util"], optional = true } [features] default = [] @@ -39,6 +40,7 @@ server = ["dep:axum", "dep:subtle", "dep:sha2"] # Build the gateway's config from the proxy YAML via an embedded Python # interpreter (links libpython; requires `litellm` importable at runtime). python-config = ["dep:pyo3"] +trace-parity = ["server", "dep:tower", "litellm-core/observability"] [dev-dependencies] futures-channel = "0.3" diff --git a/litellm-rust/crates/ai-gateway/src/lib.rs b/litellm-rust/crates/ai-gateway/src/lib.rs index 057db6457c4..a2950748afc 100644 --- a/litellm-rust/crates/ai-gateway/src/lib.rs +++ b/litellm-rust/crates/ai-gateway/src/lib.rs @@ -28,6 +28,8 @@ pub mod auth; pub mod routes; #[cfg(feature = "server")] pub mod state; +#[cfg(feature = "trace-parity")] +pub mod trace_parity; mod constants; pub mod integrations; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs index c1fb328893b..d2be17260a3 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs @@ -12,6 +12,7 @@ use litellm_core::providers::azure_ai::ocr::transformation::{ AZURE_AI_OCR_CONFIG, AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG, }; use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; +use litellm_core::providers::reducto::ocr::transformation as reducto; use litellm_core::providers::vertex_ai::ocr::transformation as vertex_ai; use litellm_core::providers::vertex_ai::ocr::transformation::{ VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG, @@ -39,6 +40,7 @@ pub(super) fn ocr_provider_config( ) -> Option<&'static dyn OcrProviderConfig> { match provider { "mistral" => Some(&MISTRAL_OCR_CONFIG), + "reducto" => reducto::config_for_model(model), "azure_ai" if is_azure_document_intelligence_model(model) => { Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG) } @@ -334,6 +336,7 @@ fn operation_status(response_json: &Value) -> Result<&str, Error> { } } +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub(super) async fn poll_document_intelligence( operation_url: &str, original_url: &str, @@ -392,9 +395,11 @@ pub(super) async fn poll_document_intelligence( #[cfg(test)] mod tests { - use super::*; + use litellm_core::ocr::transformation::OcrResponseHandling; use serde_json::json; + use super::*; + #[test] fn blocks_private_and_metadata_ips() { assert!(is_blocked_ip("127.0.0.1".parse().unwrap())); @@ -438,4 +443,87 @@ mod tests { assert_eq!(transformed, document); } + + #[test] + fn truncate_error_body_passes_short_strings_through() { + let body = "Unauthorized"; + assert_eq!(truncate_error_body(body), "Unauthorized"); + } + + #[test] + fn truncate_error_body_caps_long_payloads() { + let body = "x".repeat(306); + let truncated = truncate_error_body(&body); + + assert!(truncated.ends_with("... (truncated)")); + let prefix_chars = truncated + .strip_suffix("... (truncated)") + .expect("truncated marker present") + .chars() + .count(); + assert_eq!(prefix_chars, 256); + } + + #[test] + fn truncate_error_body_does_not_split_multibyte_chars() { + let body = "é".repeat(266); + let truncated = truncate_error_body(&body); + assert!(truncated.is_char_boundary(truncated.len())); + } + + #[test] + fn ocr_dispatch_supports_migrated_providers() { + assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some()); + assert!( + ocr_provider_config("azure_ai", "pixtral-12b-2409") + .expect("azure ai config resolves") + .requires_data_uri_document() + ); + assert_eq!( + ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read") + .expect("document intelligence config resolves") + .response_handling(), + OcrResponseHandling::AzureDocumentIntelligencePoll + ); + assert!( + ocr_provider_config("vertex_ai", "deepseek-ocr-maas") + .expect("vertex deepseek config resolves") + .supported_ocr_params() + .contains(&"temperature") + ); + assert!(ocr_provider_config("openai", "gpt-4o").is_none()); + } + + #[test] + fn string_headers_accepts_string_values() { + let headers = json!({ + "x-trace-id": "trace-1" + }) + .as_object() + .unwrap() + .clone(); + + assert_eq!( + string_headers(Some(headers)).expect("string headers accepted"), + vec![("x-trace-id".to_string(), "trace-1".to_string())] + ); + } + + #[test] + fn string_headers_rejects_non_string_values() { + let headers = json!({ + "x-retry-count": 3 + }) + .as_object() + .unwrap() + .clone(); + + let err = string_headers(Some(headers)).expect_err("non-string header rejected"); + assert_eq!( + err, + Error::InvalidRequest( + "OCR extra_headers.x-retry-count must be a string, got number".to_string() + ) + ); + } } diff --git a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs index 856d9571201..6c6e12724cd 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs @@ -50,7 +50,11 @@ pub(crate) async fn execute_ocr_provider_call( .await?; return Ok(request .config - .transform_ocr_response(&request.model, response_json)? + .transform_ocr_response_with_params( + &request.model, + response_json, + &request.optional_params, + )? .into_json()); } @@ -71,6 +75,10 @@ pub(crate) async fn execute_ocr_provider_call( Ok(request .config - .transform_ocr_response(&request.model, response_json)? + .transform_ocr_response_with_params( + &request.model, + response_json, + &request.optional_params, + )? .into_json()) } diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs index f8c4f8fe8c5..446b323db3a 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -1,11 +1,15 @@ use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; use litellm_core::error::Error; +use litellm_core::providers::reducto::ocr::transformation::{ + build_upload_request, extract_document_source, extract_upload_file_id, +}; use serde_json::{Map, Value, json}; use std::future::Future; use std::pin::Pin; -use super::common_utils::{convert_document_url_to_data_uri, string_headers}; +use super::common_utils::{convert_document_url_to_data_uri, string_headers, truncate_error_body}; use super::types::{PreparedOcrRequest, ProviderOcrRequest}; +use crate::client::http_client; use crate::integrations::custom_guardrail::{ CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest, }; @@ -89,22 +93,39 @@ impl OcrLifecycleHooks { )?; let model = request.model.clone(); let custom_llm_provider = request.custom_llm_provider.clone(); - let document = if config.requires_data_uri_document() { + let is_reducto = custom_llm_provider == "reducto"; + let document = if is_reducto { + let guarded_document = self + .run_during_call_guardrails(&model, &custom_llm_provider, &url, request.document) + .await?; + upload_reducto_document( + &guarded_document, + request.api_base.as_deref(), + request.timeout, + &upstream_headers, + ) + .await? + } else if config.requires_data_uri_document() { convert_document_url_to_data_uri(request.document).await? } else { request.document }; + let optional_params = request.optional_params; let body = config - .transform_ocr_request(&request.model, document, request.optional_params)? + .transform_ocr_request(&request.model, document, optional_params.clone())? .data; - let body = self - .run_during_call_guardrails(&model, &custom_llm_provider, &url, body) - .await?; + let body = if is_reducto { + body + } else { + self.run_during_call_guardrails(&model, &custom_llm_provider, &url, body) + .await? + }; Ok(ProviderOcrRequest { model, config, url, body, + optional_params, upstream_headers, timeout: request.timeout, }) @@ -165,6 +186,63 @@ impl OcrLifecycleHooks { } } +async fn upload_reducto_document( + document: &Value, + api_base: Option<&str>, + timeout: Option, + upstream_headers: &[(String, String)], +) -> Result { + let source = extract_document_source(document)?; + let Some(authorization) = upstream_headers + .iter() + .find(|(name, _)| name.eq_ignore_ascii_case("authorization")) + .map(|(_, value)| value.as_str()) + else { + return Err(Error::Auth( + "Reducto upload requires an Authorization header".to_string(), + )); + }; + let Some(upload) = build_upload_request(source, authorization, api_base) else { + return Ok(document.clone()); + }; + let part = reqwest::multipart::Part::bytes(upload.bytes) + .file_name(upload.file_name) + .mime_str(&upload.mime_type) + .map_err(|error| Error::InvalidRequest(error.to_string()))?; + let form = reqwest::multipart::Form::new().part("file", part); + let mut request_builder = http_client().post(upload.url).multipart(form); + for (name, value) in upstream_headers { + if !name.eq_ignore_ascii_case("content-type") + && !name.eq_ignore_ascii_case("content-length") + { + request_builder = request_builder.header(name, value); + } + } + if let Some(timeout) = timeout { + request_builder = request_builder.timeout(timeout); + } + let response = request_builder + .send() + .await + .map_err(|error| Error::Network(error.to_string()))?; + let status = response.status(); + let body = response + .text() + .await + .map_err(|error| Error::Network(error.to_string()))?; + if !status.is_success() { + return Err(Error::Http { + status: status.as_u16(), + body: truncate_error_body(&body), + }); + } + let response_json: Value = serde_json::from_str(&body).map_err(|error| { + Error::InvalidResponse(format!("invalid Reducto upload response JSON: {error}")) + })?; + let file_id = extract_upload_file_id(&response_json)?; + Ok(json!({"type": "document_url", "document_url": file_id})) +} + impl CallLifecycleHooks for OcrLifecycleHooks { type PreCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>; type DuringCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs index d9230af1c59..2acdd232c80 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs @@ -24,4 +24,151 @@ pub async fn ocr(request: OcrRequest<'_>) -> Result { } #[cfg(test)] -mod tests; +mod tests { + use serde_json::{Map, json}; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::{TcpListener, TcpStream}; + + use super::{OcrRequest, ocr}; + use crate::integrations::types::RequestMetadata; + + async fn read_http_request(socket: &mut TcpStream) -> String { + let mut request = Vec::new(); + let mut buffer = [0_u8; 1024]; + let header_end = loop { + let n = socket.read(&mut buffer).await.expect("reads request"); + if n == 0 { + break request.len(); + } + request.extend_from_slice(&buffer[..n]); + if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") { + break position + 4; + } + }; + let headers = String::from_utf8_lossy(&request[..header_end]); + let content_length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) + .unwrap_or(0); + while request.len().saturating_sub(header_end) < content_length { + let n = socket.read(&mut buffer).await.expect("reads body"); + if n == 0 { + break; + } + request.extend_from_slice(&buffer[..n]); + } + String::from_utf8(request).expect("request is utf8") + } + + fn base_ocr_request(model: &str) -> OcrRequest<'_> { + OcrRequest { + model, + document: json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }), + api_key: Some("sk-test"), + api_base: None, + custom_llm_provider: None, + extra_headers: None, + optional_params: Map::new(), + timeout: None, + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: RequestMetadata::default(), + litellm_call_id: None, + } + } + + #[tokio::test] + async fn reducto_file_upload_then_parse_maps_response() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let address = listener.local_addr().expect("listener has local address"); + let server = tokio::spawn(async move { + let (mut upload_socket, _) = listener.accept().await.expect("accepts upload request"); + let upload_request = read_http_request(&mut upload_socket).await; + let upload_body = r#"{"file_id":"reducto://uploaded.pdf"}"#; + let upload_response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + upload_body.len(), + upload_body + ); + upload_socket + .write_all(upload_response.as_bytes()) + .await + .expect("writes upload response"); + + let (mut parse_socket, _) = listener.accept().await.expect("accepts parse request"); + let parse_request = read_http_request(&mut parse_socket).await; + let parse_body = r#"{"job_id":"job_123","usage":{"num_pages":3,"credits":3},"result":{"chunks":[{"content":"Page 1 block A","blocks":[{"content":"Page 1 block A","bbox":{"page":1},"kind":"text"}]},{"content":"Page 2 block A","blocks":[{"content":"Page 2 block A","bbox":{"page":2},"kind":"table"}]},{"content":"Page 1 block B","blocks":[{"content":"Page 1 block B","bbox":{"page":1},"kind":"text"}]},{"content":"Page 3 block A","blocks":[{"content":"Page 3 block A","bbox":{"page":3},"kind":"figure"}]}]}}"#; + let parse_response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + parse_body.len(), + parse_body + ); + parse_socket + .write_all(parse_response.as_bytes()) + .await + .expect("writes parse response"); + (upload_request, parse_request) + }); + let api_base = format!("http://{address}"); + let mut request = base_ocr_request("reducto/parse-v3"); + request.api_base = Some(&api_base); + request.api_key = None; + request.extra_headers = Some(Map::from_iter([ + ("Authorization".to_string(), json!("Bearer test-key")), + ("x-trace-id".to_string(), json!("trace-1")), + ])); + request.document = json!({ + "type": "document_url", + "document_url": "data:application/pdf;base64,JVBERi0xLjQ=" + }); + request.optional_params = Map::from_iter([ + ( + "formatting".to_string(), + json!({"table_output_format": "html"}), + ), + ("retrieval".to_string(), json!({"chunk_mode": "section"})), + ("settings".to_string(), json!({"ocr_system": "standard"})), + ]); + + let response = ocr(request).await.expect("Reducto OCR succeeds"); + + assert_eq!(response["pages"].as_array().map(Vec::len), Some(3)); + assert_eq!( + response["pages"][0]["markdown"], + "Page 1 block A\n\nPage 1 block B" + ); + assert_eq!(response["pages"][1]["markdown"], "Page 2 block A"); + assert_eq!(response["pages"][2]["markdown"], "Page 3 block A"); + assert_eq!(response["usage_info"]["pages_processed"], 3); + assert_eq!(response["usage_info"]["credits"], 3); + assert_eq!(response["provider_native_response"]["job_id"], "job_123"); + let (upload_request, parse_request) = server.await.expect("server task completes"); + assert!( + upload_request + .to_ascii_lowercase() + .contains("authorization: bearer test-key") + ); + assert!(upload_request.contains("application/pdf")); + assert!(upload_request.contains("%PDF-1.4")); + assert!(upload_request.contains("x-trace-id: trace-1")); + assert!( + parse_request + .to_ascii_lowercase() + .contains("authorization: bearer test-key") + ); + assert!(parse_request.contains(r#""input":"reducto://uploaded.pdf""#)); + assert!(parse_request.contains(r#""table_output_format":"html""#)); + assert!(parse_request.contains(r#""chunk_mode":"section""#)); + assert!(parse_request.contains(r#""ocr_system":"standard""#)); + } +} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs index fedacc62760..fa9ca1a193e 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs @@ -2,6 +2,7 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{SystemTime, UNIX_EPOCH}; use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; +use serde_json::{Map, Value}; use super::common_utils::ocr_provider_config; use super::hooks::OcrLifecycleHooks; @@ -28,17 +29,33 @@ pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall { let model = provider_info.model.to_string(); let custom_llm_provider = provider_info.custom_llm_provider.to_string(); let config = ocr_provider_config(&custom_llm_provider, &model) - .ok_or_else(|| litellm_core::Error::InvalidProvider(custom_llm_provider.clone())); + .ok_or_else(|| litellm_core::Error::InvalidProvider(custom_llm_provider.clone())) + .and_then(|config| { + validate_request_format(config, &request.optional_params, &custom_llm_provider)?; + Ok(config) + }); let optional_params = match &config { Ok(config) => { let supported = config.supported_ocr_params(); - config.map_ocr_params( + let mut mapped = config.map_ocr_params( &request .optional_params - .into_iter() + .iter() .filter(|(name, _)| supported.contains(&name.as_str())) + .map(|(name, value)| (name.clone(), value.clone())) .collect(), - ) + ); + for name in [ + "vertex_project", + "vertex_ai_project", + "vertex_location", + "vertex_ai_location", + ] { + if let Some(value) = request.optional_params.get(name) { + mapped.insert(name.to_string(), value.clone()); + } + } + mapped } Err(_) => request.optional_params, }; @@ -64,6 +81,26 @@ pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall { } } +fn validate_request_format( + config: &'static dyn litellm_core::ocr::transformation::OcrProviderConfig, + optional_params: &Map, + provider: &str, +) -> Result<(), litellm_core::Error> { + let Some(format) = optional_params.get("req_format") else { + return Ok(()); + }; + match format.as_str() { + Some("litellm") => Ok(()), + Some("native") if config.supported_ocr_params().contains(&"req_format") => Ok(()), + Some("native") => Err(litellm_core::Error::InvalidRequest(format!( + "`req_format=native` is not supported for provider {provider}" + ))), + _ => Err(litellm_core::Error::InvalidRequest(format!( + "Invalid `req_format`: {format}. Expected `litellm` or `native`" + ))), + } +} + fn new_ocr_call_id() -> String { static COUNTER: AtomicU64 = AtomicU64::new(1); let sequence = COUNTER.fetch_add(1, Ordering::Relaxed); @@ -73,3 +110,54 @@ fn new_ocr_call_id() -> String { .unwrap_or(0); format!("ocr-{timestamp}-{sequence}") } + +#[cfg(test)] +mod tests { + use litellm_core::error::Error; + use serde_json::{Map, json}; + + use super::{OcrRequest, prepare_ocr_call}; + use crate::integrations::types::RequestMetadata; + + fn base_ocr_request(model: &str) -> OcrRequest<'_> { + OcrRequest { + model, + document: json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }), + api_key: Some("sk-test"), + api_base: None, + custom_llm_provider: None, + extra_headers: None, + optional_params: Map::new(), + timeout: None, + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: RequestMetadata::default(), + litellm_call_id: None, + } + } + + fn request_with_format(format: &str) -> OcrRequest<'_> { + let mut request = base_ocr_request("mistral/mistral-ocr-latest"); + request.optional_params = Map::from_iter([("req_format".to_string(), json!(format))]); + request + } + + #[test] + fn native_format_rejected_for_provider_without_support_as_bad_request() { + let prepared = prepare_ocr_call(request_with_format("native")); + assert!( + matches!(prepared.request.config, Err(Error::InvalidRequest(message)) if message.contains("not supported for provider")) + ); + } + + #[test] + fn unknown_format_rejected_for_provider_without_support_as_bad_request() { + let prepared = prepare_ocr_call(request_with_format("raw")); + assert!( + matches!(prepared.request.config, Err(Error::InvalidRequest(message)) if message.contains("Invalid `req_format`")) + ); + } +} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/types.rs b/litellm-rust/crates/ai-gateway/src/ocr/types.rs index 95e551d79ca..75a8e61ddbf 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/types.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/types.rs @@ -53,6 +53,7 @@ pub(crate) struct ProviderOcrRequest { pub(crate) config: &'static dyn OcrProviderConfig, pub(crate) url: String, pub(crate) body: Value, + pub(crate) optional_params: Map, pub(crate) upstream_headers: Vec<(String, String)>, pub(crate) timeout: Option, } diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs index e9f8c477f36..bb9f3851a77 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs @@ -21,6 +21,12 @@ pub fn router() -> Router { Router::new().route(MESSAGES_ROUTE_PATH, post(handle)) } +#[tracing::instrument( + name = "messages_gateway_route", + target = "litellm::function_trace", + level = "trace", + skip_all +)] async fn handle( _auth: RequireMasterKey, State(state): State, diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs index 4fd29db05d6..5434719987b 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs @@ -12,6 +12,12 @@ pub(crate) enum MessagesResponse { Stream(reqwest::Response), } +#[tracing::instrument( + name = "messages_gateway_service", + target = "litellm::function_trace", + level = "trace", + skip_all +)] pub async fn run( router: &Arc, body: Value, diff --git a/litellm-rust/crates/ai-gateway/src/trace_parity.rs b/litellm-rust/crates/ai-gateway/src/trace_parity.rs new file mode 100644 index 00000000000..614852c541d --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/trace_parity.rs @@ -0,0 +1,65 @@ +//! Harness-only in-process adapters. Never mounted as production routes. + +use std::sync::Arc; + +use axum::body::{Body, to_bytes}; +use axum::http::header::{AUTHORIZATION, CONTENT_TYPE}; +use axum::http::{Request, StatusCode}; +use litellm_core::Error; +use litellm_core::router::{Deployment, LiteLLMParams, Router as ModelRouter}; +use serde::Serialize; +use serde_json::Value; +use tower::ServiceExt; + +use crate::io::realtime_pool::RealtimePool; +use crate::routes; +use crate::state::AppState; + +#[derive(Debug, Serialize)] +pub struct GatewayResponse { + pub status: u16, + pub body: Value, +} + +pub async fn messages_request( + model_alias: String, + provider_model: String, + api_base: String, + body: Value, +) -> Result { + let state = AppState { + router: Arc::new(ModelRouter::new(vec![Deployment { + model_name: model_alias, + litellm_params: LiteLLMParams { + model: provider_model, + api_key: Some("trace-provider-key".to_string()), + api_base: Some(api_base), + }, + }])), + master_key: Some(Arc::from("trace-master-key")), + loggers: Arc::new(Vec::new()), + realtime_pool: RealtimePool::disabled(), + }; + let request = Request::builder() + .method("POST") + .uri("/v1/messages") + .header(AUTHORIZATION, "Bearer trace-master-key") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(body.to_string())) + .map_err(|error| Error::InvalidRequest(error.to_string()))?; + let response = routes::app(state) + .oneshot(request) + .await + .map_err(|error| match error {})?; + let status: StatusCode = response.status(); + let bytes = to_bytes(response.into_body(), usize::MAX) + .await + .map_err(|error| Error::InvalidResponse(error.to_string()))?; + let body = serde_json::from_slice(&bytes).map_err(|error| { + Error::InvalidResponse(format!("gateway returned invalid JSON: {error}")) + })?; + Ok(GatewayResponse { + status: status.as_u16(), + body, + }) +} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs b/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs similarity index 82% rename from litellm-rust/crates/ai-gateway/src/ocr/tests.rs rename to litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs index 85e4c408045..c3a89f4394d 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs +++ b/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs @@ -1,23 +1,19 @@ use std::sync::{Arc, Mutex}; use std::time::Duration; -use litellm_core::error::Error; -use litellm_core::http_utils::has_header; -use litellm_core::ocr::transformation::OcrResponseHandling; -use serde_json::{Map, Value, json}; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; -use tokio::net::{TcpListener, TcpStream}; - -use super::common_utils::{ocr_provider_config, string_headers, truncate_error_body}; -use super::{OcrRequest, ocr}; -use crate::integrations::custom_guardrail::{ +use litellm_ai_gateway::integrations::custom_guardrail::{ CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook, GuardrailFuture, GuardrailRequest, }; -use crate::integrations::custom_logger::{ +use litellm_ai_gateway::integrations::custom_logger::{ CallbackTiming, CallbackValue, CustomLogger, LogFuture, ModelCallDetails, }; -use crate::integrations::types::RequestMetadata; +use litellm_ai_gateway::integrations::types::RequestMetadata; +use litellm_ai_gateway::ocr::{OcrRequest, ocr}; +use litellm_core::error::Error; +use serde_json::{Map, Value, json}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; async fn read_http_headers(socket: &mut TcpStream) -> String { let mut request = Vec::new(); @@ -136,6 +132,7 @@ struct RecordingOcrGuardrail { hooks: Vec, events: Mutex>, block_pre_call: bool, + block_during_call: bool, } impl RecordingOcrGuardrail { @@ -144,6 +141,7 @@ impl RecordingOcrGuardrail { hooks, events: Mutex::new(Vec::new()), block_pre_call: false, + block_during_call: false, } } @@ -152,6 +150,16 @@ impl RecordingOcrGuardrail { hooks: vec![GuardrailEventHook::PreCall], events: Mutex::new(Vec::new()), block_pre_call: true, + block_during_call: false, + } + } + + fn blocking_during_call() -> Self { + Self { + hooks: vec![GuardrailEventHook::DuringCall], + events: Mutex::new(Vec::new()), + block_pre_call: false, + block_during_call: true, } } @@ -193,91 +201,95 @@ impl CustomGuardrail for RecordingOcrGuardrail { ) -> GuardrailFuture<'a> { Box::pin(async move { self.events.lock().unwrap().push("async_moderation_hook"); + if self.block_during_call { + return Ok(GuardrailDecision::Block(GuardrailError::blocked( + "blocked before provider", + ))); + } request.data["body"]["guarded_during"] = json!(true); Ok(GuardrailDecision::Mask(request)) }) } } -#[test] -fn truncate_error_body_passes_short_strings_through() { - let body = "Unauthorized"; - assert_eq!(truncate_error_body(body), "Unauthorized"); +fn base_ocr_request(model: &str) -> OcrRequest<'_> { + OcrRequest { + model, + document: json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }), + api_key: Some("sk-test"), + api_base: None, + custom_llm_provider: None, + extra_headers: None, + optional_params: Map::new(), + timeout: None, + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: RequestMetadata::default(), + litellm_call_id: None, + } } -#[test] -fn truncate_error_body_caps_long_payloads() { - let body = "x".repeat(306); - let truncated = truncate_error_body(&body); +#[tokio::test] +async fn reducto_during_call_guardrail_blocks_before_upload() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let address = listener.local_addr().expect("listener has local address"); + let api_base = format!("http://{address}"); + let guardrail = Arc::new(RecordingOcrGuardrail::blocking_during_call()); + let mut request = base_ocr_request("reducto/parse-v3"); + request.api_base = Some(&api_base); + request.document = json!({ + "type": "document_url", + "document_url": "data:application/pdf;base64,JVBERi0xLjQ=" + }); + request.guardrails = vec![guardrail.clone()]; - assert!(truncated.ends_with("... (truncated)")); - let prefix_chars = truncated - .strip_suffix("... (truncated)") - .expect("truncated marker present") - .chars() - .count(); - assert_eq!(prefix_chars, 256); + let error = ocr(request).await.expect_err("guardrail blocks upload"); + + assert!(matches!(error, Error::InvalidRequest(_))); + assert_eq!(guardrail.events(), vec!["async_moderation_hook"]); + let accepted = tokio::time::timeout(Duration::from_millis(100), listener.accept()).await; + assert!(accepted.is_err(), "upload socket should not be touched"); } -#[test] -fn truncate_error_body_does_not_split_multibyte_chars() { - let body = "é".repeat(266); - let truncated = truncate_error_body(&body); - assert!(truncated.is_char_boundary(truncated.len())); -} +#[tokio::test] +async fn reducto_upload_error_body_is_truncated() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let address = listener.local_addr().expect("listener has local address"); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts upload request"); + let _request = read_http_request(&mut socket).await; + let body = "x".repeat(300); + let response = format!( + "HTTP/1.1 500 Internal Server Error\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + body.len(), + body + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes upload response"); + }); + let api_base = format!("http://{address}"); + let mut request = base_ocr_request("reducto/parse-v3"); + request.api_base = Some(&api_base); + request.document = json!({ + "type": "document_url", + "document_url": "data:application/pdf;base64,JVBERi0xLjQ=" + }); + + let error = ocr(request).await.expect_err("upload should fail"); -#[test] -fn ocr_dispatch_supports_migrated_providers() { - assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some()); assert!( - ocr_provider_config("azure_ai", "pixtral-12b-2409") - .expect("azure ai config resolves") - .requires_data_uri_document() + matches!(error, Error::Http { status: 500, body } if body.chars().count() < 300 && body.ends_with("... (truncated)")) ); - assert_eq!( - ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read") - .expect("document intelligence config resolves") - .response_handling(), - OcrResponseHandling::AzureDocumentIntelligencePoll - ); - assert!( - ocr_provider_config("vertex_ai", "deepseek-ocr-maas") - .expect("vertex deepseek config resolves") - .supported_ocr_params() - .contains(&"temperature") - ); - assert!(ocr_provider_config("openai", "gpt-4o").is_none()); -} - -#[test] -fn string_headers_accepts_string_values() { - let headers = json!({ - "x-trace-id": "trace-1" - }) - .as_object() - .unwrap() - .clone(); - - assert_eq!( - string_headers(Some(headers)).expect("string headers accepted"), - vec![("x-trace-id".to_string(), "trace-1".to_string())] - ); -} - -#[test] -fn auth_header_detection_is_case_insensitive() { - let headers = vec![ - ("x-trace-id".to_string(), "trace-1".to_string()), - ("authorization".to_string(), "Bearer sk-test".to_string()), - ]; - - assert!(has_header(&headers, "authorization")); - - let headers = vec![("Authorization".to_string(), "Bearer sk-test".to_string())]; - assert!(has_header(&headers, "authorization")); - - let headers = vec![("x-trace-id".to_string(), "trace-1".to_string())]; - assert!(!has_header(&headers, "authorization")); + server.await.expect("server task completes"); } #[tokio::test] @@ -595,21 +607,3 @@ async fn document_intelligence_poll_uses_resolved_subscription_key() { "{poll_request}" ); } - -#[test] -fn string_headers_rejects_non_string_values() { - let headers = json!({ - "x-retry-count": 3 - }) - .as_object() - .unwrap() - .clone(); - - let err = string_headers(Some(headers)).expect_err("non-string header rejected"); - assert_eq!( - err, - Error::InvalidRequest( - "OCR extra_headers.x-retry-count must be a string, got number".to_string() - ) - ); -} diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 389dbd49505..c0de7ff3977 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -6,12 +6,14 @@ license.workspace = true repository.workspace = true [dependencies] +base64.workspace = true rand.workspace = true reqwest.workspace = true serde.workspace = true serde_json.workspace = true thiserror.workspace = true tracing.workspace = true +tracing-subscriber = { workspace = true, optional = true } sha2.workspace = true aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true } aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"], optional = true } @@ -30,6 +32,9 @@ bedrock-auth = [ "dep:aws-types", "dep:aws-smithy-runtime-api", ] +observability = ["dep:tracing-subscriber"] [dev-dependencies] +rstest.workspace = true tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } +tracing-subscriber.workspace = true diff --git a/litellm-rust/crates/core/src/constants.rs b/litellm-rust/crates/core/src/constants.rs index a73961060eb..fc81f4fa029 100644 --- a/litellm-rust/crates/core/src/constants.rs +++ b/litellm-rust/crates/core/src/constants.rs @@ -41,3 +41,5 @@ pub const CHAT_COMPLETION_OBJECT: &str = "chat.completion"; /// `litellm/litellm_core_utils/prompt_templates/factory.py`. pub const EMPTY_TEXT_PLACEHOLDER: &str = "[System: Empty message content sanitised to satisfy protocol]"; + +pub const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace"; diff --git a/litellm-rust/crates/core/src/http_utils.rs b/litellm-rust/crates/core/src/http_utils.rs index 3633130528d..cb472dd5a57 100644 --- a/litellm-rust/crates/core/src/http_utils.rs +++ b/litellm-rust/crates/core/src/http_utils.rs @@ -101,6 +101,23 @@ mod tests { assert!(!has_header(&headers, "authorization")); } + #[test] + fn auth_header_detection_is_case_insensitive() { + let headers = vec![ + ("x-trace-id".to_string(), "trace-1".to_string()), + ("authorization".to_string(), "Bearer sk-test".to_string()), + ]; + + assert!(has_header(&headers, "authorization")); + + let headers = vec![("Authorization".to_string(), "Bearer sk-test".to_string())]; + + assert!(has_header(&headers, "authorization")); + + let headers = vec![("x-trace-id".to_string(), "trace-1".to_string())]; + assert!(!has_header(&headers, "authorization")); + } + #[test] fn bearer_detection_requires_a_non_empty_token() { assert!(has_bearer_auth(&[( diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 0e18d24e5d8..b93e084f57e 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -6,6 +6,8 @@ pub mod constants; pub mod error; pub mod http_utils; pub mod messages; +#[cfg(any(feature = "observability", test))] +pub mod observability; pub mod ocr; pub mod providers; pub mod realtime; diff --git a/litellm-rust/crates/core/src/observability/function_trace.rs b/litellm-rust/crates/core/src/observability/function_trace.rs new file mode 100644 index 00000000000..2031e35901c --- /dev/null +++ b/litellm-rust/crates/core/src/observability/function_trace.rs @@ -0,0 +1,215 @@ +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; + +use serde::Serialize; +use tracing::span::{Attributes, Id}; +use tracing::{Dispatch, Subscriber}; +use tracing_subscriber::layer::Context; +use tracing_subscriber::prelude::*; +use tracing_subscriber::registry::LookupSpan; +use tracing_subscriber::{Layer, Registry}; + +use super::function_trace_filter; + +#[derive(Clone, Debug, PartialEq, Serialize)] +pub struct FunctionTraceEvent { + pub id: usize, + pub parent_id: Option, + pub function: &'static str, + pub module_path: Option<&'static str>, + pub file: Option<&'static str>, + pub line: Option, +} + +#[derive(Clone, Default)] +pub struct FunctionTrace { + events: Arc>>, + span_events: Arc>>, +} + +impl FunctionTrace { + pub fn dispatcher(&self) -> Dispatch { + Dispatch::new( + Registry::default().with( + FunctionTraceLayer { + trace: self.clone(), + } + .with_filter(function_trace_filter()), + ), + ) + } + + pub fn events(&self) -> Vec { + self.events + .lock() + .unwrap_or_else(|error| error.into_inner()) + .clone() + } +} + +struct FunctionTraceLayer { + trace: FunctionTrace, +} + +impl Layer for FunctionTraceLayer +where + S: Subscriber + for<'lookup> LookupSpan<'lookup>, +{ + fn on_new_span(&self, attributes: &Attributes<'_>, id: &Id, context: Context<'_, S>) { + let parent_id = context.span(id).and_then(|span| { + let span_events = self + .trace + .span_events + .lock() + .unwrap_or_else(|error| error.into_inner()); + span.scope() + .skip(1) + .find_map(|ancestor| span_events.get(&ancestor.id()).copied()) + }); + let mut events = self + .trace + .events + .lock() + .unwrap_or_else(|error| error.into_inner()); + let event_id = events.len(); + events.push(FunctionTraceEvent { + id: event_id, + parent_id, + function: attributes.metadata().name(), + module_path: attributes.metadata().module_path(), + file: attributes.metadata().file(), + line: attributes.metadata().line(), + }); + self.trace + .span_events + .lock() + .unwrap_or_else(|error| error.into_inner()) + .insert(id.clone(), event_id); + } +} + +#[cfg(test)] +mod tests { + use crate::constants::FUNCTION_TRACE_TARGET; + + use super::*; + + fn event( + id: usize, + parent_id: Option, + function: &'static str, + ) -> (usize, Option, &'static str) { + (id, parent_id, function) + } + + fn structural_events( + events: &[FunctionTraceEvent], + ) -> Vec<(usize, Option, &'static str)> { + events + .iter() + .map(|event| (event.id, event.parent_id, event.function)) + .collect() + } + + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] + async fn outer() { + tokio::task::yield_now().await; + inner().await; + } + + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] + async fn inner() { + tokio::task::yield_now().await; + } + + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] + async fn concurrent_parent() { + tokio::join!(inner(), inner()); + } + + #[tokio::test] + async fn concurrent_futures_keep_separate_traces_across_yields() { + use tracing::instrument::WithSubscriber; + + let first = FunctionTrace::default(); + let second = FunctionTrace::default(); + let outside = FunctionTrace::default(); + + async { + tokio::join!( + outer().with_subscriber(first.dispatcher()), + inner().with_subscriber(second.dispatcher()), + ); + inner().await; + } + .with_subscriber(outside.dispatcher()) + .await; + + assert_eq!( + structural_events(&first.events()), + vec![event(0, None, "outer"), event(1, Some(0), "inner")], + ); + assert_eq!( + structural_events(&second.events()), + vec![event(0, None, "inner")], + ); + assert_eq!( + structural_events(&outside.events()), + vec![event(0, None, "inner")], + ); + } + + #[tokio::test] + async fn concurrent_siblings_keep_the_same_parent() { + use tracing::instrument::WithSubscriber; + + let trace = FunctionTrace::default(); + concurrent_parent() + .with_subscriber(trace.dispatcher()) + .await; + + assert_eq!( + structural_events(&trace.events()), + vec![ + event(0, None, "concurrent_parent"), + event(1, Some(0), "inner"), + event(2, Some(0), "inner"), + ] + ); + } + + #[test] + fn records_matching_spans_in_creation_order() { + let trace = FunctionTrace::default(); + let dispatch = trace.dispatcher(); + + tracing::dispatcher::with_default(&dispatch, || { + let _ignored = tracing::trace_span!(target: "other", "ignored"); + let _first = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name"); + let _wrong_level = tracing::debug_span!(target: FUNCTION_TRACE_TARGET, "wrong_level"); + let _second = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name"); + }); + + assert_eq!( + structural_events(&trace.events()), + vec![event(0, None, "same_name"), event(1, None, "same_name")] + ); + } + + #[test] + fn records_matching_span_nesting_depth() { + let trace = FunctionTrace::default(); + let dispatch = trace.dispatcher(); + + tracing::dispatcher::with_default(&dispatch, || { + let outer = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "outer"); + let _outer_guard = outer.enter(); + let _inner = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "inner"); + }); + + assert_eq!( + structural_events(&trace.events()), + vec![event(0, None, "outer"), event(1, Some(0), "inner")] + ); + } +} diff --git a/litellm-rust/crates/core/src/observability/mod.rs b/litellm-rust/crates/core/src/observability/mod.rs new file mode 100644 index 00000000000..3f9da8e2bb4 --- /dev/null +++ b/litellm-rust/crates/core/src/observability/mod.rs @@ -0,0 +1,59 @@ +use tracing::span::Id; +use tracing::{Level, Metadata, Subscriber}; +use tracing_subscriber::filter::{FilterFn, LevelFilter, filter_fn}; +use tracing_subscriber::layer::Context; +use tracing_subscriber::registry::LookupSpan; + +use crate::constants::FUNCTION_TRACE_TARGET; + +pub mod function_trace; + +pub use function_trace::{FunctionTrace, FunctionTraceEvent}; + +pub fn function_trace_filter() -> FilterFn) -> bool> { + filter_fn(|metadata| { + metadata.is_span() + && metadata.target() == FUNCTION_TRACE_TARGET + && *metadata.level() == Level::TRACE + }) + .with_max_level_hint(LevelFilter::TRACE) +} + +pub fn span_depth(context: &Context<'_, S>, id: &Id) -> usize +where + S: Subscriber + for<'lookup> LookupSpan<'lookup>, +{ + context + .span(id) + .map(|span| span.scope().skip(1).count()) + .unwrap_or_default() +} + +#[cfg(test)] +mod tests { + use tracing::instrument::WithSubscriber; + + use super::*; + + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] + async fn instrumented_with_literal_target() {} + + #[tokio::test] + async fn literal_instrument_target_matches_filter_constant() { + assert_eq!(FUNCTION_TRACE_TARGET, "litellm::function_trace"); + + let trace = FunctionTrace::default(); + instrumented_with_literal_target() + .with_subscriber(trace.dispatcher()) + .await; + + let events = trace.events(); + assert_eq!(events.len(), 1); + assert_eq!(events[0].id, 0); + assert_eq!(events[0].parent_id, None); + assert_eq!(events[0].function, "instrumented_with_literal_target"); + assert_eq!(events[0].module_path, Some(module_path!())); + assert_eq!(events[0].file, Some(file!())); + assert!(events[0].line.is_some()); + } +} diff --git a/litellm-rust/crates/core/src/ocr/transformation.rs b/litellm-rust/crates/core/src/ocr/transformation.rs index ad484c8f968..62299faf9ed 100644 --- a/litellm-rust/crates/core/src/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/transformation.rs @@ -51,6 +51,15 @@ pub trait OcrProviderConfig: Sync { response_json: Value, ) -> Result; + fn transform_ocr_response_with_params( + &self, + model: &str, + response_json: Value, + _optional_params: &Map, + ) -> Result { + self.transform_ocr_response(model, response_json) + } + fn complete_url( &self, api_base: Option<&str>, diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index 1a72b8f1d66..71cdb232a87 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -1,5 +1,5 @@ use serde::{Deserialize, Serialize}; -use serde_json::Value; +use serde_json::{Map, Value}; #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct OcrRequestData { @@ -14,16 +14,25 @@ pub struct OcrResponseData { pub document_annotation: Option, pub usage_info: Option, pub object: String, + pub extra_fields: Map, + pub provider_native_response: Option, } impl OcrResponseData { pub fn into_json(self) -> Value { - serde_json::json!({ + let mut response = serde_json::json!({ "pages": self.pages, "model": self.model, "document_annotation": self.document_annotation, "usage_info": self.usage_info, "object": self.object, - }) + }); + if let Value::Object(object) = &mut response { + object.extend(self.extra_fields); + if let Some(native_response) = self.provider_native_response { + object.insert("provider_native_response".to_string(), native_response); + } + } + response } } diff --git a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs index 641a019476e..d15c032f0bc 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs @@ -14,7 +14,8 @@ const AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV: &str = "AZURE_DOCUMENT_INTELLIGE const AZURE_DOCUMENT_INTELLIGENCE_API_VERSION: &str = "2024-11-30"; const AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI: i64 = 96; -const AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS: &[&str] = &["pages", "features"]; +const AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS: &[&str] = + &["pages", "features", "req_format"]; pub struct AzureAiOcrConfig; pub struct AzureDocumentIntelligenceOcrConfig; @@ -98,9 +99,76 @@ pub fn resolve_document_intelligence_endpoint( ) } -fn encode_model_id(model: &str) -> String { +fn prepend_auth_header( + headers: Vec<(String, String)>, + name: &str, + value: String, +) -> Vec<(String, String)> { + std::iter::once((name.to_string(), value)) + .chain(headers) + .collect() +} + +pub fn validate_azure_ai_environment( + headers: Vec<(String, String)>, + api_key: Option<&str>, + azure_ad_token: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> Result, Error> { + if crate::http_utils::has_header(&headers, "Authorization") + || crate::http_utils::has_header(&headers, "Api-Key") + { + return Ok(headers); + } + if let Ok(api_key) = resolve_azure_ai_api_key(api_key, env_lookup) { + return Ok(prepend_auth_header(headers, "Api-Key", api_key)); + } + non_empty(azure_ad_token) + .map(|token| prepend_auth_header(headers, "Authorization", format!("Bearer {token}"))) + .ok_or_else(|| { + Error::Auth( + "Missing Azure AI credentials - set AZURE_AI_API_KEY or provide azure_ad_token" + .to_string(), + ) + }) +} + +pub fn validate_document_intelligence_environment( + headers: Vec<(String, String)>, + api_key: Option<&str>, + azure_ad_token: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> Result, Error> { + if crate::http_utils::has_header(&headers, "Authorization") + || crate::http_utils::has_header(&headers, "Ocp-Apim-Subscription-Key") + { + return Ok(headers); + } + if let Ok(api_key) = resolve_document_intelligence_api_key(api_key, env_lookup) { + return Ok(prepend_auth_header( + headers, + "Ocp-Apim-Subscription-Key", + api_key, + )); + } + non_empty(azure_ad_token) + .map(|token| prepend_auth_header(headers, "Authorization", format!("Bearer {token}"))) + .ok_or_else(|| { + Error::Auth( + "Missing Azure Document Intelligence credentials - set AZURE_DOCUMENT_INTELLIGENCE_API_KEY or provide azure_ad_token" + .to_string(), + ) + }) +} + +fn encode_model_id(model: &str) -> Result { let model_id = model.rsplit('/').next().unwrap_or(model); - model_id + if matches!(model_id, "." | "..") { + return Err(Error::InvalidRequest( + "model_id cannot be a dot path segment".to_string(), + )); + } + Ok(model_id .bytes() .flat_map(|byte| match byte { b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => { @@ -108,7 +176,7 @@ fn encode_model_id(model: &str) -> String { } _ => format!("%{byte:02X}").chars().collect(), }) - .collect() + .collect()) } fn pages_token_is_valid(token: &str) -> bool { @@ -147,6 +215,11 @@ fn normalize_pages_param(pages: &Value) -> Result, Error> { if values.is_empty() { return Ok(None); } + if values.iter().any(Value::is_boolean) { + return Err(Error::InvalidRequest( + "`pages` must be integers, not booleans".to_string(), + )); + } if values.iter().all(Value::is_i64) { let mut pages = BTreeSet::new(); for value in values { @@ -232,6 +305,38 @@ fn normalize_features_param(features: &Value) -> Result, Error> { } } +fn normalize_req_format(req_format: &Value) -> Result { + match req_format.as_str() { + Some(value @ ("native" | "litellm")) => Ok(value.to_string()), + _ => Err(Error::InvalidRequest(format!( + "Invalid `req_format` for Azure Document Intelligence: {req_format:?}. Expected 'native' or 'litellm'." + ))), + } +} + +pub fn map_document_intelligence_ocr_params( + non_default_params: &Map, +) -> Result, Error> { + let mut mapped = Map::new(); + if let Some(pages) = non_default_params.get("pages") + && let Some(normalized) = normalize_pages_param(pages)? + { + mapped.insert("pages".to_string(), Value::String(normalized)); + } + if let Some(features) = non_default_params.get("features") + && let Some(normalized) = normalize_features_param(features)? + { + mapped.insert("features".to_string(), Value::String(normalized)); + } + if let Some(req_format) = non_default_params.get("req_format") { + mapped.insert( + "req_format".to_string(), + Value::String(normalize_req_format(req_format)?), + ); + } + Ok(mapped) +} + pub fn complete_document_intelligence_url( api_base: Option<&str>, model: &str, @@ -242,7 +347,7 @@ pub fn complete_document_intelligence_url( let mut url = format!( "{}/documentintelligence/documentModels/{}:analyze?api-version={}", endpoint.trim_end_matches('/'), - encode_model_id(model), + encode_model_id(model)?, AZURE_DOCUMENT_INTELLIGENCE_API_VERSION ); @@ -260,6 +365,10 @@ pub fn complete_document_intelligence_url( url.push_str(&normalized); } + if let Some(req_format) = optional_params.get("req_format") { + normalize_req_format(req_format)?; + } + Ok(url) } @@ -327,11 +436,78 @@ fn page_dimensions(page: &Map) -> Value { }) } +fn transform_document_intelligence_response( + model: &str, + response_json: Value, + preserve_native_response: bool, +) -> Result { + let response = response_json + .as_object() + .ok_or_else(|| Error::InvalidType { + expected: "object", + actual: json_type_name(&response_json), + })?; + let status = response + .get("status") + .and_then(Value::as_str) + .ok_or(Error::MissingField("status"))?; + if status != "succeeded" { + return Err(Error::InvalidResponse(format!( + "Azure Document Intelligence analysis failed with status: {status}" + ))); + } + + let analyze_result = response.get("analyzeResult").and_then(Value::as_object); + let azure_pages = analyze_result + .and_then(|result| result.get("pages")) + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(); + let pages = azure_pages + .iter() + .filter_map(Value::as_object) + .map(|page| { + let page_number = page.get("pageNumber").and_then(Value::as_i64).unwrap_or(1); + json!({ + "index": page_number - 1, + "markdown": page_markdown(page), + "dimensions": page_dimensions(page), + }) + }) + .collect::>(); + let extra_fields = ["content", "tables", "keyValuePairs"] + .into_iter() + .map(|field| { + ( + field.to_string(), + analyze_result + .and_then(|result| result.get(field)) + .cloned() + .unwrap_or(Value::Null), + ) + }) + .collect(); + + Ok(OcrResponseData { + usage_info: Some(json!({ + "pages_processed": pages.len(), + "doc_size_bytes": null, + })), + pages, + model: model.to_string(), + document_annotation: None, + object: "ocr".to_string(), + extra_fields, + provider_native_response: preserve_native_response.then_some(response_json), + }) +} + impl OcrProviderConfig for AzureAiOcrConfig { fn supported_ocr_params(&self) -> &'static [&'static str] { MISTRAL_OCR_CONFIG.supported_ocr_params() } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn transform_ocr_request( &self, model: &str, @@ -349,6 +525,7 @@ impl OcrProviderConfig for AzureAiOcrConfig { MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn complete_url( &self, api_base: Option<&str>, @@ -373,10 +550,25 @@ impl OcrProviderConfig for AzureAiOcrConfig { } impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn supported_ocr_params(&self) -> &'static [&'static str] { AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] + fn map_ocr_params(&self, non_default_params: &Map) -> Map { + map_document_intelligence_ocr_params(non_default_params).unwrap_or_else(|_| { + non_default_params + .iter() + .filter(|(name, _)| { + AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS.contains(&name.as_str()) + }) + .map(|(name, value)| (name.clone(), value.clone())) + .collect() + }) + } + + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn transform_ocr_request( &self, _model: &str, @@ -402,59 +594,29 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { }) } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn transform_ocr_response( &self, model: &str, response_json: Value, ) -> Result { - let response = response_json - .as_object() - .ok_or_else(|| Error::InvalidType { - expected: "object", - actual: json_type_name(&response_json), - })?; - let status = response - .get("status") - .and_then(Value::as_str) - .ok_or(Error::MissingField("status"))?; - if status != "succeeded" { - return Err(Error::InvalidResponse(format!( - "Azure Document Intelligence analysis failed with status: {status}" - ))); - } - - let azure_pages = response - .get("analyzeResult") - .and_then(|result| result.get("pages")) - .and_then(Value::as_array) - .cloned() - .unwrap_or_default(); - - let pages = azure_pages - .iter() - .filter_map(Value::as_object) - .map(|page| { - let page_number = page.get("pageNumber").and_then(Value::as_i64).unwrap_or(1); - json!({ - "index": page_number - 1, - "markdown": page_markdown(page), - "dimensions": page_dimensions(page), - }) - }) - .collect::>(); - - Ok(OcrResponseData { - usage_info: Some(json!({ - "pages_processed": pages.len(), - "doc_size_bytes": null, - })), - pages, - model: model.to_string(), - document_annotation: None, - object: "ocr".to_string(), - }) + transform_document_intelligence_response(model, response_json, false) } + fn transform_ocr_response_with_params( + &self, + model: &str, + response_json: Value, + optional_params: &Map, + ) -> Result { + transform_document_intelligence_response( + model, + response_json, + optional_params.get("req_format").and_then(Value::as_str) == Some("native"), + ) + } + + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn complete_url( &self, api_base: Option<&str>, @@ -485,6 +647,101 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { #[cfg(test)] mod tests { use super::*; + use rstest::{fixture, rstest}; + + const ENDPOINT: &str = "https://example.cognitiveservices.azure.com"; + + #[fixture] + fn document_intelligence_config() -> AzureDocumentIntelligenceOcrConfig { + AzureDocumentIntelligenceOcrConfig + } + + fn header_value<'a>(headers: &'a [(String, String)], name: &str) -> Option<&'a str> { + headers + .iter() + .find(|(header_name, _)| header_name.eq_ignore_ascii_case(name)) + .map(|(_, value)| value.as_str()) + } + + #[fixture] + fn native_operation() -> Value { + json!({ + "status": "succeeded", + "createdDateTime": "2026-07-02T00:00:00Z", + "lastUpdatedDateTime": "2026-07-02T00:00:05Z", + "analyzeResult": { + "content": "Invoice\nInvoice No: INV-12345\nTotal: $100.00", + "pages": [{ + "pageNumber": 1, + "width": 8.5, + "height": 11, + "unit": "inch", + "angle": 0.13, + "lines": [ + {"content": "Invoice"}, + {"content": "Invoice No: INV-12345"}, + {"content": "Total: $100.00"} + ], + "words": [{"content": "Invoice", "confidence": 0.994}] + }], + "tables": [ + { + "rowCount": 2, + "columnCount": 2, + "cells": [ + {"kind": "columnHeader", "rowIndex": 0, "columnIndex": 0, "content": "Item"}, + {"kind": "columnHeader", "rowIndex": 0, "columnIndex": 1, "content": "Price"}, + {"rowIndex": 1, "columnIndex": 0, "content": "Widget"}, + {"rowIndex": 1, "columnIndex": 1, "content": "$100.00"} + ] + }, + { + "rowCount": 1, + "columnCount": 1, + "cells": [{"rowIndex": 0, "columnIndex": 0, "content": "Totals"}] + } + ], + "keyValuePairs": [ + { + "key": {"content": "Invoice No"}, + "value": {"content": "INV-12345"}, + "confidence": 0.98 + }, + { + "key": {"content": "Total"}, + "value": {"content": "$100.00"}, + "confidence": 0.95 + } + ], + "paragraphs": [{"content": "Invoice"}] + } + }) + } + + fn assert_native_fields_preserved(response: &OcrResponseData, operation: &Value) { + let analyze_result = &operation["analyzeResult"]; + + assert_eq!(response.extra_fields["content"], analyze_result["content"]); + assert_eq!(response.extra_fields["tables"], analyze_result["tables"]); + assert_eq!( + response.extra_fields["keyValuePairs"], + analyze_result["keyValuePairs"] + ); + assert_eq!(response.object, "ocr"); + assert_eq!( + response.usage_info, + Some(json!({"pages_processed": 1, "doc_size_bytes": null})) + ); + assert_eq!(response.pages[0]["index"], 0); + assert_eq!( + response.pages[0]["markdown"], + "Invoice\nInvoice No: INV-12345\nTotal: $100.00" + ); + assert_eq!( + response.pages[0]["dimensions"], + json!({"width": 816, "height": 1056, "dpi": 96}) + ); + } #[test] fn azure_ai_reuses_mistral_body_transform() { @@ -568,6 +825,11 @@ mod tests { #[test] fn document_intelligence_url_omits_empty_feature_list() { let params = serde_json::Map::from_iter([("features".to_string(), json!([]))]); + assert!( + map_document_intelligence_ocr_params(¶ms) + .expect("empty features map") + .is_empty() + ); let url = complete_document_intelligence_url( Some("https://example.cognitiveservices.azure.com"), "prebuilt-layout", @@ -582,40 +844,43 @@ mod tests { ); } - #[test] - fn document_intelligence_url_rejects_invalid_features() { - for features in [ - json!("keyValuePairs&pages=9"), - json!(""), - json!(["keyValuePairs", 1]), - json!({"feature": "keyValuePairs"}), - ] { - let params = serde_json::Map::from_iter([("features".to_string(), features.clone())]); - let error = complete_document_intelligence_url( - Some("https://example.cognitiveservices.azure.com"), - "prebuilt-layout", - ¶ms, - &|_| None, - ) - .expect_err("invalid features must fail"); + #[rstest] + #[case::query_injection(json!("keyValuePairs&pages=9"))] + #[case::spaces(json!("key value pairs"))] + #[case::empty_string(json!(""))] + #[case::integer_list(json!([1, 2]))] + #[case::nested_list(json!([["keyValuePairs"]]))] + #[case::object(json!({"feature": "keyValuePairs"}))] + #[case::number(json!(5))] + fn document_intelligence_mapping_rejects_invalid_features(#[case] features: Value) { + let params = serde_json::Map::from_iter([("features".to_string(), features)]); + let error = + map_document_intelligence_ocr_params(¶ms).expect_err("invalid features must fail"); - assert!( - matches!(error, Error::InvalidRequest(message) if message.contains("Invalid `features`")), - "features={features:?}" - ); - } + assert!(matches!( + error, + Error::InvalidRequest(message) if message.contains("Invalid `features`") + )); } - #[test] - fn document_intelligence_maps_features() { + #[rstest] + #[case::single_list(json!(["keyValuePairs"]), "keyValuePairs")] + #[case::multiple_list( + json!(["keyValuePairs", "languages"]), + "keyValuePairs,languages" + )] + #[case::single_string(json!("keyValuePairs"), "keyValuePairs")] + #[case::comma_separated(json!("keyValuePairs,languages"), "keyValuePairs,languages")] + #[case::spaces(json!("keyValuePairs, languages"), "keyValuePairs,languages")] + fn document_intelligence_maps_features(#[case] features: Value, #[case] expected: &str) { let params = Map::from_iter([ - ("features".to_string(), json!(["keyValuePairs"])), + ("features".to_string(), features), ("unsupported".to_string(), json!(true)), ]); assert_eq!( AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG.map_ocr_params(¶ms), - Map::from_iter([("features".to_string(), json!(["keyValuePairs"]))]) + Map::from_iter([("features".to_string(), json!(expected))]) ); } @@ -633,32 +898,484 @@ mod tests { assert_eq!(body, json!({"base64Source": "abc123"})); } + #[rstest] + fn document_intelligence_response_normalizes_pages(native_operation: Value) { + let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG + .transform_ocr_response("prebuilt-layout", native_operation.clone()) + .expect("response transforms"); + + assert_native_fields_preserved(&response, &native_operation); + } + #[test] - fn document_intelligence_response_normalizes_pages() { + fn azure_document_intelligence_model_id_is_encoded() { + let url = complete_document_intelligence_url( + Some(ENDPOINT), + "prebuilt-layout?x=1#frag", + &Map::new(), + &|_| None, + ) + .expect("url builds"); + + assert_eq!( + url, + "https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout%3Fx%3D1%23frag:analyze?api-version=2024-11-30" + ); + } + + #[test] + fn azure_document_intelligence_dot_segment_model_id_is_rejected() { + let error = complete_document_intelligence_url( + Some(ENDPOINT), + "azure_ai/doc-intelligence/..", + &Map::new(), + &|_| None, + ) + .expect_err("dot segment must fail"); + + assert_eq!( + error, + Error::InvalidRequest("model_id cannot be a dot path segment".to_string()) + ); + } + + #[rstest] + fn document_intelligence_async_response_preserves_normalized_fields(native_operation: Value) { let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG .transform_ocr_response( - "prebuilt-layout", + "azure_ai/doc-intelligence/prebuilt-layout", + native_operation.clone(), + ) + .expect("response transforms"); + + assert_native_fields_preserved(&response, &native_operation); + } + + #[test] + fn document_intelligence_response_tolerates_missing_native_fields() { + let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG + .transform_ocr_response( + "azure_ai/doc-intelligence/prebuilt-read", json!({ "status": "succeeded", "analyzeResult": { "pages": [{ - "pageNumber": 2, + "pageNumber": 1, "width": 8.5, "height": 11, "unit": "inch", - "lines": [{"content": "hello"}, {"content": "world"}] + "lines": [{"content": "hello"}] }] } }), ) + .expect("missing optional fields are allowed"); + + assert_eq!(response.pages[0]["markdown"], "hello"); + assert_eq!(response.extra_fields["content"], Value::Null); + assert_eq!(response.extra_fields["tables"], Value::Null); + assert_eq!(response.extra_fields["keyValuePairs"], Value::Null); + } + + #[test] + fn document_intelligence_non_succeeded_status_is_rejected() { + let error = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG + .transform_ocr_response( + "azure_ai/doc-intelligence/prebuilt-layout", + json!({"status": "failed"}), + ) + .expect_err("failed status must fail"); + + assert_eq!( + error, + Error::InvalidResponse( + "Azure Document Intelligence analysis failed with status: failed".to_string() + ) + ); + } + + #[test] + fn document_intelligence_supported_params_include_features() { + assert_eq!( + AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG.supported_ocr_params(), + &["pages", "features", "req_format"] + ); + } + + #[rstest] + fn document_intelligence_native_format_carries_raw_operation(native_operation: Value) { + let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG + .transform_ocr_response_with_params( + "azure_ai/doc-intelligence/prebuilt-layout", + native_operation.clone(), + &Map::from_iter([("req_format".to_string(), json!("native"))]), + ) + .expect("native response transforms"); + + assert_eq!( + response.provider_native_response, + Some(native_operation.clone()) + ); + assert_native_fields_preserved(&response, &native_operation); + } + + #[rstest] + fn document_intelligence_async_native_format_carries_raw_operation(native_operation: Value) { + let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG + .transform_ocr_response_with_params( + "azure_ai/doc-intelligence/prebuilt-layout", + native_operation.clone(), + &Map::from_iter([("req_format".to_string(), json!("native"))]), + ) + .expect("native response transforms"); + + assert_eq!( + response.provider_native_response, + Some(native_operation.clone()) + ); + assert_native_fields_preserved(&response, &native_operation); + } + + #[rstest] + #[case::default(Map::new())] + #[case::litellm(Map::from_iter([("req_format".to_string(), json!("litellm"))]))] + fn document_intelligence_default_format_omits_raw_operation( + #[case] optional_params: Map, + native_operation: Value, + ) { + let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG + .transform_ocr_response_with_params( + "azure_ai/doc-intelligence/prebuilt-layout", + native_operation.clone(), + &optional_params, + ) .expect("response transforms"); - assert_eq!(response.pages[0]["index"], 1); - assert_eq!(response.pages[0]["markdown"], "hello\nworld"); - assert_eq!(response.pages[0]["dimensions"]["width"], 816); + assert_eq!(response.provider_native_response, None); + assert_native_fields_preserved(&response, &native_operation); + } + + #[rstest] + #[case::native("native")] + #[case::litellm("litellm")] + fn document_intelligence_maps_req_format(#[case] req_format: &str) { + let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( + "req_format".to_string(), + json!(req_format), + )])) + .expect("req_format maps"); + assert_eq!( - response.usage_info, - Some(json!({"pages_processed": 1, "doc_size_bytes": null})) + mapped, + Map::from_iter([("req_format".to_string(), json!(req_format))]) + ); + } + + #[test] + fn document_intelligence_rejects_unknown_req_format() { + let error = map_document_intelligence_ocr_params(&Map::from_iter([( + "req_format".to_string(), + json!("azure"), + )])) + .expect_err("unknown req_format must fail"); + + assert!( + matches!(error, Error::InvalidRequest(message) if message.contains("Invalid `req_format`")) + ); + } + + #[test] + fn document_intelligence_url_omits_req_format() { + let url = complete_document_intelligence_url( + Some(ENDPOINT), + "prebuilt-layout", + &Map::from_iter([("req_format".to_string(), json!("native"))]), + &|_| None, + ) + .expect("url builds"); + + assert!(!url.contains("req_format")); + } + + #[test] + fn document_intelligence_validate_environment_uses_subscription_key() { + let headers = + validate_document_intelligence_environment(Vec::new(), Some("my-key"), None, &|_| None) + .expect("api key authenticates"); + + assert_eq!( + header_value(&headers, "Ocp-Apim-Subscription-Key"), + Some("my-key") + ); + } + + #[test] + fn document_intelligence_validate_environment_falls_back_to_entra_token() { + let headers = validate_document_intelligence_environment( + Vec::new(), + None, + Some("entra-token"), + &|_| None, + ) + .expect("Entra token authenticates"); + + assert_eq!( + header_value(&headers, "Authorization"), + Some("Bearer entra-token") + ); + assert_eq!(header_value(&headers, "Ocp-Apim-Subscription-Key"), None); + } + + #[test] + fn document_intelligence_supported_params_include_pages_features_and_req_format() { + assert_eq!( + AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG.supported_ocr_params(), + &["pages", "features", "req_format"] + ); + } + + #[test] + fn document_intelligence_maps_zero_based_page_list() { + let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( + "pages".to_string(), + json!([0, 1, 2]), + )])) + .expect("pages map"); + + assert_eq!( + mapped, + Map::from_iter([("pages".to_string(), json!("1,2,3"))]) + ); + } + + #[test] + fn document_intelligence_page_mapping_dedupes_and_sorts() { + let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( + "pages".to_string(), + json!([2, 0, 0, 1]), + )])) + .expect("pages map"); + + assert_eq!(mapped["pages"], "1,2,3"); + } + + #[test] + fn document_intelligence_page_mapping_omits_empty_list() { + let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( + "pages".to_string(), + json!([]), + )])) + .expect("empty pages map"); + + assert!(mapped.is_empty()); + } + + #[test] + fn document_intelligence_page_mapping_accepts_native_range() { + let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( + "pages".to_string(), + json!("3-9"), + )])) + .expect("range maps"); + + assert_eq!(mapped["pages"], "3-9"); + } + + #[test] + fn document_intelligence_page_mapping_strips_spaces() { + let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( + "pages".to_string(), + json!("1-3, 5"), + )])) + .expect("range maps"); + + assert_eq!(mapped["pages"], "1-3,5"); + } + + #[test] + fn document_intelligence_page_mapping_accepts_string_tokens() { + let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( + "pages".to_string(), + json!(["1", "3-5"]), + )])) + .expect("tokens map"); + + assert_eq!(mapped["pages"], "1,3-5"); + } + + #[test] + fn document_intelligence_page_mapping_rejects_invalid_string() { + let error = map_document_intelligence_ocr_params(&Map::from_iter([( + "pages".to_string(), + json!("a,b"), + )])) + .expect_err("invalid pages must fail"); + + assert!( + matches!(error, Error::InvalidRequest(message) if message.contains("Invalid `pages` string")) + ); + } + + #[test] + fn document_intelligence_page_mapping_rejects_negative_index() { + let error = map_document_intelligence_ocr_params(&Map::from_iter([( + "pages".to_string(), + json!([-1]), + )])) + .expect_err("negative pages must fail"); + + assert!( + matches!(error, Error::InvalidRequest(message) if message.contains("must be >= 0")) + ); + } + + #[test] + fn document_intelligence_page_mapping_rejects_bool_list() { + let error = map_document_intelligence_ocr_params(&Map::from_iter([( + "pages".to_string(), + json!([true, false]), + )])) + .expect_err("boolean pages must fail"); + + assert!( + matches!(error, Error::InvalidRequest(message) if message.contains("integers, not booleans")) + ); + } + + #[test] + fn document_intelligence_page_mapping_rejects_unsupported_type() { + let error = map_document_intelligence_ocr_params(&Map::from_iter([( + "pages".to_string(), + json!(5), + )])) + .expect_err("unsupported pages must fail"); + + assert!( + matches!(error, Error::InvalidRequest(message) if message.contains("Mistral-style")) + ); + } + + #[test] + fn document_intelligence_url_appends_pages_query() { + let url = complete_document_intelligence_url( + Some("https://example.cognitiveservices.azure.com/"), + "azure_ai/doc-intelligence/prebuilt-layout", + &Map::from_iter([("pages".to_string(), json!("1-3,5"))]), + &|_| None, + ) + .expect("url builds"); + + assert!(url.contains("api-version=2024-11-30")); + assert!(url.contains("pages=1-3,5")); + assert!(url.contains("/documentintelligence/documentModels/prebuilt-layout:analyze")); + } + + #[test] + fn document_intelligence_url_has_no_pages_when_params_are_empty() { + let url = complete_document_intelligence_url( + Some(ENDPOINT), + "prebuilt-layout", + &Map::new(), + &|_| None, + ) + .expect("url builds"); + + assert!(!url.contains("pages=")); + } + + #[rstest] + fn document_intelligence_request_keeps_pages_out_of_body( + document_intelligence_config: AzureDocumentIntelligenceOcrConfig, + ) { + let request = document_intelligence_config + .transform_ocr_request( + "prebuilt-layout", + json!({"type": "document_url", "document_url": "https://example.com/x.pdf"}), + Map::from_iter([("pages".to_string(), json!("1,2,3"))]), + ) + .expect("request transforms"); + + assert_eq!( + request.data, + json!({"urlSource": "https://example.com/x.pdf"}) + ); + } + + #[test] + fn document_intelligence_mistral_pages_flow_to_query_only() { + let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( + "pages".to_string(), + json!([2, 3, 4, 5, 6, 7, 8]), + )])) + .expect("pages map"); + let url = + complete_document_intelligence_url(Some(ENDPOINT), "prebuilt-layout", &mapped, &|_| { + None + }) + .expect("url builds"); + let request = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG + .transform_ocr_request( + "prebuilt-layout", + json!({"type": "document_url", "document_url": "https://example.com/x.pdf"}), + mapped, + ) + .expect("request transforms"); + + assert!(url.contains("pages=3,4,5,6,7,8,9")); + assert_eq!( + request.data, + json!({"urlSource": "https://example.com/x.pdf"}) + ); + } + + #[test] + fn document_intelligence_endpoint_ignores_generic_azure_ai_base() { + let resolved = resolve_document_intelligence_endpoint(None, &|name| match name { + AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV => Some(ENDPOINT.to_string()), + AZURE_AI_API_BASE_ENV => Some("https://generic.example.com".to_string()), + _ => None, + }) + .expect("endpoint resolves"); + + assert_eq!(resolved, ENDPOINT); + } + + #[test] + fn document_intelligence_endpoint_honors_explicit_api_base() { + let resolved = resolve_document_intelligence_endpoint( + Some("https://my-di.cognitiveservices.azure.com"), + &|name| match name { + AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV => Some(ENDPOINT.to_string()), + AZURE_AI_API_BASE_ENV => Some("https://generic.example.com".to_string()), + _ => None, + }, + ) + .expect("endpoint resolves"); + + assert_eq!(resolved, "https://my-di.cognitiveservices.azure.com"); + } + + #[test] + fn azure_ai_mistral_ocr_uses_generic_api_base() { + let resolved = resolve_azure_ai_api_base(None, &|name| match name { + AZURE_AI_API_BASE_ENV => Some("https://generic-azure-ai.example.com".to_string()), + AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV => Some(ENDPOINT.to_string()), + _ => None, + }) + .expect("api base resolves"); + + assert_eq!(resolved, "https://generic-azure-ai.example.com"); + } + + #[test] + fn azure_ai_ocr_authenticates_with_entra_token() { + let headers = + validate_azure_ai_environment(Vec::new(), None, Some("entra-token"), &|_| None) + .expect("Entra token authenticates"); + + assert_eq!( + header_value(&headers, "Authorization"), + Some("Bearer entra-token") ); } } diff --git a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs index 0125886aac1..11e8fe7db18 100644 --- a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs @@ -134,6 +134,8 @@ impl OcrProviderConfig for MistralOcrConfig { document_annotation, usage_info, object: "ocr".to_string(), + extra_fields: Map::new(), + provider_native_response: None, }) } diff --git a/litellm-rust/crates/core/src/providers/mod.rs b/litellm-rust/crates/core/src/providers/mod.rs index 805600d6dbe..c0c2c69831b 100644 --- a/litellm-rust/crates/core/src/providers/mod.rs +++ b/litellm-rust/crates/core/src/providers/mod.rs @@ -4,4 +4,5 @@ pub mod azure_ai; pub mod bedrock; pub mod mistral; pub mod openai; +pub mod reducto; pub mod vertex_ai; diff --git a/litellm-rust/crates/core/src/providers/reducto/mod.rs b/litellm-rust/crates/core/src/providers/reducto/mod.rs new file mode 100644 index 00000000000..3621ff6a2fd --- /dev/null +++ b/litellm-rust/crates/core/src/providers/reducto/mod.rs @@ -0,0 +1 @@ +pub mod ocr; diff --git a/litellm-rust/crates/core/src/providers/reducto/ocr/mod.rs b/litellm-rust/crates/core/src/providers/reducto/ocr/mod.rs new file mode 100644 index 00000000000..8acee8f770c --- /dev/null +++ b/litellm-rust/crates/core/src/providers/reducto/ocr/mod.rs @@ -0,0 +1,4 @@ +pub mod transformation; + +#[cfg(test)] +mod tests; diff --git a/litellm-rust/crates/core/src/providers/reducto/ocr/tests.rs b/litellm-rust/crates/core/src/providers/reducto/ocr/tests.rs new file mode 100644 index 00000000000..2b66d058b5d --- /dev/null +++ b/litellm-rust/crates/core/src/providers/reducto/ocr/tests.rs @@ -0,0 +1,202 @@ +use rstest::{fixture, rstest}; +use serde_json::{Value, json}; + +use super::transformation::*; +use crate::ocr::transformation::OcrProviderConfig; + +#[fixture] +fn parse_response() -> Value { + json!({ + "job_id": "job_123", + "usage": {"num_pages": 3, "credits": 3}, + "result": { + "chunks": [ + { + "content": "Page 1 block A", + "blocks": [{ + "content": "Page 1 block A", + "bbox": {"page": 1}, + "kind": "text", + }], + }, + { + "content": "Page 2 block A", + "blocks": [{ + "content": "Page 2 block A", + "bbox": {"page": 2}, + "kind": "table", + }], + }, + { + "content": "Page 1 block B", + "blocks": [{ + "content": "Page 1 block B", + "bbox": {"page": 1}, + "kind": "text", + }], + }, + { + "content": "Page 3 block A", + "blocks": [{ + "content": "Page 3 block A", + "bbox": {"page": 3}, + "kind": "figure", + }], + }, + ], + }, + }) +} + +#[rstest] +fn test_parse_v3_file_upload_and_response_mapping(parse_response: Value) { + let source = classify_document_source("data:application/pdf;base64,JVBERi0xLjQ=") + .expect("PDF data URI should be valid"); + let upload = build_upload_request( + source, + "Bearer test-key", + Some("https://platform.reducto.ai"), + ) + .expect("data URI should require upload"); + assert_eq!(upload.url, "https://platform.reducto.ai/upload"); + assert_eq!(upload.authorization, "Bearer test-key"); + assert_eq!(upload.file_name, "document"); + assert_eq!(upload.mime_type, "application/pdf"); + assert_eq!(upload.bytes, b"%PDF-1.4"); + + let optional_params = json!({ + "formatting": {"table_output_format": "html"}, + "retrieval": {"chunk_mode": "section"}, + "settings": {"ocr_system": "standard"}, + }) + .as_object() + .expect("params should be an object") + .clone(); + let request = build_parse_v3_request("reducto://uploaded.pdf", optional_params); + assert_eq!( + request.data, + json!({ + "input": "reducto://uploaded.pdf", + "formatting": {"table_output_format": "html"}, + "retrieval": {"chunk_mode": "section"}, + "settings": {"ocr_system": "standard"}, + }) + ); + + let transformed = transform_reducto_response("parse-v3", parse_response.clone()) + .expect("response should transform"); + assert_eq!( + transformed.usage_info, + Some(json!({"pages_processed": 3, "credits": 3})) + ); + assert_eq!(transformed.pages.len(), 3); + assert_eq!( + transformed.pages[0], + json!({ + "index": 0, + "markdown": "Page 1 block A\n\nPage 1 block B", + "blocks": [ + {"content": "Page 1 block A", "bbox": {"page": 1}, "kind": "text"}, + {"content": "Page 1 block B", "bbox": {"page": 1}, "kind": "text"}, + ], + }) + ); + assert_eq!(transformed.pages[1]["markdown"], "Page 2 block A"); + assert_eq!(transformed.pages[2]["markdown"], "Page 3 block A"); + assert_eq!(transformed.provider_native_response, Some(parse_response)); +} + +#[rstest] +fn test_parse_v3_reducto_id_passthrough_skips_upload(parse_response: Value) { + let document = json!({ + "type": "document_url", + "document_url": "reducto://already-uploaded.pdf", + }); + let source = extract_document_source(&document).expect("Reducto ID should be valid"); + assert!(build_upload_request(source.clone(), "Bearer test-key", None).is_none()); + assert_eq!( + source, + ReductoDocumentSource::FileId("reducto://already-uploaded.pdf".to_string()) + ); + + let request = REDUCTO_PARSE_V3_CONFIG + .transform_ocr_request( + "parse-v3", + document, + json!({"retrieval": {"chunk_mode": "section"}}) + .as_object() + .expect("params should be object") + .clone(), + ) + .expect("direct ID should transform"); + assert_eq!(request.data["input"], "reducto://already-uploaded.pdf"); + assert_eq!(request.data["retrieval"]["chunk_mode"], "section"); + + let response = REDUCTO_PARSE_V3_CONFIG + .transform_ocr_response("parse-v3", parse_response) + .expect("response should transform"); + assert!( + response.pages[0]["markdown"] + .as_str() + .expect("markdown should be string") + .starts_with("Page 1 block A") + ); +} + +#[rstest] +fn test_parse_legacy_wraps_enhance_under_options() { + let request = build_parse_legacy_request( + "reducto://legacy.pdf", + json!({"enhance": {"agentic": [{"type": "table"}]}}) + .as_object() + .expect("params should be object"), + ); + assert_eq!( + request.data, + json!({ + "document_url": "reducto://legacy.pdf", + "options": {"enhance": {"agentic": [{"type": "table"}]}}, + }) + ); +} + +#[rstest] +fn test_parse_v3_image_data_uri_upload_uses_image_mime() { + let source = classify_document_source("data:image/png;base64,iVBORw0KGgo=") + .expect("PNG data URI should be valid"); + let upload = build_upload_request( + source, + "Bearer programmatic-key", + Some("https://custom.reducto.test/"), + ) + .expect("data URI should require upload"); + assert_eq!(upload.url, "https://custom.reducto.test/upload"); + assert_eq!(upload.authorization, "Bearer programmatic-key"); + assert_eq!(upload.mime_type, "image/png"); + assert_eq!(upload.bytes, b"\x89PNG\r\n\x1a\n"); +} + +#[rstest] +#[case::http("http://example.com/document.pdf")] +#[case::https("https://example.com/document.pdf")] +fn test_parse_v3_rejects_plain_http_urls(#[case] source: &str) { + let error = classify_document_source(source).expect_err("plain URL should be rejected"); + assert!(error.to_string().contains("upload the file first")); +} + +#[rstest] +fn test_parse_v3_uses_programmatic_api_key_over_env() { + let key = resolve_api_key(Some("passed-key"), &|_| Some("env-reducto-key".to_string())) + .expect("explicit key should resolve"); + assert_eq!(key, "passed-key"); + + let headers = REDUCTO_PARSE_V3_CONFIG + .validate_environment(Vec::new(), Some("passed-key"), &|_| { + Some("env-reducto-key".to_string()) + }) + .expect("headers should validate"); + assert_eq!( + headers, + vec![("Authorization".to_string(), "Bearer passed-key".to_string())] + ); +} diff --git a/litellm-rust/crates/core/src/providers/reducto/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/reducto/ocr/transformation.rs new file mode 100644 index 00000000000..b8507541e18 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/reducto/ocr/transformation.rs @@ -0,0 +1,407 @@ +use std::collections::BTreeMap; + +use base64::Engine; +use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; +use serde_json::{Map, Value, json}; + +use crate::error::{Error, json_type_name}; +use crate::ocr::transformation::OcrProviderConfig; +use crate::ocr::types::{OcrRequestData, OcrResponseData}; + +pub const REDUCTO_API_BASE: &str = "https://platform.reducto.ai"; +pub const REDUCTO_API_KEY_ENV: &str = "REDUCTO_API_KEY"; +pub const REDUCTO_ID_PREFIX: &str = "reducto://"; + +const PARSE_V3_SUPPORTED_OCR_PARAMS: &[&str] = &["formatting", "retrieval", "settings"]; +const PARSE_LEGACY_SUPPORTED_OCR_PARAMS: &[&str] = &["enhance"]; +const MISSING_KEY_MESSAGE: &str = "Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()"; +const DATA_URI_UPLOAD_REQUIRED: &str = + "Reducto data URI upload must complete before OCR request transformation"; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum ReductoDocumentSource { + FileId(String), + Upload { bytes: Vec, mime_type: String }, +} + +#[derive(Clone, PartialEq, Eq)] +pub struct ReductoUploadRequest { + pub url: String, + pub authorization: String, + pub file_name: &'static str, + pub bytes: Vec, + pub mime_type: String, +} + +pub struct ReductoParseV3Config; +pub struct ReductoParseLegacyConfig; + +pub const REDUCTO_PARSE_V3_CONFIG: ReductoParseV3Config = ReductoParseV3Config; +pub const REDUCTO_PARSE_LEGACY_CONFIG: ReductoParseLegacyConfig = ReductoParseLegacyConfig; + +pub fn config_for_model(model: &str) -> Option<&'static dyn OcrProviderConfig> { + match model { + "parse-v3" => Some(&REDUCTO_PARSE_V3_CONFIG), + "parse-legacy" => Some(&REDUCTO_PARSE_LEGACY_CONFIG), + _ => None, + } +} + +pub fn normalize_api_base(api_base: Option<&str>) -> String { + api_base + .map(str::trim) + .filter(|base| !base.is_empty()) + .unwrap_or(REDUCTO_API_BASE) + .trim_end_matches('/') + .to_string() +} + +pub fn parse_url(api_base: Option<&str>) -> String { + format!("{}/parse", normalize_api_base(api_base)) +} + +pub fn upload_url(api_base: Option<&str>) -> String { + format!("{}/upload", normalize_api_base(api_base)) +} + +pub fn resolve_api_key( + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> Result { + api_key + .map(str::trim) + .filter(|key| !key.is_empty()) + .map(str::to_string) + .or_else(|| { + env_lookup(REDUCTO_API_KEY_ENV) + .map(|key| key.trim().to_string()) + .filter(|key| !key.is_empty()) + }) + .ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string())) +} + +pub fn extract_document_source(document: &Value) -> Result { + let document = document.as_object().ok_or_else(|| Error::InvalidType { + expected: "object", + actual: json_type_name(document), + })?; + let source = document + .get("document_url") + .and_then(Value::as_str) + .filter(|source| !source.is_empty()) + .or_else(|| document.get("image_url").and_then(Value::as_str)) + .ok_or_else(|| { + Error::InvalidRequest( + "Reducto expected OCR preprocessing to produce document_url or image_url" + .to_string(), + ) + })?; + classify_document_source(source) +} + +pub fn classify_document_source(source: &str) -> Result { + if source.starts_with(REDUCTO_ID_PREFIX) { + return Ok(ReductoDocumentSource::FileId(source.to_string())); + } + if source.starts_with("http://") || source.starts_with("https://") { + return Err(Error::InvalidRequest( + "Reducto requires type='file' (auto-uploaded) or a reducto:// id. Plain http(s) URLs are not supported; upload the file first." + .to_string(), + )); + } + if !source.starts_with("data:") { + return Err(Error::InvalidRequest( + "Reducto requires a reducto:// id or a base64 data URI after OCR preprocessing." + .to_string(), + )); + } + + let (header, encoded) = source + .split_once(',') + .ok_or_else(|| Error::InvalidRequest("Invalid Reducto data URI provided.".to_string()))?; + if !header.split(';').any(|part| part == "base64") { + return Err(Error::InvalidRequest( + "Reducto only supports base64-encoded data URIs.".to_string(), + )); + } + + let mime_type = header + .strip_prefix("data:") + .and_then(|header| header.split(';').next()) + .filter(|mime| !mime.is_empty()) + .unwrap_or("application/octet-stream") + .to_string(); + let bytes = BASE64_STANDARD.decode(encoded).map_err(|_| { + Error::InvalidRequest("Invalid Reducto base64 payload provided.".to_string()) + })?; + + Ok(ReductoDocumentSource::Upload { bytes, mime_type }) +} + +pub fn build_upload_request( + source: ReductoDocumentSource, + authorization: &str, + api_base: Option<&str>, +) -> Option { + let ReductoDocumentSource::Upload { bytes, mime_type } = source else { + return None; + }; + + Some(ReductoUploadRequest { + url: upload_url(api_base), + authorization: authorization.to_string(), + file_name: "document", + bytes, + mime_type, + }) +} + +pub fn extract_upload_file_id(response_json: &Value) -> Result<&str, Error> { + response_json + .as_object() + .and_then(|response| response.get("file_id")) + .and_then(Value::as_str) + .filter(|file_id| !file_id.is_empty()) + .ok_or_else(|| { + Error::InvalidResponse(format!( + "Reducto /upload returned 200 without a file_id; got payload={response_json}" + )) + }) +} + +pub fn build_parse_v3_request( + file_id: &str, + optional_params: Map, +) -> OcrRequestData { + let data = std::iter::once(("input".to_string(), Value::String(file_id.to_string()))) + .chain(optional_params) + .collect(); + OcrRequestData { + data: Value::Object(data), + files: None, + } +} + +pub fn build_parse_legacy_request( + file_id: &str, + optional_params: &Map, +) -> OcrRequestData { + let options = optional_params + .get("enhance") + .filter(|enhance| !enhance.is_null()) + .map(|enhance| json!({"options": {"enhance": enhance}})); + let data = match options { + Some(Value::Object(options)) => std::iter::once(( + "document_url".to_string(), + Value::String(file_id.to_string()), + )) + .chain(options) + .collect(), + _ => Map::from_iter([( + "document_url".to_string(), + Value::String(file_id.to_string()), + )]), + }; + OcrRequestData { + data: Value::Object(data), + files: None, + } +} + +fn source_file_id(document: &Value) -> Result { + match extract_document_source(document)? { + ReductoDocumentSource::FileId(file_id) => Ok(file_id), + ReductoDocumentSource::Upload { .. } => Err(Error::Unsupported(DATA_URI_UPLOAD_REQUIRED)), + } +} + +fn page_number(block: &Map) -> Option { + let page = block.get("bbox")?.as_object()?.get("page")?; + page.as_i64() + .or_else(|| page.as_u64().and_then(|page| i64::try_from(page).ok())) + .or_else(|| page.as_str().and_then(|page| page.parse().ok())) +} + +fn chunks(result: &Map) -> &[Value] { + result + .get("chunks") + .and_then(Value::as_array) + .map(Vec::as_slice) + .unwrap_or_default() +} + +fn build_pages(result: &Map) -> Vec { + let blocks_by_page = chunks(result) + .iter() + .filter_map(Value::as_object) + .filter_map(|chunk| chunk.get("blocks").and_then(Value::as_array)) + .flatten() + .filter_map(|block| block.as_object().map(|object| (block, object))) + .filter_map(|(block, object)| page_number(object).map(|page| (page, block.clone()))) + .fold( + BTreeMap::>::new(), + |mut pages, (page, block)| { + pages.entry(page).or_default().push(block); + pages + }, + ); + + if blocks_by_page.is_empty() { + let markdown = chunks(result) + .iter() + .filter_map(Value::as_object) + .filter_map(|chunk| chunk.get("content").and_then(Value::as_str)) + .filter(|content| !content.is_empty()) + .collect::>() + .join("\n\n"); + return if markdown.is_empty() { + Vec::new() + } else { + vec![json!({"index": 0, "markdown": markdown})] + }; + } + + blocks_by_page + .into_iter() + .map(|(page, blocks)| { + let markdown = blocks + .iter() + .filter_map(Value::as_object) + .filter_map(|block| block.get("content").and_then(Value::as_str)) + .filter(|content| !content.is_empty()) + .collect::>() + .join("\n\n"); + json!({ + "index": page.saturating_sub(1).max(0), + "markdown": markdown, + "blocks": blocks, + }) + }) + .collect() +} + +pub fn transform_reducto_response( + model: &str, + response_json: Value, +) -> Result { + let response = response_json + .as_object() + .ok_or_else(|| Error::InvalidType { + expected: "object", + actual: json_type_name(&response_json), + })?; + let empty_result = Map::new(); + let result = match response.get("result") { + Some(Value::Object(result)) => result, + Some(Value::Null) => &empty_result, + Some(_) => { + return Err(Error::InvalidResponse( + "Reducto result must be an object".to_string(), + )); + } + None => response, + }; + let usage = response + .get("usage") + .and_then(Value::as_object) + .cloned() + .unwrap_or_default(); + let usage_info = Some(json!({ + "pages_processed": usage.get("num_pages").cloned().unwrap_or(Value::Null), + "credits": usage.get("credits").cloned().unwrap_or(Value::Null), + })); + + Ok(OcrResponseData { + pages: build_pages(result), + model: model.to_string(), + document_annotation: None, + usage_info, + object: "ocr".to_string(), + extra_fields: Map::new(), + provider_native_response: Some(response_json), + }) +} + +impl OcrProviderConfig for ReductoParseV3Config { + fn supported_ocr_params(&self) -> &'static [&'static str] { + PARSE_V3_SUPPORTED_OCR_PARAMS + } + + fn transform_ocr_request( + &self, + _model: &str, + document: Value, + optional_params: Map, + ) -> Result { + let file_id = source_file_id(&document)?; + Ok(build_parse_v3_request(&file_id, optional_params)) + } + + fn transform_ocr_response( + &self, + model: &str, + response_json: Value, + ) -> Result { + transform_reducto_response(model, response_json) + } + + fn complete_url( + &self, + api_base: Option<&str>, + _model: &str, + _optional_params: &Map, + _env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + Ok(parse_url(api_base)) + } + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + resolve_api_key(api_key, env_lookup) + } +} + +impl OcrProviderConfig for ReductoParseLegacyConfig { + fn supported_ocr_params(&self) -> &'static [&'static str] { + PARSE_LEGACY_SUPPORTED_OCR_PARAMS + } + + fn transform_ocr_request( + &self, + _model: &str, + document: Value, + optional_params: Map, + ) -> Result { + let file_id = source_file_id(&document)?; + Ok(build_parse_legacy_request(&file_id, &optional_params)) + } + + fn transform_ocr_response( + &self, + model: &str, + response_json: Value, + ) -> Result { + transform_reducto_response(model, response_json) + } + + fn complete_url( + &self, + api_base: Option<&str>, + _model: &str, + _optional_params: &Map, + _env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + Ok(parse_url(api_base)) + } + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + resolve_api_key(api_key, env_lookup) + } +} diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs index ee095447028..c324de8cb45 100644 --- a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs @@ -212,6 +212,7 @@ impl OcrProviderConfig for VertexAiOcrConfig { MISTRAL_OCR_CONFIG.supported_ocr_params() } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn transform_ocr_request( &self, model: &str, @@ -229,6 +230,7 @@ impl OcrProviderConfig for VertexAiOcrConfig { MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn complete_url( &self, api_base: Option<&str>, @@ -253,10 +255,21 @@ impl OcrProviderConfig for VertexAiOcrConfig { } impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn supported_ocr_params(&self) -> &'static [&'static str] { DEEPSEEK_SUPPORTED_OCR_PARAMS } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] + fn map_ocr_params(&self, non_default_params: &Map) -> Map { + non_default_params + .iter() + .filter(|(name, _)| DEEPSEEK_SUPPORTED_OCR_PARAMS.contains(&name.as_str())) + .map(|(name, value)| (name.clone(), value.clone())) + .collect() + } + + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn transform_ocr_request( &self, model: &str, @@ -283,6 +296,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { }) } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn transform_ocr_response( &self, model: &str, @@ -335,9 +349,12 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { document_annotation: object.get("document_annotation").cloned(), usage_info, object: "ocr".to_string(), + extra_fields: Map::new(), + provider_native_response: None, }) } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn complete_url( &self, api_base: Option<&str>, @@ -360,6 +377,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { #[cfg(test)] mod tests { use super::*; + use rstest::rstest; #[test] fn vertex_mistral_url_uses_project_location_and_model() { @@ -411,6 +429,22 @@ mod tests { ); } + #[rstest] + #[case::bare_model("deepseek-ocr-maas")] + #[case::namespaced_model("deepseek-ai/deepseek-ocr-maas")] + fn vertex_deepseek_request_uses_single_provider_namespace(#[case] model: &str) { + let body = VERTEX_AI_DEEPSEEK_OCR_CONFIG + .transform_ocr_request( + model, + json!({"type": "image_url", "image_url": "data:image/png;base64,AA=="}), + Map::new(), + ) + .expect("request transforms") + .data; + + assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); + } + #[test] fn vertex_deepseek_response_wraps_markdown_content() { let response = VERTEX_AI_DEEPSEEK_OCR_CONFIG diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 637e5580170..bda09a7d840 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -14,11 +14,15 @@ default = ["abi3"] abi3 = ["pyo3/abi3-py310"] extension-module = ["pyo3/extension-module"] panic-test = [] +trace-parity = [ + "dep:tracing", + "litellm-core/observability", + "litellm-ai-gateway/trace-parity", +] [dependencies] futures-util.workspace = true -tracing.workspace = true -tracing-subscriber.workspace = true +tracing = { workspace = true, optional = true } litellm-core = { workspace = true, features = ["bedrock-auth"] } litellm-ai-gateway = { workspace = true, default-features = false } litellm-python-interop.workspace = true @@ -31,6 +35,7 @@ tokio.workspace = true [dev-dependencies] criterion = "0.8.2" tokio-tungstenite.workspace = true +tracing.workspace = true [[bench]] name = "serialization" diff --git a/litellm-rust/crates/python-bridge/src/constants.rs b/litellm-rust/crates/python-bridge/src/constants.rs deleted file mode 100644 index 07b2836b838..00000000000 --- a/litellm-rust/crates/python-bridge/src/constants.rs +++ /dev/null @@ -1 +0,0 @@ -pub(crate) const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace"; diff --git a/litellm-rust/crates/python-bridge/src/function_trace.rs b/litellm-rust/crates/python-bridge/src/function_trace.rs index 420d237c79d..ea7d9f4993e 100644 --- a/litellm-rust/crates/python-bridge/src/function_trace.rs +++ b/litellm-rust/crates/python-bridge/src/function_trace.rs @@ -1,216 +1,22 @@ use std::future::Future; -use std::sync::{Arc, Mutex}; +use litellm_core::observability::{FunctionTrace, FunctionTraceEvent}; use serde::Serialize; use tracing::instrument::WithSubscriber; -use tracing::span::{Attributes, Id}; -use tracing::{Dispatch, Level, Subscriber}; -use tracing_subscriber::filter::{LevelFilter, filter_fn}; -use tracing_subscriber::layer::Context; -use tracing_subscriber::prelude::*; -use tracing_subscriber::registry::LookupSpan; -use tracing_subscriber::{Layer, Registry}; - -use crate::constants::FUNCTION_TRACE_TARGET; #[derive(Serialize)] -#[serde(untagged)] -pub(crate) enum TraceResponse { - Plain(T), - Traced { - response: T, - trace: Vec, - }, +pub(crate) struct TracedResponse { + response: T, + trace: Vec, } -pub(crate) async fn trace_call( +pub(crate) async fn capture( future: impl Future>, - enabled: bool, -) -> Result, E> { - if !enabled { - return future.await.map(TraceResponse::Plain); - } +) -> Result, E> { let trace = FunctionTrace::default(); let response = future.with_subscriber(trace.dispatcher()).await?; - Ok(TraceResponse::Traced { + Ok(TracedResponse { response, trace: trace.events(), }) } - -#[derive(Clone, Debug, PartialEq, Serialize)] -pub struct FunctionTraceEvent { - pub function: &'static str, - pub depth: usize, -} - -#[derive(Clone, Default)] -pub struct FunctionTrace { - events: Arc>>, -} - -impl FunctionTrace { - pub fn dispatcher(&self) -> Dispatch { - let filter = filter_fn(|metadata| { - metadata.is_span() - && metadata.target() == FUNCTION_TRACE_TARGET - && *metadata.level() == Level::TRACE - }) - .with_max_level_hint(LevelFilter::TRACE); - Dispatch::new( - Registry::default().with( - FunctionTraceLayer { - trace: self.clone(), - } - .with_filter(filter), - ), - ) - } - - pub fn events(&self) -> Vec { - self.events - .lock() - .unwrap_or_else(|error| error.into_inner()) - .clone() - } -} - -struct FunctionTraceLayer { - trace: FunctionTrace, -} - -impl Layer for FunctionTraceLayer -where - S: Subscriber + for<'lookup> LookupSpan<'lookup>, -{ - fn on_new_span(&self, attributes: &Attributes<'_>, id: &Id, context: Context<'_, S>) { - let depth = context - .span(id) - .map(|span| span.scope().skip(1).count()) - .unwrap_or_default(); - self.trace - .events - .lock() - .unwrap_or_else(|error| error.into_inner()) - .push(FunctionTraceEvent { - function: attributes.metadata().name(), - depth, - }); - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - async fn outer() { - tokio::task::yield_now().await; - inner().await; - } - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - async fn inner() { - tokio::task::yield_now().await; - } - - #[tokio::test] - async fn concurrent_futures_keep_separate_traces_across_yields() { - use tracing::instrument::WithSubscriber; - - let first = FunctionTrace::default(); - let second = FunctionTrace::default(); - let outside = FunctionTrace::default(); - - async { - tokio::join!( - outer().with_subscriber(first.dispatcher()), - inner().with_subscriber(second.dispatcher()), - ); - inner().await; - } - .with_subscriber(outside.dispatcher()) - .await; - - assert_eq!( - first.events(), - vec![ - FunctionTraceEvent { - function: "outer", - depth: 0 - }, - FunctionTraceEvent { - function: "inner", - depth: 1 - }, - ], - ); - assert_eq!( - second.events(), - vec![FunctionTraceEvent { - function: "inner", - depth: 0 - }], - ); - assert_eq!( - outside.events(), - vec![FunctionTraceEvent { - function: "inner", - depth: 0 - }], - ); - } - - #[test] - fn records_matching_spans_in_creation_order() { - let trace = FunctionTrace::default(); - let dispatch = trace.dispatcher(); - - tracing::dispatcher::with_default(&dispatch, || { - let _ignored = tracing::trace_span!(target: "other", "ignored"); - let _first = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name"); - let _wrong_level = tracing::debug_span!(target: FUNCTION_TRACE_TARGET, "wrong_level"); - let _second = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name"); - }); - - assert_eq!( - trace.events(), - vec![ - FunctionTraceEvent { - function: "same_name", - depth: 0, - }, - FunctionTraceEvent { - function: "same_name", - depth: 0, - }, - ] - ); - } - - #[test] - fn records_matching_span_nesting_depth() { - let trace = FunctionTrace::default(); - let dispatch = trace.dispatcher(); - - tracing::dispatcher::with_default(&dispatch, || { - let outer = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "outer"); - let _outer_guard = outer.enter(); - let _inner = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "inner"); - }); - - assert_eq!( - trace.events(), - vec![ - FunctionTraceEvent { - function: "outer", - depth: 0, - }, - FunctionTraceEvent { - function: "inner", - depth: 1, - }, - ] - ); - } -} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 5f36a22370a..384f0be5a1b 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,8 +1,8 @@ -mod constants; mod diagnostics; mod errors; mod execution; -pub mod function_trace; +#[cfg(feature = "trace-parity")] +mod function_trace; mod marshal; mod routes; @@ -115,9 +115,43 @@ mod tests { .extract::>() .expect("module names should be strings") .into_iter() - .filter(|name| !name.starts_with("__")) + .filter(|name| !name.starts_with('_')) .collect(); assert_eq!(public_names, expected); + + #[cfg(not(feature = "trace-parity"))] + assert!(!module.hasattr("_trace").expect("module lookup should work")); + + #[cfg(feature = "trace-parity")] + { + let trace = module + .getattr("_trace") + .expect("trace build should expose its diagnostic namespace"); + let trace_names: Vec = trace + .cast::() + .expect("trace namespace should be a module") + .dict() + .keys() + .extract::>() + .expect("trace names should be strings") + .into_iter() + .filter(|name| !name.starts_with("__")) + .collect(); + assert_eq!( + trace_names, + [ + "ocr", + "aocr", + "transcription", + "atranscription", + "messages", + "amessages", + "chat_completions", + "achat_completions", + "gateway_messages", + ] + ); + } }); } diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs index 10b86132be7..af60515b0e2 100644 --- a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs +++ b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs @@ -54,16 +54,16 @@ bridge_route! { required = { model: String, #[pyo3(from_py_with = litellm_python_interop::from_py)] - audio: Value, + audio: serde_json::Value, }, optional = { api_key: Option, api_base: Option, custom_llm_provider: Option, #[pyo3(from_py_with = litellm_python_interop::from_py)] - extra_headers: Option, + extra_headers: Option, #[pyo3(from_py_with = litellm_python_interop::from_py)] - optional_params: Option, + optional_params: Option, timeout_seconds: Option, }, prepare = prepare_transcription, diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs index 68b7762cb10..08ab476005c 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs @@ -73,16 +73,16 @@ bridge_route! { required = { model: String, #[pyo3(from_py_with = litellm_python_interop::from_py)] - messages: Value, + messages: serde_json::Value, }, optional = { #[pyo3(from_py_with = litellm_python_interop::from_py)] - optional_params: Option, + optional_params: Option, api_key: Option, api_base: Option, custom_llm_provider: Option, #[pyo3(from_py_with = litellm_python_interop::from_py)] - extra_headers: Option, + extra_headers: Option, timeout_seconds: Option, }, prepare = prepare_chat_completions, diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs index 21a7fd5a766..3285da14d5f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/definition.rs +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -20,43 +20,33 @@ macro_rules! bridge_route { } #[pyfunction] - #[pyo3(signature = ($($required_name),*, $($optional_name=None,)* trace=false))] + #[pyo3(signature = ($($required_name),*, $($optional_name=None),*))] #[allow(clippy::too_many_arguments)] fn $sync_name( py: pyo3::Python<'_>, $($(#[$required_attr])* $required_name: $required_type,)* $($(#[$optional_attr])* $optional_name: $optional_type,)* - trace: bool, ) -> pyo3::PyResult> { let future = $prepare($inputs { $($required_name,)* $($optional_name),* })?; - $crate::execution::run_sync( - py, - $crate::function_trace::trace_call(future, trace), - $map_error, - ) + $crate::execution::run_sync(py, future, $map_error) } #[pyfunction] - #[pyo3(signature = ($($required_name),*, $($optional_name=None,)* trace=false))] + #[pyo3(signature = ($($required_name),*, $($optional_name=None),*))] #[allow(clippy::too_many_arguments)] fn $async_name( py: pyo3::Python<'_>, $($(#[$required_attr])* $required_name: $required_type,)* $($(#[$optional_attr])* $optional_name: $optional_type,)* - trace: bool, ) -> pyo3::PyResult> { let future = $prepare($inputs { $($required_name,)* $($optional_name),* })?; - $crate::execution::run_async( - py, - $crate::function_trace::trace_call(future, trace), - $map_error, - ) + $crate::execution::run_async(py, future, $map_error) } pub(super) fn register( @@ -67,6 +57,71 @@ macro_rules! bridge_route { $crate::routes::definition::add_function(module, pyo3::wrap_pyfunction!($async_name, module)?)?; Ok(()) } + + #[cfg(feature = "trace-parity")] + mod trace { + use pyo3::prelude::*; + use super::{$inputs, $map_error, $prepare}; + + #[pyfunction] + #[pyo3(signature = ($($required_name),*, $($optional_name=None),*))] + #[allow(clippy::too_many_arguments)] + fn $sync_name( + py: pyo3::Python<'_>, + $($(#[$required_attr])* $required_name: $required_type,)* + $($(#[$optional_attr])* $optional_name: $optional_type,)* + ) -> pyo3::PyResult> { + let future = $prepare($inputs { + $($required_name,)* + $($optional_name),* + })?; + $crate::execution::run_sync( + py, + $crate::function_trace::capture(future), + $map_error, + ) + } + + #[pyfunction] + #[pyo3(signature = ($($required_name),*, $($optional_name=None),*))] + #[allow(clippy::too_many_arguments)] + fn $async_name( + py: pyo3::Python<'_>, + $($(#[$required_attr])* $required_name: $required_type,)* + $($(#[$optional_attr])* $optional_name: $optional_type,)* + ) -> pyo3::PyResult> { + let future = $prepare($inputs { + $($required_name,)* + $($optional_name),* + })?; + $crate::execution::run_async( + py, + $crate::function_trace::capture(future), + $map_error, + ) + } + + pub(super) fn register( + module: &pyo3::Bound<'_, pyo3::types::PyModule>, + ) -> pyo3::PyResult<()> { + $crate::routes::definition::add_function( + module, + pyo3::wrap_pyfunction!($sync_name, module)?, + )?; + $crate::routes::definition::add_function( + module, + pyo3::wrap_pyfunction!($async_name, module)?, + )?; + Ok(()) + } + } + + #[cfg(feature = "trace-parity")] + pub(super) fn register_trace( + module: &pyo3::Bound<'_, pyo3::types::PyModule>, + ) -> pyo3::PyResult<()> { + trace::register(module) + } }; } @@ -130,20 +185,26 @@ mod tests { ) -> PyResult> + Send + 'static> { FUTURE_DROPPED.store(false, Ordering::SeqCst); let drop_guard = (inputs.value == "pending").then_some(DropGuard); - Ok(async move { - let _drop_guard = drop_guard; - tokio::task::yield_now().await; - match inputs.value.as_str() { - "error" => Err(Error::InvalidRequest("synthetic error".to_string())), - "map_panic" => Err(Error::InvalidRequest("panic in mapper".to_string())), - "panic" => panic!("synthetic panic"), - "pending" => { - pending::<()>().await; - unreachable!() - } - _ => Ok(inputs.value), + Ok(execute_echo(inputs, drop_guard)) + } + + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] + async fn execute_echo( + inputs: EchoInputs, + drop_guard: Option, + ) -> Result { + let _drop_guard = drop_guard; + tokio::task::yield_now().await; + match inputs.value.as_str() { + "error" => Err(Error::InvalidRequest("synthetic error".to_string())), + "map_panic" => Err(Error::InvalidRequest("panic in mapper".to_string())), + "panic" => panic!("synthetic panic"), + "pending" => { + pending::<()>().await; + unreachable!() } - }) + _ => Ok(inputs.value), + } } fn map_error(error: Error) -> PyErr { @@ -164,22 +225,22 @@ mod tests { ( "ocr", "aocr", - "(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=False)", + "(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None)", ), ( "transcription", "atranscription", - "(model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=False)", + "(model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None)", ), ( "messages", "amessages", - "(model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=False)", + "(model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None)", ), ( "chat_completions", "achat_completions", - "(model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=False)", + "(model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None)", ), ]; @@ -411,6 +472,32 @@ asyncio.run(exercise()) }); } + #[cfg(feature = "trace-parity")] + #[test] + fn diagnostic_route_returns_the_response_and_filtered_trace() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::new(py, "synthetic").expect("module should be created"); + synthetic::register_trace(&module).expect("trace routes should register"); + let locals = PyDict::new(py); + locals + .set_item("routes", &module) + .expect("module should enter Python locals"); + let code = CString::new( + r#" +result = routes.echo("traced") +assert result == { + "response": "traced", + "trace": [{"function": "execute_echo", "depth": 0}], +} +"#, + ) + .expect("Python source should not contain null bytes"); + py.run(&code, Some(&locals), Some(&locals)) + .expect("diagnostic route should return its response and trace"); + }); + } + #[test] fn route_registration_rejects_duplicate_python_names() { Python::initialize(); diff --git a/litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs b/litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs new file mode 100644 index 00000000000..97ff93f299a --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs @@ -0,0 +1,29 @@ +use pyo3::prelude::*; +use serde_json::Value; + +use crate::errors::core_error_to_pyerr; + +#[pyfunction] +fn gateway_messages<'py>( + py: Python<'py>, + model_alias: String, + provider_model: String, + api_base: String, + #[pyo3(from_py_with = litellm_python_interop::from_py)] body: Value, +) -> PyResult> { + let future = litellm_ai_gateway::trace_parity::messages_request( + model_alias, + provider_model, + api_base, + body, + ); + crate::execution::run_async( + py, + crate::function_trace::capture(future), + core_error_to_pyerr, + ) +} + +pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> { + super::definition::add_function(module, wrap_pyfunction!(gateway_messages, module)?) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages.rs index 2bb64a7a763..f69b5e9251d 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages.rs @@ -50,14 +50,14 @@ bridge_route! { required = { model: String, #[pyo3(from_py_with = litellm_python_interop::from_py)] - body: Value, + body: serde_json::Value, }, optional = { api_key: Option, api_base: Option, custom_llm_provider: Option, #[pyo3(from_py_with = litellm_python_interop::from_py)] - extra_headers: Option, + extra_headers: Option, timeout_seconds: Option, }, prepare = prepare_messages, diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index bf611c26d44..7e81f2ffe9b 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -3,6 +3,9 @@ use pyo3::prelude::*; #[macro_use] mod definition; +#[cfg(feature = "trace-parity")] +mod gateway_messages; + mod audio_transcription; mod chat_completions; mod messages; @@ -12,5 +15,16 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { ocr::register(module)?; audio_transcription::register(module)?; messages::register(module)?; - chat_completions::register(module) + chat_completions::register(module)?; + #[cfg(feature = "trace-parity")] + { + let trace = PyModule::new(module.py(), "_trace")?; + ocr::register_trace(&trace)?; + audio_transcription::register_trace(&trace)?; + messages::register_trace(&trace)?; + chat_completions::register_trace(&trace)?; + gateway_messages::register_trace(&trace)?; + module.add_submodule(&trace)?; + } + Ok(()) } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index 5cc8804238b..cc2f8e43cea 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -56,16 +56,16 @@ bridge_route! { required = { model: String, #[pyo3(from_py_with = litellm_python_interop::from_py)] - document: Value, + document: serde_json::Value, }, optional = { api_key: Option, api_base: Option, custom_llm_provider: Option, #[pyo3(from_py_with = litellm_python_interop::from_py)] - extra_headers: Option, + extra_headers: Option, #[pyo3(from_py_with = litellm_python_interop::from_py)] - optional_params: Option, + optional_params: Option, timeout_seconds: Option, }, prepare = prepare_ocr, diff --git a/litellm-rust/crates/python-bridge/src/routes/runtime.rs b/litellm-rust/crates/python-bridge/src/routes/runtime.rs deleted file mode 100644 index 87a0c3e0104..00000000000 --- a/litellm-rust/crates/python-bridge/src/routes/runtime.rs +++ /dev/null @@ -1,423 +0,0 @@ -use std::future::Future; -use std::panic::AssertUnwindSafe; -use std::time::Duration; - -use futures_util::FutureExt; -use litellm_core::error::Error; -use litellm_python_interop::{Pythonized, panic_to_pyerr, release_gil}; -use pyo3::exceptions::PyRuntimeError; -use pyo3::prelude::*; -use serde::Serialize; -use tokio::runtime::{Handle, Runtime}; -use tokio::time::{self, MissedTickBehavior}; - -pub(super) fn run_sync( - py: Python<'_>, - future: F, - map_error: fn(Error) -> PyErr, -) -> PyResult> -where - T: Serialize + Send + 'static, - F: Future> + Send + 'static, -{ - run_sync_on( - py, - pyo3_async_runtimes::tokio::get_runtime(), - future, - map_error, - ) -} - -fn run_sync_on( - py: Python<'_>, - runtime: &Runtime, - future: F, - map_error: fn(Error) -> PyErr, -) -> PyResult> -where - T: Serialize + Send + 'static, - F: Future> + Send + 'static, -{ - if Handle::try_current().is_ok() { - return Err(PyRuntimeError::new_err( - "synchronous native routes cannot run from a Tokio context; use the async route", - )); - } - - let result = release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?; - let result = map_core_result(result, map_error)?; - Pythonized(result).into_pyobject(py).map(Bound::unbind) -} - -pub(super) fn run_async( - py: Python<'_>, - future: F, - map_error: fn(Error) -> PyErr, -) -> PyResult> -where - T: Serialize + Send + 'static, - F: Future> + Send + 'static, -{ - pyo3_async_runtimes::tokio::future_into_py(py, async move { - let result = catch_route_panic(future).await?; - let result = map_core_result(result, map_error)?; - Ok(Pythonized(result)) - }) -} - -fn map_core_result(result: Result, map_error: fn(Error) -> PyErr) -> PyResult { - match result { - Ok(value) => Ok(value), - Err(error) => Err( - std::panic::catch_unwind(AssertUnwindSafe(|| map_error(error))) - .map_err(panic_to_pyerr)?, - ), - } -} - -async fn catch_route_panic(future: F) -> PyResult> -where - F: Future>, -{ - AssertUnwindSafe(future) - .catch_unwind() - .await - .map_err(panic_to_pyerr) -} - -async fn wait_for_sync_result(future: F) -> PyResult> -where - F: Future>, -{ - let future = catch_route_panic(future); - tokio::pin!(future); - - let signal_interval = Duration::from_millis(50); - let mut signal_checks = - time::interval_at(time::Instant::now() + signal_interval, signal_interval); - signal_checks.set_missed_tick_behavior(MissedTickBehavior::Delay); - loop { - tokio::select! { - result = &mut future => return result, - _ = signal_checks.tick() => Python::attach(|py| py.check_signals())?, - } - } -} - -#[cfg(test)] -mod tests { - use std::ffi::CString; - use std::future::poll_fn; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::{Arc, mpsc}; - use std::task::Poll; - use std::thread; - use std::time::Instant; - - use pyo3::panic::PanicException; - use pyo3::types::{PyDict, PyModule}; - use serde::Serializer; - use tokio::runtime::Builder; - - use super::*; - - fn runtime_error(error: Error) -> PyErr { - PyRuntimeError::new_err(error.to_string()) - } - - fn panicking_error_mapper(_error: Error) -> PyErr { - panic!("error mapper panicked") - } - - struct PanickingOutput; - - static ASYNC_PROBE_COMPLETED: AtomicUsize = AtomicUsize::new(0); - - impl Serialize for PanickingOutput { - fn serialize(&self, _serializer: S) -> Result - where - S: Serializer, - { - panic!("serializer panicked") - } - } - - #[pyfunction] - fn async_serialization_panic(py: Python<'_>) -> PyResult> { - run_async(py, async { Ok(PanickingOutput) }, runtime_error) - } - - #[pyfunction] - fn async_runtime_probe(py: Python<'_>) -> PyResult> { - run_async( - py, - async { - ASYNC_PROBE_COMPLETED.fetch_add(1, Ordering::SeqCst); - Ok(true) - }, - runtime_error, - ) - } - - #[pyfunction] - fn runtime_worker_count() -> usize { - pyo3_async_runtimes::tokio::get_runtime() - .metrics() - .num_workers() - } - - #[pyfunction] - fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> bool { - let completion_deadline = Instant::now() + Duration::from_secs(2); - while ASYNC_PROBE_COMPLETED.load(Ordering::SeqCst) < expected_completions { - if Instant::now() >= completion_deadline { - return false; - } - thread::sleep(Duration::from_millis(1)); - } - - let (heartbeat_tx, heartbeat_rx) = mpsc::sync_channel(1); - pyo3_async_runtimes::tokio::get_runtime().spawn(async move { - let _ = heartbeat_tx.send(()); - }); - heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok() - } - - fn extract_bool(py: Python<'_>, result: PyResult>) -> bool { - result - .expect("route should complete") - .bind(py) - .extract() - .expect("result should convert") - } - - #[test] - fn sync_runner_polls_future_on_the_caller_thread() { - Python::initialize(); - Python::attach(|py| { - let caller_thread = std::thread::current().id(); - let result = run_sync( - py, - async move { Ok(std::thread::current().id() == caller_thread) }, - runtime_error, - ); - - assert!(extract_bool(py, result)); - }); - } - - #[test] - fn sync_runner_releases_gil_while_waiting() { - Python::initialize(); - Python::attach(|py| { - let result = run_sync( - py, - async { - let gil_acquired = tokio::time::timeout( - Duration::from_secs(2), - tokio::task::spawn_blocking(|| Python::attach(|_| true)), - ) - .await; - Ok(matches!(gil_acquired, Ok(Ok(true)))) - }, - runtime_error, - ); - - assert!(extract_bool(py, result)); - }); - } - - #[test] - fn sync_runner_rejects_calls_from_a_tokio_context() { - Python::initialize(); - let runtime = Builder::new_current_thread() - .enable_all() - .build() - .expect("runtime should build"); - - let error = runtime.block_on(async { - Python::attach(|py| { - run_sync::(py, async { Ok(true) }, runtime_error) - .expect_err("sync route should reject a nested Tokio runtime") - }) - }); - - assert_eq!( - error.to_string(), - "RuntimeError: synchronous native routes cannot run from a Tokio context; use the async route" - ); - } - - #[test] - fn sync_runner_can_drive_a_current_thread_runtime() { - Python::initialize(); - let runtime = Builder::new_current_thread() - .enable_all() - .build() - .expect("runtime should build"); - Python::attach(|py| { - let result = run_sync_on( - py, - &runtime, - async { - tokio::task::yield_now().await; - Ok(true) - }, - runtime_error, - ); - assert!(extract_bool(py, result)); - }); - } - - #[test] - fn sync_runner_maps_a_panicked_future() { - Python::initialize(); - Python::attach(|py| { - let error = run_sync::( - py, - poll_fn(|_| -> Poll> { panic!("route future panicked") }), - runtime_error, - ) - .expect_err("panicked route should become a Python exception"); - - assert!(error.is_instance_of::(py)); - assert_eq!(error.to_string(), "PanicException: route future panicked"); - }); - } - - #[test] - fn sync_runner_maps_a_panicked_error_mapper() { - Python::initialize(); - Python::attach(|py| { - let error = run_sync::( - py, - async { Err(Error::InvalidRequest("invalid".to_string())) }, - panicking_error_mapper, - ) - .expect_err("panicked mapper should become a Python exception"); - - assert!(error.is_instance_of::(py)); - assert_eq!(error.to_string(), "PanicException: error mapper panicked"); - }); - } - - #[test] - fn sync_runner_surfaces_serializer_panics() { - Python::initialize(); - Python::attach(|py| { - let error = run_sync(py, async { Ok(PanickingOutput) }, runtime_error) - .expect_err("serializer panic should become a Python exception"); - - assert!(error.is_instance_of::(py)); - assert_eq!(error.to_string(), "PanicException: serializer panicked"); - }); - } - - #[test] - fn sync_runner_supports_concurrent_callers_on_the_shared_runtime() { - Python::initialize(); - let barrier = Arc::new(tokio::sync::Barrier::new(2)); - let callers: Vec<_> = (0..2) - .map(|_| { - let barrier = Arc::clone(&barrier); - thread::spawn(move || { - Python::attach(|py| { - extract_bool( - py, - run_sync( - py, - async move { - Ok(tokio::time::timeout(Duration::from_secs(2), barrier.wait()) - .await - .is_ok()) - }, - runtime_error, - ), - ) - }) - }) - }) - .collect(); - let results: Vec<_> = callers - .into_iter() - .map(|caller| caller.join().expect("caller should not panic")) - .collect(); - - assert_eq!(results, vec![true, true]); - } - - #[test] - fn async_runner_surfaces_serializer_panics() { - Python::initialize(); - Python::attach(|py| { - let module = PyModule::new(py, "runtime").expect("module should be created"); - module - .add_function( - wrap_pyfunction!(async_serialization_panic, &module) - .expect("function should wrap"), - ) - .expect("function should register"); - let locals = PyDict::new(py); - locals - .set_item("runtime", &module) - .expect("module should enter Python locals"); - let code = CString::new( - r#" -import asyncio - -async def exercise(): - try: - await runtime.async_serialization_panic() - except BaseException as error: - assert type(error).__name__ == "PanicException" - assert str(error) == "serializer panicked" - else: - raise AssertionError("serializer panic was not raised") - -asyncio.run(exercise()) -"#, - ) - .expect("Python source should not contain null bytes"); - py.run(&code, Some(&locals), Some(&locals)) - .expect("serializer panic should reach the Python awaiter"); - }); - } - - #[test] - fn async_result_delivery_does_not_stall_tokio_workers() { - Python::initialize(); - ASYNC_PROBE_COMPLETED.store(0, Ordering::SeqCst); - Python::attach(|py| { - let module = PyModule::new(py, "runtime").expect("module should be created"); - for function in [ - wrap_pyfunction!(async_runtime_probe, &module).expect("function should wrap"), - wrap_pyfunction!(runtime_worker_count, &module).expect("function should wrap"), - wrap_pyfunction!(runtime_is_responsive, &module).expect("function should wrap"), - ] { - module - .add_function(function) - .expect("function should register"); - } - let locals = PyDict::new(py); - locals - .set_item("runtime", &module) - .expect("module should enter Python locals"); - let code = CString::new( - r#" -import asyncio - -async def exercise(): - worker_count = runtime.runtime_worker_count() - awaitables = [runtime.async_runtime_probe() for _ in range(worker_count)] - assert runtime.runtime_is_responsive(worker_count) - assert await asyncio.gather(*awaitables) == [True] * worker_count - -asyncio.run(exercise()) -"#, - ) - .expect("Python source should not contain null bytes"); - py.run(&code, Some(&locals), Some(&locals)) - .expect("result delivery should leave Tokio workers responsive"); - }); - } -} diff --git a/litellm/anthropic_interface/exceptions/exceptions.py b/litellm/anthropic_interface/exceptions/exceptions.py index ae333d1f4ad..91bcf82f455 100644 --- a/litellm/anthropic_interface/exceptions/exceptions.py +++ b/litellm/anthropic_interface/exceptions/exceptions.py @@ -1,8 +1,9 @@ """Anthropic error format type definitions.""" +from collections.abc import Mapping from typing import Literal -from typing_extensions import Required, TypedDict +from typing_extensions import NotRequired, ReadOnly, Required, TypedDict # Known Anthropic error types # Source: https://docs.anthropic.com/en/api/errors @@ -23,6 +24,7 @@ class AnthropicErrorDetail(TypedDict): type: AnthropicErrorType message: str + provider_specific_fields: NotRequired[ReadOnly[Mapping[str, object]]] class AnthropicErrorResponse(TypedDict, total=False): diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 754815fce47..884d095793c 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -100,6 +100,7 @@ class Cache: qdrant_semantic_cache_vector_size: int | None = None, semantic_cache_embedding_max_input_tokens: int | None = None, semantic_cache_embedding_timeout: float | None = None, + semantic_cache_scope: str = SemanticCacheScope.KEY.value, # GCP IAM authentication parameters gcp_service_account: str | None = None, gcp_ssl_ca_certs: str | None = None, @@ -127,6 +128,7 @@ class Cache: similarity_threshold (float, optional): The similarity threshold for semantic-caching, Required if type is "redis-semantic" or "qdrant-semantic". semantic_cache_embedding_max_input_tokens (int, optional): Truncate prompts to this many tokens before embedding them for semantic caching. Defaults to the embedding deployment's configured max_input_tokens. semantic_cache_embedding_timeout (float, optional): Seconds a semantic-cache lookup may spend embedding the prompt before it gives up and lets the request continue to the LLM. Defaults to SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS. + semantic_cache_scope (str, optional): "key" isolates semantic-cache buckets per key/team/org. "end_user" additionally isolates per end user (falls back to the key scope when the request carries no end-user id). Defaults to "key". # Disk Cache Args disk_cache_dir (str, optional): The directory for the disk cache. Defaults to None. @@ -274,6 +276,7 @@ class Cache: self.redis_flush_size = redis_flush_size self.ttl = ttl self.mode: CacheMode = mode or CacheMode.default_on + self.semantic_cache_scope: str = SemanticCacheScope(semantic_cache_scope).value if self.type == LiteLLMCacheType.LOCAL and default_in_memory_ttl is not None: self.ttl = default_in_memory_ttl @@ -301,6 +304,7 @@ class Cache: "user_api_key_team_id", "user_api_key_org_id", ) + _SEMANTIC_CACHE_END_USER_SCOPE_FIELD: Final = "user_api_key_end_user_id" def _is_semantic_cache(self) -> bool: return self.type in ( @@ -309,19 +313,21 @@ class Cache: LiteLLMCacheType.VALKEY_SEMANTIC, ) - def _get_semantic_cache_tenant_scope(self, kwargs: dict) -> str: - metadata: Final[dict] = kwargs.get("metadata") or {} - litellm_params: Final[dict] = kwargs.get("litellm_params") or {} - metadata_in_litellm_params: Final[dict] = litellm_params.get("metadata") or {} + def _semantic_cache_scope_fields(self) -> tuple[str, ...]: + if self.semantic_cache_scope == SemanticCacheScope.END_USER: + return (*self._SEMANTIC_CACHE_TENANT_SCOPE_FIELDS, self._SEMANTIC_CACHE_END_USER_SCOPE_FIELD) + return self._SEMANTIC_CACHE_TENANT_SCOPE_FIELDS - scope = "" - for field in self._SEMANTIC_CACHE_TENANT_SCOPE_FIELDS: - value = metadata.get(field) - if value is None: - value = metadata_in_litellm_params.get(field) - if value is not None: - scope += f"{field}: {value}" - return scope + def _get_semantic_cache_tenant_scope(self, kwargs: dict) -> str: + litellm_params: Final[dict] = kwargs.get("litellm_params") or {} + metadata_sources: Final[tuple[dict, ...]] = tuple( + source.get(key) or {} for source in (kwargs, litellm_params) for key in ("metadata", "litellm_metadata") + ) + scope_values: Final = ( + (field, next((source[field] for source in metadata_sources if source.get(field) is not None), None)) + for field in self._semantic_cache_scope_fields() + ) + return "".join(f"{field}: {value}" for field, value in scope_values if value is not None) def get_cache_key(self, **kwargs) -> str: """ diff --git a/litellm/caching/redis_cluster_node_isolation.py b/litellm/caching/redis_cluster_node_isolation.py index ae8c78709d9..0035801018b 100644 --- a/litellm/caching/redis_cluster_node_isolation.py +++ b/litellm/caching/redis_cluster_node_isolation.py @@ -19,13 +19,13 @@ connections untouched. Every other branch (MOVED, ASK, CLUSTERDOWN, slot-not-cov retry-exhaustion) is unchanged from upstream, since those already carry real evidence the topology changed. -redis-py 8.x fixed this upstream with gentler machinery than this override's -``node.disconnect()`` (which also kills connections other coroutines are mid-operation -on, so one timeout cascades into a reconnect storm and, with TLS, a fresh handshake per -killed connection): it marks in-use connections for reconnect only after their current -operation completes, disconnects only the idle pooled ones, and defers reinitialization -to the outer retry loop. When the installed ``ClusterNode`` has that per-connection -recovery API, the factory returns the base ``RedisCluster`` unmodified. +redis-py 8.x recovers connections per-connection, so the copied override is not used. Upstream +still flips the shared ``_initialize`` flag on any node's timeout, funneling every concurrent +caller through the reinit lock and, if ``CLUSTER SLOTS`` lands on the slow node, into a full +teardown. For those versions the factory returns a thin wrapper around upstream's +``_execute_command`` that clears the flag again after an isolated timeout (a ConnectionError, +a third consecutive timeout on the same node, or a concurrent request from any other command +or ``aclose()`` still reinits). """ import asyncio @@ -44,6 +44,8 @@ class _ClusterNodeAttrs(Protocol): mode; typing ``target_node`` as this Protocol at the one boundary keeps the override's own logic fully typed without a banned ``typing.cast``.""" + name: str + async def execute_command( self, *args: object, @@ -78,18 +80,20 @@ class _ClusterAttrs(Protocol): #: this override can't see (Python won't error -- it'll just run our now-stale copy), so #: construction logs a loud warning rather than silently trusting an unverified copy. _VERIFIED_REDIS_VERSIONS: Final = frozenset({"5.3.1"}) +_CONSECUTIVE_TIMEOUTS_BEFORE_REINIT: Final = 3 -def get_litellm_async_redis_cluster_class( +def get_litellm_async_redis_cluster_class( # noqa: C901 # supports redis-py version-specific cluster implementations cluster_node_class: type | None = None, + base_cluster_class: type | None = None, ) -> type["_AsyncRedisClusterType"]: - """Returns the base ``RedisCluster`` when the installed redis-py already recovers a - node-level connection error per-connection (8.x+), else builds the ``RedisCluster`` - subclass with the per-node isolation fix for older versions whose upstream branch - tears down the whole cluster client. + """Returns a timeout-tolerant ``RedisCluster`` subclass when installed redis-py already + recovers node-level connections per-connection (8.x+), else builds the ``RedisCluster`` + subclass with the per-node isolation fix for older versions whose upstream branch tears + down the whole cluster client. - ``cluster_node_class`` exists for dependency injection in tests; production callers - leave it unset and the installed ``ClusterNode`` is used. + ``cluster_node_class`` and ``base_cluster_class`` exist for dependency injection in tests; + production callers leave them unset and the installed redis-py classes are used. Imported lazily because this module is reachable from a base ``import litellm`` while redis is not a base dependency. Cheap to call repeatedly: the underlying redis @@ -118,13 +122,68 @@ def get_litellm_async_redis_cluster_class( from redis.exceptions import TimeoutError as _RedisTimeoutError node_class: Final = cluster_node_class if cluster_node_class is not None else _AsyncClusterNode + base_class: Final = base_cluster_class if base_cluster_class is not None else _BaseAsyncRedisCluster if hasattr(node_class, "update_active_connections_for_reconnect"): verbose_logger.debug( - "redis-py %s recovers a node-level connection error per-connection upstream; " - "using the base RedisCluster without litellm's node-isolation override.", + "redis-py %s recovers node connections per-connection upstream; using " + "LiteLLM's timeout-tolerant RedisCluster wrapper.", redis.__version__, ) - return _BaseAsyncRedisCluster + + class LiteLLMAsyncRedisClusterTimeoutTolerant( + base_class # pyright: ignore[reportGeneralTypeIssues, reportUntypedBaseClass] # the injected base class is selected at runtime + ): + def __init__( + self, + *args: object, + **kwargs: object, # kwargs-ok: passes redis-py's constructor kwargs through untouched + ) -> None: + self._litellm_initialize = False + self._litellm_reinit_requests = 0 + self._litellm_tolerated_timeouts = 0 + super().__init__(*args, **kwargs) + self._litellm_consecutive_timeouts: dict[ # mutable-ok: per-node counter updated on the command hot path + str, int + ] = {} + + @property + def _initialize(self) -> bool: + return self._litellm_initialize + + @_initialize.setter + def _initialize(self, value: bool) -> None: + if value: + self._litellm_reinit_requests += 1 + self._litellm_initialize = value + + async def _execute_command( + self, + target_node: _ClusterNodeAttrs, + *args: object, + **kwargs: object, # kwargs-ok: matches redis-py's own command dispatch signature + ) -> object: + outstanding_before: Final = self._litellm_reinit_requests - self._litellm_tolerated_timeouts + pending_before: Final = self._litellm_initialize + try: + result: Final = await super()._execute_command(target_node, *args, **kwargs) + except _RedisTimeoutError: + timeouts: Final = self._litellm_consecutive_timeouts.get(target_node.name, 0) + 1 + if timeouts >= _CONSECUTIVE_TIMEOUTS_BEFORE_REINIT: + self._litellm_consecutive_timeouts.pop(target_node.name, None) + raise + self._litellm_consecutive_timeouts[target_node.name] = timeouts + self._litellm_tolerated_timeouts += 1 + if ( + not pending_before + and self._litellm_reinit_requests - self._litellm_tolerated_timeouts == outstanding_before + ): + self._initialize = False + raise + if self._litellm_consecutive_timeouts: + self._litellm_consecutive_timeouts.pop(target_node.name, None) + return result + + return LiteLLMAsyncRedisClusterTimeoutTolerant if redis.__version__ not in _VERIFIED_REDIS_VERSIONS: verbose_logger.warning( diff --git a/litellm/constants.py b/litellm/constants.py index cd72adc3db5..da731cb5eb2 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -151,6 +151,7 @@ DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD = float(os.getenv("DEFAULT_SEMANTIC_ MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS: Final = int(os.getenv("MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS", "60")) MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE: Final = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE", "200")) MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL: Final = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL", "3600")) +MCP_SSO_ASSERTION_CACHE_TTL_SECONDS: Final = int(os.getenv("MCP_SSO_ASSERTION_CACHE_TTL_SECONDS", "60")) # Default npm cache directory for STDIO MCP servers. # npm/npx needs a writable cache dir; in containers the default (~/.npm) diff --git a/litellm/integrations/azure_storage/azure_storage.py b/litellm/integrations/azure_storage/azure_storage.py index 16ef6920114..13058bf4f22 100644 --- a/litellm/integrations/azure_storage/azure_storage.py +++ b/litellm/integrations/azure_storage/azure_storage.py @@ -36,7 +36,7 @@ AZURE_STORAGE_TOKEN_SCOPE: Final = "https://storage.azure.com/.default" def _cached_credential_chain_token_provider() -> Callable[[], str]: return get_azure_ad_token_provider( azure_scope=AZURE_STORAGE_TOKEN_SCOPE, - azure_credential=AzureCredentialType.DefaultAzureCredential, + azure_credential=AzureCredentialType.DeploymentIdentityCredential, ) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 51d6858b7dc..17a19f05fa3 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -10,7 +10,7 @@ import subprocess import sys import time import traceback -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence from datetime import datetime as dt_object from functools import lru_cache from types import MappingProxyType, TracebackType @@ -5878,14 +5878,28 @@ def _get_status_fields( ######################################################### # Map - guardrail_information.guardrail_status to guardrail_status ######################################################### - guardrail_status: GuardrailStatus = "not_run" - if guardrail_information and isinstance(guardrail_information, list): - for information in guardrail_information: - if isinstance(information, dict): - raw_status = information.get("guardrail_status", "not_run") - if raw_status != "not_run": - guardrail_status = GUARDRAIL_STATUS_MAP.get(raw_status, "not_run") - break + # Severity order, least severe first. The status aggregates across ALL + # guardrail entries rather than taking the first non-"not_run" one: a + # pre_call guardrail that passed (e.g. a mask) records its entry before a + # later guardrail's block, and first-wins would report a blocked request + # as "success". + GUARDRAIL_STATUS_SEVERITY: Final[tuple[GuardrailStatus, ...]] = ( + "not_run", + "success", + "guardrail_failed_to_respond", + "guardrail_intervened", + ) + entries: Final[Sequence[object]] = guardrail_information if isinstance(guardrail_information, list) else () + raw_statuses: Final[Iterator[object]] = ( + entry.get("guardrail_status", "not_run") for entry in entries if isinstance(entry, dict) + ) + # A guardrail is free to write any value here, and an unhashable one would + # raise TypeError on the mapping lookup and drop the whole payload. + guardrail_status: Final[GuardrailStatus] = max( + (GUARDRAIL_STATUS_MAP.get(raw_status, "not_run") for raw_status in raw_statuses if isinstance(raw_status, str)), + key=GUARDRAIL_STATUS_SEVERITY.index, + default="not_run", + ) return StandardLoggingPayloadStatusFields(llm_api_status=llm_api_status, guardrail_status=guardrail_status) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py index 9deff950724..242300c7b6d 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py @@ -6,6 +6,7 @@ from litellm.types.llms.anthropic import AnthropicMessagesRequestOptionalParams from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) +from litellm.types.llms.openai import ChatCompletionSystemMessage if TYPE_CHECKING: from litellm.exceptions import ContentPolicyViolationError @@ -36,6 +37,16 @@ def safeguard_refusal_error(model: str, stop_details: Mapping[str, object]) -> " ) +def anthropic_system_to_openai_message(system: object) -> ChatCompletionSystemMessage | None: + """ + Return the Anthropic Messages top-level ``system`` (a string or a list of text + blocks) as an OpenAI-style system message, or None when the request has none. + """ + if not isinstance(system, (str, list)) or not system: + return None + return ChatCompletionSystemMessage(role="system", content=system) + + @lru_cache(maxsize=1) def _anthropic_messages_optional_param_keys() -> frozenset[str]: """ diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index 0968185b084..c64fc583edc 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/litellm/llms/snowflake/chat/transformation.py @@ -8,23 +8,31 @@ Routes to native Cortex REST API endpoints based on model: Ref: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-rest-api """ +import copy import json +import re from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict import httpx from typing_extensions import ReadOnly -from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk +from litellm.litellm_core_utils.prompt_templates.factory import ( + anthropic_process_openai_file_message, + convert_to_anthropic_tool_result, + create_anthropic_image_param, + select_anthropic_content_block_type_for_file, +) +from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicStreamParser +from litellm.llms.anthropic.chat.transformation import AnthropicConfig +from litellm.llms.anthropic.common_utils import normalize_cache_control_in_anthropic_payload +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk, ChatCompletionToolMessage from litellm.types.utils import ( - ChatCompletionMessageToolCall, - ChatCompletionUsageBlock, Choices, - Function, GenericStreamingChunk, Message, ModelResponse, - Usage, + ModelResponseStream, ) from ...base_llm.base_model_iterator import BaseModelResponseIterator @@ -93,6 +101,103 @@ def _is_claude_model(model: str) -> bool: return any(name.startswith(p) for p in _CLAUDE_MODEL_PREFIXES) +def _convert_image_url_to_anthropic(block: Mapping[str, object]) -> object: + """One OpenAI ``image_url`` block in the native shape Cortex accepts. + + Cortex documents base64 sources only, so remote URLs are inlined the way every + other base64-only Anthropic dialect (Bedrock invoke, Vertex) inlines them, and + pdf/text data URIs become document blocks rather than malformed image blocks. + """ + image_url: Final = block.get("image_url") + url: Final = image_url if isinstance(image_url, str) else _image_url_field(image_url, "url") + if not url: + return block + + converted: Final = ( + anthropic_process_openai_file_message({"type": "file", "file": {"file_data": url}}) + if select_anthropic_content_block_type_for_file(_data_uri_media_type(url)) == "document" + else create_anthropic_image_param( + image_url if isinstance(image_url, dict) else url, # mutable-ok: caller's JSON block + format=_image_url_field(image_url, "format"), + is_bedrock_invoke=True, + ) + ) + cache_control: Final = block.get("cache_control") + if cache_control is None: + return converted + return {**converted, "cache_control": cache_control} # mutable-ok: JSON wire block + + +def _image_url_field(image_url: object, key: str) -> str | None: + value: Final = image_url.get(key) if isinstance(image_url, dict) else None + return value if isinstance(value, str) else None + + +def _data_uri_media_type(url: str) -> str: + match: Final = re.match(r"data:([^;,]+)", url) + return match.group(1) if match else "" + + +def _convert_image_url_blocks_to_anthropic(content: object) -> object: + if not isinstance(content, list): + return content + return [ # mutable-ok: JSON wire blocks + _convert_image_url_to_anthropic(block) + if isinstance(block, Mapping) and block.get("type") == "image_url" + else block + for block in content + ] + + +def _convert_tool_result_to_anthropic( + content: object, tool_call_id: str, cache_control: object +) -> Mapping[str, object]: + """The Anthropic ``tool_result`` block for one OpenAI tool message. + + Delegating to the shared converter keeps image, document and per-block cache + breakpoints identical to every other Anthropic dialect; only the plain-string + and non-list shapes it does not model are handled here. + """ + if not isinstance(content, list): + plain: Final[dict[str, object]] = { # mutable-ok: JSON wire block + "type": "tool_result", + "tool_use_id": tool_call_id, + "content": content if isinstance(content, str) else json.dumps(content), + } + return {**plain, "cache_control": cache_control} if cache_control is not None else plain + converted: Final = convert_to_anthropic_tool_result( + ChatCompletionToolMessage(role="tool", tool_call_id=tool_call_id, content=content), + force_base64=True, + ) + if cache_control is None: + return converted + return {**converted, "cache_control": cache_control} # mutable-ok: JSON wire block + + +def _signed_thinking_blocks(msg: object) -> list[dict[str, object]]: # mutable-ok: JSON wire blocks + """The assistant turn's thinking blocks that can legally be echoed back. + + Only signed blocks round-trip: Cortex rejects a thinking block whose signature is + missing, which is what an unsigned block from a non-thinking turn would produce. + """ + blocks: Final = msg.get("thinking_blocks") if isinstance(msg, dict) else getattr(msg, "thinking_blocks", None) + if not isinstance(blocks, list): + return [] # mutable-ok: JSON wire blocks + return [ # mutable-ok: JSON wire blocks + dict(block) + for block in blocks + if isinstance(block, Mapping) and (block.get("signature") or block.get("type") == "redacted_thinking") + ] + + +def _clean_input_schema(schema: object) -> object: # mutable-ok: JSON schema copy + return ( + {key: value for key, value in schema.items() if key != "$schema"} + if isinstance(schema, Mapping) + else schema # mutable-ok: JSON schema copy + ) # mutable-ok: JSON schema copy + + class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): """ Snowflake Cortex REST API — unified provider. @@ -178,7 +283,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): if "description" in func: anthropic_tool["description"] = func["description"] if "parameters" in func: - anthropic_tool["input_schema"] = func["parameters"] + anthropic_tool["input_schema"] = _clean_input_schema(func["parameters"]) else: anthropic_tool["input_schema"] = { "type": "object", @@ -186,10 +291,16 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): } anthropic_tools.append(anthropic_tool) else: - anthropic_tools.append(tool) + anthropic_tools.append( + {**tool, "input_schema": _clean_input_schema(tool["input_schema"])} # mutable-ok: JSON wire tool + if "input_schema" in tool + else tool + ) return anthropic_tools - def _extract_system_and_messages(self, messages: list[AllMessageValues]) -> tuple[str | None, list[dict]]: + def _extract_system_and_messages( # mutable-ok: JSON wire messages + self, messages: list[AllMessageValues] + ) -> tuple[list[dict] | None, list[dict]]: """ Split messages into system prompt and conversation turns for Anthropic format. @@ -197,26 +308,39 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): - assistant messages with tool_calls → tool_use content blocks - tool role messages → user role with tool_result content blocks """ - system_parts: Final[list[str]] = [] - conversation: Final[list[dict]] = [] + system_parts: Final[list[dict]] = [] # mutable-ok: JSON wire messages + conversation: Final[list[dict]] = [] # mutable-ok: JSON wire messages for msg in messages: if isinstance(msg, dict): role = msg.get("role", "") content: Any = msg.get("content", "") + msg_cache_control: object = msg.get("cache_control") else: role = getattr(msg, "role", "") content = getattr(msg, "content", "") + msg_cache_control = getattr(msg, "cache_control", None) if role == "system": if isinstance(content, str) and content: - system_parts.append(content) + system_parts.append({"type": "text", "text": content}) # mutable-ok: JSON wire system block elif isinstance(content, list): - system_parts.append("\n".join(b.get("text", "") for b in content if b.get("type") == "text")) + system_parts.extend( + { # mutable-ok: JSON wire system block + "type": "text", + "text": block.get("text", ""), + **( + {"cache_control": block["cache_control"]} if "cache_control" in block else {} + ), # mutable-ok: JSON wire block + } + for block in content + if isinstance(block, Mapping) and block.get("type") == "text" + ) elif role == "assistant": tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else getattr(msg, "tool_calls", None) + thinking_blocks = _signed_thinking_blocks(msg) if tool_calls: - content_blocks: list[dict[str, object]] = [] + content_blocks: list[dict[str, object]] = list(thinking_blocks) # mutable-ok: JSON wire blocks if content: content_blocks.append({"type": "text", "text": content}) for tc in tool_calls: @@ -239,18 +363,26 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): } ) conversation.append({"role": "assistant", "content": content_blocks}) + elif thinking_blocks: + thinking_content = ( + [ + *thinking_blocks, + *copy.deepcopy(content), + ] + if isinstance(content, list) + else [*thinking_blocks, *([{"type": "text", "text": content}] if content else [])] + ) # rebind-ok: loop-local normalized content + conversation.append({"role": "assistant", "content": thinking_content}) else: conversation.append({"role": "assistant", "content": content}) elif role == "tool": - tool_call_id = ( + tool_call_id_value = ( msg.get("tool_call_id", "") if isinstance(msg, dict) else getattr(msg, "tool_call_id", "") ) - tool_content = content if isinstance(content, str) else json.dumps(content) - tool_result_block = { - "type": "tool_result", - "tool_use_id": tool_call_id, - "content": tool_content, - } + tool_call_id = ( + tool_call_id_value if isinstance(tool_call_id_value, str) else "" + ) # rebind-ok: normalized loop value + tool_result_block = _convert_tool_result_to_anthropic(content, tool_call_id, msg_cache_control) if ( conversation and conversation[-1]["role"] == "user" @@ -260,11 +392,18 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): ): conversation[-1]["content"].append(tool_result_block) else: - conversation.append({"role": "user", "content": [tool_result_block]}) + conversation.append( + {"role": "user", "content": [tool_result_block]} # mutable-ok: JSON wire message + ) # mutable-ok: JSON wire message else: - conversation.append({"role": role, "content": content}) + conversation.append( # mutable-ok: JSON wire message + { # mutable-ok: JSON wire message + "role": role, + "content": _convert_image_url_blocks_to_anthropic(content), + } # mutable-ok: JSON wire message + ) - system: Final[str | None] = "\n\n".join(system_parts) if system_parts else None + system: Final[list[dict] | None] = system_parts if system_parts else None # mutable-ok: JSON wire messages return system, conversation def transform_request( @@ -339,7 +478,9 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): extra_body: dict, ) -> dict: """Anthropic Messages format for /messages endpoint.""" - system, conversation = self._extract_system_and_messages(messages) + passthrough_system: Final = optional_params.pop("system", None) + extracted_system, conversation = self._extract_system_and_messages(messages) + system: Final = passthrough_system if passthrough_system is not None else extracted_system if "tools" in optional_params: optional_params["tools"] = self._transform_tools_to_anthropic(optional_params["tools"]) @@ -353,16 +494,19 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): model_name: Final = model.removeprefix("snowflake/") - body: Final[dict[str, object]] = { - "model": model_name, - "messages": conversation, - "stream": stream, - **optional_params, - **extra_body, - } - + body: Final[dict[str, object]] = normalize_cache_control_in_anthropic_payload( # mutable-ok: JSON wire body + { # mutable-ok: JSON wire body + "model": model_name, + "messages": conversation, + "stream": stream, + **optional_params, + **extra_body, # mutable-ok: JSON wire body + } + ) if system is not None: - body["system"] = system + body["system"] = normalize_cache_control_in_anthropic_payload( # mutable-ok: JSON wire payload + {"system": system} # mutable-ok: JSON wire payload + )["system"] if "max_tokens" not in body: body["max_tokens"] = 4096 # reasonable default; Anthropic API max varies by model @@ -435,23 +579,10 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): additional_args={"complete_input_dict": request_data}, ) - text_content = "" - tool_calls: Final = [] - - for block in response_json.get("content", []): - if block.get("type") == "text": - text_content += block.get("text", "") - elif block.get("type") == "tool_use": - tool_calls.append( - ChatCompletionMessageToolCall( - id=block.get("id", ""), - type="function", - function=Function( - name=block.get("name", ""), - arguments=json.dumps(block.get("input", {})), - ), - ) - ) + anthropic_config: Final = AnthropicConfig() + text_content, _, thinking_blocks, reasoning_content, tool_calls, _, _, _ = ( + anthropic_config.extract_response_content(completion_response=dict(response_json)) + ) _stop_reason_map: Final = { "end_turn": "stop", @@ -461,9 +592,13 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): } finish_reason: Final = _stop_reason_map.get(response_json.get("stop_reason", "end_turn"), "stop") - message: Final = Message(content=text_content or None, role="assistant") - if tool_calls: - message.tool_calls = tool_calls + message: Final = Message( + content=text_content or None, + role="assistant", + tool_calls=tool_calls or None, + thinking_blocks=thinking_blocks, + reasoning_content=reasoning_content, + ) choice: Final = Choices( finish_reason=finish_reason, @@ -471,11 +606,13 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): message=message, ) - usage_data: Final = response_json.get("usage", {}) - usage: Final = Usage( - prompt_tokens=usage_data.get("input_tokens", 0), - completion_tokens=usage_data.get("output_tokens", 0), - total_tokens=usage_data.get("input_tokens", 0) + usage_data.get("output_tokens", 0), + # Cortex reports prompt-cache creation/read counts alongside input_tokens; the + # shared calculator folds them into prompt_tokens_details so cached input is + # visible and billed at its own rate. + usage: Final = anthropic_config.calculate_usage( + usage_object=response_json.get("usage", {}), + reasoning_content=reasoning_content, + completion_response=dict(response_json), ) model_response.choices = [choice] @@ -516,15 +653,19 @@ class SnowflakeStreamingHandler(BaseModelResponseIterator): json_mode: bool | None = False, ): super().__init__(streaming_response=streaming_response, sync_stream=sync_stream) - self._tool_index = 0 - self._tool_id = "" - self._tool_name = "" - self._input_tokens = 0 + # Cortex streams the Anthropic SSE dialect on /messages, so its events are parsed + # by Anthropic's own parser: thinking deltas, signatures and prompt-cache usage + # all arrive the way they do on every other Anthropic-dialect provider. + self._anthropic_parser: Final = AnthropicStreamParser( + streaming_response=streaming_response, + sync_stream=sync_stream, + json_mode=json_mode, + ) - def chunk_parser(self, chunk: dict) -> GenericStreamingChunk: + def chunk_parser(self, chunk: dict) -> GenericStreamingChunk | ModelResponseStream: if "choices" in chunk: return self._parse_openai_chunk(chunk) - return self._parse_anthropic_chunk(chunk) + return self._anthropic_parser.chunk_parser(chunk) def _parse_openai_chunk(self, chunk: dict) -> GenericStreamingChunk: choices: Final = chunk.get("choices", []) @@ -566,117 +707,3 @@ class SnowflakeStreamingHandler(BaseModelResponseIterator): index=choice.get("index", 0), tool_use=tool_use, ) - - def _parse_anthropic_chunk(self, chunk: dict) -> GenericStreamingChunk: - event_type: Final = chunk.get("type", "") - - if event_type == "message_start": - message: Final = chunk.get("message", {}) - usage_data = message.get("usage", {}) - self._input_tokens = usage_data.get("input_tokens", 0) - return GenericStreamingChunk( - text="", - is_finished=False, - finish_reason="", - usage=None, - index=0, - tool_use=None, - ) - - elif event_type == "content_block_delta": - delta = chunk.get("delta", {}) - delta_type: Final = delta.get("type", "") - - if delta_type == "text_delta": - return GenericStreamingChunk( - text=delta.get("text", ""), - is_finished=False, - finish_reason="", - usage=None, - index=chunk.get("index", 0), - tool_use=None, - ) - elif delta_type == "input_json_delta": - return GenericStreamingChunk( - text="", - is_finished=False, - finish_reason="", - usage=None, - index=chunk.get("index", 0), - tool_use=ChatCompletionToolCallChunk( - id=self._tool_id, - type="function", - function={ - "name": self._tool_name, - "arguments": delta.get("partial_json", ""), - }, - index=self._tool_index, - ), - ) - - elif event_type == "content_block_start": - content_block: Final = chunk.get("content_block", {}) - if content_block.get("type") == "tool_use": - self._tool_id = content_block.get("id", "") - self._tool_name = content_block.get("name", "") - self._tool_index = chunk.get("index", 0) - return GenericStreamingChunk( - text="", - is_finished=False, - finish_reason="", - usage=None, - index=chunk.get("index", 0), - tool_use=ChatCompletionToolCallChunk( - id=self._tool_id, - type="function", - function={"name": self._tool_name, "arguments": ""}, - index=self._tool_index, - ), - ) - - elif event_type == "message_delta": - delta = chunk.get("delta", {}) - stop_reason: Final = delta.get("stop_reason", "") - usage_data = chunk.get("usage", {}) - _stop_map: Final = { - "end_turn": "stop", - "max_tokens": "length", - "tool_use": "tool_calls", - "stop_sequence": "stop", - } - usage = None - if usage_data or self._input_tokens: - output_t: Final = usage_data.get("output_tokens", 0) - input_t: Final = self._input_tokens or usage_data.get("input_tokens", 0) - usage = ChatCompletionUsageBlock( - prompt_tokens=input_t, - completion_tokens=output_t, - total_tokens=input_t + output_t, - ) - return GenericStreamingChunk( - text="", - is_finished=True, - finish_reason=_stop_map.get(stop_reason, "stop"), - usage=usage, - index=0, - tool_use=None, - ) - - elif event_type == "message_stop": - return GenericStreamingChunk( - text="", - is_finished=True, - finish_reason="stop", - usage=None, - index=0, - tool_use=None, - ) - - return GenericStreamingChunk( - text="", - is_finished=False, - finish_reason="", - usage=None, - index=0, - tool_use=None, - ) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 0c3932d2ba1..ce4928ff83d 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -3638,30 +3638,48 @@ class MCPServerManager: user_api_key_auth: UserAPIKeyAuth | None, raw_headers: Mapping[str, str] | None = None, ) -> None: - """Run the OBO exchange for a caller-supplied subject at the transport edge. + """Mint an exchange-backed server's upstream credential at the transport edge. Single-server routes call this before the MCP session opens, where an HTTP status and ``WWW-Authenticate`` still reach the client. A rejected subject raises the RFC 9728 challenge and any other ``CredError`` maps onto its public HTTP status, so an exchange failure surfaces as a failure instead of the session continuing into an empty tool list. A successful exchange is cached by the exchanger, so the session's list/call reuses it. + + Each mode pre-flights only where it would resolve the subject the session goes on to use, + which is what keeps the pre-flight from reaching a verdict the session would contradict. + ``oauth2_token_exchange`` mints from the caller's inbound bearer, so without one there is + nothing to exchange and the missing-subject case stays the preemptive challenge's job. + ``oauth2_id_jag`` is the mirror image: tool listing resolves it from the identity assertion + captured for this user at SSO login and never from the inbound bearer, so the pre-flight is + faithful exactly when no identity bearer was sent (a LiteLLM key in ``Authorization`` is not one), + and a caller that did send one is passed through + untouched rather than judged against a subject the listing will not use. That store-sourced + case is the one whose missing-assertion 412 and store-outage 503 the session cannot report. + Only OBO has a discovery challenge to raise; ID-JAG's failures are plain statuses whose body + already names what the user has to do, so they map through ``raise_public`` as at egress. """ - if server.auth_type != MCPAuth.oauth2_token_exchange: - return - if not self._extract_bearer_token(oauth2_headers, None): - return - resolved_server: Final = await self.ensure_oauth_metadata_discovered(server) - spec: Final = to_server_spec(resolved_server) - if spec is None or not isinstance(spec.config, TokenExchangeConfig): - return subject_token: Final = self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth) - if subject_token is None: + match server.auth_type: + case MCPAuth.oauth2_token_exchange: + if not self._extract_bearer_token(oauth2_headers, None): + return + case MCPAuth.oauth2_id_jag: + if subject_token is not None: + return + case _: + return + resolved_server: Final = await self.ensure_oauth_metadata_discovered(server) + spec: Final = _to_server_spec_fail_closed(resolved_server) + if spec is None or not isinstance(spec.config, (TokenExchangeConfig, IdJagConfig)): + return + if subject_token is None and isinstance(spec.config, TokenExchangeConfig): raise_token_exchange_challenge(resolved_server, root_path=get_server_root_path()) match await self._cred_provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec): case Ok(_): return case Error(err): - if err.tag == "unauthorized": + if err.tag == "unauthorized" and isinstance(spec.config, TokenExchangeConfig): raise_token_exchange_challenge( resolved_server, root_path=get_server_root_path(), diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py index 8a0d41584f9..6552008ca54 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py @@ -11,7 +11,8 @@ being registered, so a gateway with no EMA upstream never stores bearer material The row is one encrypted payload per user, latest login wins. ``expires_at`` mirrors the id_token ``exp`` claim and is judged by the reader, never enforced by deletion here: an expired assertion with a refresh token is still renewable, and the DB row is the source of -truth, the same contract as the per-user OAuth credential store. +truth, the same contract as the per-user OAuth credential store. Reads use a per-process cache with +TTL ``MCP_SSO_ASSERTION_CACHE_TTL_SECONDS``; invalidation also guards against stale in-flight reads. """ from __future__ import annotations @@ -24,6 +25,8 @@ import jwt from pydantic import BaseModel, ConfigDict, SecretStr, TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.constants import MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE, MCP_SSO_ASSERTION_CACHE_TTL_SECONDS if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient @@ -45,6 +48,46 @@ class SSOIdentityAssertion(BaseModel): expires_at: datetime | None = None +class SSOAssertionCache: + """Process-local read cache. ``invalidate`` bumps a process-wide epoch so a fetch that started + before a login cannot repopulate the old assertion after it.""" + + def __init__(self, ttl_seconds: int = MCP_SSO_ASSERTION_CACHE_TTL_SECONDS) -> None: + self._entries = InMemoryCache( + max_size_in_memory=MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE, + default_ttl=ttl_seconds, + ) + self._epoch: int = 0 + + def epoch(self) -> int: + return self._epoch + + def get(self, user_id: str) -> SSOIdentityAssertion | None: + cached: Final = self._entries.get_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped + user_id + ) + return cached if isinstance(cached, SSOIdentityAssertion) else None + + def set_if_unchanged(self, user_id: str, assertion: SSOIdentityAssertion, seen_epoch: int) -> None: + if self._epoch != seen_epoch: + return + self._entries.set_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped + user_id, assertion + ) + + def invalidate(self, user_id: str) -> None: + self._epoch += 1 + self._entries.delete_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped + user_id + ) + + def flush(self) -> None: + self._entries.flush_cache() # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped + + +_ASSERTION_CACHE: Final = SSOAssertionCache() + + class _IdTokenClaims(BaseModel): exp: float | None = None iss: str | None = None @@ -107,7 +150,9 @@ async def ema_assertion_retention_enabled() -> bool: return row is not None -async def persist_sso_identity_assertion(user_id: str, assertion: SSOIdentityAssertion) -> None: +async def persist_sso_identity_assertion( + user_id: str, assertion: SSOIdentityAssertion, cache: SSOAssertionCache = _ASSERTION_CACHE +) -> None: from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper # noqa: PLC0415 # runtime global from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # runtime global @@ -127,11 +172,10 @@ async def persist_sso_identity_assertion(user_id: str, assertion: SSOIdentityAss "update": {"assertion_b64": encoded}, }, ) + cache.invalidate(user_id) -async def fetch_sso_identity_assertion(user_id: str) -> SSOIdentityAssertion | None: - """The stored assertion for ``user_id``, or ``None`` when absent, undecryptable (salt-key - rotation), or unparseable. Expiry is not judged here; the reader owns that policy.""" +async def _read_assertion_from_db(user_id: str) -> SSOIdentityAssertion | None: from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper # noqa: PLC0415 # runtime global from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # runtime global @@ -160,6 +204,21 @@ async def fetch_sso_identity_assertion(user_id: str) -> SSOIdentityAssertion | N ) +async def fetch_sso_identity_assertion( + user_id: str, cache: SSOAssertionCache = _ASSERTION_CACHE +) -> SSOIdentityAssertion | None: + """The stored assertion for ``user_id``, or ``None`` when absent, undecryptable (salt-key + rotation), or unparseable. Expiry is not judged here; the reader owns that policy.""" + cached: Final = cache.get(user_id) + if cached is not None: + return cached + seen_epoch: Final = cache.epoch() + assertion: Final = await _read_assertion_from_db(user_id) + if assertion is not None: + cache.set_if_unchanged(user_id, assertion, seen_epoch) + return assertion + + class AssertionStoreUnavailable(Exception): """Raised by ``fetch`` when the backing store is unreachable (e.g. the DB is down). @@ -189,9 +248,12 @@ class DbSSOAssertionStore: from credential resolution and from the upstream-401 retry. """ + def __init__(self, cache: SSOAssertionCache = _ASSERTION_CACHE) -> None: + self._cache = cache + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: try: - return await fetch_sso_identity_assertion(user_id) + return await fetch_sso_identity_assertion(user_id, cache=self._cache) except Exception as exc: # noqa: BLE001 # any driver/storage failure is an outage, not an absence raise AssertionStoreUnavailable(str(exc)) from exc diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index bb075220530..3d7c947a913 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -3851,15 +3851,15 @@ if MCP_AVAILABLE: raise_token_exchange_challenge(server, root_path=get_server_root_path()) - # token_exchange (OBO) with a subject present: run the exchange here at the transport - # edge, so a rejected subject raises the RFC 9728 challenge (and a gateway fault its - # public status) instead of the session opening and list_tools masking the failure as - # an empty tool list. Gated to single-server routes; the multi-server aggregate keeps - # absorbing per-server auth failures so one bad server cannot 401 the whole connect. + # Exchange-backed modes (token_exchange's OBO mint, id_jag's stored-assertion mint): run + # the exchange here at the transport edge, so a rejected subject raises the RFC 9728 + # challenge and any other failure its public status, instead of the session opening and + # list_tools masking it as an empty tool list. The manager owns which modes pre-flight + # and what each mints from. Gated to single-server routes the key may reach; the + # multi-server aggregate keeps absorbing per-server auth failures so one bad server + # cannot 401 the whole connect. if ( server - and server.auth_type == MCPAuth.oauth2_token_exchange - and oauth2_headers and len(mcp_servers or []) == 1 and server.server_id in frozenset( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5d5a25e7cd6..832d941f5b5 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -737,6 +737,9 @@ class LiteLLMRoutes(enum.Enum): "/.well-known/litellm-ui-config", "/public/model_hub", "/public/v1/model_hub", + "/public/v1/model_hub/providers", + "/public/v1/model_hub/modes", + "/public/v1/model_hub/features", "/public/model_hub/info", "/public/agent_hub", "/public/mcp_hub", @@ -2404,6 +2407,15 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): """ completion_model: str | None = Field(None, description="proxy level default model for all chat completion calls") + max_in_flight_requests_per_worker: int | None = Field( + None, gt=0, description="maximum concurrent requests handled by each worker" + ) + max_queued_requests_per_worker: int | None = Field( + None, ge=0, description="maximum requests waiting for a worker slot" + ) + admission_queue_timeout_seconds: float = Field( + 1.0, gt=0, description="maximum time a request waits for a worker slot" + ) plugins: list[PluginConfig] | None = Field( None, description="external services registered as embeddable UI plugins" ) diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 7f0045c1d93..b243b737b0a 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -9,7 +9,7 @@ from fastapi.responses import JSONResponse import litellm from litellm._logging import verbose_proxy_logger -from litellm.anthropic_interface.exceptions import AnthropicExceptionMapping +from litellm.anthropic_interface.exceptions import AnthropicErrorResponse, AnthropicExceptionMapping from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.llms.anthropic.experimental_pass_through.context_management import ( AnthropicContextManagementError, @@ -30,6 +30,27 @@ from litellm.types.utils import TokenCountResponse router: Final = APIRouter() +def _anthropic_error_json_response(exc: ProxyException, request: Request) -> JSONResponse: + from litellm.proxy.proxy_server import ( + _close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; error JSONResponses returned by the route must stamp the OTel server span like the global ProxyException handler does + ) + + status_code: Final = int(exc.code) if exc.code is not None and exc.code.isdigit() else 500 + _close_dangling_otel_server_span(request, status_code, exc=exc) + envelope: Final = AnthropicExceptionMapping.transform_to_anthropic_error( + status_code=status_code, + raw_message=exc.message, + request_id=request.headers.get("x-request-id"), + ) + if not exc.provider_specific_fields: + return JSONResponse(status_code=status_code, content=envelope, headers=exc.headers) + content: Final[AnthropicErrorResponse] = { + **envelope, + "error": {**envelope["error"], "provider_specific_fields": exc.provider_specific_fields}, + } + return JSONResponse(status_code=status_code, content=content, headers=exc.headers) + + def _strip_total_tokens_from_anthropic_response(response: Any) -> None: """Remove the OpenAI-flavored `usage.total_tokens` field that LiteLLM injects into Anthropic /v1/messages responses. @@ -195,7 +216,7 @@ async def anthropic_response( verbose_proxy_logger.exception("litellm.proxy.proxy_server.anthropic_response(): Exception occured - %s", e) if isinstance(e, ProxyException): - raise + return _anthropic_error_json_response(e, request) # Extract model_id from request metadata (same as success path) litellm_metadata: Final = data.get("litellm_metadata", {}) or {} @@ -216,15 +237,18 @@ async def anthropic_response( ) if isinstance(e, HTTPException): - raise proxy_exception_from_http_exception(e, headers) + return _anthropic_error_json_response(proxy_exception_from_http_exception(e, headers), request) error_msg: Final = f"{e}" - raise ProxyException( - message=getattr(e, "message", error_msg), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), - headers=headers, + return _anthropic_error_json_response( + ProxyException( + message=getattr(e, "message", error_msg), + type=getattr(e, "type", "None"), + param=getattr(e, "param", "None"), + code=getattr(e, "status_code", 500), + headers=headers, + ), + request, ) diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index 1625198892f..de2ca4762f1 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -387,33 +387,22 @@ def _get_wildcard_models( all_wildcard_models: Final = [] for model in unique_models: if _check_wildcard_routing(model=model): - if return_wildcard_routes: # will add the wildcard route to the list eg: anthropic/*. + if return_wildcard_routes: all_wildcard_models.append(model) - ## get litellm params from model - if llm_router is not None: - model_list = llm_router.get_model_list(model_name=model, team_id=team_id) - if model_list: - for router_model in model_list: - wildcard_models = get_known_models_from_wildcard( + models_to_remove.add(model) + + model_list = llm_router.get_model_list(model_name=model, team_id=team_id) if llm_router else None + if model_list: + for router_model in model_list: + all_wildcard_models.extend( + get_known_models_from_wildcard( wildcard_model=model, litellm_params=LiteLLM_Params(**router_model["litellm_params"]), ) - all_wildcard_models.extend(wildcard_models) - else: - # Router has no deployment for this wildcard (e.g., BYOK team models) - # Fall back to expanding from known provider models - wildcard_models = get_known_models_from_wildcard(wildcard_model=model, litellm_params=None) - if wildcard_models: - models_to_remove.add(model) - all_wildcard_models.extend(wildcard_models) + ) else: - # get all known provider models - wildcard_models = get_known_models_from_wildcard(wildcard_model=model, litellm_params=None) - - if wildcard_models: - models_to_remove.add(model) - all_wildcard_models.extend(wildcard_models) + all_wildcard_models.extend(get_known_models_from_wildcard(wildcard_model=model, litellm_params=None)) for model in models_to_remove: unique_models.remove(model) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 8b57bdca2fe..65d0ec8c0dc 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -50,6 +50,9 @@ from litellm.proxy.health_check import ( perform_health_check, run_with_timeout, ) +from litellm.proxy.middleware.admission_control_middleware import ( + get_admission_control_stats, +) from litellm.proxy.middleware.in_flight_requests_middleware import ( get_in_flight_requests, ) @@ -63,6 +66,13 @@ from litellm.secret_managers.main import get_secret_bool #### Health ENDPOINTS #### +class _HealthBacklogResponse(TypedDict): + in_flight_requests: ReadOnly[int] + admitted_requests: ReadOnly[int] + queued_requests: ReadOnly[int] + rejected_requests: ReadOnly[int] + + def _reject_os_environ_references(params: dict) -> None: """ Validate that the provided params do not contain any ``os.environ/`` @@ -1759,7 +1769,14 @@ async def health_backlog(): for the event loop to get to them, adding latency before LiteLLM even starts its own timer. """ - return {"in_flight_requests": get_in_flight_requests()} + stats: Final = get_admission_control_stats() + response: Final[_HealthBacklogResponse] = { + "in_flight_requests": get_in_flight_requests(), + "admitted_requests": stats.admitted, + "queued_requests": stats.queued, + "rejected_requests": stats.rejected_total, + } + return response @router.get( diff --git a/litellm/proxy/list_api/in_memory.py b/litellm/proxy/list_api/in_memory.py index bada8ea0a35..3f3f173ba1a 100644 --- a/litellm/proxy/list_api/in_memory.py +++ b/litellm/proxy/list_api/in_memory.py @@ -141,3 +141,15 @@ class InMemoryListExecutor(Generic[TRow]): async def find_many(self, plan: QueryPlan) -> Sequence[TRow]: page: Final = _ordered(self._matching(plan.where), plan.order)[plan.skip : plan.skip + plan.take] return await self.enrich_page(tuple(row for _, row in page)) + + async def distinct(self, field: str, where: tuple[Predicate, ...]) -> Sequence[str]: + """A repeated field contributes each of its elements, so a facet over `providers` + lists providers rather than the tuples rows happen to carry.""" + cells: Final = (cells.get(field) for cells, _ in self._matching(where)) + values: Final = ( + value + for cell in cells + for value in (cell if isinstance(cell, tuple) else (cell,)) + if isinstance(value, str) and value + ) + return tuple(sorted(frozenset(values))) diff --git a/litellm/proxy/list_api/list_framework.py b/litellm/proxy/list_api/list_framework.py index 21ee4e6860f..af422b7650a 100644 --- a/litellm/proxy/list_api/list_framework.py +++ b/litellm/proxy/list_api/list_framework.py @@ -28,12 +28,15 @@ from litellm.proxy.list_api.common import ( PROBLEM_TYPE_BASE, ManagementProblem, build_list_links, + build_page_links, escape_like, unknown_query_param_problem, ) from litellm.types.proxy.management_endpoints.management_v1 import ( + FacetListResponse, ListMeta, ListResponse, + PageMeta, ProblemDetail, ) @@ -186,6 +189,13 @@ class ListExecutor(Protocol[TRow_co]): async def find_many(self, plan: QueryPlan) -> Sequence[TRow_co]: ... +class FacetExecutor(Protocol): + """The half of a facet that knows the rows. Separate from `ListExecutor` so a SQL + executor is not forced to implement `distinct` to keep serving entity lists.""" + + async def distinct(self, field: str, where: tuple[Predicate, ...]) -> Sequence[str]: ... + + def order_by_sql(order: tuple[SortKey, ...]) -> str: """`ORDER BY` body for a plan, NULLS LAST in both directions. @@ -515,6 +525,78 @@ def build_query_plan( ) +def _facet_allowed_params(spec: ListSpec[TRow, TOut]) -> tuple[str, ...]: + """A facet's values are always ascending, so `sort` is not one of its parameters.""" + return tuple(name for name in _allowed_params(spec) if name != SORT_PARAM) + + +def _facet_where( + spec: ListSpec[TRow, TOut], + params: Mapping[str, str], + caller: UserAPIKeyAuth, +) -> tuple[Predicate, ...] | ProblemDetail: + scope_predicates: Final = _scope_predicates(spec.scope(caller)) + if isinstance(scope_predicates, ProblemDetail): + return scope_predicates + filters: Final = _parse_filters(spec, params) + if isinstance(filters, ProblemDetail): + return filters + search: Final = _search_predicate(spec, params) + return scope_predicates + filters + ((search,) if search is not None else ()) + + +async def handle_facet( + spec: ListSpec[TRow, TOut], + executor: FacetExecutor, + request: Request, + caller: UserAPIKeyAuth, + field: str, +) -> FacetListResponse: + """The distinct values one column takes over a filtered query on a resource. + + Carries the parent's parameters so a filter dropdown offers exactly the values the + table can show, and `has_more` rather than a total, which would cost a COUNT(*) over + the whole match set on every keystroke. + """ + params: Final = request.query_params + unknown: Final = tuple(sorted(name for name in params if name == SORT_PARAM or not _is_known_param(spec, name))) + if unknown: + raise ManagementProblem(unknown_query_param_problem(unknown=unknown, allowed=_facet_allowed_params(spec))) + + duplicates: Final = _duplicate_params(request) + if duplicates: + raise ManagementProblem( + _problem( + "duplicate-query-parameter", + "Duplicate query parameter", + 400, + f"Repeated query parameter(s): {', '.join(duplicates)}. Each may appear once; " + f"use a comma-separated list for multiple filter values.", + ) + ) + + page: Final = _parse_page(params) + if isinstance(page, ProblemDetail): + raise ManagementProblem(page) + page_size: Final = _parse_page_size(spec, params) + if isinstance(page_size, ProblemDetail): + raise ManagementProblem(page_size) + + where: Final = _facet_where(spec, params, caller) + if isinstance(where, ProblemDetail): + raise ManagementProblem(where) + + values: Final = await executor.distinct(field, where) + skip: Final = (page - 1) * page_size + window: Final = values[skip : skip + page_size + 1] + has_more: Final = len(window) > page_size + return FacetListResponse( + data=tuple(window[:page_size]), + meta=PageMeta(page=page, page_size=page_size, has_more=has_more), + links=build_page_links(request=request, page=page, has_more=has_more), + ) + + def _duplicate_params(request: Request) -> tuple[str, ...]: names: Final = tuple(name for name, _ in request.query_params.multi_items()) return tuple(sorted(frozenset(name for name in names if names.count(name) > 1))) diff --git a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py index f58f3722741..69356922ea1 100644 --- a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py +++ b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py @@ -12,7 +12,7 @@ All /policy management endpoints import copy import json import os -from collections.abc import AsyncIterator +from collections.abc import AsyncGenerator, AsyncIterator from typing import TYPE_CHECKING, Final, Literal, cast from fastapi import APIRouter, Depends, HTTPException, Request @@ -20,6 +20,7 @@ from fastapi.responses import Response, StreamingResponse from pydantic import BaseModel, Field from typing_extensions import TypedDict +import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import ( COMPETITOR_LLM_TEMPERATURE, @@ -32,6 +33,10 @@ from litellm.llms.openai.chat.guardrail_translation.handler import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.sse_keepalive import ( + SSE_COMMENT_PING, + wrap_sse_stream_with_keepalive_pings, +) from litellm.proxy.guardrails.guardrail_hooks.custom_code import ( RESPONSE_REJECTION_GUARDRAIL_CODE, CustomCodeGuardrail, @@ -811,7 +816,7 @@ async def _stream_competitor_events( llm_enrichment: dict, brand_name: str, model: str, -) -> AsyncIterator[str]: +) -> AsyncGenerator[str, None]: """Stream competitor names as SSE events, then emit a final 'done' event.""" competitors: Final[list[str]] = list(data.competitors or []) @@ -883,7 +888,11 @@ async def enrich_policy_template_stream( model: Final = data.model or DEFAULT_COMPETITOR_DISCOVERY_MODEL return StreamingResponse( - _stream_competitor_events(data, template, llm_enrichment, brand_name, model), + wrap_sse_stream_with_keepalive_pings( + _stream_competitor_events(data, template, llm_enrichment, brand_name, model), + ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds, + ping_chunk=SSE_COMMENT_PING, + ), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, ) diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index 069f86c852c..ceb67e3eee8 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -1123,13 +1123,15 @@ async def _create_user_if_not_exists(user_id: str, created_via: str = "scim_grou user_id=user_id, user_email=user_id, # We don't have email from group membership user_alias=None, - teams=[], # Teams will be added separately metadata={"created_via": created_via}, auto_create_key=False, user_role=default_role, ) - created_user: Final = await new_user(data=new_user_request) + created_user: Final = await new_user( + data=new_user_request, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) verbose_proxy_logger.info("Created user %s via %s", user_id, created_via) return created_user @@ -1699,7 +1701,7 @@ async def create_user( user_id=user_id, user_email=user_data["user_email"], user_alias=user_data["user_alias"], - teams=user_data["teams"], + teams=user_data["teams"] or None, metadata=metadata, auto_create_key=False, user_role=resolved_role if admin_group is not None else default_role, @@ -1717,6 +1719,7 @@ async def create_user( created_user: Final = await new_user( data=new_user_request, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), ) scim_user: Final = await ScimTransformations.transform_litellm_user_to_scim_user(user=created_user) @@ -1771,22 +1774,25 @@ async def update_user( roles=user_data["roles"], ) + # SCIM User.groups is readOnly (RFC 7643 4.1.2): IdPs sync membership via /Groups and send + # no groups or `groups: []` on profile PUTs, so empty means unspecified, not "remove from every team" + target_teams: Final = user_data["teams"] or existing_user.teams await _handle_team_membership_changes( user_id=user_id, - existing_teams=existing_user.teams or [], - new_teams=user_data["teams"], + existing_teams=existing_user.teams, + new_teams=target_teams, ) update_data: Final = { "user_email": user_data["user_email"], "user_alias": user_data["user_alias"], "sso_user_id": user_data["sso_user_id"], - "teams": user_data["teams"], + "teams": target_teams, "metadata": safe_dumps(metadata), } admin_group: Final = await _get_scim_admin_group() - if admin_group is not None: + if admin_group is not None and user_data["teams"]: update_data["user_role"] = _resolve_scim_user_role( user.groups or [], admin_group, _default_scim_user_role() ) diff --git a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py index 9d5ddda017a..7ef496f94e3 100644 --- a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py +++ b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py @@ -4,7 +4,7 @@ usage/spend data by querying the aggregated daily activity endpoints. """ import json -from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, Sequence from datetime import date from typing import Any, Final, Literal, Protocol, cast, overload @@ -543,7 +543,7 @@ async def stream_usage_ai_chat( model: str | None = None, user_id: str | None = None, is_admin: bool = False, -) -> AsyncIterator[str]: +) -> AsyncGenerator[str, None]: """Stream SSE events: status → tool_call → chunk → done.""" resolved_model: Final = (model or "").strip() or DEFAULT_COMPETITOR_DISCOVERY_MODEL truncated: Final = messages[-MAX_CHAT_MESSAGES:] if len(messages) > MAX_CHAT_MESSAGES else messages diff --git a/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py b/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py index b7b0ae2d8e5..d1e92d0c7e4 100644 --- a/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py +++ b/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py @@ -10,8 +10,13 @@ from fastapi import APIRouter, Depends, Request from fastapi.responses import StreamingResponse from pydantic import BaseModel, Field +import litellm from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.sse_keepalive import ( + SSE_COMMENT_PING, + wrap_sse_stream_with_keepalive_pings, +) router: Final = APIRouter() @@ -56,11 +61,15 @@ async def usage_ai_chat( messages: Final = [{"role": m.role, "content": m.content} for m in data.messages] return StreamingResponse( - stream_usage_ai_chat( - messages=messages, - model=data.model, - user_id=user_id, - is_admin=is_admin, + wrap_sse_stream_with_keepalive_pings( + stream_usage_ai_chat( + messages=messages, + model=data.model, + user_id=user_id, + is_admin=is_admin, + ), + ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds, + ping_chunk=SSE_COMMENT_PING, ), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, diff --git a/litellm/proxy/middleware/admission_control_middleware.py b/litellm/proxy/middleware/admission_control_middleware.py new file mode 100644 index 00000000000..aa62ef9e3bf --- /dev/null +++ b/litellm/proxy/middleware/admission_control_middleware.py @@ -0,0 +1,315 @@ +import asyncio +import os +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from functools import lru_cache +from typing import Annotated, Final, Protocol, TypeAlias, runtime_checkable + +from pydantic import Field, TypeAdapter, ValidationError +from starlette.responses import JSONResponse +from starlette.types import ASGIApp, Receive, Scope, Send + +from litellm._logging import verbose_proxy_logger + +_EXEMPT_PATHS: Final[frozenset[str]] = frozenset( + { + "/health/liveliness", + "/health/liveness", + "/health/readiness", + "/health/readiness/details", + "/health/backlog", + "/health/drain", + "/metrics", + "/metrics/", + } +) + + +@dataclass(frozen=True, slots=True) +class AdmissionControlSettings: + max_in_flight_requests: int + max_queued_requests: int + queue_timeout_seconds: float + + +AdmissionControlSettingsGetter: TypeAlias = Callable[[], AdmissionControlSettings | None] # mutable-ok: Callable params + + +@dataclass(frozen=True, slots=True) +class AdmissionControlStats: + admitted: int + queued: int + rejected_total: int + + +@runtime_checkable +class _Gauge(Protocol): + def inc(self, amount: float = 1) -> None: ... + + def dec(self, amount: float = 1) -> None: ... + + +@runtime_checkable +class _CounterChild(Protocol): + def inc(self, amount: float = 1) -> None: ... + + +@runtime_checkable +class _Counter(Protocol): + def labels(self, reason: str) -> _CounterChild: ... + + +@dataclass(frozen=True, slots=True) +class AdmissionControlMetrics: + admitted_gauge: _Gauge + queued_gauge: _Gauge + rejected_counter: _Counter + + +AdmissionControlMetricsFactory: TypeAlias = Callable[[], AdmissionControlMetrics | None] # mutable-ok: Callable params + + +class AdmissionControlState: + """Per-process admission counters and the in-flight semaphore shared by one worker's requests.""" + + def __init__(self, metrics_factory: AdmissionControlMetricsFactory) -> None: + self._metrics_factory = metrics_factory + self._metrics: AdmissionControlMetrics | None = None + self._metrics_init_attempted = False + self._admitted = 0 + self._queued = 0 + self._rejected_total = 0 + self._semaphore: asyncio.Semaphore | None = None + self._semaphore_loop: asyncio.AbstractEventLoop | None = None + + def get_stats(self) -> AdmissionControlStats: + return AdmissionControlStats( + admitted=self._admitted, + queued=self._queued, + rejected_total=self._rejected_total, + ) + + def get_semaphore(self, max_in_flight_requests: int) -> asyncio.Semaphore: + loop: Final = asyncio.get_running_loop() + if self._semaphore_loop is not loop: + self._semaphore = asyncio.Semaphore(max_in_flight_requests) + self._semaphore_loop = loop + semaphore: Final = self._semaphore + if semaphore is None: + raise RuntimeError("Admission control semaphore was not initialized") + return semaphore + + def record_admission(self) -> None: + self._admitted += 1 + metrics: Final = self._get_metrics() + if metrics is not None: + metrics.admitted_gauge.inc() + + def record_release(self) -> None: + self._admitted -= 1 + metrics: Final = self._get_metrics() + if metrics is not None: + metrics.admitted_gauge.dec() + + def record_queue(self) -> None: + self._queued += 1 + metrics: Final = self._get_metrics() + if metrics is not None: + metrics.queued_gauge.inc() + + def record_dequeue(self) -> None: + self._queued -= 1 + metrics: Final = self._get_metrics() + if metrics is not None: + metrics.queued_gauge.dec() + + def record_rejection(self, reason: str) -> None: + self._rejected_total += 1 + metrics: Final = self._get_metrics() + if metrics is not None: + metrics.rejected_counter.labels(reason=reason).inc() + + def _get_metrics(self) -> AdmissionControlMetrics | None: + if not self._metrics_init_attempted: + self._metrics_init_attempted = True + self._metrics = self._metrics_factory() + return self._metrics + + +class AdmissionControlMiddleware: + def __init__( + self, + app: ASGIApp, + get_settings: AdmissionControlSettingsGetter, + state: AdmissionControlState, + ) -> None: + self.app = app + self.get_settings = get_settings + self.state = state + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] != "http": + await self.app(scope, receive, send) + return + + settings: Final = self.get_settings() + if settings is None or _get_route_path(scope) in _EXEMPT_PATHS: + await self.app(scope, receive, send) + return + + state: Final = self.state + semaphore: Final = state.get_semaphore(settings.max_in_flight_requests) + if not semaphore.locked(): + await semaphore.acquire() + state.record_admission() + elif state.get_stats().queued >= settings.max_queued_requests: + state.record_rejection("queue_full") + await _overloaded_response(state)(scope, receive, send) + return + else: + state.record_queue() + try: + await asyncio.wait_for( + semaphore.acquire(), + timeout=settings.queue_timeout_seconds, + ) + except asyncio.TimeoutError: + state.record_dequeue() + state.record_rejection("queue_timeout") + await _overloaded_response(state)(scope, receive, send) + return + except asyncio.CancelledError: + state.record_dequeue() + raise + state.record_dequeue() + state.record_admission() + + try: + await self.app(scope, receive, send) + finally: + semaphore.release() + state.record_release() + + +def _get_route_path(scope: Scope) -> str: + """Strip the ASGI root_path (SERVER_ROOT_PATH) the same way Starlette does before route matching.""" + path: Final[str] = scope["path"] + root_path: Final[str] = scope.get("root_path", "") + if not root_path or not path.startswith(root_path): + return path + if path == root_path: + return "" + if path[len(root_path)] == "/": + return path[len(root_path) :] + return path + + +def _create_gauge(gauge_type: Callable[..., object], name: str, description: str) -> _Gauge: + metric: Final = ( + gauge_type(name, description, multiprocess_mode="livesum") + if "PROMETHEUS_MULTIPROC_DIR" in os.environ + else gauge_type(name, description) + ) + if not isinstance(metric, _Gauge): + raise TypeError("Admission gauge has an unexpected type") + return metric + + +def create_prometheus_admission_metrics() -> AdmissionControlMetrics | None: + try: + from prometheus_client import Counter, Gauge + + return AdmissionControlMetrics( + admitted_gauge=_create_gauge( + Gauge, + "litellm_admission_admitted_requests", + "Number of requests admitted by this worker", + ), + queued_gauge=_create_gauge( + Gauge, + "litellm_admission_queued_requests", + "Number of requests queued by this worker", + ), + rejected_counter=Counter( # mutable-ok: Prometheus requires runtime Counter construction + "litellm_admission_rejected_requests_total", + "Number of requests rejected by this worker", + labelnames=("reason",), + ), + ) + except (ImportError, ValueError): + return None + + +admission_control_state: Final = AdmissionControlState(create_prometheus_admission_metrics) + + +def get_admission_control_stats() -> AdmissionControlStats: + return admission_control_state.get_stats() + + +_PositiveInt: TypeAlias = Annotated[int, Field(gt=0)] +_NonNegativeInt: TypeAlias = Annotated[int, Field(ge=0)] +_PositiveFloat: TypeAlias = Annotated[float, Field(gt=0)] +_AdmissionControlRaw: TypeAlias = int | float | str | None + + +def _hashable(value: object) -> _AdmissionControlRaw: + return value if value is None or isinstance(value, (int, float, str)) else repr(value) + + +_POSITIVE_INT_ADAPTER: Final[TypeAdapter[int]] = TypeAdapter(_PositiveInt) +_NON_NEGATIVE_INT_ADAPTER: Final[TypeAdapter[int]] = TypeAdapter(_NonNegativeInt) +_POSITIVE_FLOAT_ADAPTER: Final[TypeAdapter[float]] = TypeAdapter(_PositiveFloat) + + +@lru_cache(maxsize=16) +def _parse_admission_control_settings( + max_in_flight_raw: _AdmissionControlRaw, + max_queued_raw: _AdmissionControlRaw, + queue_timeout_raw: _AdmissionControlRaw, +) -> AdmissionControlSettings | None: + try: + max_in_flight: Final = _POSITIVE_INT_ADAPTER.validate_python(max_in_flight_raw) + max_queued: Final = ( + max_in_flight if max_queued_raw is None else _NON_NEGATIVE_INT_ADAPTER.validate_python(max_queued_raw) + ) + queue_timeout: Final = _POSITIVE_FLOAT_ADAPTER.validate_python(queue_timeout_raw) + except ValidationError as exc: + verbose_proxy_logger.error( + "Ignoring invalid admission control settings, per-worker admission control is disabled: %s", + exc, + ) + return None + return AdmissionControlSettings( + max_in_flight_requests=max_in_flight, + max_queued_requests=max_queued, + queue_timeout_seconds=queue_timeout, + ) + + +def get_admission_control_settings(settings: Mapping[str, object]) -> AdmissionControlSettings | None: + max_in_flight_raw: Final = settings.get("max_in_flight_requests_per_worker") + if max_in_flight_raw is None: + return None + return _parse_admission_control_settings( + _hashable(max_in_flight_raw), + _hashable(settings.get("max_queued_requests_per_worker")), + _hashable(settings.get("admission_queue_timeout_seconds", 1.0)), + ) + + +def _overloaded_response(state: AdmissionControlState) -> JSONResponse: + stats: Final = state.get_stats() + return JSONResponse( + status_code=503, + headers={"retry-after": "1"}, # mutable-ok: Starlette expects a plain headers mapping + content={ # mutable-ok: Starlette serializes a plain response mapping + "error": { # mutable-ok: nested response mapping + "message": ( + f"Worker at capacity: {stats.admitted} in-flight, {stats.queued} queued requests. Retry later." + ), + "type": "overloaded_error", + "code": "503", + } + }, + ) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 29f216fd450..6b1d6405a6a 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -9,10 +9,11 @@ Use litellm with Anthropic SDK, Vertex AI SDK, Cohere SDK, etc. from __future__ import annotations import hmac +import inspect import json import os import re -from collections.abc import Callable, Mapping +from collections.abc import AsyncGenerator, Callable, Mapping from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, cast @@ -32,6 +33,7 @@ from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +from litellm.passthrough.main import AsyncPassthroughStreamingResponse from litellm.proxy._types import * from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.route_checks import RouteChecks @@ -40,6 +42,7 @@ from litellm.proxy.auth.user_api_key_auth import ( user_api_key_auth, user_api_key_auth_websocket, ) +from litellm.proxy.common_request_processing import open_sse_before_first_byte from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, @@ -47,6 +50,9 @@ from litellm.proxy.common_utils.http_parsing_utils import ( get_form_data, get_request_body, ) +from litellm.proxy.common_utils.sse_keepalive import ( + wrap_passthrough_sse_bytes_with_keepalive_pings, +) from litellm.proxy.pass_through_endpoints.common_utils import get_litellm_virtual_key from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( HttpPassThroughEndpointHelpers, @@ -1478,6 +1484,74 @@ def is_azure_ai_search_service_level_index_create(method: str, endpoint: str) -> return path == "indexes" or path.endswith("/indexes") +async def _relay_upstream_bytes(upstream: AsyncGenerator[bytes, bytes]) -> AsyncGenerator[bytes, None]: + try: + async for chunk in upstream: + yield chunk + finally: + await upstream.aclose() + + +async def _relay_azure_router_model( + llm_router: litellm.Router, + model: str, + endpoint: str, + request: Request, + request_body: Mapping[str, object], + is_streaming_request: bool, + user_api_key_dict: UserAPIKeyAuth, +) -> Response: + result: Final = await llm_router.allm_passthrough_route( + model=model, + method=request.method, + endpoint=endpoint, + request_query_params=request.query_params, + request_headers=_safe_get_request_headers(request), + stream=is_streaming_request, + content=None, + data=None, + files=None, + json=(request_body if request.headers.get("content-type") == "application/json" else None), + params=None, + headers=None, + cookies=None, + litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict), + ) + + if not is_streaming_request: + upstream: Final = cast(httpx.Response, result) + return Response( + content=await upstream.aread(), + status_code=upstream.status_code, + headers=HttpPassThroughEndpointHelpers.get_response_headers(headers=upstream.headers, custom_headers=None), + ) + + if inspect.isasyncgen(result): + sse_headers: Final = {"content-type": "text/event-stream"} + return StreamingResponse( + content=wrap_passthrough_sse_bytes_with_keepalive_pings( + stream=_relay_upstream_bytes(result), + ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds, + upstream_headers=sse_headers, + ), + status_code=200, + headers=sse_headers, + ) + + upstream_stream: Final = cast(AsyncPassthroughStreamingResponse, result) + return StreamingResponse( + content=wrap_passthrough_sse_bytes_with_keepalive_pings( + stream=_relay_upstream_bytes(upstream_stream), + ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds, + upstream_headers=upstream_stream.headers, + ), + status_code=upstream_stream.status_code, + headers=HttpPassThroughEndpointHelpers.get_response_headers( + headers=upstream_stream.headers, custom_headers=None + ), + ) + + @router.api_route( "/azure_ai/{endpoint:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"], @@ -1528,55 +1602,18 @@ async def azure_proxy_route( if is_router_model: request_body = await get_request_body(request) is_streaming_request = is_passthrough_request_streaming(request_body) - result = await llm_router.allm_passthrough_route( - model=part, - method=request.method, - endpoint=endpoint, - request_query_params=request.query_params, - request_headers=_safe_get_request_headers(request), - stream=is_streaming_request, - content=None, - data=None, - files=None, - json=(request_body if request.headers.get("content-type") == "application/json" else None), - params=None, - headers=None, - cookies=None, - litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict), - ) - - if is_streaming_request: - # Check if result is an async generator (from _async_streaming) - import inspect - - if inspect.isasyncgen(result): - # Result is already an async generator, use it directly - return StreamingResponse( - content=result, - status_code=200, - headers={"content-type": "text/event-stream"}, - ) - else: - # Result is an httpx.Response, use aiter_bytes() - result = cast(httpx.Response, result) - return StreamingResponse( - content=result.aiter_bytes(), - status_code=result.status_code, - headers=HttpPassThroughEndpointHelpers.get_response_headers( - headers=result.headers, - custom_headers=None, - ), - ) - - # Non-streaming response - result = cast(httpx.Response, result) - content = await result.aread() - return Response( - content=content, - status_code=result.status_code, - headers=HttpPassThroughEndpointHelpers.get_response_headers( - headers=result.headers, - custom_headers=None, + return await open_sse_before_first_byte( + _relay_azure_router_model( + llm_router=llm_router, + model=part, + endpoint=endpoint, + request=request, + request_body=request_body, + is_streaming_request=is_streaming_request, + user_api_key_dict=user_api_key_dict, + ), + ping_interval_seconds=( + litellm.sse_keepalive_ping_interval_seconds if is_streaming_request else None ), ) elif is_vector_store_index: @@ -1659,6 +1696,12 @@ async def azure_proxy_route( from abc import ABC, abstractmethod +_VERTEX_LOCATION_REQUIRED_DETAIL: Final = ( + "No Vertex AI location for this request. Include /projects//locations// in the " + "route, set vertex_location in default_vertex_config (or DEFAULT_VERTEXAI_LOCATION), or add the " + "model to model_list with use_in_pass_through: true." +) + class BaseVertexAIPassThroughHandler(ABC): @staticmethod @@ -1666,29 +1709,18 @@ class BaseVertexAIPassThroughHandler(ABC): def get_default_base_target_url(vertex_location: str | None) -> str: pass - @staticmethod - @abstractmethod - def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: str | None) -> str: - pass - class VertexAIDiscoveryPassThroughHandler(BaseVertexAIPassThroughHandler): @staticmethod def get_default_base_target_url(vertex_location: str | None) -> str: return "https://discoveryengine.googleapis.com/" - @staticmethod - def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: str | None) -> str: - return base_target_url - class VertexAIPassThroughHandler(BaseVertexAIPassThroughHandler): @staticmethod def get_default_base_target_url(vertex_location: str | None) -> str: - return get_vertex_base_url(vertex_location) - - @staticmethod - def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: str | None) -> str: + if vertex_location is None: + raise HTTPException(status_code=400, detail=_VERTEX_LOCATION_REQUIRED_DETAIL) return get_vertex_base_url(vertex_location) @@ -1911,10 +1943,8 @@ async def _prepare_vertex_auth_headers( router_credentials: LiteLLM_ManagedVectorStore | None, vertex_project: str | None, vertex_location: str | None, - base_target_url: str | None, - get_vertex_pass_through_handler: BaseVertexAIPassThroughHandler, user_api_key_dict: UserAPIKeyAuth, -) -> tuple[Mapping[str, str], str | None, bool, str | None, str | None]: +) -> tuple[Mapping[str, str], bool, str | None, str | None]: """ Prepare authentication headers for Vertex AI pass-through requests. @@ -1924,15 +1954,12 @@ async def _prepare_vertex_auth_headers( router_credentials: Optional vector store credentials from registry vertex_project: Vertex project ID vertex_location: Vertex location - base_target_url: Base URL for the Vertex AI service - get_vertex_pass_through_handler: Handler for the specific Vertex AI service user_api_key_dict: The caller's resolved authentication, so only the secret that authenticated them is stripped on the credential-less branch Returns: tuple containing: - headers: dict - Authentication headers to use - - base_target_url: str | None - Updated base target URL - headers_passed_through: bool - Whether headers were passed through from request - vertex_project: str | None - Updated vertex project ID - vertex_location: str | None - Updated vertex location @@ -1985,14 +2012,8 @@ async def _prepare_vertex_auth_headers( # Add the Authorization header with vendor credentials headers["Authorization"] = f"Bearer {auth_header}" - if base_target_url is not None: - base_target_url = get_vertex_pass_through_handler.update_base_target_url_with_credential_location( - base_target_url, vertex_location - ) - return ( headers, - base_target_url, headers_passed_through, vertex_project, vertex_location, @@ -2085,12 +2106,9 @@ async def _base_vertex_proxy_route( location=vertex_location, ) - base_target_url = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location) - # Prepare authentication headers ( headers, - base_target_url, headers_passed_through, vertex_project, vertex_location, @@ -2100,13 +2118,10 @@ async def _base_vertex_proxy_route( router_credentials=router_credentials, vertex_project=vertex_project, vertex_location=vertex_location, - base_target_url=base_target_url, - get_vertex_pass_through_handler=get_vertex_pass_through_handler, user_api_key_dict=user_api_key_dict, ) - if base_target_url is None: - base_target_url = get_vertex_base_url(vertex_location) + base_target_url: Final = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location) request_route: Final = encoded_endpoint verbose_proxy_logger.debug("request_route %s", request_route) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1d3455615fd..6091225af15 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -583,6 +583,11 @@ try: except ImportError: build_billing_metrics_recorder = None shutdown_billing_metrics_recorder = None +from litellm.proxy.middleware.admission_control_middleware import ( + AdmissionControlMiddleware, + admission_control_state, + get_admission_control_settings, +) from litellm.proxy.middleware.in_flight_requests_middleware import ( InFlightRequestsMiddleware, ) @@ -15265,20 +15270,33 @@ async def async_queue_request( if llm_router is None: raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) - - response: Final = await llm_router.schedule_acompletion(**data) + router: Final = llm_router if "stream" in data and data["stream"] is True: # use generate_responses to stream responses - return StreamingResponse( - async_data_generator( - user_api_key_dict=user_api_key_dict, - response=response, - request_data=data, - request=request, - ), - media_type="text/event-stream", + + async def produce_queue_stream() -> StreamingResponse: + return StreamingResponse( + async_data_generator( + user_api_key_dict=user_api_key_dict, + response=await router.schedule_acompletion(**data), + request_data=data, + request=request, + ), + media_type="text/event-stream", + ) + + async def audit_late_failure(exc: Exception) -> HTTPException | None: + return await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, original_exception=exc, request_data=data + ) + + return await open_sse_before_first_byte( + produce_queue_stream(), + ping_interval_seconds=ttft_keepalive_interval(data, router), + on_late_failure=audit_late_failure, ) + response: Final = await router.schedule_acompletion(**data) fastapi_response.headers.update({"x-litellm-priority": str(data["priority"])}) return response except Exception as e: @@ -16534,6 +16552,9 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro { "max_parallel_requests": "Integer", "global_max_parallel_requests": "Integer", + "max_in_flight_requests_per_worker": "Integer", + "max_queued_requests_per_worker": "Integer", + "admission_queue_timeout_seconds": "Float", "max_request_size_mb": "Integer", "max_batch_file_size_mb": "Integer", "max_file_size_mb": "Integer", @@ -18209,6 +18230,11 @@ app.add_middleware( get_max_request_size_mb=lambda: general_settings.get("max_request_size_mb"), is_request_size_limit_enabled=lambda: premium_user is True, ) +app.add_middleware( + AdmissionControlMiddleware, + get_settings=lambda: get_admission_control_settings(general_settings), + state=admission_control_state, +) async def _stream_mcp_asgi_response(handle_fn, scope: dict, receive) -> "StreamingResponse": diff --git a/litellm/proxy/public_endpoints/public_v1/model_hub.py b/litellm/proxy/public_endpoints/public_v1/model_hub.py index 5a2d8068af7..b0e688740e4 100644 --- a/litellm/proxy/public_endpoints/public_v1/model_hub.py +++ b/litellm/proxy/public_endpoints/public_v1/model_hub.py @@ -3,7 +3,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import Annotated, Final, Protocol +from typing import Annotated, Final, Literal, Protocol from fastapi import APIRouter, Depends, Request from typing_extensions import ReadOnly, TypedDict @@ -20,10 +20,12 @@ from litellm.proxy.list_api.list_framework import ( Scope, ScopeAll, SortKey, + handle_facet, handle_list, ) from litellm.proxy.utils import PrismaClient from litellm.types.proxy.management_endpoints.management_v1 import ( + FacetListResponse, ListResponse, ProblemDetail, ) @@ -95,16 +97,37 @@ class HealthEnricher: return tuple(_with_health(row, health.get(row.model_group)) for row in rows) +FEATURE_PREFIX: Final = "supports_" + + +def _features(row: ModelGroupInfoProxy) -> tuple[str, ...]: + """A row's capabilities as one repeated field, so selecting two of them matches either. + + The hub's feature control has always been a multi-select over the `supports_*` flags. + One boolean filter per flag would AND them, which is the opposite of what it does. + """ + return tuple( + sorted( + name.removeprefix(FEATURE_PREFIX) + for name, value in row.model_dump().items() + if name.startswith(FEATURE_PREFIX) and value is True + ) + ) + + def _cells(row: ModelGroupInfoProxy) -> Cells: return MappingProxyType( { "model_group": row.model_group, "mode": row.mode, "providers": tuple(row.providers), + "features": _features(row), "max_input_tokens": row.max_input_tokens, "max_output_tokens": row.max_output_tokens, "input_cost_per_token": row.input_cost_per_token, "output_cost_per_token": row.output_cost_per_token, + "rpm": row.rpm, + "tpm": row.tpm, } ) @@ -126,20 +149,28 @@ def _scope(_caller: UserAPIKeyAuth) -> Scope: MODEL_HUB_FILTERS: Final[Mapping[str, FilterSpec]] = MappingProxyType( { "mode": FilterSpec(type=str, ops=frozenset(("eq", "in"))), - "providers": FilterSpec(type=str, ops=frozenset(("contains",))), + "providers": FilterSpec(type=str, ops=frozenset(("contains", "in"))), + "features": FilterSpec(type=str, ops=frozenset(("in",))), } ) +MODEL_HUB_FACETS: Final[Mapping[str, str]] = MappingProxyType( + {"providers": "providers", "modes": "mode", "features": "features"} +) + MODEL_HUB_LIST_SPEC: Final[ListSpec[ModelGroupInfoProxy, ModelGroupInfoProxy]] = ListSpec( resource="model groups", sortable=frozenset( ( "model_group", "mode", + "providers", "max_input_tokens", "max_output_tokens", "input_cost_per_token", "output_cost_per_token", + "rpm", + "tpm", ) ), searchable=frozenset(("model_group",)), @@ -153,6 +184,32 @@ MODEL_HUB_LIST_SPEC: Final[ListSpec[ModelGroupInfoProxy, ModelGroupInfoProxy]] = ) +def _published_rows() -> Sequence[ModelGroupInfoProxy]: + from litellm.proxy.proxy_server import ( + _get_model_group_info, # pyright: ignore[reportPrivateUsage] # /public/model_hub imports it the same way + llm_router, + ) + + if llm_router is None: + raise ManagementProblem( + ProblemDetail( + type=f"{PROBLEM_TYPE_BASE}no-llm-router", + title="No models configured", + status=400, + detail=CommonProxyErrors.no_llm_router.value, + ) + ) + if litellm.public_model_groups is None: + return () + return tuple( + _get_model_group_info( + llm_router=llm_router, + all_models_str=litellm.public_model_groups, + model_group=None, + ) + ) + + def _executor( rows: Sequence[ModelGroupInfoProxy], prisma_client: PrismaClient | None, @@ -191,37 +248,11 @@ async def public_model_hub_list( ``` """ try: - from litellm.proxy.proxy_server import ( - _get_model_group_info, # pyright: ignore[reportPrivateUsage] # /public/model_hub imports it the same way - llm_router, - prisma_client, - ) - - if llm_router is None: - raise ManagementProblem( - ProblemDetail( - type=f"{PROBLEM_TYPE_BASE}no-llm-router", - title="No models configured", - status=400, - detail=CommonProxyErrors.no_llm_router.value, - ) - ) - - rows: Final[Sequence[ModelGroupInfoProxy]] = ( - () - if litellm.public_model_groups is None - else tuple( - _get_model_group_info( - llm_router=llm_router, - all_models_str=litellm.public_model_groups, - model_group=None, - ) - ) - ) + from litellm.proxy.proxy_server import prisma_client return await handle_list( spec=MODEL_HUB_LIST_SPEC, - executor=_executor(rows, prisma_client), + executor=_executor(_published_rows(), prisma_client), request=request, caller=user_api_key_dict, ) @@ -240,3 +271,53 @@ async def public_model_hub_list( detail="Failed to list public model groups.", ) ) + + +@router.get( + "/model_hub/{facet}", + tags=["public", "model management"], # mutable-ok: fastapi types tags as list[str | Enum] + dependencies=(Depends(user_api_key_auth),), + response_model=FacetListResponse, +) +async def public_model_hub_facet( + request: Request, + facet: Literal["providers", "modes", "features"], + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> FacetListResponse: + """ + The distinct providers, modes or features across the published model groups, for the + Model Hub's filter dropdowns. No authentication. + + Carries the same filters and search as the list route, so a dropdown offers exactly + the values the table can show: asking for providers under `filter[mode][in]=chat` + lists only the providers that serve a chat model. + + Example curl: + ``` + curl --location --globoff \ + 'http://0.0.0.0:4000/public/v1/model_hub/providers?filter[mode][in]=chat&page_size=50' + ``` + """ + try: + return await handle_facet( + spec=MODEL_HUB_LIST_SPEC, + executor=InMemoryListExecutor(rows=_published_rows(), cells=_cells), + request=request, + caller=user_api_key_dict, + field=MODEL_HUB_FACETS[facet], + ) + + except ManagementProblem: + raise + except Exception as e: # noqa: BLE001 # a router error answers as a problem document, not the OpenAI error shape + verbose_proxy_logger.exception( + "litellm.proxy.public_endpoints.public_v1.model_hub.public_model_hub_facet(): Exception occured - %s", e + ) + raise ManagementProblem( + ProblemDetail( + type=f"{PROBLEM_TYPE_BASE}internal-server-error", + title="Internal server error", + status=500, + detail="Failed to list public model group values.", + ) + ) diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index c8c6c505375..d6a402e1860 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -23,11 +23,14 @@ from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, ) -from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.proxy._types import * from litellm.proxy.auth.auth_utils import is_request_body_safe from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth -from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + open_sse_before_first_byte, + ttft_keepalive_interval, +) from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, @@ -48,6 +51,7 @@ from litellm.proxy.vector_store_endpoints.utils import ( assert_user_can_access_vector_store_id, ) from litellm.repositories.table_repositories import ManagedVectorStoresRepository +from litellm.types.utils import ModelResponse if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient @@ -756,43 +760,53 @@ async def rag_query( merged_retrieval_config.get("custom_llm_provider"), ) - # Call query - response: Final = await litellm.aquery( - model=model, - messages=messages, - retrieval_config=merged_retrieval_config, - vector_store_params=store_data, - rerank=rerank, - stream=stream, - router=llm_router, - **request_data, - ) - - hidden_params: Final = getattr(response, "_hidden_params", {}) or {} - custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers( - user_api_key_dict=user_api_key_dict, - call_id=hidden_params.get("litellm_call_id", None) or "", - model_id=hidden_params.get("model_id", None) or "", - cache_key=hidden_params.get("cache_key", None) or "", - api_base=hidden_params.get("api_base", None) or "", - version=version, - response_cost=hidden_params.get("response_cost", None), - request_data=request_data, - ) - - if isinstance(response, CustomStreamWrapper): - return StreamingResponse( - select_data_generator( - response=response, - user_api_key_dict=user_api_key_dict, - request_data=request_data, - request=request, - ), - media_type="text/event-stream", - headers=custom_headers, + async def query() -> ModelResponse: + return await litellm.aquery( + model=model, + messages=messages, + retrieval_config=merged_retrieval_config, + vector_store_params=store_data, + rerank=rerank, + stream=stream, + router=llm_router, + **request_data, ) - fastapi_response.headers.update(custom_headers) + def custom_headers_for(response: ModelResponse) -> Mapping[str, str]: + hidden_params: Final = getattr(response, "_hidden_params", {}) or {} + return ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + call_id=hidden_params.get("litellm_call_id", None) or "", + model_id=hidden_params.get("model_id", None) or "", + cache_key=hidden_params.get("cache_key", None) or "", + api_base=hidden_params.get("api_base", None) or "", + version=version, + response_cost=hidden_params.get("response_cost", None), + request_data=request_data, + ) + + if stream: + + async def produce_stream() -> StreamingResponse: + response: Final = await query() + return StreamingResponse( + select_data_generator( + response=response, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + request=request, + ), + media_type="text/event-stream", + headers=custom_headers_for(response), + ) + + return await open_sse_before_first_byte( + produce_stream(), + ping_interval_seconds=ttft_keepalive_interval(data, llm_router), + ) + + response: Final = await query() + fastapi_response.headers.update(custom_headers_for(response)) return response except HTTPException: diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index b86a877e8f9..dff100bdea7 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3,7 +3,8 @@ import collections import json import os from collections.abc import Mapping, Sequence -from datetime import datetime, timedelta, timezone +from datetime import date, datetime, timedelta, timezone +from itertools import groupby from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -16,7 +17,6 @@ from typing import ( TypeAlias, TypedDict, TypeVar, - cast, # noqa: TID251 # prisma group_by returns untyped aggregate mappings ) import fastapi @@ -201,16 +201,12 @@ class _SessionSpendStats(NamedTuple): _SessionSpendMap: TypeAlias = Mapping[tuple[str, str], _SessionSpendStats] -class _SpendSumAggregate(TypedDict, total=False): - spend: ReadOnly[float] - - -class _SpendGroupByRow(TypedDict): +class _SpendDailySummaryRow(TypedDict): + day: ReadOnly[str] api_key: ReadOnly[str] user: ReadOnly[str | None] model: ReadOnly[str] - startTime: ReadOnly[object] - _sum: ReadOnly[_SpendSumAggregate] + spend: ReadOnly[float] async def _query_raw(prisma_client: PrismaClient, sql_query: str, *args: object) -> Sequence[_RowT]: @@ -251,6 +247,66 @@ def _verification_token_table(prisma_client: PrismaClient) -> _VerificationToken return VerificationTokenRepository(prisma_client).table +def _spend_logs_daily_summary_sql( + *, + start_date_iso: str, + end_date_iso: str, + api_key: str | None, + request_id: str | None, + user_id: str | None, +) -> tuple[str, tuple[object, ...]]: + filter_params: Final[tuple[tuple[str, object], ...]] = tuple( + (column, value) + for column, value in ( + ("api_key", api_key), + ("request_id", request_id), + ('"user"', user_id), + ) + if value is not None + ) + filter_clauses: Final[tuple[str, ...]] = tuple( + f"AND {column} = ${index}" for index, (column, _) in enumerate(filter_params, start=3) + ) + filter_sql: Final = "\n".join(filter_clauses) + sql_query: Final = f""" +SELECT + to_char(date_trunc('day', "startTime"), 'YYYY-MM-DD') AS day, + api_key, + "user", + model, + SUM(spend) AS spend +FROM "LiteLLM_SpendLogs" +WHERE "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') AND "startTime" <= ($2::timestamptz AT TIME ZONE 'UTC') +{filter_sql} +GROUP BY 1, 2, 3, 4 +ORDER BY 1 +""" + params: Final[tuple[object, ...]] = ( + start_date_iso, + end_date_iso, + *(value for _, value in filter_params), + ) + return sql_query, params + + +def _sum_spend_by( + rows: Sequence[_SpendDailySummaryRow], column: Literal["api_key", "user", "model"] +) -> Mapping[str | None, float]: + keys: Final = frozenset(row[column] for row in rows) + return {key: sum(float(row["spend"]) for row in rows if row[column] == key) for key in keys} + + +def _daily_summary_item(summary_date: date, rows: Sequence[_SpendDailySummaryRow]) -> Mapping[str, object]: + api_key_spend: Final = {key: value for key, value in _sum_spend_by(rows, "api_key").items() if key is not None} + return { + **api_key_spend, + "startTime": summary_date, + "spend": sum(float(row["spend"]) for row in rows), + "users": _sum_spend_by(rows, "user"), + "models": _sum_spend_by(rows, "model"), + } + + async def _find_spend_logs( prisma_client: PrismaClient, where: Mapping[str, object], @@ -3266,18 +3322,22 @@ async def view_spend_logs( start_date_iso: Final = start_date_obj.isoformat() end_date_iso: Final = end_date_obj.isoformat() - filter_query: Final = { + filter_query: Final[ + dict[str, object] + ] = { # mutable-ok: legacy filters are extended for optional parameters "startTime": { "gte": start_date_iso, # Greater than or equal to Start Date "lte": end_date_iso, # Less than or equal to End Date } } + summary_api_key: Final[str | None] = ( + prisma_client.hash_token(token=api_key) + if api_key is not None and api_key.startswith("sk-") + else api_key + ) if api_key is not None and isinstance(api_key, str): - if api_key.startswith("sk-"): - filter_query["api_key"] = prisma_client.hash_token(token=api_key) - else: - filter_query["api_key"] = api_key + filter_query["api_key"] = summary_api_key if request_id is not None and isinstance(request_id, str): filter_query["request_id"] = request_id if user_id is not None and isinstance(user_id, str): @@ -3296,58 +3356,34 @@ async def view_spend_logs( return data # Legacy behavior: return summarized data (when summarize=true) - # SQL query - response: Final = await SpendLogsRepository(prisma_client).table.group_by( - by=["api_key", "user", "model", "startTime"], - where=filter_query, - sum={ - "spend": True, - }, + summary_sql_and_params: Final = _spend_logs_daily_summary_sql( + start_date_iso=start_date_iso, + end_date_iso=end_date_iso, + api_key=summary_api_key, + request_id=request_id, + user_id=user_id, ) + sql_query, params = summary_sql_and_params + rows: Final[Sequence[_SpendDailySummaryRow]] = await _query_raw(prisma_client, sql_query, *params) + if len(rows) == 0: + return [] # pyright: ignore[reportUnknownVariableType] # empty summary has no element type - if isinstance(response, list) and len(response) > 0 and isinstance(response[0], dict): - spend_rows: Final = cast(Sequence[_SpendGroupByRow], response) # cast-ok: by/sum fix the shape - result: Final[dict] = {} - for record in spend_rows: - dt_object = datetime.strptime(str(record["startTime"]), "%Y-%m-%dT%H:%M:%S.%fZ") - date = dt_object.date() - if date not in result: - result[date] = {"users": {}, "models": {}} - api_key = record["api_key"] - user_id = record["user"] - model = record["model"] - result[date]["spend"] = result[date].get("spend", 0) + record.get("_sum", {}).get("spend", 0) - result[date][api_key] = result[date].get(api_key, 0) + record.get("_sum", {}).get("spend", 0) - result[date]["users"][user_id] = result[date]["users"].get(user_id, 0) + record.get("_sum", {}).get( - "spend", 0 - ) - result[date]["models"][model] = result[date]["models"].get(model, 0) + record.get("_sum", {}).get( - "spend", 0 - ) - return_list: Final = [] - final_date = None - for k, v in sorted(result.items()): - return_list.append({**v, "startTime": k}) - final_date = k - - end_date_date: Final = end_date_obj.date() - if final_date is not None and final_date < end_date_date: - current_date = final_date + timedelta(days=1) - while current_date <= end_date_date: - # Represent current_date as string because original response has it this way - return_list.append( - { - "startTime": current_date, - "spend": 0, - "users": {}, - "models": {}, - } - ) # If no data, will stay as zero - current_date += timedelta(days=1) # Move on to the next day - - return return_list - - return response + summary_items: Final = tuple( + _daily_summary_item(date.fromisoformat(day), tuple(day_rows)) + for day, day_rows in groupby(rows, key=lambda row: row["day"]) + ) + final_date: Final = date.fromisoformat(rows[-1]["day"]) + end_date_date: Final = end_date_obj.date() + padding: Final[tuple[Mapping[str, object], ...]] = tuple( + { + "startTime": final_date + timedelta(days=offset), + "spend": 0, + "users": {}, + "models": {}, + } + for offset in range(1, (end_date_date - final_date).days + 1) + ) + return [*summary_items, *padding] else: scoped_filter: Final[dict[str, str]] = {} diff --git a/litellm/router.py b/litellm/router.py index f33dfbba7bf..6da201725b6 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -197,6 +197,7 @@ from litellm.router_utils.router_callbacks.track_deployment_metrics import ( from litellm.scheduler import FlowItem, Scheduler from litellm.types.llms.openai import ( AllMessageValues, + ChatCompletionToolParam, FileTypes, OpenAIFileObject, OpenAIFilesPurpose, @@ -9545,6 +9546,7 @@ class Router: public_model_name for _, public_model_name in self.team_model_to_deployment_indices ) + self.pattern_router.remove_deployment(model_id) for team_id in list(self.team_pattern_routers.keys()): team_pattern_router = self.team_pattern_routers[team_id] team_pattern_router.remove_deployment(model_id) @@ -11762,7 +11764,7 @@ class Router: self, messages: list[dict[str, str]] | None, input: str | list | None, - instructions: str | None = None, + request_kwargs: Mapping[str, object] | None = None, ) -> int: """ Count input tokens for context-window pre-call checks. @@ -11772,9 +11774,28 @@ class Router: The Responses payload is normalized to chat messages via the shared LiteLLMCompletionResponsesConfig transform so the same token_counter path covers both API surfaces and `instructions` tokens are included in the count. + + Prompt content the message list never carries is read from `request_kwargs`: + `tools` (Chat Completions, Responses and Anthropic Messages shapes) and the + Anthropic Messages top-level `system` block. """ + from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( + anthropic_system_to_openai_message, + ) + + extras: Final = request_kwargs if request_kwargs is not None else MappingProxyType({}) + raw_instructions: Final = extras.get("instructions") + instructions: Final = raw_instructions if isinstance(raw_instructions, str) else None + raw_tools: Final = extras.get("tools") + tools: Final = ( + cast(list[ChatCompletionToolParam], raw_tools) # cast-ok: token_counter formats any tool dict shape + if isinstance(raw_tools, list) and raw_tools + else None + ) + system_message: Final = anthropic_system_to_openai_message(extras.get("system")) if messages is not None: - return litellm.token_counter(messages=messages) + counted_messages: Final = (system_message, *messages) if system_message is not None else messages + return litellm.token_counter(messages=counted_messages, tools=tools) if input is not None: from openai.types.responses.response_create_params import ResponseInputParam @@ -11787,7 +11808,10 @@ class Router: input=typed_input, responses_api_request={"instructions": instructions} if instructions is not None else {}, ) - return litellm.token_counter(messages=cast(list, input_messages)) # cast-ok: transformed chat messages + return litellm.token_counter( + messages=cast(list, input_messages), # cast-ok: transformed chat messages + tools=tools, + ) raise ValueError("Either messages or input must be provided to count tokens") def _deployment_max_input_tokens(self, model: str, deployment: Mapping[str, object]) -> int | None: @@ -11833,14 +11857,13 @@ class Router: """ if messages is None and input is None: return None - raw_instructions: Final = request_kwargs.get("instructions") if request_kwargs else None try: if not self._pre_call_checks_need_token_count(model, healthy_deployments): return None return await asyncify(self._count_pre_call_check_tokens)( messages=cast(list[dict[str, str]] | None, messages), # cast-ok: forwarded to the sync counter input=cast(str | list | None, input), # cast-ok: forwarded to the sync counter - instructions=raw_instructions if isinstance(raw_instructions, str) else None, + request_kwargs=request_kwargs, ) except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request verbose_router_logger.error( @@ -11887,8 +11910,6 @@ class Router: _rate_limit_error = False parent_otel_span: Final = _get_parent_otel_span_from_kwargs(request_kwargs) - raw_instructions: Final = request_kwargs.get("instructions") if request_kwargs else None - instructions: Final = raw_instructions if isinstance(raw_instructions, str) else None has_countable_input: Final = messages is not None or input is not None ## get model group RPM ## @@ -11919,7 +11940,7 @@ class Router: return _returned_deployments try: input_tokens = self._count_pre_call_check_tokens( - messages=messages, input=input, instructions=instructions + messages=messages, input=input, request_kwargs=request_kwargs ) except Exception as e: verbose_router_logger.error( diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index e3da70f50fe..afa27719064 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -190,6 +190,9 @@ model_list: # Let that replacement also override a kept session pin, for image turns only (default: false) modality_pin_override: true + + # Refreshes on every pin reuse, so this is idle time rather than total session length (default: 3600) + session_affinity_ttl_seconds: 300 ``` ## Usage @@ -240,6 +243,10 @@ affinity write happens upstream of the gate and stores the session's own model, turn replays the original pin and the override is never pinned in its place. It does nothing unless `modality_routing` is also on. +### Session pin retention + +`session_affinity_ttl_seconds` is the idle window for both the model pin selected by session affinity and the deployment pin. Every request that reuses a pin refreshes its TTL, so a session actively sending requests stays pinned. After the window passes with no pin reuse, the next request classifies again and creates a fresh pin. Omit the setting to track the default of 3600 seconds. + ### Heuristic-first chaining `classifier_type: heuristic_first` runs the local scorer on every request and only calls the LLM diff --git a/litellm/rust_bridge/__init__.py b/litellm/rust_bridge/__init__.py index e6d8ffef48c..9e8558bbf7d 100644 --- a/litellm/rust_bridge/__init__.py +++ b/litellm/rust_bridge/__init__.py @@ -4,6 +4,7 @@ from litellm.rust_bridge.configuration import use_litellm_rust from litellm.rust_bridge.loader import ( get_native_bridge, native_bridge_available, + reset_native_bridge_cache, ) -__all__ = ["get_native_bridge", "native_bridge_available", "use_litellm_rust"] +__all__ = ["get_native_bridge", "native_bridge_available", "reset_native_bridge_cache", "use_litellm_rust"] diff --git a/litellm/rust_bridge/loader.py b/litellm/rust_bridge/loader.py index 1c11d6435d8..022c38f5a85 100644 --- a/litellm/rust_bridge/loader.py +++ b/litellm/rust_bridge/loader.py @@ -24,6 +24,12 @@ def get_native_bridge() -> ModuleType | None: return _native +def reset_native_bridge_cache() -> None: + """Forget the cached extension so the next lookup reimports it from disk.""" + global _cached_bridge + _cached_bridge = _BRIDGE_SENTINEL + + def native_bridge_available() -> bool: """Whether the packaged Rust extension is importable.""" return get_native_bridge() is not None diff --git a/litellm/secret_managers/get_azure_ad_token_provider.py b/litellm/secret_managers/get_azure_ad_token_provider.py index c2dc09bc65d..5d056ea3fe0 100644 --- a/litellm/secret_managers/get_azure_ad_token_provider.py +++ b/litellm/secret_managers/get_azure_ad_token_provider.py @@ -57,9 +57,11 @@ def get_azure_ad_token_provider( from azure import identity from azure.identity import ( CertificateCredential, + ChainedTokenCredential, ClientSecretCredential, DefaultAzureCredential, ManagedIdentityCredential, + WorkloadIdentityCredential, get_bearer_token_provider, ) @@ -101,6 +103,28 @@ def get_azure_ad_token_provider( # DefaultAzureCredential doesn't require explicit environment variables # It automatically discovers credentials from the environment (managed identity, CLI, etc.) credential = DefaultAzureCredential() + elif cred == AzureCredentialType.DeploymentIdentityCredential: + # DefaultAzureCredential cannot express this: excluding its developer credentials still + # leaves one managed identity link, which AZURE_CLIENT_ID pins to a user assigned identity, + # so a host running as a system assigned identity never gets asked + workload_client_id: Final = os.environ.get("AZURE_CLIENT_ID") + workload_tenant_id: Final = os.environ.get("AZURE_TENANT_ID") + workload_token_file: Final = os.environ.get("AZURE_FEDERATED_TOKEN_FILE") + credential = ChainedTokenCredential( + *( + ( + WorkloadIdentityCredential( + client_id=workload_client_id, + tenant_id=workload_tenant_id, + token_file_path=workload_token_file, + ), + ) + if workload_client_id and workload_tenant_id and workload_token_file + else () + ), + *((ManagedIdentityCredential(client_id=workload_client_id),) if workload_client_id else ()), + ManagedIdentityCredential(), + ) else: cred_cls: Final = getattr(identity, cred) credential = cred_cls() diff --git a/litellm/types/caching.py b/litellm/types/caching.py index 10c83376a6c..747f202d35e 100644 --- a/litellm/types/caching.py +++ b/litellm/types/caching.py @@ -18,6 +18,11 @@ class LiteLLMCacheType(str, Enum): GCS = "gcs" +class SemanticCacheScope(str, Enum): + KEY = "key" + END_USER = "end_user" + + CachingSupportedCallTypes = Literal[ "completion", "acompletion", diff --git a/litellm/types/management_endpoints/cache_settings_endpoints.py b/litellm/types/management_endpoints/cache_settings_endpoints.py index 2e6d0be9071..32b32991449 100644 --- a/litellm/types/management_endpoints/cache_settings_endpoints.py +++ b/litellm/types/management_endpoints/cache_settings_endpoints.py @@ -187,6 +187,19 @@ CACHE_SETTINGS_FIELDS: Final[list[CacheSettingsField]] = [ ui_field_name="Embedding Model", redis_type="semantic", ), + CacheSettingsField( + field_name="semantic_cache_scope", + field_type="String", + field_value=None, + field_description=( + "Isolation granularity for semantic cache hits. 'key' shares hits between all end users of a key/team/org." + " 'end_user' also isolates per end user; requests without an end user fall back to the key scope." + ), + field_default="key", + options=["key", "end_user"], + ui_field_name="Semantic Cache Scope", + redis_type="semantic", + ), # GCP IAM authentication fields CacheSettingsField( field_name="gcp_service_account", diff --git a/litellm/types/proxy/management_endpoints/management_v1.py b/litellm/types/proxy/management_endpoints/management_v1.py index b2244f6eb9b..aa82138110d 100644 --- a/litellm/types/proxy/management_endpoints/management_v1.py +++ b/litellm/types/proxy/management_endpoints/management_v1.py @@ -1,5 +1,6 @@ """Shared response shapes for the `/management/v1` control-plane surface.""" +from collections.abc import Sequence from typing import Generic, TypeVar from pydantic import BaseModel, ConfigDict, Field @@ -38,7 +39,7 @@ class PageMeta(BaseModel): class FacetListResponse(BaseModel): """The distinct values one column takes over a filtered query. `data` holds bare values, not entity rows.""" - data: list[str] + data: Sequence[str] meta: PageMeta links: PageLinks diff --git a/litellm/types/secret_managers/get_azure_ad_token_provider.py b/litellm/types/secret_managers/get_azure_ad_token_provider.py index 5d2f7409f95..6b4700d081d 100644 --- a/litellm/types/secret_managers/get_azure_ad_token_provider.py +++ b/litellm/types/secret_managers/get_azure_ad_token_provider.py @@ -6,3 +6,4 @@ class AzureCredentialType(str, Enum): ManagedIdentityCredential = "ManagedIdentityCredential" CertificateCredential = "CertificateCredential" DefaultAzureCredential = "DefaultAzureCredential" + DeploymentIdentityCredential = "DeploymentIdentityCredential" diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 24c0ff6b181..8763318b4eb 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -96,7 +96,7 @@ "limit": 10 }, "DTZ007": { - "limit": 17 + "limit": 6 }, "DTZ011": { "limit": 3 diff --git a/tests/e2e/ui/tests/tables/tableScrolling.spec.ts b/tests/e2e/ui/tests/tables/tableScrolling.spec.ts new file mode 100644 index 00000000000..5c7438cfd11 --- /dev/null +++ b/tests/e2e/ui/tests/tables/tableScrolling.spec.ts @@ -0,0 +1,218 @@ +import { test, expect, type APIRequestContext, type Locator, type Page as PlaywrightPage } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; +import { CHAT_MODEL_A, masterKey, sendChatCompletion, waitForSpendLog } from "../../helpers/traffic"; + +const VIEWPORT = { width: 1280, height: 720 }; +const SEED_ROWS = 40; +const LOG_ROWS = 20; +const BODY_SCROLL_PX = 500; +const MAX_FOOTER_GAP_PX = 40; + +interface GeneratedKey { + key: string; +} + +interface CreatedTeam { + team_id: string; +} + +interface CreatedModel { + model_info: { id: string }; +} + +interface BoxMetrics { + top: number; + bottom: number; + scrollHeight: number; + clientHeight: number; + scrollWidth: number; + clientWidth: number; +} + +const uniqueSuffix = (): string => `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; + +const adminHeaders = (): Record => ({ Authorization: `Bearer ${masterKey()}` }); + +const appShellMain = (page: PlaywrightPage): Locator => page.locator("main").first(); + +const visibleTestId = (page: PlaywrightPage, id: string): Locator => page.getByTestId(id).filter({ visible: true }); + +const visibleDataTable = (page: PlaywrightPage): Locator => visibleTestId(page, "data-table-root").first(); + +const visibleRows = (page: PlaywrightPage): Locator => visibleDataTable(page).locator("tbody tr"); + +const metrics = (locator: Locator): Promise => + locator.evaluate((el) => { + const rect = el.getBoundingClientRect(); + return { + top: rect.top, + bottom: rect.bottom, + scrollHeight: el.scrollHeight, + clientHeight: el.clientHeight, + scrollWidth: el.scrollWidth, + clientWidth: el.clientWidth, + }; + }); + +async function postOk(request: APIRequestContext, path: string, data: Record): Promise { + const res = await request.post(path, { headers: adminHeaders(), data }); + expect(res.ok(), `POST ${path} failed (${res.status()}): ${await res.text()}`).toBe(true); + return (await res.json()) as T; +} + +const oneAtATime = (count: number, call: (index: number) => Promise): Promise => + Array.from({ length: count }, (_, i) => i).reduce>( + async (previous, i) => [...(await previous), await call(i)], + Promise.resolve([]), + ); + +async function expectRowsAtLeast(page: PlaywrightPage, count: number): Promise { + await expect.poll(() => visibleRows(page).count(), { timeout: 30_000 }).toBeGreaterThanOrEqual(count); +} + +async function setRowsPerPage(page: PlaywrightPage, size: "25" | "50" | "100"): Promise { + await visibleTestId(page, "pagination-page-size").click(); + await page.getByRole("option", { name: size, exact: true }).click(); +} + +async function expectBodyIsTheOnlyScroller(page: PlaywrightPage): Promise { + const scroller = await metrics(appShellMain(page)); + const body = visibleTestId(page, "data-table-scroller"); + const bodyBefore = await metrics(body); + const headBefore = await metrics(visibleTestId(page, "data-table-head")); + const footer = await metrics(visibleDataTable(page)); + + expect(scroller.scrollHeight, "page scroller must not overflow vertically").toBe(scroller.clientHeight); + expect(scroller.scrollWidth, "page scroller must not overflow horizontally").toBe(scroller.clientWidth); + expect(bodyBefore.scrollHeight, "table body must be the element that scrolls").toBeGreaterThan( + bodyBefore.clientHeight, + ); + expect(footer.bottom, "pagination footer must be inside the page").toBeLessThanOrEqual(scroller.bottom); + expect(scroller.bottom - footer.bottom, "pagination footer must sit at the bottom of the page").toBeLessThanOrEqual( + MAX_FOOTER_GAP_PX, + ); + + await body.evaluate((el, px) => { + el.scrollTop = px; + }, BODY_SCROLL_PX); + await expect.poll(() => body.evaluate((el) => el.scrollTop)).toBeGreaterThan(0); + const headAfter = await metrics(visibleTestId(page, "data-table-head")); + expect(Math.round(headAfter.top), "header must stay put while the body scrolls").toBe(Math.round(headBefore.top)); +} + +const rowsPaintingPastAnAncestor = (page: PlaywrightPage): Promise => + visibleDataTable(page) + .locator("table") + .evaluate((table) => { + const scrollsVertically = (el: Element): boolean => + /auto|scroll/.test(getComputedStyle(el).overflowY) && el.scrollHeight > el.clientHeight + 1; + const boxesUpToTheScroller = (el: Element | null): Element[] => + el === null || el === document.body || scrollsVertically(el) + ? [] + : [el, ...boxesUpToTheScroller(el.parentElement)]; + const describe = (el: Element): string => + `<${el.tagName.toLowerCase()} class="${el.getAttribute("class") ?? ""}">`; + return Array.from(table.querySelectorAll("tbody tr")).flatMap((row, index) => { + const rowBottom = row.getBoundingClientRect().bottom; + return boxesUpToTheScroller(row.parentElement) + .filter((box) => rowBottom > box.getBoundingClientRect().bottom + 1) + .map( + (box) => + `row ${index} bottom ${Math.round(rowBottom)} past ${describe(box)} bottom ${Math.round(box.getBoundingClientRect().bottom)}`, + ); + }); + }); + +test.describe("Admin tables scroll inside the page", () => { + test.use({ storageState: ADMIN_STORAGE_PATH, viewport: VIEWPORT }); + + test("Virtual Keys: rows scroll under a sticky header and the page itself never scrolls", async ({ + page, + request, + }) => { + const suffix = uniqueSuffix(); + const keys = await oneAtATime(SEED_ROWS, (i) => + postOk(request, "/key/generate", { key_alias: `e2e-scroll-key-${suffix}-${i}` }), + ); + try { + await navigateToPage(page, Page.ApiKeys); + await expectRowsAtLeast(page, SEED_ROWS); + await expectBodyIsTheOnlyScroller(page); + } finally { + await request.post("/key/delete", { headers: adminHeaders(), data: { keys: keys.map((k) => k.key) } }); + } + }); + + test("Teams: rows scroll under a sticky header and the page itself never scrolls", async ({ page, request }) => { + const suffix = uniqueSuffix(); + const teams = await oneAtATime(SEED_ROWS, (i) => + postOk(request, "/team/new", { team_alias: `e2e-scroll-team-${suffix}-${i}` }), + ); + try { + await navigateToPage(page, Page.Teams); + await expectRowsAtLeast(page, SEED_ROWS); + await expectBodyIsTheOnlyScroller(page); + } finally { + await request.post("/team/delete", { headers: adminHeaders(), data: { team_ids: teams.map((t) => t.team_id) } }); + } + }); + + test("Request Logs: rows scroll under a sticky header and the page itself never scrolls", async ({ + page, + request, + }) => { + const suffix = uniqueSuffix(); + const ids = await oneAtATime(LOG_ROWS, (i) => + sendChatCompletion(request, { model: CHAT_MODEL_A, prompt: `scroll ${suffix} ${i}` }), + ); + await waitForSpendLog(request, ids[ids.length - 1]); + + await navigateToPage(page, Page.Logs); + await expect(visibleTestId(page, "datatable-search")).toBeVisible({ timeout: 20_000 }); + await setRowsPerPage(page, "25"); + await expectRowsAtLeast(page, LOG_ROWS); + await expectBodyIsTheOnlyScroller(page); + }); + + test("Tags: no row paints past the box it lives in", async ({ page, request }) => { + const suffix = uniqueSuffix(); + const names = Array.from({ length: SEED_ROWS }, (_, i) => `e2e-scroll-tag-${suffix}-${i}`); + await oneAtATime(SEED_ROWS, (i) => + postOk(request, "/tag/new", { name: names[i], description: "LIT-4738 scroll" }), + ); + try { + await navigateToPage(page, Page.TagManagement); + await expectRowsAtLeast(page, SEED_ROWS); + expect(await rowsPaintingPastAnAncestor(page)).toEqual([]); + } finally { + await oneAtATime(SEED_ROWS, (i) => + request.post("/tag/delete", { headers: adminHeaders(), data: { name: names[i] } }), + ); + } + }); + + test("Model Hub: no row paints past the box it lives in", async ({ page, request }) => { + const suffix = uniqueSuffix(); + const models = await oneAtATime(SEED_ROWS, (i) => + postOk(request, "/model/new", { + model_name: `e2e-scroll-model-${suffix}-${i}`, + litellm_params: { + model: "openai/fake-gpt-4", + api_base: `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`, + api_key: "fake-key", + }, + }), + ); + try { + await navigateToPage(page, Page.ModelHubTable); + await expectRowsAtLeast(page, SEED_ROWS); + expect(await rowsPaintingPastAnAncestor(page)).toEqual([]); + } finally { + await oneAtATime(SEED_ROWS, (i) => + request.post("/model/delete", { headers: adminHeaders(), data: { id: models[i].model_info.id } }), + ); + } + }); +}); diff --git a/tests/e2e/ui/tests/users/searchUsers.spec.ts b/tests/e2e/ui/tests/users/searchUsers.spec.ts index fa8f32764e8..64fce24bf2a 100644 --- a/tests/e2e/ui/tests/users/searchUsers.spec.ts +++ b/tests/e2e/ui/tests/users/searchUsers.spec.ts @@ -17,7 +17,7 @@ test.describe("Internal Users Search", () => { test("narrows the table to the matching email, and restores it when cleared", async ({ page }) => { await goToInternalUsers(page); - const search = page.getByPlaceholder("Search by email…"); + const search = page.getByPlaceholder("Search by email or ID…"); await expect(search).toBeVisible(); await search.fill("noteam@"); diff --git a/tests/guardrails_tests/test_tracing_guardrails.py b/tests/guardrails_tests/test_tracing_guardrails.py index bd8b7bad33f..ac85803ba39 100644 --- a/tests/guardrails_tests/test_tracing_guardrails.py +++ b/tests/guardrails_tests/test_tracing_guardrails.py @@ -806,3 +806,91 @@ def test_guardrail_status_fields_computation(): ) assert status_fields_no_guardrail.get("llm_api_status") == "success" assert status_fields_no_guardrail.get("guardrail_status") == "not_run" + + +@pytest.mark.parametrize( + "status, guardrail_information, expected_guardrail_status", + [ + pytest.param( + "failure", + [ + {"guardrail_status": "success"}, + {"guardrail_status": "guardrail_intervened"}, + ], + "guardrail_intervened", + id="pre_call_success_before_blocker", + ), + pytest.param( + "failure", + [ + {"guardrail_status": "guardrail_intervened"}, + {"guardrail_status": "success"}, + ], + "guardrail_intervened", + id="blocker_before_success", + ), + pytest.param( + "failure", + [ + {"guardrail_status": "success"}, + {"guardrail_status": "guardrail_failed_to_respond"}, + ], + "guardrail_failed_to_respond", + id="failure_outranks_success", + ), + pytest.param( + "failure", + [ + {"guardrail_status": "guardrail_failed_to_respond"}, + {"guardrail_status": "guardrail_intervened"}, + ], + "guardrail_intervened", + id="intervention_outranks_failure", + ), + pytest.param( + "success", + [ + {"guardrail_status": "success"}, + {"guardrail_status": "success"}, + ], + "success", + id="all_success_stays_success", + ), + pytest.param( + "failure", + [ + {"guardrail_status": "some_new_status"}, + {"guardrail_status": "blocked"}, + ], + "guardrail_intervened", + id="unknown_status_does_not_mask_blocker", + ), + pytest.param( + "failure", + [ + {"guardrail_status": {"unhashable": True}}, + {"guardrail_status": "guardrail_intervened"}, + ], + "guardrail_intervened", + id="unhashable_status_is_skipped", + ), + ], +) +def test_guardrail_status_fields_severity_across_entries( + status, guardrail_information, expected_guardrail_status +): + """ + A blocked request must never be reported as a guardrail success. + + With multiple guardrails on one request (e.g. a pre_call mask that passes, + then a post_call guardrail that blocks), entries are recorded in execution + order, so the earlier "success" entry must not shadow the later + "guardrail_intervened" entry: the aggregate takes the most severe status, + regardless of entry order. + """ + from litellm.litellm_core_utils.litellm_logging import _get_status_fields + + fields = _get_status_fields( + status=status, guardrail_information=guardrail_information, error_str=None + ) + assert fields.get("guardrail_status") == expected_guardrail_status diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index 9ffb57924b6..3fab20a28ad 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -763,7 +763,7 @@ class _MigrateDeployHarness: "_resolve_specific_migration", staticmethod(self.resolved.append), ) - monkeypatch.setattr(utils_module.subprocess, "run", self._fake_run) + monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", self._fake_run) monkeypatch.setattr(utils_module.time, "sleep", lambda seconds: None) self.baseline_succeeds = True diff --git a/tests/load_tests/test_granian_admission_saturation.py b/tests/load_tests/test_granian_admission_saturation.py new file mode 100644 index 00000000000..b42c06037c2 --- /dev/null +++ b/tests/load_tests/test_granian_admission_saturation.py @@ -0,0 +1,150 @@ +import asyncio +import os +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Final + +import httpx +import pytest + +pytestmark = pytest.mark.skipif( + os.environ.get("LITELLM_RUN_SATURATION_BENCHMARK") != "1", + reason="set LITELLM_RUN_SATURATION_BENCHMARK=1 to run the saturation benchmark", +) + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener: + listener.bind(("127.0.0.1", 0)) + return int(listener.getsockname()[1]) + + +def _percentile(values: list[float], percentile: float) -> float: + return sorted(values)[min(int(len(values) * percentile), len(values) - 1)] + + +@pytest.mark.asyncio +async def test_granian_admission_control_saturation(tmp_path: Path) -> None: + fake_port: Final = _free_port() + proxy_port: Final = _free_port() + fake_script: Final = Path(__file__).parents[1] / "_fake_openai_endpoint_server.py" + config_path: Final = tmp_path / "saturation_config.yaml" + config_path.write_text( + f"""model_list: + - model_name: slow-endpoint + litellm_params: + model: openai/slow-endpoint + api_base: http://127.0.0.1:{fake_port}/v1 +general_settings: + master_key: sk-saturation + max_in_flight_requests_per_worker: 8 + max_queued_requests_per_worker: 8 + admission_queue_timeout_seconds: 0.5 +""" + ) + fake_process: Final = subprocess.Popen( + [sys.executable, str(fake_script), "--host", "127.0.0.1", "--port", str(fake_port)], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + try: + proxy_process: Final = subprocess.Popen( + [ + sys.executable, + "-m", + "litellm.proxy.proxy_cli", + "--config", + str(config_path), + "--run_granian", + "--num_workers", + "1", + "--port", + str(proxy_port), + ], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + try: + async with httpx.AsyncClient(base_url=f"http://127.0.0.1:{proxy_port}") as client: + deadline: Final = time.monotonic() + 60 + while time.monotonic() < deadline: + try: + response: Final = await client.get("/health/liveliness", timeout=2) + if response.status_code == 200: + break + except httpx.HTTPError: + pass + await asyncio.sleep(0.25) + else: + raise AssertionError("Granian proxy did not become healthy") + + liveness_latencies: Final[list[float]] = [] + stop_sampling: Final = asyncio.Event() + + async def sample_liveness() -> None: + while not stop_sampling.is_set(): + start: Final = time.perf_counter() + try: + response = await client.get("/health/liveliness", timeout=2) + response.raise_for_status() + liveness_latencies.append(time.perf_counter() - start) + except httpx.HTTPError: + pass + await asyncio.sleep(0.05) + + async def send_completion() -> tuple[int, float, bool]: + start: Final = time.perf_counter() + response = await client.post( + "/chat/completions", + headers={"Authorization": "Bearer sk-saturation"}, + json={ + "model": "slow-endpoint", + "messages": [{"role": "user", "content": "hello"}], + }, + timeout=10, + ) + return response.status_code, time.perf_counter() - start, "retry-after" in response.headers + + sampler: Final = asyncio.create_task(sample_liveness()) + results: Final = await asyncio.gather(*(send_completion() for _ in range(200))) + stop_sampling.set() + await sampler + + statuses: Final = [result[0] for result in results] + latencies: Final = [result[1] for result in results] + rejected: Final = [result for result in results if result[0] == 503] + assert set(statuses) <= {200, 503} + assert rejected + assert all(result[2] for result in rejected) + assert _percentile(latencies, 0.99) < 5 + assert liveness_latencies + assert _percentile(liveness_latencies, 0.95) < 0.5 + + duration: Final = max(latencies) + print( + "\nmetric value\n" + f"rps {len(results) / duration:.2f}\n" + f"200 count {statuses.count(200)}\n" + f"503 count {statuses.count(503)}\n" + f"p50 {_percentile(latencies, 0.50):.3f}s\n" + f"p95 {_percentile(latencies, 0.95):.3f}s\n" + f"p99 {_percentile(latencies, 0.99):.3f}s\n" + f"liveness p95 {_percentile(liveness_latencies, 0.95):.3f}s" + ) + finally: + proxy_process.terminate() + try: + proxy_process.wait(timeout=10) + except subprocess.TimeoutExpired: + proxy_process.kill() + proxy_process.wait() + finally: + fake_process.terminate() + try: + fake_process.wait(timeout=10) + except subprocess.TimeoutExpired: + fake_process.kill() + fake_process.wait() diff --git a/tests/rust-python-harness/AGENTS.md b/tests/rust-python-harness/AGENTS.md index e9b17027ddc..017668d4289 100644 --- a/tests/rust-python-harness/AGENTS.md +++ b/tests/rust-python-harness/AGENTS.md @@ -3,42 +3,75 @@ ```text tests/rust-python-harness/ ├── __main__.py +├── cli/ +│ ├── __init__.py +│ ├── catalog.py +│ └── commands.py │ ├── strategies/ │ ├── e2e_parity/ -│ │ ├── runner.py +│ │ ├── __init__.py +│ │ ├── reporting.py │ │ ├── sdk/ -│ │ │ ├── ocr/ -│ │ │ ├── messages/ -│ │ │ ├── chat_completions/ -│ │ │ └── responses/ -│ │ └── gateway/ +│ │ │ └── ocr/ │ │ │ ├── trace_parity/ -│ │ ├── runner.py -│ │ ├── sdk/ -│ │ └── gateway/ +│ │ ├── __init__.py +│ │ ├── models.py +│ │ ├── reporting.py +│ │ └── sdk/ +│ │ ├── chat_completions/ +│ │ ├── messages/ +│ │ ├── ocr/ +│ │ └── transcription/ │ │ -│ └── unit_tests/ -│ ├── runner.py -│ ├── mapping_validator.py -│ ├── python_runner.py -│ └── rust_runner.py +│ ├── unit_tests_mapping/ +│ │ ├── __init__.py +│ │ ├── contracts.py +│ │ ├── cases/ +│ │ │ └── ocr.py +│ │ ├── mapping_report.py +│ │ ├── mappings.py +│ │ ├── mapping_validator.py +│ │ ├── reporting.py +│ │ └── runner.py +│ │ +│ ├── unit_tests_parity/ +│ │ ├── __init__.py +│ │ ├── reporting.py +│ │ └── runner.py +│ │ +│ └── unit_tests_rust/ +│ ├── __init__.py +│ ├── reporting.py +│ └── runner.py │ └── shared/ ├── parity/ ├── tracing/ - └── reporting/ + ├── reporting/ + │ └── strategy.py + └── unit_runners/ + └── suite_runner.py ``` +- A strategy is a folder under `strategies/` with a one-line `AGENTS.md` and an `__init__.py` exporting exactly one `STRATEGY: StrategyDefinition`; its id must equal the folder name +- `shared/reporting/strategy.py` is the contract: runnable module/suite specs, not-implemented/skipped specs, the runner protocol, and `StrategyDefinition` +- Every `STRATEGY` explicitly classifies every SDK function; surface-aware strategies declare their surfaces and classify the complete surface-by-function matrix - Run locally only; no CI integration -- `__main__.py` selects strategies and combines their reports; each strategy also runs independently +- `python -m tests.rust-python-harness run |all` runs the selected strategy; `--function` is common, while each strategy exposes only its supported options +- Examples: `run e2e_parity --surface sdk --function ocr`, `run unit_tests_parity --function ocr --pytest-arg=-x`, or `run all --function ocr` +- `cli/catalog.py` discovers strategies, validates their Python definitions, and orders them; `cli/__init__.py` builds the Click command tree; `cli/commands.py` runs selected cases - `e2e_parity/` compares SDK objects, exceptions, callbacks, and streams, or gateway HTTP responses -- `trace_parity/` compares mapped operations, call counts, and required execution ordering -- E2E and trace runners share orchestration across `sdk/` and `gateway/`; surface-specific execution lives in those folders -- `unit_tests/runner.py` combines mapping validation, Python test runs, and native Rust test runs -- `mapping_validator.py` matches Python/Rust tests by agreed names or annotations and reports missing or ambiguous counterparts -- `python_runner.py` runs existing Python tests with Rust disabled and enabled in separate processes, verifies backend selection, and compares results -- `rust_runner.py` runs Cargo tests; native Rust unit tests stay beside their implementation -- `shared/` contains reusable parity, tracing, and reporting machinery +- `trace_parity/` compares mapped operations, call counts, and required execution ordering; before running it rebuilds the native bridge with the `trace-parity` feature whenever `litellm-rust` sources are newer than the installed extension (`shared/native_build.py`) +- E2E and trace strategies load their registered module cases and run surface-specific execution from their folders +- `unit_tests_mapping/contracts.py` owns typed harness-side mapping contracts, per-function contracts live below `cases/`, and `mappings.py` exports the registry; live test discovery derives unmapped Python and Rust-only tests without an exhaustive manifest +- `unit_tests_mapping/runner.py` validates confirmed mappings against the live Python and Rust inventories and attaches the derived status report +- `unit_tests_parity/runner.py` runs each contract's `unit_parity_scope` with `LITELLM_RUST=0` and `LITELLM_RUST=1` in separate processes and requires matching outcomes, including failures; exclusions require a reason in the contract +- `unit_tests_rust/runner.py` runs each contract's focused Cargo test suite; native Rust unit tests stay beside their implementation +- `shared/unit_runners/suite_runner.py` runs typed suites registered in code with nodeids of the form `suite:::` +- Every strategy declares its report sections and presentation in its own `reporting.py`; shared reporting code only provides reusable models and cell-formatting primitives +- `shared/` contains reusable parity, tracing, reporting primitives, and unit-runner machinery - Keep fixtures with their owning API and existing Python tests in their current locations +- Each strategy folder carries an `AGENTS.md` one-liner stating what it should be doing +- Run the harness's own checks with `uv run pytest -o consider_namespace_packages=true tests/rust-python-harness/shared tests/rust-python-harness/cli tests/rust-python-harness/strategies/unit_tests_mapping tests/rust-python-harness/strategies/unit_tests_parity tests/rust-python-harness/strategies/unit_tests_rust tests/test_rust_python_harness.py -q` diff --git a/tests/rust-python-harness/README.md b/tests/rust-python-harness/README.md deleted file mode 100644 index 77df4a24dd3..00000000000 --- a/tests/rust-python-harness/README.md +++ /dev/null @@ -1,105 +0,0 @@ -# Rust/Python migration harness - -This local harness follows [the agreed structure](AGENTS.md). The root command selects strategies and combines their reports. Each strategy has an independent entry point - -```text -strategies/ - e2e_parity/runner.py - sdk/ocr/fixtures/ - sdk/messages/ - sdk/chat_completions/ - sdk/responses/ - gateway/ - existing_e2e_test_sdk/runner.py - trace_parity/runner.py - sdk/ - gateway/ - unit_tests/ - runner.py - mapping_validator.py - python_runner.py - rust_runner.py -shared/ - parity/ - tracing/ - reporting/ -``` - -## Run locally - -```bash -uv run python -m tests.rust-python-harness --list -uv run python -m tests.rust-python-harness --function ocr --plain -uv run python -m tests.rust-python-harness --strategy e2e_parity --surface sdk --function ocr --plain -uv run python -m tests.rust-python-harness.strategies.e2e_parity.runner --function ocr --plain -uv run python -m tests.rust-python-harness.strategies.trace_parity.runner --plain -uv run python -m tests.rust-python-harness.strategies.unit_tests.runner --plain -uv run python -m tests.rust-python-harness.strategies.existing_e2e_test_sdk.runner --function transcription --plain -``` - -Use `--interactive` for strategy and function selection, `--pytest-arg=-x` to stop pytest on its first failure, and `--coverage` to write Python coverage under `target/rust-python-harness/`. The harness enables pytest namespace-package discovery only for its own invocations - -This harness has no CI execution. A configured test that fails or disappears makes the command fail. An unconfigured strategy cell remains planned and contributes no passing evidence. Interruptions and collection errors stop execution; ordinary test failures remain in the combined report while later strategies run - -## Strategy responsibilities - -E2E parity compares SDK objects, exceptions, callbacks, streams, and provider requests. Gateway tests compare HTTP responses. Both surfaces use the same strategy runner and keep execution details and fixtures in their own folders. OCR has recorded sync/async SDK coverage; the existing Messages and Responses bridge checks remain partial - -Trace parity compares operation names through an explicit Python/Rust mapping, call counts, and required completion-before-start ordering with `shared/tracing/compare.py`. Surface tests supply captured operation intervals. No production trace instrumentation or trace case is configured yet - -Unit testing combines test mapping validation, separate Python processes with Rust disabled and enabled, backend verification, result comparison, and native Cargo tests. Native tests stay beside their Rust implementation. Existing Python tests stay at their original paths. No complete Python/native unit mapping is configured yet, so these cells remain planned - -The existing E2E SDK strategy retains the live provider tests configured upstream. It runs OCR, Chat Completions, and Transcription checks from their existing paths and reports them separately from parity tests. These tests require provider credentials - -## Configure cases - -Each strategy has a `strategy.json`. Its `functions` object defines SDK cases for OCR, Messages, Responses, Count Tokens, Chat Completions, and Transcription. E2E and trace manifests also accept a `gateway` object keyed by API name. A case has `coverage`, `selectors`, and an optional `note` - -```json -{ - "coverage": "partial", - "selectors": ["tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py"] -} -``` - -Selectors use pytest file or node syntax. A selector ending in `/` includes tests recursively from that directory - -Use `planned` with no selectors until an executable contract exists, `partial` for incomplete coverage, `complete` for the full contract, and `not_applicable` when a strategy does not apply. The dashboard shows passing evidence separately from coverage completeness and LOC coverage - -Unit cases use `unit_suite` instead of `selectors`, pointing to a repository-relative JSON file with this shape: - -```json -{ - "python_selectors": ["tests/test_api.py::test_decode"], - "cargo_manifest": "litellm-rust/Cargo.toml", - "cargo_package": "litellm-core", - "cargo_filter": "ocr::", - "backend": { - "environment_variable": "LITELLM_USE_RUST_OCR", - "probe": "tests.rust-python-harness.strategies.unit_tests.python_runner:ocr_backend" - }, - "mappings": [{"python": "tests/test_api.py::test_decode", "rust": "ocr::test_decode"}] -} -``` - -Names match automatically when the collected Python and Rust test names agree. Explicit `mappings` handle different names, class names, and parametrized cases. Missing or ambiguous counterparts fail validation in either direction. The Cargo filter must select the same behavior as the Python selectors - -The backend probe returns `python` or `rust` and runs at startup and before every test call, after fixtures have run. The OCR probe verifies the dispatch flag and native extension availability. Surface tests must also assert that calls reach their intended implementation to catch per-call fallback. Python outcomes must agree, and failed runs remain failures even if both backends fail identically - -## OCR fixtures - -Fixtures, provider configuration, input strategies, and recording commands live in [the OCR package](strategies/e2e_parity/sdk/ocr/fixtures/README.md). Record with provider credentials: - -```bash -uv run python -m tests.rust-python-harness.strategies.e2e_parity.sdk.ocr.fixtures.record --examples 1000 -``` - -`LITELLM_OCR_FIXTURE_DIR` and `--fixture-dir` override the default directory. Shared recording, replay, comparison, streaming, and cassette persistence live in `shared/parity/` - -Run the harness's own checks locally: - -```bash -uv run pytest -o consider_namespace_packages=true tests/rust-python-harness/shared tests/rust-python-harness/strategies/unit_tests tests/test_rust_python_harness.py -q -``` - -Existing OCR parity gaps remain visible: invalid-model provider errors differ, Reducto lacks a native contract, and the expanded Azure corpus exposes duplicate Content-Type headers. Moving the harness does not change provider responses or weaken assertions diff --git a/tests/rust-python-harness/__init__.py b/tests/rust-python-harness/__init__.py index 70362674d2b..448e24de03b 100644 --- a/tests/rust-python-harness/__init__.py +++ b/tests/rust-python-harness/__init__.py @@ -1,5 +1,4 @@ -"""Interactive Rust/Python SDK parity test harness.""" +from .cli import main +from .cli.catalog import load_catalog -from .catalog import load_catalog - -__all__ = ["load_catalog"] +__all__ = ["load_catalog", "main"] diff --git a/tests/rust-python-harness/catalog.py b/tests/rust-python-harness/catalog.py deleted file mode 100644 index f40fd5fc6b0..00000000000 --- a/tests/rust-python-harness/catalog.py +++ /dev/null @@ -1,75 +0,0 @@ -from __future__ import annotations - -import json -from pathlib import Path -from typing import Final - -from pydantic import BaseModel, ConfigDict, ValidationError - -from .shared.reporting.models import Coverage, HarnessCase, SDK_FUNCTIONS, Strategy - -STRATEGIES_ROOT: Final = Path(__file__).parent / "strategies" - - -class CaseSpec(BaseModel): - model_config = ConfigDict(frozen=True, extra="forbid") - - coverage: Coverage - selectors: tuple[str, ...] = () - note: str = "" - unit_suite: str | None = None - - -class StrategySpec(BaseModel): - model_config = ConfigDict(frozen=True, extra="forbid") - - order: int - id: str - label: str - description: str - functions: dict[str, CaseSpec] - gateway: dict[str, CaseSpec] = {} - - -def _load_strategy(source: Path) -> Strategy: - data: Final = StrategySpec.model_validate_json(source.read_text(encoding="utf-8")) - if set(data.functions) != set(SDK_FUNCTIONS): - raise ValueError(f"{source}: functions must exactly match {SDK_FUNCTIONS}") - cases: Final = tuple( - HarnessCase( - strategy_id=data.id, - strategy_label=data.label, - sdk_function=name, - coverage=case.coverage, - selectors=case.selectors, - note=case.note, - surface=surface, - unit_suite=case.unit_suite, - ) - for surface, functions in (("sdk", data.functions), ("gateway", data.gateway)) - for name in (SDK_FUNCTIONS if surface == "sdk" else functions) - for case in (functions[name],) - ) - for case in cases: - if case.coverage in {Coverage.PLANNED, Coverage.NOT_APPLICABLE} and (case.selectors or case.unit_suite): - raise ValueError(f"{source}: {case.coverage.value} case {case.key} cannot configure tests") - if any(not selector.strip() for selector in case.selectors): - raise ValueError(f"{source}: empty selector in {case.key}") - if data.id == "unit_tests" and case.selectors: - raise ValueError(f"{source}: unit_tests must configure unit_suite instead of pytest selectors") - if data.id != "unit_tests" and case.unit_suite: - raise ValueError(f"{source}: unit_suite is only valid for unit_tests") - return Strategy(data.order, data.id, data.label, data.description, source.parent, cases) - - -def load_catalog(root: Path = STRATEGIES_ROOT) -> tuple[Strategy, ...]: - sources: Final = tuple(sorted(root.glob("*/strategy.json"))) - if not sources: - raise ValueError(f"No strategy manifests found below {root}") - try: - strategies: Final = tuple(sorted((_load_strategy(source) for source in sources), key=lambda item: item.order)) - except (ValidationError, json.JSONDecodeError) as error: - raise ValueError(str(error)) from error - if len({strategy.id for strategy in strategies}) != len(strategies): - raise ValueError(f"Duplicate strategy id in {root}") - return strategies diff --git a/tests/rust-python-harness/cli.py b/tests/rust-python-harness/cli.py deleted file mode 100644 index d266a12ce92..00000000000 --- a/tests/rust-python-harness/cli.py +++ /dev/null @@ -1,245 +0,0 @@ -from __future__ import annotations - -import argparse -import importlib.util -from collections.abc import Sequence -from pathlib import Path - -from .catalog import load_catalog -from .shared.reporting.models import SDK_FUNCTIONS, HarnessCase, Strategy -from .shared.reporting.orchestration import StrategyRunner, run_strategies -from .shared.reporting.ui import make_dashboard -from .strategies.e2e_parity.runner import run as run_e2e -from .strategies.existing_e2e_test_sdk.runner import run as run_existing -from .strategies.trace_parity.runner import run as run_trace -from .strategies.unit_tests.mapping_validator import FunctionReport, build_function_report -from .strategies.unit_tests.runner import run as run_units - -REPO_ROOT = Path(__file__).resolve().parents[2] -COVERAGE_ROOT = REPO_ROOT / "target" / "rust-python-harness" - - -def _parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser( - prog="rust-python-harness", - description="Run Rust/Python parity tests with a live strategy-by-SDK-function dashboard.", - ) - parser.add_argument( - "-i", - "--interactive", - action="store_true", - help="pick strategies and SDK functions in a guided terminal menu", - ) - parser.add_argument( - "--list", action="store_true", help="show the catalog without running tests" - ) - parser.add_argument( - "--strategy", - action="append", - default=[], - metavar="ID", - help="run only this strategy", - ) - parser.add_argument( - "--function", - action="append", - default=[], - dest="sdk_functions", - choices=SDK_FUNCTIONS, - help="run only this SDK function", - ) - parser.add_argument("--surface", choices=("sdk", "gateway"), help="run only this API surface") - parser.add_argument( - "--validate-ledger", - action="store_true", - help=( - "report Python<->Rust test-parity ledger gaps and drift instead of " - "running the dashboard; narrow with --function" - ), - ) - parser.add_argument( - "--plain", - action="store_true", - help="disable the interactive terminal dashboard", - ) - parser.add_argument( - "--coverage", - action="store_true", - help="write Python reference LOC reports (HTML, JSON, and XML)", - ) - parser.add_argument( - "--pytest-arg", - action="append", - default=[], - metavar="ARG", - help="append an argument to pytest (repeatable, for example --pytest-arg=-x)", - ) - return parser - - -def _coverage_pytest_args(output_root: Path = COVERAGE_ROOT) -> tuple[str, ...]: - output_root.mkdir(parents=True, exist_ok=True) - return ( - "--cov=litellm", - "--cov-context=test", - f"--cov-report=json:{output_root / 'python.json'}", - f"--cov-report=xml:{output_root / 'python.xml'}", - f"--cov-report=html:{output_root / 'python-html'}", - ) - - -def _pick_values( - title: str, options: Sequence[tuple[str, str]], input_fn=input -) -> set[str]: - print(f"\n{title} (Enter = all)") - for index, (value, label) in enumerate(options, start=1): - print(f" {index:>2}. {label} [{value}]") - while True: - answer = input_fn("Choose numbers, comma-separated: ").strip() - if not answer: - return set() - try: - indexes = {int(part.strip()) for part in answer.split(",")} - except ValueError: - print("Please enter numbers separated by commas.") - continue - if indexes and all(1 <= index <= len(options) for index in indexes): - return {options[index - 1][0] for index in indexes} - print(f"Choose values from 1 to {len(options)}.") - - -def _interactive_filters(strategies: Sequence[Strategy]) -> tuple[set[str], set[str]]: - strategy_ids = _pick_values( - "Testing strategies", [(strategy.id, strategy.label) for strategy in strategies] - ) - sdk_functions = _pick_values( - "SDK functions", - [(name, name) for name in SDK_FUNCTIONS], - ) - return strategy_ids, sdk_functions - - -def _select( - strategies: Sequence[Strategy], strategy_ids: set[str], sdk_functions: set[str] -) -> tuple[HarnessCase, ...]: - known_ids = {strategy.id for strategy in strategies} - unknown = strategy_ids - known_ids - if unknown: - raise ValueError(f"Unknown strategy: {', '.join(sorted(unknown))}") - return tuple( - case - for strategy in strategies - if not strategy_ids or strategy.id in strategy_ids - for case in strategy.cases - if not sdk_functions or case.sdk_function in sdk_functions - ) - - -def _print_catalog(strategies: Sequence[Strategy]) -> None: - for strategy in strategies: - print(f"{strategy.id:20} {strategy.label}") - for case in strategy.cases: - selectors = ( - ", ".join(case.selectors) if case.selectors else case.unit_suite or "no test configured" - ) - print(f" {case.surface}/{case.sdk_function:12} {case.coverage.value:14} {selectors}") - - -def _print_function_report(report: FunctionReport) -> None: - print(f"\n{report.sdk_function}") - if report.ledger is None or report.audit is None: - print(" no ledger yet") - return - ledger, audit = report.ledger, report.audit - print( - f" {ledger.mapped_count}/{ledger.total_count} python tests mapped to rust " - f"({ledger.percentage}%)" - ) - print(f" {len(ledger.rust_only_tests)} rust-only tests with no python counterpart") - if audit.is_clean: - print(" ledger is in sync with the live test files") - return - for label, items in ( - ("ledger references a python test that no longer exists", audit.missing_python_tests), - ("python test exists but is not tracked in the ledger", audit.stale_python_tests), - ("ledger references a rust test that no longer exists", audit.missing_rust_tests), - ("rust test exists but is not tracked in the ledger", audit.stale_rust_tests), - ): - for item in items: - print(f" {label}: {item}") - - -def _validate_ledger(sdk_functions: set[str]) -> int: - functions = sdk_functions or set(SDK_FUNCTIONS) - reports = tuple(build_function_report(function) for function in sorted(functions)) - for report in reports: - _print_function_report(report) - return 0 if all(report.is_clean for report in reports) else 1 - - -def _resolve_runner(strategy_id: str) -> StrategyRunner: - match strategy_id: - case "e2e_parity": - return run_e2e - case "trace_parity": - return run_trace - case "unit_tests": - return run_units - case "existing_e2e_test_sdk": - return run_existing - case _: - raise ValueError(f"Unknown strategy: {strategy_id}") - - -def main(argv: Sequence[str] | None = None, *, strategy_id: str | None = None) -> int: - args = _parser().parse_args(argv) - if args.coverage and importlib.util.find_spec("pytest_cov") is None: - _parser().error( - "--coverage requires the project's pytest-cov dependency; run with " - "`poetry run python -m tests.rust-python-harness --coverage`" - ) - if args.validate_ledger: - return _validate_ledger(set(args.sdk_functions)) - catalog = load_catalog() - strategies = tuple(strategy for strategy in catalog if strategy_id is None or strategy.id == strategy_id) - if args.list: - _print_catalog(strategies) - return 0 - - strategy_ids = set(args.strategy) - sdk_functions = set(args.sdk_functions) - if args.interactive: - picked_strategies, picked_functions = _interactive_filters(strategies) - strategy_ids = strategy_ids or picked_strategies - sdk_functions = sdk_functions or picked_functions - - try: - selected = _select(strategies, strategy_ids, sdk_functions) - cases = tuple(case for case in selected if args.surface is None or case.surface == args.surface) - except ValueError as exc: - _parser().error(str(exc)) - selected_strategy_ids = {case.strategy_id for case in cases} - visible_strategies = tuple( - strategy for strategy in strategies if strategy.id in selected_strategy_ids - ) - dashboard = make_dashboard( - visible_strategies, - plain=args.plain, - confidence_strategies=strategies, - ) - pytest_args = [*args.pytest_arg] - if args.coverage: - pytest_args.extend(_coverage_pytest_args()) - with dashboard: - exit_code, run = run_strategies( - cases=cases, - repo_root=REPO_ROOT, - on_update=dashboard.update, - pytest_args=pytest_args, - resolve_runner=_resolve_runner, - ) - dashboard.finish(run, exit_code) - if args.coverage and (COVERAGE_ROOT / "python.json").exists(): - print(f"Python LOC heatmap: {COVERAGE_ROOT / 'python-html' / 'index.html'}") - print(f"Machine-readable coverage: {COVERAGE_ROOT / 'python.json'}") - return exit_code diff --git a/tests/rust-python-harness/cli/__init__.py b/tests/rust-python-harness/cli/__init__.py new file mode 100644 index 00000000000..13b995825dd --- /dev/null +++ b/tests/rust-python-harness/cli/__init__.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +import sys +from collections.abc import Sequence +from typing import Final, cast + +import click + +from ..shared.reporting.models import SDK_FUNCTIONS, SdkFunction, Strategy, Surface +from .catalog import load_catalog +from .commands import run_command, select_cases + +__all__ = ["load_catalog", "main"] + +_INTERRUPTED_EXIT_CODE: Final = 130 + + +def _function_option() -> click.Option: + return click.Option( + ("--function", "sdk_functions"), + type=click.Choice(SDK_FUNCTIONS), + multiple=True, + help="run only this SDK function; repeat to select more than one", + ) + + +def _run_all_command(strategies: Sequence[Strategy]) -> click.Command: + def run_all(sdk_functions: tuple[str, ...]) -> int: + selected_functions: Final = cast(frozenset[SdkFunction], frozenset(sdk_functions)) + cases: Final = select_cases(strategies, selected_functions) + return run_command(strategies, cases) + + return click.Command( + "all", + params=[_function_option()], + callback=run_all, + help="run every strategy", + ) + + +def _strategy_command(strategy: Strategy) -> click.Command: + params: list[click.Parameter] = [_function_option()] + if strategy.definition.surfaces: + params.append( + click.Option( + ("--surface",), + type=click.Choice(strategy.definition.surfaces), + help="run only this API surface; omit to run every surface", + ) + ) + runner_argument: Final = strategy.definition.runner_argument + if runner_argument is not None: + params.append( + click.Option( + (runner_argument.option, "runner_args"), + multiple=True, + metavar=runner_argument.metavar, + help=runner_argument.help, + ) + ) + + def run_strategy( + sdk_functions: tuple[str, ...], + surface: str | None = None, + runner_args: tuple[str, ...] = (), + ) -> int: + selected_functions: Final = cast(frozenset[SdkFunction], frozenset(sdk_functions)) + selected_surface: Final = cast(Surface | None, surface) + cases: Final = select_cases((strategy,), selected_functions, selected_surface) + return run_command((strategy,), cases, runner_args) + + return click.Command( + strategy.id, + params=params, + callback=run_strategy, + help=strategy.description, + ) + + +def _build_cli(strategies: Sequence[Strategy]) -> click.Group: + root: Final = click.Group( + "rust-python-harness", + help="Run Rust/Python parity tests with raw progress and strategy reports.", + ) + run: Final = click.Group("run", help="run one strategy or the complete harness") + run.add_command(_run_all_command(strategies)) + for strategy in strategies: + run.add_command(_strategy_command(strategy)) + root.add_command(run) + return root + + +def main(argv: Sequence[str] | None = None) -> int: + try: + strategies: Final = load_catalog() + result: Final = _build_cli(strategies).main( + args=None if argv is None else list(argv), + prog_name="rust-python-harness", + standalone_mode=False, + ) + exit_code: Final = result if isinstance(result, int) else 0 + except click.ClickException as error: + error.show() + return error.exit_code + except click.Abort: + click.echo("Aborted!", err=True) + return 1 + except KeyboardInterrupt: + sys.stderr.write("\nInterrupted\n") + return _INTERRUPTED_EXIT_CODE + if exit_code == _INTERRUPTED_EXIT_CODE: + sys.stderr.write("Interrupted\n") + return exit_code diff --git a/tests/rust-python-harness/cli/catalog.py b/tests/rust-python-harness/cli/catalog.py new file mode 100644 index 00000000000..03eb032d9c6 --- /dev/null +++ b/tests/rust-python-harness/cli/catalog.py @@ -0,0 +1,116 @@ +from __future__ import annotations + +import hashlib +import importlib +import importlib.util +import pkgutil +import sys +from pathlib import Path +from types import ModuleType +from typing import Final + +from .. import strategies as _strategies_package +from ..shared.reporting.models import SDK_FUNCTIONS, SURFACES, CaseDisposition, HarnessCase, Strategy +from ..shared.reporting.strategy import StrategyDefinition + +_STRATEGIES_PACKAGE: Final = _strategies_package +STRATEGIES_ROOT: Final = Path(_STRATEGIES_PACKAGE.__path__[0]) + + +def _load_strategy_module(name: str, folder: Path, prefix: str | None) -> ModuleType: + if prefix is not None: + return importlib.import_module(f"{prefix}.{name}") + module_name: Final = _synthetic_module_name(folder) + spec: Final = importlib.util.spec_from_file_location( + module_name, folder / "__init__.py" + ) + if spec is None or spec.loader is None: + raise ValueError(f"{folder}: cannot load strategy package") + module: Final = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + try: + spec.loader.exec_module(module) + except Exception as error: + del sys.modules[module_name] + raise ValueError(f"{folder}: cannot import strategy package: {error}") from error + return module + + +def _synthetic_module_name(folder: Path) -> str: + digest: Final = hashlib.sha1(str(folder.resolve()).encode()).hexdigest()[:8] + return f"_harness_strategy_{folder.name}_{digest}" + + +def _load_strategy(name: str, folder: Path, prefix: str | None) -> Strategy: + module: Final = _load_strategy_module(name, folder, prefix) + definition: Final = getattr(module, "STRATEGY", None) + if not isinstance(definition, StrategyDefinition): + raise ValueError(f"{folder}: __init__.py must export STRATEGY: StrategyDefinition") + if definition.id != name: + raise ValueError(f"{folder}: strategy id {definition.id!r} must match folder name {name!r}") + if definition.directory.resolve() != folder.resolve(): + raise ValueError(f"{folder}: strategy directory must be {folder}") + if len(set(definition.surfaces)) != len(definition.surfaces) or any( + surface not in SURFACES for surface in definition.surfaces + ): + raise ValueError(f"{folder}: invalid strategy surfaces: {definition.surfaces}") + keys: Final = tuple((case.surface, case.sdk_function) for case in definition.cases) + duplicates: Final = tuple(sorted(key for key in set(keys) if keys.count(key) > 1)) + if duplicates: + raise ValueError(f"{folder}: duplicate strategy cases: {duplicates}") + expected: Final = frozenset( + (surface, function) + for surface in (definition.surfaces or (None,)) + for function in SDK_FUNCTIONS + ) + actual: Final = frozenset(keys) + if actual != expected: + missing: Final = tuple(sorted(expected - actual)) + extra: Final = tuple(sorted(actual - expected)) + raise ValueError( + f"{folder}: strategy cases must exactly match its declared matrix; missing={missing}, extra={extra}" + ) + incompatible: Final = tuple( + (case.surface, case.sdk_function) + for case in definition.cases + if case.spec.disposition is CaseDisposition.RUNNABLE + and not isinstance(case.spec, definition.runnable_spec) + ) + if incompatible: + raise ValueError(f"{folder}: runnable cases do not match {definition.runnable_spec.__name__}: {incompatible}") + cases: Final = tuple( + HarnessCase( + strategy_id=definition.id, + strategy_label=definition.label, + sdk_function=case.sdk_function, + spec=case.spec, + surface=case.surface, + ) + for case in definition.cases + ) + return Strategy( + definition.order, + definition.id, + definition.label, + definition.description, + definition.directory, + cases, + definition, + ) + + +def load_catalog(root: Path | None = None) -> tuple[Strategy, ...]: + resolved: Final = STRATEGIES_ROOT if root is None else root + prefix: Final = _STRATEGIES_PACKAGE.__name__ if resolved == STRATEGIES_ROOT else None + folders: Final = tuple( + info.name for info in pkgutil.iter_modules([str(resolved)]) if info.ispkg + ) + if not folders: + raise ValueError(f"No strategy packages found below {resolved}") + strategies: Final = tuple( + _load_strategy(name, resolved / name, prefix) for name in sorted(folders) + ) + ids: Final = [strategy.id for strategy in strategies] + if len(set(ids)) != len(ids): + raise ValueError(f"Duplicate strategy id in {resolved}") + return tuple(sorted(strategies, key=lambda strategy: (strategy.order, strategy.id))) diff --git a/tests/rust-python-harness/cli/commands.py b/tests/rust-python-harness/cli/commands.py new file mode 100644 index 00000000000..f94c3277dc2 --- /dev/null +++ b/tests/rust-python-harness/cli/commands.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +from collections.abc import Sequence, Set +from dataclasses import replace +from pathlib import Path +from typing import Final + +from ..shared.reporting.models import HarnessCase, SdkFunction, Strategy, Surface +from ..shared.reporting.orchestration import run_strategies +from ..shared.reporting.ui import make_dashboard + +REPO_ROOT: Final = Path(__file__).resolve().parents[3] + + +def select_cases( + strategies: Sequence[Strategy], + sdk_functions: Set[SdkFunction], + surface: Surface | None = None, +) -> tuple[HarnessCase, ...]: + return tuple( + case + for strategy in strategies + for case in strategy.cases + if (not sdk_functions or case.sdk_function in sdk_functions) + and (surface is None or case.surface == surface) + ) + + +def run_command( + strategies: Sequence[Strategy], + cases: Sequence[HarnessCase], + runner_args: Sequence[str] = (), +) -> int: + grouped: Final = { + strategy.id: tuple(case for case in cases if case.strategy_id == strategy.id) + for strategy in strategies + } + visible: Final = tuple(strategy for strategy in strategies if grouped[strategy.id]) + runners: Final = tuple(replace(strategy, cases=grouped[strategy.id]) for strategy in visible) + dashboard: Final = make_dashboard(visible) + with dashboard: + exit_code, run = run_strategies(runners, REPO_ROOT, dashboard.update, runner_args) + if exit_code != 130: + dashboard.finish(run, exit_code) + return exit_code diff --git a/tests/rust-python-harness/cli/test_cli.py b/tests/rust-python-harness/cli/test_cli.py new file mode 100644 index 00000000000..226c89843d0 --- /dev/null +++ b/tests/rust-python-harness/cli/test_cli.py @@ -0,0 +1,458 @@ +from __future__ import annotations + +import importlib +from collections.abc import Callable, Sequence +from dataclasses import replace +from pathlib import Path +from typing import Final + +import pytest + +from ..shared.reporting.models import ( + SDK_FUNCTIONS, + SURFACES, + CaseDisposition, + HarnessCase, + HarnessRun, + RunStatus, + Strategy, +) +from ..shared.reporting.strategy import NotImplementedCaseSpec, SkippedCaseSpec, StrategyDefinition +from ..shared.reporting.ui import PlainDashboard, final_report, make_dashboard +from ..strategies.unit_tests_mapping.mappings import UNIT_TEST_CONTRACTS +from ..strategies.unit_tests_parity import UNIT_PARITY_SUITES +from ..strategies.unit_tests_rust import RUST_SUITES +from . import main +from .catalog import STRATEGIES_ROOT, load_catalog +from .commands import REPO_ROOT, select_cases + + +def _strategy_source( + *, + strategy_id: str = "example", + surfaces: tuple[str, ...] = (), + drop: tuple[str | None, str] | None = None, + duplicate: tuple[str | None, str] | None = None, + incompatible: tuple[str | None, str] | None = None, +) -> str: + cells: Final = tuple( + (surface, function) + for surface in (surfaces or (None,)) + for function in SDK_FUNCTIONS + if (surface, function) != drop + ) + definitions: Final = tuple( + ( + f"strategy.CaseDefinition({function!r}, " + "strategy.ModuleCaseSpec(coverage=models.Coverage.COMPLETE, module='tests.example'), " + f"surface={surface!r})" + if (surface, function) == incompatible + else ( + f"strategy.CaseDefinition({function!r}, " + "strategy.NotImplementedCaseSpec(reason='Not implemented yet'), " + f"surface={surface!r})" + ) + ) + for surface, function in (*cells, *((duplicate,) if duplicate is not None else ())) + ) + return ( + "import importlib\n" + "from pathlib import Path\n" + "strategy = importlib.import_module('tests.rust-python-harness.shared.reporting.strategy')\n" + "models = importlib.import_module('tests.rust-python-harness.shared.reporting.models')\n" + "runner = importlib.import_module('tests.rust-python-harness.strategies.trace_parity.runner')\n" + "rendering = importlib.import_module('tests.rust-python-harness.shared.reporting.rendering')\n" + "def render(results):\n" + " return (rendering.ReportSection('Example outcomes', " + "tuple(rendering.render_case_outcome(r) for r in results)),)\n" + f"CASES = ({','.join(definitions)},)\n" + "STRATEGY = strategy.StrategyDefinition(" + f"id={strategy_id!r}, order=1, label='Example strategy', description='Example description', " + "directory=Path(__file__).parent, runnable_spec=strategy.SuiteCaseSpec, cases=CASES, " + f"run=runner.run_trace_cases, render=render, surfaces={surfaces!r})\n" + ) + + +def _write_strategy_folder( + root: Path, + name: str = "example", + *, + init_source: str | None = None, +) -> Path: + folder: Final = root / name + folder.mkdir(parents=True) + (folder / "__init__.py").write_text(init_source or _strategy_source(), encoding="utf-8") + return folder + + +def test_should_load_surface_aware_and_function_only_strategies() -> None: + strategies: Final = load_catalog() + + assert [strategy.id for strategy in strategies] == [ + "e2e_parity", + "trace_parity", + "unit_tests_mapping", + "unit_tests_parity", + "unit_tests_rust", + ] + for strategy in strategies: + expected: Final = tuple( + (surface, function) for surface in (strategy.definition.surfaces or (None,)) for function in SDK_FUNCTIONS + ) + assert tuple((case.surface, case.sdk_function) for case in strategy.cases) == expected + + +def test_unit_strategies_use_function_only_cases() -> None: + strategies: Final = { + strategy.id: strategy + for strategy in load_catalog() + if strategy.id in {"unit_tests_mapping", "unit_tests_parity", "unit_tests_rust"} + } + + for sdk_function in SDK_FUNCTIONS: + cases: Final = tuple( + case for strategy in strategies.values() for case in strategy.cases if case.sdk_function == sdk_function + ) + assert len(cases) == 3 + assert all(case.surface is None for case in cases) + expected_mapping: Final = ( + CaseDisposition.RUNNABLE if sdk_function in UNIT_TEST_CONTRACTS else CaseDisposition.NOT_IMPLEMENTED + ) + assert cases[0].spec.disposition is expected_mapping + expected_parity: Final = ( + CaseDisposition.RUNNABLE if sdk_function in UNIT_PARITY_SUITES else CaseDisposition.NOT_IMPLEMENTED + ) + expected_rust: Final = ( + CaseDisposition.RUNNABLE if sdk_function in RUST_SUITES else CaseDisposition.NOT_IMPLEMENTED + ) + assert cases[1].spec.disposition is expected_parity + assert cases[2].spec.disposition is expected_rust + + +def test_raw_dashboard_is_always_the_default() -> None: + assert isinstance(make_dashboard(load_catalog()), PlainDashboard) + + +def test_every_strategy_folder_complies() -> None: + strategies: Final = load_catalog() + folders: Final = { + path.name for path in STRATEGIES_ROOT.iterdir() if path.is_dir() and (path / "__init__.py").exists() + } + + assert folders == {strategy.id for strategy in strategies} + for strategy in strategies: + definition: Final = strategy.definition + assert isinstance(definition, StrategyDefinition) + assert definition.directory == strategy.directory + assert not (strategy.directory / "strategy.json").exists() + assert (strategy.directory / "AGENTS.md").exists() + for case in strategy.cases: + if case.spec.disposition is CaseDisposition.RUNNABLE: + assert isinstance(case.spec, definition.runnable_spec) + + +@pytest.mark.parametrize("surfaces", ((), SURFACES)) +def test_should_reject_a_registry_missing_a_declared_matrix_cell(tmp_path: Path, surfaces: tuple[str, ...]) -> None: + surface: Final = surfaces[0] if surfaces else None + _write_strategy_folder( + tmp_path, + init_source=_strategy_source(surfaces=surfaces, drop=(surface, "count_tokens")), + ) + + with pytest.raises(ValueError, match="must exactly match its declared matrix"): + load_catalog(tmp_path) + + +def test_should_reject_a_duplicate_matrix_cell(tmp_path: Path) -> None: + _write_strategy_folder(tmp_path, init_source=_strategy_source(duplicate=(None, "ocr"))) + + with pytest.raises(ValueError, match="duplicate strategy cases"): + load_catalog(tmp_path) + + +def test_should_reject_invalid_declared_surfaces(tmp_path: Path) -> None: + _write_strategy_folder(tmp_path, init_source=_strategy_source(surfaces=("sdk", "sdk"))) + + with pytest.raises(ValueError, match="invalid strategy surfaces"): + load_catalog(tmp_path) + + +def test_should_reject_a_folder_without_a_strategy_definition(tmp_path: Path) -> None: + folder: Final = tmp_path / "example" + folder.mkdir() + (folder / "__init__.py").write_text("VALUE = 1\n", encoding="utf-8") + + with pytest.raises(ValueError, match="STRATEGY"): + load_catalog(tmp_path) + + +def test_should_reject_a_strategy_id_that_differs_from_its_folder(tmp_path: Path) -> None: + _write_strategy_folder(tmp_path, init_source=_strategy_source(strategy_id="other")) + + with pytest.raises(ValueError, match="must match folder name"): + load_catalog(tmp_path) + + +def test_should_reject_a_runnable_case_incompatible_with_the_strategy(tmp_path: Path) -> None: + _write_strategy_folder(tmp_path, init_source=_strategy_source(incompatible=(None, "ocr"))) + + with pytest.raises(ValueError, match="runnable cases do not match SuiteCaseSpec"): + load_catalog(tmp_path) + + +@pytest.mark.parametrize("case_type", (NotImplementedCaseSpec, SkippedCaseSpec)) +def test_should_reject_an_unavailable_case_with_a_blank_reason( + case_type: type[NotImplementedCaseSpec] | type[SkippedCaseSpec], +) -> None: + with pytest.raises(ValueError, match="at least 1 character"): + case_type(reason=" ") + + +def test_should_select_functions_and_surfaces() -> None: + strategy: Final = next(strategy for strategy in load_catalog() if strategy.id == "e2e_parity") + + assert tuple(case.key for case in select_cases((strategy,), {"messages"})) == ( + "e2e_parity:messages", + "e2e_parity:gateway:messages", + ) + assert tuple(case.display_name for case in select_cases((strategy,), {"ocr"}, "gateway")) == ("gateway/ocr",) + + +def _assert_unavailable_cell(strategy: Strategy, case: HarnessCase, section_title: str) -> None: + spec: Final = case.spec + assert isinstance(spec, (NotImplementedCaseSpec, SkippedCaseSpec)) + scoped: Final = replace(strategy, cases=(case,)) + exit_code, run = strategy.definition.run((case,), REPO_ROOT, lambda _: None) + result: Final = run.results[case.key] + expected: Final = ( + RunStatus.NOT_IMPLEMENTED if spec.disposition is CaseDisposition.NOT_IMPLEMENTED else RunStatus.SKIPPED + ) + report: Final = final_report(run, exit_code, (scoped,)) + + assert exit_code == 0 + assert result.status is expected + assert spec.reason in report + assert section_title in report + expected_result: Final = "NOT RUN" if expected is RunStatus.NOT_IMPLEMENTED else "SKIPPED" + expected_implemented: Final = 0 if expected is RunStatus.NOT_IMPLEMENTED else 1 + assert f"Result: {expected_result}" in report + assert f"Harness support: {expected_implemented}/1 cases implemented" in report + + +def test_every_unavailable_case_finishes_and_explains_itself() -> None: + section_titles: Final = { + "e2e_parity": "End-to-end parity outcomes", + "trace_parity": "trace comparisons", + "unit_tests_mapping": "Python/Rust unit-test mappings", + "unit_tests_parity": "Python backend parity outcomes", + "unit_tests_rust": "Native Rust unit-test outcomes", + } + unavailable: Final = tuple( + (strategy, case) + for strategy in load_catalog() + for case in strategy.cases + if case.spec.disposition is not CaseDisposition.RUNNABLE + ) + + for strategy, case in unavailable: + _assert_unavailable_cell(strategy, case, section_titles[strategy.id]) + + +@pytest.mark.parametrize( + ("strategy_id", "present", "absent"), + ( + ("e2e_parity", "--surface", "--pytest-arg"), + ("trace_parity", "--surface", "--pytest-arg"), + ("unit_tests_parity", "--pytest-arg", "--surface"), + ("unit_tests_mapping", "--detail", "--surface"), + ("unit_tests_rust", "--function", "--surface"), + ), +) +def test_strategy_help_only_lists_supported_options( + strategy_id: str, + present: str, + absent: str, + capsys: pytest.CaptureFixture[str], +) -> None: + exit_code: Final = main(["run", strategy_id, "--help"]) + captured: Final = capsys.readouterr() + + assert exit_code == 0 + assert present in captured.out + assert absent not in captured.out + + +def test_run_help_lists_all_and_every_strategy(capsys: pytest.CaptureFixture[str]) -> None: + exit_code: Final = main(["run", "--help"]) + captured: Final = capsys.readouterr() + + assert exit_code == 0 + for command in ( + "all", + "e2e_parity", + "trace_parity", + "unit_tests_mapping", + "unit_tests_parity", + "unit_tests_rust", + ): + assert command in captured.out + + +@pytest.mark.parametrize( + "argv", + ( + ("list",), + ("check",), + ("run", "--strategy", "unit_tests_parity"), + ("run", "unit_tests_parity", "--surface", "sdk"), + ("run", "unit_tests_parity", "--plain"), + ("run", "unit_tests_parity", "--runner-arg=-x"), + ("run", "all", "--pytest-arg=-x"), + ), +) +def test_removed_commands_and_options_are_rejected(argv: tuple[str, ...], capsys: pytest.CaptureFixture[str]) -> None: + exit_code: Final = main(argv) + captured: Final = capsys.readouterr() + + assert exit_code == 2 + assert captured.err + + +def test_strategy_command_forwards_repeated_filters_and_runner_arguments( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cli: Final = importlib.import_module("tests.rust-python-harness.cli") + captured: list[tuple[tuple[str, ...], tuple[str, ...], tuple[str, ...]]] = [] + + def capture_run( + strategies: Sequence[Strategy], + cases: Sequence[HarnessCase], + runner_args: Sequence[str] = (), + ) -> int: + captured.append( + ( + tuple(strategy.id for strategy in strategies), + tuple(case.display_name for case in cases), + tuple(runner_args), + ) + ) + return 0 + + monkeypatch.setattr(cli, "run_command", capture_run) + + assert ( + main( + [ + "run", + "unit_tests_parity", + "--function", + "ocr", + "--function", + "messages", + "--pytest-arg=-x", + ] + ) + == 0 + ) + assert captured == [ + (("unit_tests_parity",), ("ocr", "messages"), ("-x",)), + ] + + +def test_omitted_surface_selects_every_strategy_surface(monkeypatch: pytest.MonkeyPatch) -> None: + cli: Final = importlib.import_module("tests.rust-python-harness.cli") + selected: list[str] = [] + + def capture_run( + strategies: Sequence[Strategy], + cases: Sequence[HarnessCase], + runner_args: Sequence[str] = (), + ) -> int: + del strategies, runner_args + selected.extend(case.display_name for case in cases) + return 0 + + monkeypatch.setattr(cli, "run_command", capture_run) + + assert main(["run", "e2e_parity", "--function", "ocr"]) == 0 + assert selected == ["sdk/ocr", "gateway/ocr"] + + +def test_run_all_selects_every_declared_case_once(monkeypatch: pytest.MonkeyPatch) -> None: + cli: Final = importlib.import_module("tests.rust-python-harness.cli") + selected: list[HarnessCase] = [] + + def capture_run( + strategies: Sequence[Strategy], + cases: Sequence[HarnessCase], + runner_args: Sequence[str] = (), + ) -> int: + del strategies, runner_args + selected.extend(cases) + return 0 + + monkeypatch.setattr(cli, "run_command", capture_run) + + assert main(["run", "all", "--function", "ocr"]) == 0 + assert len(selected) == 7 + assert sum(case.surface is None for case in selected) == 3 + assert sum(case.surface is not None for case in selected) == 4 + + +def test_run_reports_not_implemented_surface_as_not_run( + capsys: pytest.CaptureFixture[str], +) -> None: + exit_code: Final = main(["run", "trace_parity", "--surface", "gateway", "--function", "ocr"]) + captured: Final = capsys.readouterr() + + assert exit_code == 0 + assert "Result: NOT RUN" in captured.out + assert "Harness support: 0/1 cases implemented" in captured.out + assert "Cases: 1 selected, 1 not implemented, 0 skipped" in captured.out + assert "Not implemented" in captured.out + assert "No gateway OCR trace-parity case is registered." in captured.out + + +def test_keyboard_interrupt_exits_cleanly( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + cli: Final = importlib.import_module("tests.rust-python-harness.cli") + + def interrupt() -> tuple[object, ...]: + raise KeyboardInterrupt + + monkeypatch.setattr(cli, "load_catalog", interrupt) + + exit_code: Final = main(["run", "all"]) + captured: Final = capsys.readouterr() + + assert exit_code == 130 + assert captured.out == "" + assert captured.err == "\nInterrupted\n" + + +def test_runner_interrupt_skips_the_completion_report( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + commands: Final = importlib.import_module("tests.rust-python-harness.cli.commands") + + def interrupt_run( + strategies: Sequence[Strategy], + repo_root: Path, + on_update: Callable[[HarnessRun], None], + runner_args: Sequence[str] = (), + ) -> tuple[int, HarnessRun]: + del repo_root, on_update, runner_args + run: Final = HarnessRun.from_cases(case for strategy in strategies for case in strategy.cases) + return 130, run + + monkeypatch.setattr(commands, "run_strategies", interrupt_run) + + exit_code: Final = main(["run", "trace_parity", "--surface", "gateway"]) + captured: Final = capsys.readouterr() + + assert exit_code == 130 + assert "Rust <-> Python parity report" not in captured.out + assert captured.err == "Interrupted\n" diff --git a/tests/rust-python-harness/conftest.py b/tests/rust-python-harness/conftest.py new file mode 100644 index 00000000000..d50d0fa4204 --- /dev/null +++ b/tests/rust-python-harness/conftest.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +from collections.abc import Callable +from pathlib import Path +from typing import Final + +import pytest + +HARNESS_ROOT: Final = Path(__file__).resolve().parents[2] + + +@pytest.fixture(autouse=True) +def subprocess_test_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("PYTHONPATH", str(HARNESS_ROOT)) + monkeypatch.setenv("PYTEST_DISABLE_PLUGIN_AUTOLOAD", "1") + + +@pytest.fixture +def cargo_project(tmp_path: Path) -> Callable[[str, str], Path]: + def create(package: str, source: str) -> Path: + manifest: Final = tmp_path / "Cargo.toml" + manifest.write_text( + f'[package]\nname = "{package}"\nversion = "0.1.0"\nedition = "2021"\n[workspace]\n' + ) + (tmp_path / "src").mkdir() + (tmp_path / "src/lib.rs").write_text(source) + return manifest + + return create diff --git a/tests/rust-python-harness/shared/native_build.py b/tests/rust-python-harness/shared/native_build.py new file mode 100644 index 00000000000..2ca7131c2c1 --- /dev/null +++ b/tests/rust-python-harness/shared/native_build.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +import importlib.util +import os +import subprocess +import sys +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +from litellm.rust_bridge import get_native_bridge, reset_native_bridge_cache + +MATURIN_SPEC: Final = "maturin==1.15.0" +BRIDGE_FEATURE: Final = "trace-parity" +_RUST_ROOT: Final = "litellm-rust" +_LOCKFILE: Final = "Cargo.lock" +_SOURCE_SUFFIXES: Final = frozenset({".rs", ".toml"}) +_FAILURE_OUTPUT_LINES: Final = 15 + + +def needs_rebuild(native_mtime: float | None, newest_source_mtime: float | None) -> bool: + if native_mtime is None: + return True + if newest_source_mtime is None: + return False + return newest_source_mtime > native_mtime + + +def _source_files(rust_root: Path) -> Iterator[Path]: + for path in rust_root.rglob("*"): + relative: Final = path.relative_to(rust_root) + if "target" in relative.parts or not path.is_file(): + continue + if path.name == _LOCKFILE or path.suffix in _SOURCE_SUFFIXES: + yield path + + +def _newest_source_mtime(repo_root: Path) -> float | None: + rust_root: Final = repo_root / _RUST_ROOT + if not rust_root.is_dir(): + return None + return max((path.stat().st_mtime for path in _source_files(rust_root)), default=None) + + +def _native_module_path() -> Path | None: + try: + spec: Final = importlib.util.find_spec("litellm.rust_bridge._native") + except (ImportError, ValueError): + return None + origin: Final = getattr(spec, "origin", None) + return Path(origin) if origin else None + + +def _drop_imported_bridge() -> None: + reset_native_bridge_cache() + for name in tuple(sys.modules): + if name.startswith("litellm.rust_bridge._native"): + del sys.modules[name] + + +def _rebuild(repo_root: Path) -> tuple[bool, str]: + command: Final = ("uvx", "--from", MATURIN_SPEC, "maturin", "develop", "--features", BRIDGE_FEATURE) + completed: Final = subprocess.run( + command, + cwd=repo_root, + env={**os.environ, "VIRTUAL_ENV": sys.prefix}, + capture_output=True, + text=True, + check=False, + ) + output: Final = f"{completed.stdout}\n{completed.stderr}".strip() + lines: Final = tuple(output.splitlines()) + return completed.returncode == 0, "\n".join(lines[-_FAILURE_OUTPUT_LINES:]) + + +def ensure_trace_bridge(repo_root: Path) -> str | None: + native_path: Final = _native_module_path() + native_mtime: Final = native_path.stat().st_mtime if native_path is not None and native_path.exists() else None + if needs_rebuild(native_mtime, _newest_source_mtime(repo_root)): + print(f"Rebuilding native Rust bridge ({BRIDGE_FEATURE} feature)...", flush=True) + succeeded: Final + output: Final + succeeded, output = _rebuild(repo_root) + if not succeeded: + return f"native Rust bridge rebuild failed:\n{output}" + _drop_imported_bridge() + bridge: Final = get_native_bridge() + if bridge is None: + return "native Rust bridge is not importable" + if getattr(bridge, "_trace", None) is None: + return f"native Rust bridge does not expose _trace; it must be built with the {BRIDGE_FEATURE} feature" + return None diff --git a/tests/rust-python-harness/shared/parity/__init__.py b/tests/rust-python-harness/shared/parity/__init__.py index f18197acd2a..e69de29bb2d 100644 --- a/tests/rust-python-harness/shared/parity/__init__.py +++ b/tests/rust-python-harness/shared/parity/__init__.py @@ -1,3 +0,0 @@ -import pytest - -pytest.register_assert_rewrite("tests.rust-python-harness.shared.parity.compare") diff --git a/tests/rust-python-harness/shared/parity/fixtures/__init__.py b/tests/rust-python-harness/shared/parity/fixtures/__init__.py index 9d48db4f9f8..6d7f8b4f048 100644 --- a/tests/rust-python-harness/shared/parity/fixtures/__init__.py +++ b/tests/rust-python-harness/shared/parity/fixtures/__init__.py @@ -1 +1,7 @@ from __future__ import annotations + +from typing import Final + +from pydantic import TypeAdapter + +JSON_OBJECT_ADAPTER: Final = TypeAdapter(dict[str, object]) diff --git a/tests/rust-python-harness/shared/parity/fixtures/cassette.py b/tests/rust-python-harness/shared/parity/fixtures/cassette.py index 03a5fc9f416..79e04d63b22 100644 --- a/tests/rust-python-harness/shared/parity/fixtures/cassette.py +++ b/tests/rust-python-harness/shared/parity/fixtures/cassette.py @@ -5,11 +5,10 @@ from datetime import datetime from itertools import accumulate from typing import Final, Literal -from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, TypeAdapter +from pydantic import AwareDatetime, BaseModel, ConfigDict, Field from vcr.serialize import serialize from vcr.serializers import yamlserializer -from .recording import RecordedInteraction from ..recorded_http import ( HttpHeader, RecordedHttpResponse, @@ -17,8 +16,8 @@ from ..recorded_http import ( RecordedResponse, RecordedStreamChunk, ) - -_OBJECT: Final = TypeAdapter(dict[str, object]) +from . import JSON_OBJECT_ADAPTER +from .recording import RecordedInteraction class _CassetteModel(BaseModel): @@ -118,7 +117,7 @@ def serialize_cassette( recorded_at: datetime, request_source: Literal["recorded", "python_replay"], ) -> str: - normalized: Final = _OBJECT.validate_python( + normalized: Final = JSON_OBJECT_ADAPTER.validate_python( yamlserializer.deserialize( serialize( { diff --git a/tests/rust-python-harness/shared/parity/fixtures/pipeline.py b/tests/rust-python-harness/shared/parity/fixtures/pipeline.py index ab5c3437c2a..8a6f72b4a4d 100644 --- a/tests/rust-python-harness/shared/parity/fixtures/pipeline.py +++ b/tests/rust-python-harness/shared/parity/fixtures/pipeline.py @@ -8,11 +8,11 @@ from types import MappingProxyType from typing import Final, Generic, Literal, Protocol, TypeVar from hypothesis.strategies import SearchStrategy -from pydantic import BaseModel from .inputs import generate_case_inputs from .recording import UpstreamEndpoint, record_upstream_interactions from .store import ( + CaseT, FixtureInput, canonical_json, fixture_cache_key, @@ -25,7 +25,6 @@ from .store import ( LOGGER: Final = logging.getLogger(__name__) InputT = TypeVar("InputT", bound=FixtureInput) InputT_contra = TypeVar("InputT_contra", bound=FixtureInput, contravariant=True) -CaseT = TypeVar("CaseT", bound=BaseModel) class RecordingInvocation(Protocol[InputT_contra]): diff --git a/tests/rust-python-harness/shared/parity/fixtures/pytest_support.py b/tests/rust-python-harness/shared/parity/fixtures/pytest_support.py index 04c097f318c..844417d1c6b 100644 --- a/tests/rust-python-harness/shared/parity/fixtures/pytest_support.py +++ b/tests/rust-python-harness/shared/parity/fixtures/pytest_support.py @@ -3,14 +3,12 @@ from __future__ import annotations import os from collections.abc import Callable from pathlib import Path -from typing import Final, TypeVar +from typing import Final import pytest -from pydantic import BaseModel, ValidationError +from pydantic import ValidationError -from .store import recorded_fixtures - -CaseT = TypeVar("CaseT", bound=BaseModel) +from .store import CaseT, fixture_directory, recorded_fixtures def parametrize_recorded_fixtures( @@ -27,9 +25,10 @@ def parametrize_recorded_fixtures( if fixture_name not in metafunc.fixturenames: return configured: Final = os.environ.get(env_var) - if configured == "": - raise pytest.UsageError(f"{env_var} is set but empty") - directory: Final = Path(configured).expanduser() if configured is not None else default_directory + try: + directory: Final = fixture_directory(None, configured, default_directory) + except ValueError as error: + raise pytest.UsageError(f"{env_var} is set but empty") from error try: fixtures: Final = recorded_fixtures(directory, case_type) except (ValidationError, ValueError) as error: diff --git a/tests/rust-python-harness/shared/parity/fixtures/recording.py b/tests/rust-python-harness/shared/parity/fixtures/recording.py index 6c23e36f20c..ee0d4b6de7b 100644 --- a/tests/rust-python-harness/shared/parity/fixtures/recording.py +++ b/tests/rust-python-harness/shared/parity/fixtures/recording.py @@ -1,23 +1,24 @@ from __future__ import annotations import queue -import threading from collections.abc import Callable, Generator, Iterable -from contextlib import contextmanager +from contextlib import AbstractContextManager from dataclasses import dataclass -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import Final, TypeVar, cast -from urllib.parse import urlsplit, urlunsplit import httpx from vcr.filters import remove_query_parameters from vcr.request import Request from ..http import ( + PARITY_PROVIDER_HOST, dropped_request_headers, dropped_response_headers, is_streaming_response, + local_response_header, + normalized_response_header, ) +from ..local_server import LocalHttpHandler, LocalHttpServer, serve_in_thread from ..recorded_http import ( HttpHeader, RecordedHttpResponse, @@ -26,7 +27,6 @@ from ..recorded_http import ( RecordedStreamChunk, ) -_PARITY_PROVIDER_HOST: Final = "parity-provider.invalid" _SECRET_HEADERS: Final = frozenset( { "authorization", @@ -61,42 +61,18 @@ def _end_to_end_headers(headers: httpx.Headers) -> tuple[HttpHeader, ...]: decoded: Final = tuple((name.decode("ascii"), value.decode("latin-1")) for name, value in headers.raw) excluded: Final = dropped_response_headers(decoded) return tuple( - HttpHeader(name=name, value=_normalized_response_header(name, value)) + HttpHeader(name=name, value=normalized_response_header(name, value)) for name, value in decoded if name.lower() not in excluded ) -def _normalized_response_header(name: str, value: str) -> str: - if name.lower() not in {"location", "operation-location"}: - return value - parsed: Final = urlsplit(value) - if not parsed.netloc: - return value - return urlunsplit(("http", _PARITY_PROVIDER_HOST, parsed.path, parsed.query, parsed.fragment)) - - -def local_response_header(name: str, value: str, provider_url: str) -> str: - if name.lower() not in {"location", "operation-location"}: - return value - parsed: Final = urlsplit(value) - if parsed.hostname != _PARITY_PROVIDER_HOST: - return value - return f"{provider_url}{parsed.path}{'?' + parsed.query if parsed.query else ''}" - - -class _RecordingProvider(ThreadingHTTPServer): - daemon_threads = True - +class _RecordingProvider(LocalHttpServer): def __init__(self, spec: UpstreamEndpoint) -> None: super().__init__(("127.0.0.1", 0), _RecordingHandler) self.spec: Final = spec self.interactions: queue.Queue[RecordedInteraction] = queue.Queue() - @property - def url(self) -> str: - return f"http://127.0.0.1:{self.server_address[1]}" - def take_interactions(self) -> tuple[RecordedInteraction, ...]: try: first: Final = self.interactions.get(timeout=5) @@ -106,9 +82,7 @@ class _RecordingProvider(ThreadingHTTPServer): return (first, *remaining) -class _RecordingHandler(BaseHTTPRequestHandler): - protocol_version = "HTTP/1.1" - +class _RecordingHandler(LocalHttpHandler): def do_POST(self) -> None: self._forward() @@ -151,7 +125,7 @@ class _RecordingHandler(BaseHTTPRequestHandler): recorded_request: Final = remove_query_parameters( Request( self.command, - f"http://{_PARITY_PROVIDER_HOST}{self.path}", + f"http://{PARITY_PROVIDER_HOST}{self.path}", request_body, {name: value for name, value in forwarded_headers if name.lower() not in _SECRET_HEADERS}, ), @@ -191,8 +165,10 @@ class _RecordingHandler(BaseHTTPRequestHandler): self.send_header("transfer-encoding", "chunked") self.end_headers() chunks: Final = tuple(self._relay_chunks(upstream.iter_bytes())) - self.wfile.write(b"0\r\n\r\n") - self.wfile.flush() + try: + self.finish_chunked() + except (BrokenPipeError, ConnectionResetError): + pass return RecordedHttpStreamResponse( kind="http_stream", status_code=upstream.status_code, @@ -202,10 +178,7 @@ class _RecordingHandler(BaseHTTPRequestHandler): def _relay_chunks(self, chunks: Iterable[bytes]) -> Generator[RecordedStreamChunk, None, None]: for chunk in chunks: - self.wfile.write(f"{len(chunk):X}\r\n".encode("ascii")) - self.wfile.write(chunk) - self.wfile.write(b"\r\n") - self.wfile.flush() + self.write_chunk(chunk) yield RecordedStreamChunk.from_bytes(chunk) def _send_response(self, status_code: int, headers: tuple[HttpHeader, ...], body: bytes) -> None: @@ -218,21 +191,8 @@ class _RecordingHandler(BaseHTTPRequestHandler): self.end_headers() self.wfile.write(body) - def log_message(self, format: str, *args: object) -> None: - return - - -@contextmanager -def _recording_provider(spec: UpstreamEndpoint) -> Generator[_RecordingProvider]: - server: Final = _RecordingProvider(spec) - thread: Final = threading.Thread(target=server.serve_forever, daemon=True) - thread.start() - try: - yield server - finally: - server.shutdown() - server.server_close() - thread.join(timeout=5) +def _recording_provider(spec: UpstreamEndpoint) -> AbstractContextManager[_RecordingProvider]: + return serve_in_thread(_RecordingProvider(spec)) def _invoke_and_take_interactions( diff --git a/tests/rust-python-harness/shared/parity/fixtures/store.py b/tests/rust-python-harness/shared/parity/fixtures/store.py index 7a10c5c6c5d..270af2a7625 100644 --- a/tests/rust-python-harness/shared/parity/fixtures/store.py +++ b/tests/rust-python-harness/shared/parity/fixtures/store.py @@ -8,15 +8,13 @@ from datetime import datetime, timezone from pathlib import Path from typing import Final, Literal, Protocol, TypeVar, cast -from pydantic import AwareDatetime, BaseModel, ConfigDict, TypeAdapter, ValidationError +from pydantic import AwareDatetime, BaseModel, ConfigDict, ValidationError +from . import JSON_OBJECT_ADAPTER from .cassette import deserialize_cassette, serialize_cassette from .recording import RecordedInteraction FIXTURE_SCHEMA_VERSION: Final = 1 -JSON_OBJECT: Final = TypeAdapter(dict[str, object]) - - class FixtureInput(Protocol): def canonical_input(self) -> dict[str, object]: ... @@ -40,10 +38,13 @@ def fixture_cache_key(case_input: FixtureInput) -> dict[str, object]: return case_input.canonical_input() -def fixture_path(directory: Path, case_input: FixtureInput) -> Path: +def _fixture_digest(case_input: FixtureInput) -> str: input_json: Final = canonical_json(fixture_cache_key(case_input)) - digest: Final = hashlib.sha256(input_json.encode("utf-8")).hexdigest() - return directory / f"{digest}.yaml" + return hashlib.sha256(input_json.encode("utf-8")).hexdigest() + + +def fixture_path(directory: Path, case_input: FixtureInput) -> Path: + return directory / f"{_fixture_digest(case_input)}.yaml" def load_fixture(directory: Path, case_input: FixtureInput, case_type: type[CaseT]) -> CaseT | None: @@ -87,7 +88,7 @@ def save_fixture( def read_fixture(path: Path, case_type: type[CaseT]) -> CaseT: contents: Final = path.read_text(encoding="utf-8") if path.suffix == ".json": - return _load_fixture(JSON_OBJECT.validate_json(contents), path, case_type) + return _load_fixture(JSON_OBJECT_ADAPTER.validate_json(contents), path, case_type) try: cassette: Final = deserialize_cassette(contents) return case_type.model_validate(cassette.case_data()) @@ -117,10 +118,12 @@ def recorded_fixtures(directory: Path, case_type: type[CaseT]) -> tuple[CaseT, . def fixture_directory(configured: Path | None, env_value: str | None, default: Path) -> Path: - return (configured or Path(env_value or default)).expanduser() + if configured is not None: + return configured.expanduser() + if env_value == "": + raise ValueError("fixture directory environment variable is set but empty") + return Path(env_value).expanduser() if env_value is not None else default.expanduser() def fixture_id(case_input: FixtureInput, prefix: str) -> str: - input_json: Final = canonical_json(case_input.canonical_input()) - digest: Final = hashlib.sha256(input_json.encode("utf-8")).hexdigest()[:8] - return f"{prefix}-{digest}" + return f"{prefix}-{_fixture_digest(case_input)[:8]}" diff --git a/tests/rust-python-harness/shared/parity/fixtures/test_cassette.py b/tests/rust-python-harness/shared/parity/fixtures/test_cassette.py index e668223782b..7850f0863f8 100644 --- a/tests/rust-python-harness/shared/parity/fixtures/test_cassette.py +++ b/tests/rust-python-harness/shared/parity/fixtures/test_cassette.py @@ -10,9 +10,6 @@ from vcr import VCR from vcr.request import Request from ..fixture_models import ParityCase, SdkInputBase -from .cassette import deserialize_cassette -from .recording import RecordedInteraction -from .store import load_fixture, save_fixture from ..recorded_http import ( HttpHeader, RecordedHttpResponse, @@ -21,6 +18,9 @@ from ..recorded_http import ( RecordedStreamChunk, ) from ..replay import replay_server +from .cassette import deserialize_cassette +from .recording import RecordedInteraction +from .store import load_fixture, save_fixture _URI: Final = "http://parity-provider.invalid/operation?api-version=1" @@ -93,3 +93,18 @@ def test_cassette_preserves_duplicate_response_headers(tmp_path: Path) -> None: save_fixture(tmp_path, sdk_input, case, (RecordedInteraction(Request("POST", _URI, b"", {}), response),)) assert load_fixture(tmp_path, sdk_input, ParityCase[_Input]) == case + + +def test_local_replay_skips_recorded_retry_delay() -> None: + response: Final = RecordedHttpResponse.from_bytes( + 200, + (HttpHeader(name="retry-after", value="5"),), + b"{}", + ) + + with replay_server() as server: + server.enqueue_response(response) + replayed: Final = httpx.post(f"{server.url}/operation", content=b"{}") + server.take_requests(1) + + assert replayed.headers["retry-after"] == "0" diff --git a/tests/rust-python-harness/shared/parity/fixtures/test_pipeline.py b/tests/rust-python-harness/shared/parity/fixtures/test_pipeline.py index 66b6b4bffdc..4535ba05bf6 100644 --- a/tests/rust-python-harness/shared/parity/fixtures/test_pipeline.py +++ b/tests/rust-python-harness/shared/parity/fixtures/test_pipeline.py @@ -2,10 +2,8 @@ from __future__ import annotations import logging import threading -from collections.abc import Generator -from contextlib import contextmanager +from contextlib import AbstractContextManager from dataclasses import dataclass -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from typing import Final, Literal @@ -14,6 +12,8 @@ import pytest from hypothesis import strategies as st from pydantic import BaseModel, ConfigDict +from ..local_server import LocalHttpHandler, LocalHttpServer, serve_in_thread +from ..recorded_http import RecordedResponse from .pipeline import ( RecordingInvocation, RecordingTarget, @@ -22,7 +22,6 @@ from .pipeline import ( ) from .recording import UpstreamEndpoint from .store import fixture_path -from ..recorded_http import RecordedResponse class _FixtureInput(BaseModel): @@ -41,20 +40,12 @@ class _ParityCase(BaseModel): provider_responses: tuple[RecordedResponse, ...] -class _Upstream(ThreadingHTTPServer): - daemon_threads = True - +class _Upstream(LocalHttpServer): def __init__(self, status: int = 200) -> None: super().__init__(("127.0.0.1", 0), _UpstreamHandler) self.response_status: Final = status - @property - def url(self) -> str: - return f"http://127.0.0.1:{self.server_address[1]}" - - -class _UpstreamHandler(BaseHTTPRequestHandler): - protocol_version = "HTTP/1.1" +class _UpstreamHandler(LocalHttpHandler): def do_POST(self) -> None: length: Final = int(self.headers.get("content-length") or "0") @@ -68,21 +59,8 @@ class _UpstreamHandler(BaseHTTPRequestHandler): self.end_headers() self.wfile.write(body) - def log_message(self, format: str, *args: object) -> None: - return - - -@contextmanager -def _upstream(status: int = 200) -> Generator[_Upstream]: - server: Final = _Upstream(status) - thread: Final = threading.Thread(target=server.serve_forever, daemon=True) - thread.start() - try: - yield server - finally: - server.shutdown() - server.server_close() - thread.join(timeout=5) +def _upstream(status: int = 200) -> AbstractContextManager[_Upstream]: + return serve_in_thread(_Upstream(status)) @dataclass(frozen=True, slots=True) diff --git a/tests/rust-python-harness/shared/parity/fixtures/test_recording.py b/tests/rust-python-harness/shared/parity/fixtures/test_recording.py index e3da59ac4d8..6181f18e89a 100644 --- a/tests/rust-python-harness/shared/parity/fixtures/test_recording.py +++ b/tests/rust-python-harness/shared/parity/fixtures/test_recording.py @@ -3,10 +3,9 @@ from __future__ import annotations import asyncio import queue import threading -from collections.abc import AsyncIterator, Callable, Generator, Iterator -from contextlib import contextmanager +from collections.abc import AsyncIterator, Callable, Iterator +from contextlib import AbstractContextManager from dataclasses import dataclass -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from typing import Final, Literal @@ -17,19 +16,8 @@ from openai._streaming import SSEDecoder from pydantic import BaseModel, ConfigDict from ..compare import assert_request_parity -from .pipeline import RecordingTarget, record_fixtures -from .recording import ( - UpstreamEndpoint, - record_upstream_interactions, - record_upstream_responses, -) -from .store import ( - FIXTURE_SCHEMA_VERSION, - fixture_path, - load_fixture, - recorded_fixtures, -) from ..inprocess import InProcessExecution, run_in_process, run_in_process_async +from ..local_server import LocalHttpHandler, LocalHttpServer, serve_in_thread from ..recorded_http import ( HttpHeader, RecordedHttpStreamResponse, @@ -45,6 +33,18 @@ from ..stream import ( consume_async_stream, consume_sync_stream, ) +from .pipeline import RecordingTarget, record_fixtures +from .recording import ( + UpstreamEndpoint, + record_upstream_interactions, + record_upstream_responses, +) +from .store import ( + FIXTURE_SCHEMA_VERSION, + fixture_path, + load_fixture, + recorded_fixtures, +) _SSE_CHUNKS: Final = ( b'data: {"choices":[{"delta":{"content":"hello"}}]}\n\n', @@ -146,9 +146,7 @@ class _Invocation: self.sdk_call(provider_url, case_input) -class _ControlledUpstream(ThreadingHTTPServer): - daemon_threads = True - +class _ControlledUpstream(LocalHttpServer): def __init__(self, stream_chunks: tuple[bytes, ...]) -> None: super().__init__(("127.0.0.1", 0), _ControlledUpstreamHandler) self.stream_chunks: Final = stream_chunks @@ -158,10 +156,6 @@ class _ControlledUpstream(ThreadingHTTPServer): self.max_active_requests: int = 0 self.request_count: int = 0 - @property - def url(self) -> str: - return f"http://127.0.0.1:{self.server_address[1]}" - def start_request(self) -> None: with self.lock: self.active_requests += 1 @@ -176,9 +170,7 @@ class _ControlledUpstream(ThreadingHTTPServer): self.active_requests -= 1 -class _ControlledUpstreamHandler(BaseHTTPRequestHandler): - protocol_version = "HTTP/1.1" - +class _ControlledUpstreamHandler(LocalHttpHandler): def do_POST(self) -> None: upstream: Final = self.server assert isinstance(upstream, _ControlledUpstream) @@ -207,13 +199,7 @@ class _ControlledUpstreamHandler(BaseHTTPRequestHandler): self.send_header("content-type", "text/event-stream") self.send_header("transfer-encoding", "chunked") self.end_headers() - for chunk in upstream.stream_chunks: - self.wfile.write(f"{len(chunk):X}\r\n".encode("ascii")) - self.wfile.write(chunk) - self.wfile.write(b"\r\n") - self.wfile.flush() - self.wfile.write(b"0\r\n\r\n") - self.wfile.flush() + self.write_chunked(upstream.stream_chunks) return if self.path == "/error": self._send_json(429, b'{"error":{"message":"rate limited"}}') @@ -252,21 +238,10 @@ class _ControlledUpstreamHandler(BaseHTTPRequestHandler): self.end_headers() self.wfile.write(body) - def log_message(self, format: str, *args: object) -> None: - return - - -@contextmanager -def _controlled_upstream(stream_chunks: tuple[bytes, ...] = _SSE_CHUNKS) -> Generator[_ControlledUpstream]: - server: Final = _ControlledUpstream(stream_chunks) - thread: Final = threading.Thread(target=server.serve_forever, daemon=True) - thread.start() - try: - yield server - finally: - server.shutdown() - server.server_close() - thread.join(timeout=5) +def _controlled_upstream( + stream_chunks: tuple[bytes, ...] = _SSE_CHUNKS, +) -> AbstractContextManager[_ControlledUpstream]: + return serve_in_thread(_ControlledUpstream(stream_chunks)) def _case(identifier: str) -> _FixtureInput: diff --git a/tests/rust-python-harness/shared/parity/http.py b/tests/rust-python-harness/shared/parity/http.py index c164e3c6549..46cdbeedb59 100644 --- a/tests/rust-python-harness/shared/parity/http.py +++ b/tests/rust-python-harness/shared/parity/http.py @@ -2,6 +2,9 @@ from __future__ import annotations from collections.abc import Iterable from typing import Final +from urllib.parse import urlsplit, urlunsplit + +PARITY_PROVIDER_HOST: Final = "parity-provider.invalid" HOP_BY_HOP_HEADERS: Final[frozenset[str]] = frozenset( { @@ -52,3 +55,21 @@ def dropped_response_headers(headers: Iterable[tuple[str, str]]) -> frozenset[st def is_streaming_response(content_type: str) -> bool: return "text/event-stream" in content_type.lower() + + +def normalized_response_header(name: str, value: str) -> str: + if name.lower() not in {"location", "operation-location"}: + return value + parsed: Final = urlsplit(value) + if not parsed.netloc: + return value + return urlunsplit(("http", PARITY_PROVIDER_HOST, parsed.path, parsed.query, parsed.fragment)) + + +def local_response_header(name: str, value: str, provider_url: str) -> str: + if name.lower() not in {"location", "operation-location"}: + return value + parsed: Final = urlsplit(value) + if parsed.hostname != PARITY_PROVIDER_HOST: + return value + return f"{provider_url}{parsed.path}{'?' + parsed.query if parsed.query else ''}" diff --git a/tests/rust-python-harness/shared/parity/ledger.py b/tests/rust-python-harness/shared/parity/ledger.py deleted file mode 100644 index 40dfed583ae..00000000000 --- a/tests/rust-python-harness/shared/parity/ledger.py +++ /dev/null @@ -1,136 +0,0 @@ -from __future__ import annotations - -import json -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -@dataclass(frozen=True, slots=True) -class LedgerEntry: - python_file: str - python_test: str - status: str - rust_file: str - rust_test: str - justification: str - reason: str - - -@dataclass(frozen=True, slots=True) -class RustOnlyEntry: - rust_file: str - rust_test: str - reason: str - - -@dataclass(frozen=True, slots=True) -class TestLedger: - sdk_function: str - python_scope: tuple[str, ...] - rust_scope: tuple[str, ...] - entries: tuple[LedgerEntry, ...] - rust_only_tests: tuple[RustOnlyEntry, ...] - - @property - def mapped_count(self) -> int: - return sum(1 for entry in self.entries if entry.status == "mapped") - - @property - def total_count(self) -> int: - return len(self.entries) - - @property - def percentage(self) -> float: - if self.total_count == 0: - return 0.0 - return round(100.0 * self.mapped_count / self.total_count, 1) - - -def _require_string(value: Any, field: str, source: Path) -> str: - if not isinstance(value, str) or not value.strip(): - raise ValueError(f"{source}: {field} must be a non-empty string") - return value - - -def _require_string_list(value: Any, field: str, source: Path) -> tuple[str, ...]: - if not isinstance(value, list) or not all(isinstance(item, str) and item for item in value): - raise ValueError(f"{source}: {field} must be a list of non-empty strings") - return tuple(value) - - -def _load_entry(data: Any, index: int, source: Path) -> LedgerEntry: - if not isinstance(data, dict): - raise ValueError(f"{source}: entries[{index}] must be an object") - python_file = _require_string(data.get("python_file"), f"entries[{index}].python_file", source) - python_test = _require_string(data.get("python_test"), f"entries[{index}].python_test", source) - status = data.get("status") - if status not in ("mapped", "unmapped"): - raise ValueError(f"{source}: entries[{index}].status must be 'mapped' or 'unmapped'") - - if status == "mapped": - rust_file = _require_string(data.get("rust_file"), f"entries[{index}].rust_file", source) - rust_test = _require_string(data.get("rust_test"), f"entries[{index}].rust_test", source) - justification = _require_string( - data.get("justification"), f"entries[{index}].justification", source - ) - return LedgerEntry( - python_file=python_file, - python_test=python_test, - status=status, - rust_file=rust_file, - rust_test=rust_test, - justification=justification, - reason="", - ) - - reason = _require_string(data.get("reason"), f"entries[{index}].reason", source) - return LedgerEntry( - python_file=python_file, - python_test=python_test, - status=status, - rust_file="", - rust_test="", - justification="", - reason=reason, - ) - - -def _load_rust_only_entry(data: Any, index: int, source: Path) -> RustOnlyEntry: - if not isinstance(data, dict): - raise ValueError(f"{source}: rust_only_tests[{index}] must be an object") - return RustOnlyEntry( - rust_file=_require_string(data.get("rust_file"), f"rust_only_tests[{index}].rust_file", source), - rust_test=_require_string(data.get("rust_test"), f"rust_only_tests[{index}].rust_test", source), - reason=_require_string(data.get("reason"), f"rust_only_tests[{index}].reason", source), - ) - - -def load_ledger(path: Path) -> TestLedger: - with path.open(encoding="utf-8") as stream: - data = json.load(stream) - - sdk_function = _require_string(data.get("sdk_function"), "sdk_function", path) - python_scope = _require_string_list(data.get("python_scope"), "python_scope", path) - rust_scope = _require_string_list(data.get("rust_scope"), "rust_scope", path) - - entries_data = data.get("entries") - if not isinstance(entries_data, list): - raise ValueError(f"{path}: entries must be a list") - entries = tuple( - _load_entry(entry, index, path) for index, entry in enumerate(entries_data) - ) - - rust_only_data = data.get("rust_only_tests") - if not isinstance(rust_only_data, list): - raise ValueError(f"{path}: rust_only_tests must be a list") - rust_only_tests = tuple( - _load_rust_only_entry(entry, index, path) for index, entry in enumerate(rust_only_data) - ) - - return TestLedger( - sdk_function=sdk_function, - python_scope=python_scope, - rust_scope=rust_scope, - entries=entries, - rust_only_tests=rust_only_tests, - ) diff --git a/tests/rust-python-harness/shared/parity/local_server.py b/tests/rust-python-harness/shared/parity/local_server.py new file mode 100644 index 00000000000..6cf2377d123 --- /dev/null +++ b/tests/rust-python-harness/shared/parity/local_server.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import threading +from collections.abc import Generator, Iterable +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Final, TypeVar + + +class LocalHttpServer(ThreadingHTTPServer): + daemon_threads = True + + @property + def url(self) -> str: + return f"http://127.0.0.1:{self.server_address[1]}" + + +class LocalHttpHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def write_chunk(self, chunk: bytes) -> None: + self.wfile.write(f"{len(chunk):X}\r\n".encode("ascii")) + self.wfile.write(chunk) + self.wfile.write(b"\r\n") + self.wfile.flush() + + def finish_chunked(self) -> None: + self.wfile.write(b"0\r\n\r\n") + self.wfile.flush() + + def write_chunked(self, chunks: Iterable[bytes]) -> None: + for chunk in chunks: + self.write_chunk(chunk) + self.finish_chunked() + + def log_message(self, format: str, *args: object) -> None: + return + + +ServerT = TypeVar("ServerT", bound=LocalHttpServer) + + +@contextmanager +def serve_in_thread(server: ServerT, poll_interval: float = 0.5) -> Generator[ServerT]: + thread: Final = threading.Thread( + target=server.serve_forever, + kwargs={"poll_interval": poll_interval}, + daemon=True, + ) + thread.start() + try: + yield server + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) diff --git a/tests/rust-python-harness/shared/parity/replay.py b/tests/rust-python-harness/shared/parity/replay.py index bde84aba3c4..c7bf76895ad 100644 --- a/tests/rust-python-harness/shared/parity/replay.py +++ b/tests/rust-python-harness/shared/parity/replay.py @@ -2,15 +2,13 @@ from __future__ import annotations import base64 import queue -import threading -from collections.abc import Generator -from contextlib import contextmanager -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from contextlib import AbstractContextManager from typing import Final from pydantic import JsonValue, TypeAdapter -from .fixtures.recording import local_response_header +from .http import local_response_header +from .local_server import LocalHttpHandler, LocalHttpServer, serve_in_thread from .models import CapturedRequest from .recorded_http import RecordedHttpResponse, RecordedHttpStreamResponse, RecordedResponse @@ -28,18 +26,18 @@ EXCLUDED_REQUEST_HEADERS: Final = frozenset( EXCLUDED_RESPONSE_HEADERS: Final = frozenset({"content-length", "transfer-encoding", "connection"}) -class ReplayServer(ThreadingHTTPServer): - daemon_threads = True +def _replay_response_header(name: str, value: str, provider_url: str) -> str: + if name.lower() == "retry-after": + return "0" + return local_response_header(name, value, provider_url) + +class ReplayServer(LocalHttpServer): def __init__(self) -> None: super().__init__(("127.0.0.1", 0), _ReplayHandler) self.responses: queue.Queue[RecordedResponse] = queue.Queue() self.requests: queue.Queue[CapturedRequest] = queue.Queue() - @property - def url(self) -> str: - return f"http://127.0.0.1:{self.server_address[1]}" - def enqueue_response(self, response: RecordedResponse) -> None: self.responses.put(response) @@ -56,9 +54,7 @@ class ReplayServer(ThreadingHTTPServer): self.requests.get_nowait() -class _ReplayHandler(BaseHTTPRequestHandler): - protocol_version = "HTTP/1.1" - +class _ReplayHandler(LocalHttpHandler): def do_POST(self) -> None: self._replay() @@ -111,7 +107,7 @@ class _ReplayHandler(BaseHTTPRequestHandler): self.send_response_only(response.status_code) for header in response.headers: if header.name.lower() not in EXCLUDED_RESPONSE_HEADERS: - self.send_header(header.name, local_response_header(header.name, header.value, provider.url)) + self.send_header(header.name, _replay_response_header(header.name, header.value, provider.url)) if isinstance(response, RecordedHttpResponse): response_body: Final = response.body_bytes() self.send_header("content-length", str(len(response_body))) @@ -121,27 +117,8 @@ class _ReplayHandler(BaseHTTPRequestHandler): assert isinstance(response, RecordedHttpStreamResponse) self.send_header("transfer-encoding", "chunked") self.end_headers() - for chunk in response.chunks: - data = chunk.data_bytes() - self.wfile.write(f"{len(data):X}\r\n".encode("ascii")) - self.wfile.write(data) - self.wfile.write(b"\r\n") - self.wfile.flush() - self.wfile.write(b"0\r\n\r\n") - self.wfile.flush() - - def log_message(self, format: str, *args: object) -> None: - return + self.write_chunked(chunk.data_bytes() for chunk in response.chunks) -@contextmanager -def replay_server() -> Generator[ReplayServer]: - server: Final = ReplayServer() - thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.01}, daemon=True) - thread.start() - try: - yield server - finally: - server.shutdown() - server.server_close() - thread.join(timeout=5) +def replay_server() -> AbstractContextManager[ReplayServer]: + return serve_in_thread(ReplayServer(), poll_interval=0.01) diff --git a/tests/rust-python-harness/shared/parity/runner.py b/tests/rust-python-harness/shared/parity/runner.py index 5add5177113..43a583382cb 100644 --- a/tests/rust-python-harness/shared/parity/runner.py +++ b/tests/rust-python-harness/shared/parity/runner.py @@ -26,6 +26,7 @@ from .replay import ReplayServer, replay_server WORKER_RESULT_PREFIX: Final = "LITELLM_PARITY_RESULT " WORKER_RESULT_ADAPTER: Final[TypeAdapter[WorkerResult]] = TypeAdapter(WorkerResult) +PROJECT_ROOT: Final = Path(__file__).resolve().parents[4] @dataclass(frozen=True, slots=True) @@ -39,7 +40,7 @@ class SubprocessRunner: sys.executable, "-m", ".".join( - self.entrypoint.resolve().relative_to(Path(__file__).resolve().parents[4]).with_suffix("").parts + self.entrypoint.resolve().relative_to(PROJECT_ROOT).with_suffix("").parts ), "--parity-worker", provider_url, @@ -54,7 +55,7 @@ class ExecutionVariant: class SubprocessWorker: def __init__(self, runner: SubprocessRunner, provider: ReplayServer, variant: ExecutionVariant) -> None: - project_root: Final = str(Path(__file__).resolve().parents[4]) + project_root: Final = str(PROJECT_ROOT) existing_pythonpath: Final = os.environ.get("PYTHONPATH") env: Final = { **os.environ, @@ -165,15 +166,6 @@ def execution_worker( worker.close() -def run_execution( - worker: SubprocessWorker, - case_file: Path, - route: str, - responses: tuple[RecordedResponse, ...], -) -> Execution: - return worker.execute(case_file, route, responses) - - @contextmanager def execution_worker_pair( runner: SubprocessRunner, diff --git a/tests/rust-python-harness/shared/parity/stream.py b/tests/rust-python-harness/shared/parity/stream.py index 72e00d3b2bd..9da7bd42d93 100644 --- a/tests/rust-python-harness/shared/parity/stream.py +++ b/tests/rust-python-harness/shared/parity/stream.py @@ -81,80 +81,61 @@ def _failed(phase: Literal["creation", "iteration"], error: Exception) -> Stream ) +def _creation_failure(error: Exception) -> StreamOutcome: + return StreamOutcome( + wrapper_type=None, + supports_sync_iteration=None, + supports_async_iteration=None, + chunks=(), + chunk_types=(), + terminal=_failed("creation", error), + ) + + +def _stream_outcome( + stream: object, + chunks: Iterable[object], + terminal: StreamTerminal, +) -> StreamOutcome: + recorded: Final = tuple(chunks) + return StreamOutcome( + wrapper_type=type(stream), + supports_sync_iteration=hasattr(stream, "__iter__"), + supports_async_iteration=hasattr(stream, "__aiter__"), + chunks=recorded, + chunk_types=tuple(type(chunk) for chunk in recorded), + terminal=terminal, + ) + + def consume_sync_stream(create: Callable[[], Iterable[object]]) -> StreamOutcome: try: stream: Final = create() except Exception as error: - return StreamOutcome( - wrapper_type=None, - supports_sync_iteration=None, - supports_async_iteration=None, - chunks=(), - chunk_types=(), - terminal=_failed("creation", error), - ) + return _creation_failure(error) chunks: list[object] = [] # mutable-ok: iterator consumption builds an ordered trace try: for chunk in stream: chunks.append(chunk) # noqa: PERF402 # partial trace is required if iteration raises except Exception as error: - recorded: Final = tuple(chunks) - return StreamOutcome( - wrapper_type=type(stream), - supports_sync_iteration=hasattr(stream, "__iter__"), - supports_async_iteration=hasattr(stream, "__aiter__"), - chunks=recorded, - chunk_types=tuple(type(chunk) for chunk in recorded), - terminal=_failed("iteration", error), - ) - completed_chunks: Final = tuple(chunks) - return StreamOutcome( - wrapper_type=type(stream), - supports_sync_iteration=hasattr(stream, "__iter__"), - supports_async_iteration=hasattr(stream, "__aiter__"), - chunks=completed_chunks, - chunk_types=tuple(type(chunk) for chunk in completed_chunks), - terminal=StreamCompleted(), - ) + return _stream_outcome(stream, chunks, _failed("iteration", error)) + return _stream_outcome(stream, chunks, StreamCompleted()) async def consume_async_stream(create: Callable[[], Awaitable[AsyncIterable[object]]]) -> StreamOutcome: try: stream: Final = await create() except Exception as error: - return StreamOutcome( - wrapper_type=None, - supports_sync_iteration=None, - supports_async_iteration=None, - chunks=(), - chunk_types=(), - terminal=_failed("creation", error), - ) + return _creation_failure(error) chunks: list[object] = [] # mutable-ok: iterator consumption builds an ordered trace try: async for chunk in stream: chunks.append(chunk) except Exception as error: - recorded: Final = tuple(chunks) - return StreamOutcome( - wrapper_type=type(stream), - supports_sync_iteration=hasattr(stream, "__iter__"), - supports_async_iteration=hasattr(stream, "__aiter__"), - chunks=recorded, - chunk_types=tuple(type(chunk) for chunk in recorded), - terminal=_failed("iteration", error), - ) - completed_chunks: Final = tuple(chunks) - return StreamOutcome( - wrapper_type=type(stream), - supports_sync_iteration=hasattr(stream, "__iter__"), - supports_async_iteration=hasattr(stream, "__aiter__"), - chunks=completed_chunks, - chunk_types=tuple(type(chunk) for chunk in completed_chunks), - terminal=StreamCompleted(), - ) + return _stream_outcome(stream, chunks, _failed("iteration", error)) + return _stream_outcome(stream, chunks, StreamCompleted()) def normalize_chunk(chunk: object) -> object: diff --git a/tests/rust-python-harness/shared/reporting/models.py b/tests/rust-python-harness/shared/reporting/models.py index 4ffacdce9ab..f78141fa552 100644 --- a/tests/rust-python-harness/shared/reporting/models.py +++ b/tests/rust-python-harness/shared/reporting/models.py @@ -1,17 +1,25 @@ from __future__ import annotations +from collections.abc import Iterable from dataclasses import dataclass, field from enum import Enum from pathlib import Path from time import monotonic -from typing import Iterable +from typing import TYPE_CHECKING, Final, Literal, TypeAlias, assert_never + +if TYPE_CHECKING: + from .strategy import CaseSpec, StrategyDefinition class Coverage(str, Enum): COMPLETE = "complete" PARTIAL = "partial" - PLANNED = "planned" - NOT_APPLICABLE = "not_applicable" + + +class CaseDisposition(str, Enum): + RUNNABLE = "runnable" + NOT_IMPLEMENTED = "not_implemented" + SKIPPED = "skipped" class RunStatus(str, Enum): @@ -23,33 +31,45 @@ class RunStatus(str, Enum): SKIPPED = "skipped" ERROR = "error" MISSING = "missing" - PLANNED = "planned" - NOT_APPLICABLE = "not_applicable" + NOT_IMPLEMENTED = "not_implemented" -class ConfidenceLevel(str, Enum): - HIGH = "HIGH" - MEDIUM = "MEDIUM" - LOW = "LOW" - - -SDK_FUNCTIONS = ("ocr", "messages", "responses", "count_tokens", "chat_completions", "transcription") +SdkFunction: TypeAlias = Literal["ocr", "messages", "responses", "count_tokens", "chat_completions", "transcription"] +Surface: TypeAlias = Literal["sdk", "gateway"] +SURFACES: Final[tuple[Surface, ...]] = ("sdk", "gateway") +SDK_FUNCTIONS: Final[tuple[SdkFunction, ...]] = ( + "ocr", + "messages", + "responses", + "count_tokens", + "chat_completions", + "transcription", +) @dataclass(frozen=True) class HarnessCase: strategy_id: str strategy_label: str - sdk_function: str - coverage: Coverage - selectors: tuple[str, ...] - note: str = "" - surface: str = "sdk" - unit_suite: str | None = None + sdk_function: SdkFunction + spec: CaseSpec + surface: Surface | None = None @property def key(self) -> str: - return f"{self.strategy_id}:{self.sdk_function}" if self.surface == "sdk" else f"{self.strategy_id}:gateway:{self.sdk_function}" + return ( + f"{self.strategy_id}:{self.sdk_function}" + if self.surface in {None, "sdk"} + else f"{self.strategy_id}:gateway:{self.sdk_function}" + ) + + @property + def display_name(self) -> str: + return self.sdk_function if self.surface is None else f"{self.surface}/{self.sdk_function}" + + @property + def coverage(self) -> Coverage | None: + return self.spec.coverage @dataclass(frozen=True) @@ -60,6 +80,7 @@ class Strategy: description: str directory: Path cases: tuple[HarnessCase, ...] + definition: StrategyDefinition @dataclass @@ -74,6 +95,7 @@ class CaseResult: errors: int = 0 outcomes: dict[str, RunStatus] = field(default_factory=dict) durations: dict[str, float] = field(default_factory=dict) + artifacts: dict[str, tuple[ResultArtifact, ...]] = field(default_factory=dict) @property def total(self) -> int: @@ -83,10 +105,18 @@ class CaseResult: def duration(self) -> float: return sum(self.durations.values()) - def record(self, nodeid: str, status: RunStatus, duration: float = 0.0) -> None: + def record( + self, + nodeid: str, + status: RunStatus, + duration: float = 0.0, + artifacts: tuple[ResultArtifact, ...] = (), + ) -> None: """Record a terminal outcome, allowing teardown errors to replace a pass.""" self.outcomes[nodeid] = status - self.durations[nodeid] = self.durations.get(nodeid, 0.0) + duration + self.add_duration(nodeid, duration) + if artifacts: + self.artifacts[nodeid] = artifacts self.completed = set(self.outcomes) values = tuple(self.outcomes.values()) self.passed = values.count(RunStatus.PASSED) @@ -95,16 +125,25 @@ class CaseResult: self.errors = values.count(RunStatus.ERROR) self.finalize() + def add_duration(self, nodeid: str, duration: float) -> None: + self.durations[nodeid] = self.durations.get(nodeid, 0.0) + duration + def set_initial_status(self) -> None: - if self.case.coverage is Coverage.NOT_APPLICABLE: - self.status = RunStatus.NOT_APPLICABLE - elif not self.case.selectors and not self.case.unit_suite: - self.status = RunStatus.PLANNED - else: - self.status = RunStatus.QUEUED + disposition: Final = self.case.spec.disposition + match disposition: + case CaseDisposition.RUNNABLE: + self.status = RunStatus.QUEUED + return + case CaseDisposition.NOT_IMPLEMENTED: + self.status = RunStatus.NOT_IMPLEMENTED + return + case CaseDisposition.SKIPPED: + self.status = RunStatus.SKIPPED + return + assert_never(disposition) def finalize(self) -> None: - if self.status in {RunStatus.NOT_APPLICABLE, RunStatus.PLANNED}: + if self.status in {RunStatus.NOT_IMPLEMENTED, RunStatus.SKIPPED} and not self.collected: return if not self.collected: self.status = RunStatus.MISSING @@ -118,11 +157,18 @@ class CaseResult: self.status = RunStatus.SKIPPED +@dataclass(frozen=True, slots=True) +class ResultArtifact: + kind: str + body: str + + @dataclass class HarnessRun: results: dict[str, CaseResult] current_nodeid: str | None = None failures: list[tuple[str, str]] = field(default_factory=list) + strategy_durations: dict[str, float] = field(default_factory=dict) started_at: float = field(default_factory=monotonic) finished_at: float | None = None @@ -131,92 +177,20 @@ class HarnessRun: return (self.finished_at or monotonic()) - self.started_at @property - def unique_tests(self) -> int: + def unique_checks(self) -> int: return len( {nodeid for result in self.results.values() for nodeid in result.collected} ) @property - def completed_tests(self) -> int: + def completed_checks(self) -> int: return len( {nodeid for result in self.results.values() for nodeid in result.completed} ) @classmethod - def from_cases(cls, cases: Iterable[HarnessCase]) -> "HarnessRun": + def from_cases(cls, cases: Iterable[HarnessCase]) -> HarnessRun: results = {case.key: CaseResult(case=case) for case in cases} for result in results.values(): result.set_initial_status() return cls(results=results) - - -@dataclass(frozen=True) -class SectionConfidence: - sdk_function: str - verified_strategies: int - required_strategies: int - level: ConfidenceLevel - details: tuple[str, ...] - - @property - def percentage(self) -> int: - if not self.required_strategies: - return 0 - return round(100 * self.verified_strategies / self.required_strategies) - - -def section_confidence( - run: HarnessRun, strategies: Iterable[Strategy] -) -> tuple[SectionConfidence, ...]: - strategy_list = tuple(strategies) - scores: list[SectionConfidence] = [] - sections = tuple(dict.fromkeys((case.surface, case.sdk_function) for strategy in strategy_list for case in strategy.cases)) - for surface, sdk_function in sections: - cases = tuple( - case - for strategy in strategy_list - for case in strategy.cases - if case.sdk_function == sdk_function and case.surface == surface - and case.coverage is not Coverage.NOT_APPLICABLE - ) - verified = 0 - details: list[str] = [] - for case in cases: - result = run.results.get(case.key) - status = result.status if result is not None else RunStatus.NOT_RUN - if status is RunStatus.PASSED: - verified += 1 - details.append( - f"{STATUS_LABELS[status]} {case.strategy_id} ({case.coverage.value})" - ) - required = len(cases) - if required and verified == required: - level = ConfidenceLevel.HIGH - elif verified: - level = ConfidenceLevel.MEDIUM - else: - level = ConfidenceLevel.LOW - scores.append( - SectionConfidence( - sdk_function=sdk_function if surface == "sdk" else f"gateway/{sdk_function}", - verified_strategies=verified, - required_strategies=required, - level=level, - details=tuple(details), - ) - ) - return tuple(scores) - - -STATUS_LABELS = { - RunStatus.NOT_RUN: "·", - RunStatus.QUEUED: "○", - RunStatus.RUNNING: "◉", - RunStatus.PASSED: "✓", - RunStatus.FAILED: "✗", - RunStatus.SKIPPED: "↷", - RunStatus.ERROR: "!", - RunStatus.MISSING: "?", - RunStatus.PLANNED: "—", - RunStatus.NOT_APPLICABLE: "n/a", -} diff --git a/tests/rust-python-harness/shared/reporting/orchestration.py b/tests/rust-python-harness/shared/reporting/orchestration.py index 5e1c0ff57d8..4740f25dc9c 100644 --- a/tests/rust-python-harness/shared/reporting/orchestration.py +++ b/tests/rust-python-harness/shared/reporting/orchestration.py @@ -1,29 +1,31 @@ from __future__ import annotations -from collections.abc import Callable, Sequence +from collections.abc import Sequence from pathlib import Path from time import monotonic -from typing import Final, Protocol +from typing import Final -from .models import HarnessCase, HarnessRun -from .pytest_runner import UpdateCallback +from .models import HarnessRun, Strategy +from .strategy import StrategyRunner, UpdateCallback + +__all__ = ["StrategyRunner", "run_strategies"] -class StrategyRunner(Protocol): - def __call__( - self, - cases: Sequence[HarnessCase], - repo_root: Path, - on_update: UpdateCallback, - pytest_args: Sequence[str] = (), - ) -> tuple[int, HarnessRun]: ... - - -def combine_reports(reports: Sequence[HarnessRun]) -> HarnessRun: +def combine_reports( + reports: Sequence[HarnessRun], + *, + timed_reports: Sequence[HarnessRun] | None = None, +) -> HarnessRun: + duration_sources: Final = reports if timed_reports is None else timed_reports return HarnessRun( results={key: result for report in reports for key, result in report.results.items()}, current_nodeid=next((report.current_nodeid for report in reversed(reports) if report.current_nodeid), None), failures=[failure for report in reports for failure in report.failures], + strategy_durations={ + strategy_id: report.duration + for report in duration_sources + for strategy_id in {result.case.strategy_id for result in report.results.values()} + }, started_at=min((report.started_at for report in reports), default=monotonic()), finished_at=( max((report.finished_at for report in reports if report.finished_at is not None), default=None) @@ -34,32 +36,38 @@ def combine_reports(reports: Sequence[HarnessRun]) -> HarnessRun: def run_strategies( - cases: Sequence[HarnessCase], + strategies: Sequence[Strategy], repo_root: Path, on_update: UpdateCallback, - pytest_args: Sequence[str], - resolve_runner: Callable[[str], StrategyRunner], + runner_args: Sequence[str] = (), ) -> tuple[int, HarnessRun]: - strategy_ids: Final = tuple(dict.fromkeys(case.strategy_id for case in cases)) + cases: Final = tuple(case for strategy in strategies for case in strategy.cases) def execute( - remaining: tuple[str, ...], reports: tuple[HarnessRun, ...], codes: tuple[int, ...] + remaining: tuple[Strategy, ...], reports: tuple[HarnessRun, ...], codes: tuple[int, ...] ) -> tuple[int, HarnessRun]: if not remaining: combined: Final = combine_reports(reports) on_update(combined) return next((code for code in codes if code), 0), combined - strategy_id, *tail = remaining - selected: Final = tuple(case for case in cases if case.strategy_id == strategy_id) - pending: Final = HarnessRun.from_cases(case for case in cases if case.strategy_id in tail) - code, report = resolve_runner(strategy_id)( + strategy, *tail = remaining + selected: Final = tuple(case for case in cases if case.strategy_id == strategy.id) + pending: Final = HarnessRun.from_cases( + case for case in cases if case.strategy_id in {later.id for later in tail} + ) + code, report = strategy.definition.run( selected, repo_root, - lambda current: on_update(combine_reports((*reports, current, pending))), - pytest_args, + lambda current: on_update( + combine_reports( + (*reports, current, pending), + timed_reports=(*reports, current), + ) + ), + runner_args, ) if code in {2, 3, 4}: return code, combine_reports((*reports, report, pending)) return execute(tuple(tail), (*reports, report), (*codes, code)) - return execute(strategy_ids, (), ()) + return execute(tuple(strategies), (), ()) diff --git a/tests/rust-python-harness/shared/reporting/pytest_runner.py b/tests/rust-python-harness/shared/reporting/pytest_runner.py deleted file mode 100644 index a7e73308f30..00000000000 --- a/tests/rust-python-harness/shared/reporting/pytest_runner.py +++ /dev/null @@ -1,172 +0,0 @@ -from __future__ import annotations - -import os -from collections.abc import Callable, Sequence -from pathlib import Path -from time import monotonic -from typing import Final - -import pytest - -from .models import CaseResult, HarnessCase, HarnessRun, RunStatus - -UpdateCallback = Callable[[HarnessRun], None] - - -def selector_matches_node(selector: str, nodeid: str) -> bool: - normalized_selector = selector.replace("\\", "/") - normalized_nodeid = nodeid.replace("\\", "/") - if normalized_selector.endswith("/"): - return normalized_nodeid.startswith(normalized_selector) - if "::" in normalized_selector: - return normalized_nodeid == normalized_selector or normalized_nodeid.startswith( - f"{normalized_selector}[" - ) - return normalized_nodeid == normalized_selector or normalized_nodeid.startswith( - f"{normalized_selector}::" - ) - - -def selector_path(selector: str) -> Path: - return Path(selector.split("::", 1)[0]) - - -def runnable_selectors( - cases: Sequence[HarnessCase], repo_root: Path -) -> tuple[str, ...]: - selectors = { - selector - for case in cases - for selector in case.selectors - if (repo_root / selector_path(selector)).exists() - } - return tuple(sorted(selectors)) - - -class HarnessPytestPlugin: - def __init__(self, run: HarnessRun, on_update: UpdateCallback) -> None: - self.run = run - self.on_update = on_update - self.node_to_results: dict[str, list[CaseResult]] = {} - - def _notify(self) -> None: - self.on_update(self.run) - - def pytest_collection_modifyitems(self, items: list[pytest.Item]) -> None: - for item in items: - matched_results: list[CaseResult] = [] - for result in self.run.results.values(): - if any( - selector_matches_node(selector, item.nodeid) - for selector in result.case.selectors - ): - result.collected.add(item.nodeid) - matched_results.append(result) - if matched_results: - self.node_to_results[item.nodeid] = matched_results - for result in self.run.results.values(): - if result.status is RunStatus.QUEUED and not result.collected: - result.status = RunStatus.MISSING - self._notify() - - def pytest_runtest_logstart( - self, nodeid: str, location: tuple[str, int | None, str] - ) -> None: - del location - self.run.current_nodeid = nodeid - for result in self.node_to_results.get(nodeid, []): - if result.status not in {RunStatus.FAILED, RunStatus.ERROR}: - result.status = RunStatus.RUNNING - self._notify() - - def pytest_runtest_logreport(self, report: pytest.TestReport) -> None: - if report.when not in {"setup", "call", "teardown"}: - return - results = self.node_to_results.get(report.nodeid, []) - if not results: - return - - terminal = report.when == "call" or report.failed or report.skipped - if not terminal: - for result in results: - result.durations[report.nodeid] = ( - result.durations.get(report.nodeid, 0.0) + report.duration - ) - return - for result in results: - if report.when == "teardown" and not report.failed: - result.durations[report.nodeid] = ( - result.durations.get(report.nodeid, 0.0) + report.duration - ) - continue - if report.skipped: - status = RunStatus.SKIPPED - elif report.failed and report.when in {"setup", "teardown"}: - status = RunStatus.ERROR - elif report.failed: - status = RunStatus.FAILED - else: - status = RunStatus.PASSED - result.record(report.nodeid, status, report.duration) - if report.failed: - failure = (report.nodeid, str(report.longrepr)) - if failure not in self.run.failures: - self.run.failures.append(failure) - self._notify() - - def pytest_sessionfinish( - self, session: pytest.Session, exitstatus: int | pytest.ExitCode - ) -> None: - del session, exitstatus - self.run.current_nodeid = None - self.run.finished_at = monotonic() - for result in self.run.results.values(): - result.finalize() - self._notify() - - -def run_pytest( - cases: Sequence[HarnessCase], - repo_root: Path, - on_update: UpdateCallback, - pytest_args: Sequence[str] = (), -) -> tuple[int, HarnessRun]: - run = HarnessRun.from_cases(cases) - selectors = runnable_selectors(cases, repo_root) - if not selectors: - for result in run.results.values(): - result.finalize() - run.finished_at = monotonic() - on_update(run) - has_missing_test = any( - result.status is RunStatus.MISSING for result in run.results.values() - ) - exit_code = ( - int(pytest.ExitCode.TESTS_FAILED) - if has_missing_test - else int(pytest.ExitCode.OK) - ) - return exit_code, run - - plugin = HarnessPytestPlugin(run=run, on_update=on_update) - args: Final = (*selectors, "-q", "--tb=no", "--no-summary", "-o", "consider_namespace_packages=true", *pytest_args) - previous_directory = Path.cwd() - try: - os.chdir(repo_root) - exit_code = int(pytest.main(list(args), plugins=[plugin])) - finally: - os.chdir(previous_directory) - for result in run.results.values(): - missing = tuple( - selector for selector in result.case.selectors - if not any(selector_matches_node(selector, node) for node in result.collected) - ) - if missing: - result.status = RunStatus.MISSING - run.failures.extend((selector, "Configured selector collected no tests") for selector in missing) - on_update(run) - if exit_code == 0 and any( - result.status is RunStatus.MISSING for result in run.results.values() - ): - exit_code = int(pytest.ExitCode.TESTS_FAILED) - return exit_code, run diff --git a/tests/rust-python-harness/shared/reporting/rendering.py b/tests/rust-python-harness/shared/reporting/rendering.py new file mode 100644 index 00000000000..217caa48526 --- /dev/null +++ b/tests/rust-python-harness/shared/reporting/rendering.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Final, Protocol, assert_never + +from .models import CaseDisposition, CaseResult + + +@dataclass(frozen=True, slots=True) +class ReportSection: + title: str + blocks: tuple[str, ...] + + +class StrategyRenderer(Protocol): + def __call__(self, results: Sequence[CaseResult]) -> tuple[ReportSection, ...]: ... + + +def render_case_outcome(result: CaseResult) -> str: + prefix: Final = f"- {result.case.display_name}: {result.status.value}" + spec: Final = result.case.spec + match spec.disposition: + case CaseDisposition.RUNNABLE: + progress: Final = f", {len(result.completed)}/{result.total} checks" if result.total else "" + return f"{prefix}{progress}, {spec.coverage.value} coverage" + case CaseDisposition.NOT_IMPLEMENTED | CaseDisposition.SKIPPED: + return f"{prefix}, {spec.reason}" + assert_never(spec.disposition) diff --git a/tests/rust-python-harness/shared/reporting/strategy.py b/tests/rust-python-harness/shared/reporting/strategy.py new file mode 100644 index 00000000000..7e76f035e20 --- /dev/null +++ b/tests/rust-python-harness/shared/reporting/strategy.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Annotated, Literal, Protocol, TypeAlias + +from pydantic import BaseModel, ConfigDict, StringConstraints + +from .models import CaseDisposition, Coverage, HarnessCase, HarnessRun, SdkFunction, Surface +from .rendering import StrategyRenderer + +UpdateCallback: TypeAlias = Callable[[HarnessRun], None] +NonBlankString: TypeAlias = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)] + + +class SuiteCaseSpec(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + disposition: Literal[CaseDisposition.RUNNABLE] = CaseDisposition.RUNNABLE + coverage: Coverage + suite: NonBlankString + note: str = "" + + +class ModuleCaseSpec(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + disposition: Literal[CaseDisposition.RUNNABLE] = CaseDisposition.RUNNABLE + coverage: Coverage + module: NonBlankString + note: str = "" + + +class NotImplementedCaseSpec(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + disposition: Literal[CaseDisposition.NOT_IMPLEMENTED] = CaseDisposition.NOT_IMPLEMENTED + coverage: None = None + reason: NonBlankString + + +class SkippedCaseSpec(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + disposition: Literal[CaseDisposition.SKIPPED] = CaseDisposition.SKIPPED + coverage: None = None + reason: NonBlankString + + +RunnableCaseSpec: TypeAlias = SuiteCaseSpec | ModuleCaseSpec +UnavailableCaseSpec: TypeAlias = NotImplementedCaseSpec | SkippedCaseSpec +CaseSpec: TypeAlias = RunnableCaseSpec | UnavailableCaseSpec + + +@dataclass(frozen=True, slots=True) +class CaseDefinition: + sdk_function: SdkFunction + spec: CaseSpec + surface: Surface | None = None + + +@dataclass(frozen=True, slots=True) +class RunnerArgumentDefinition: + option: str + help: str + metavar: str = "ARG" + + +class StrategyRunner(Protocol): + def __call__( + self, + cases: Sequence[HarnessCase], + repo_root: Path, + on_update: UpdateCallback, + runner_args: Sequence[str] = (), + ) -> tuple[int, HarnessRun]: ... + + +@dataclass(frozen=True, slots=True) +class StrategyDefinition: + id: str + order: int + label: str + description: str + directory: Path + runnable_spec: type[SuiteCaseSpec] | type[ModuleCaseSpec] + cases: tuple[CaseDefinition, ...] + run: StrategyRunner + render: StrategyRenderer + surfaces: tuple[Surface, ...] = () + runner_argument: RunnerArgumentDefinition | None = None diff --git a/tests/rust-python-harness/shared/reporting/test_orchestration.py b/tests/rust-python-harness/shared/reporting/test_orchestration.py index 8aea6d67e6e..7696ab781e2 100644 --- a/tests/rust-python-harness/shared/reporting/test_orchestration.py +++ b/tests/rust-python-harness/shared/reporting/test_orchestration.py @@ -1,47 +1,161 @@ from __future__ import annotations +import logging +from collections.abc import Sequence from pathlib import Path +from time import monotonic from typing import Final -from .models import Coverage, HarnessCase, RunStatus +from .models import CaseResult, Coverage, HarnessCase, HarnessRun, RunStatus, Strategy from .orchestration import run_strategies -from .pytest_runner import run_pytest +from .rendering import ReportSection, StrategyRenderer, render_case_outcome +from .strategy import ( + CaseDefinition, + ModuleCaseSpec, + NotImplementedCaseSpec, + StrategyDefinition, + UpdateCallback, +) +from .ui import HarnessOutputFilter, final_report -def test_combines_independent_strategy_reports_and_keeps_failures(tmp_path: Path) -> None: - (tmp_path / "test_first.py").write_text("def test_first():\n assert 1 == 2\n") - (tmp_path / "test_second.py").write_text("def test_second():\n assert True\n") - cases: Final = tuple( - HarnessCase( - strategy_id=name, - strategy_label=name, - sdk_function="ocr", - coverage=Coverage.COMPLETE, - selectors=(f"test_{name}.py",), - ) - for name in ("first", "second") +def _run_cases( + cases: Sequence[HarnessCase], + repo_root: Path, + on_update: UpdateCallback, + runner_args: Sequence[str] = (), +) -> tuple[int, HarnessRun]: + del repo_root, runner_args + run: Final = HarnessRun.from_cases(cases) + for case in cases: + _record_case(run, case, on_update) + run.finished_at = monotonic() + return int(bool(run.failures)), run + + +def _record_case(run: HarnessRun, case: HarnessCase, on_update: UpdateCallback) -> None: + result: Final = run.results[case.key] + nodeid: Final = f"check:{case.strategy_id}:{case.sdk_function}" + result.collected.add(nodeid) + failed: Final = isinstance(case.spec, ModuleCaseSpec) and case.spec.module == "fail" + result.record(nodeid, RunStatus.FAILED if failed else RunStatus.PASSED) + if failed: + run.failures.append((nodeid, "comparison failed")) + on_update(run) + + +def _render_test_results(results: Sequence[CaseResult]) -> tuple[ReportSection, ...]: + return (ReportSection("Test outcomes", tuple(render_case_outcome(result) for result in results)),) + + +def _strategy(name: str, module: str, *, render: StrategyRenderer = _render_test_results) -> Strategy: + case_definition: Final = CaseDefinition("ocr", ModuleCaseSpec(coverage=Coverage.COMPLETE, module=module)) + definition: Final = StrategyDefinition( + id=name, + order=1, + label=name, + description="Example strategy", + directory=Path.cwd(), + runnable_spec=ModuleCaseSpec, + cases=(case_definition,), + run=_run_cases, + render=render, ) - code, report = run_strategies(cases, tmp_path, lambda _: None, (), lambda _: run_pytest) + case: Final = HarnessCase( + strategy_id=name, + strategy_label=name, + sdk_function="ocr", + spec=case_definition.spec, + ) + return Strategy(1, name, name, "", Path.cwd(), (case,), definition) + + +def test_combines_strategy_reports_and_delegates_rendering() -> None: + strategies: Final = (_strategy("first", "fail"), _strategy("second", "pass")) + + code, report = run_strategies(strategies, Path.cwd(), lambda _: None) + assert code == 1 assert report.results["first:ocr"].status is RunStatus.FAILED assert report.results["second:ocr"].status is RunStatus.PASSED - assert report.completed_tests == 2 - assert len(report.failures) == 1 - assert "assert 1 == 2" in report.failures[0][1] - assert "terminalreporter" not in report.failures[0][1] + assert report.completed_checks == 2 + rendered: Final = final_report(report, code, strategies) + assert "Result: FAILED" in rendered + assert rendered.count("Test outcomes") == 2 + assert "- ocr: failed, 1/1 checks, complete coverage" in rendered + assert "- ocr: passed, 1/1 checks, complete coverage" in rendered + assert "Failures (showing 1 of 1)" in rendered + assert "Port confidence" not in rendered + assert "Slowest tests" not in rendered -def test_missing_selector_cannot_hide_behind_a_passing_surface(tmp_path: Path) -> None: - (tmp_path / "test_present.py").write_text("def test_present():\n assert True\n") - case: Final = HarnessCase( - strategy_id="e2e_parity", - strategy_label="End-to-end parity", - sdk_function="ocr", - surface="gateway", - coverage=Coverage.PARTIAL, - selectors=("test_present.py", "test_missing.py"), +def test_strategy_can_replace_the_generic_result_view() -> None: + def render_custom(results: Sequence[CaseResult]) -> tuple[ReportSection, ...]: + del results + return (ReportSection("Custom comparison", ("domain-owned diff",)),) + + strategy: Final = _strategy("custom", "pass", render=render_custom) + code, report = run_strategies((strategy,), Path.cwd(), lambda _: None) + + rendered: Final = final_report(report, code, (strategy,)) + assert "Custom comparison\ndomain-owned diff" in rendered + assert "sdk/ocr" not in rendered + + +def test_report_separates_successful_execution_from_incomplete_coverage() -> None: + runnable: Final = _strategy("mixed", "pass") + unavailable: Final = HarnessCase( + strategy_id="mixed", + strategy_label="mixed", + sdk_function="messages", + spec=NotImplementedCaseSpec(reason="No Messages case is registered."), ) - code, report = run_pytest((case,), tmp_path, lambda _: None) - assert code == 1 - assert report.results["e2e_parity:gateway:ocr"].status is RunStatus.MISSING - assert ("test_missing.py", "Configured selector collected no tests") in report.failures + code, executed = run_strategies((runnable,), Path.cwd(), lambda _: None) + unavailable_run: Final = HarnessRun.from_cases((unavailable,)) + combined: Final = HarnessRun( + results={**executed.results, **unavailable_run.results}, + started_at=executed.started_at, + finished_at=executed.finished_at, + ) + + rendered: Final = final_report(combined, code, (runnable,)) + + assert code == 0 + assert "Result: PASSED" in rendered + assert "Harness support: 1/2 cases implemented" in rendered + assert "Cases: 2 selected, 1 not implemented, 0 skipped" in rendered + + +def test_harness_output_filter_suppresses_expected_harness_warnings() -> None: + output_filter: Final = HarnessOutputFilter() + ocr_cost_warning: Final = logging.LogRecord( + "LiteLLM", + logging.WARNING, + "/repo/litellm/cost_calculator.py", + 1953, + "OCR cost: model=%s has no pricing", + ("example",), + None, + ) + other_warning: Final = logging.LogRecord( + "LiteLLM", + logging.WARNING, + "/repo/litellm/main.py", + 1, + "Provider warning", + (), + None, + ) + loop_warning: Final = logging.LogRecord( + "LiteLLM", + logging.WARNING, + "/repo/litellm/litellm_core_utils/logging_worker.py", + 129, + "LoggingWorker: event loop changed; carried %d pending and revived %d dequeued logging task(s) onto the new loop", + (1, 0), + None, + ) + + assert output_filter.filter(ocr_cost_warning) is False + assert output_filter.filter(loop_warning) is False + assert output_filter.filter(other_warning) is True diff --git a/tests/rust-python-harness/shared/reporting/ui.py b/tests/rust-python-harness/shared/reporting/ui.py index 3807af8c53b..e4b1b7b0442 100644 --- a/tests/rust-python-harness/shared/reporting/ui.py +++ b/tests/rust-python-harness/shared/reporting/ui.py @@ -1,45 +1,39 @@ from __future__ import annotations -import os -import shlex -import sys +import logging from collections.abc import Sequence from contextlib import AbstractContextManager -from pathlib import Path -from typing import Any +from textwrap import indent +from types import TracebackType +from typing import TYPE_CHECKING, Final -from .models import ( - Coverage, - HarnessRun, - RunStatus, - Strategy, - section_confidence, -) +from litellm._logging import handler as litellm_log_handler -STATUS_GLYPHS = { - RunStatus.NOT_RUN: "·", - RunStatus.QUEUED: "○", - RunStatus.RUNNING: "◉", - RunStatus.PASSED: "✓", - RunStatus.FAILED: "✗", - RunStatus.SKIPPED: "↷", - RunStatus.ERROR: "!", - RunStatus.MISSING: "?", - RunStatus.PLANNED: "—", - RunStatus.NOT_APPLICABLE: "n/a", -} +from .models import HarnessRun, RunStatus, Strategy +from .rendering import ReportSection -STATUS_STYLES = { - RunStatus.QUEUED: "dim", - RunStatus.RUNNING: "bold cyan", - RunStatus.PASSED: "bold green", - RunStatus.FAILED: "bold red", - RunStatus.SKIPPED: "yellow", - RunStatus.ERROR: "bold red", - RunStatus.MISSING: "magenta", - RunStatus.PLANNED: "dim", - RunStatus.NOT_APPLICABLE: "dim", -} +if TYPE_CHECKING: + from rich.live import Live + + +class HarnessOutputFilter(logging.Filter): + def filter(self, record: logging.LogRecord) -> bool: + noisy_prefixes: Final = ( + "OCR cost:", + "LoggingWorker: event loop changed;", + ) + return not (record.name == "LiteLLM" and record.getMessage().startswith(noisy_prefixes)) + + +_HARNESS_OUTPUT_FILTER: Final = HarnessOutputFilter() + + +def _start_output_filtering() -> None: + litellm_log_handler.addFilter(_HARNESS_OUTPUT_FILTER) + + +def _stop_output_filtering() -> None: + litellm_log_handler.removeFilter(_HARNESS_OUTPUT_FILTER) def _format_duration(seconds: float) -> str: @@ -50,250 +44,221 @@ def _format_duration(seconds: float) -> str: return f"{int(seconds // 60)}m {seconds % 60:.0f}s" -def _rerun_command(nodeid: str) -> str: - if nodeid.startswith("unit-suite:"): - return "uv run python -m tests.rust-python-harness.strategies.unit_tests.runner --plain" - return f"poetry run pytest {shlex.quote(nodeid)} -q -o consider_namespace_packages=true" - - def _summary(run: HarnessRun) -> tuple[int, int, int, int]: outcomes: dict[str, RunStatus] = {} for result in run.results.values(): outcomes.update(result.outcomes) + values: Final = tuple(outcomes.values()) return ( - list(outcomes.values()).count(RunStatus.PASSED), - list(outcomes.values()).count(RunStatus.FAILED), - list(outcomes.values()).count(RunStatus.ERROR), - list(outcomes.values()).count(RunStatus.SKIPPED), + values.count(RunStatus.PASSED), + values.count(RunStatus.FAILED), + values.count(RunStatus.ERROR), + values.count(RunStatus.SKIPPED), ) -def _cell_text(run: HarnessRun, strategy_id: str, sdk_function: str, surface: str = "sdk") -> tuple[str, str]: - key = f"{strategy_id}:{sdk_function}" if surface == "sdk" else f"{strategy_id}:gateway:{sdk_function}" - result = run.results.get(key) - if result is None: - return "", "" - counts = "" - if result.total: - counts = f" {len(result.completed)}/{result.total}" - coverage = " ◐" if result.case.coverage is Coverage.PARTIAL else "" - return f"{STATUS_GLYPHS[result.status]}{counts}{coverage}", STATUS_STYLES.get( - result.status, "" +def _strategy_state(statuses: tuple[RunStatus, ...], outcomes: tuple[RunStatus, ...]) -> str: + for status in ( + RunStatus.ERROR, + RunStatus.FAILED, + RunStatus.MISSING, + RunStatus.RUNNING, + RunStatus.QUEUED, + ): + if status in statuses: + return status.value + if RunStatus.NOT_IMPLEMENTED in statuses: + return RunStatus.NOT_IMPLEMENTED.value + if outcomes and all(outcome is RunStatus.SKIPPED for outcome in outcomes): + return RunStatus.SKIPPED.value + if statuses and all(status is RunStatus.SKIPPED for status in statuses): + return RunStatus.SKIPPED.value + for status in (RunStatus.PASSED, RunStatus.SKIPPED): + if status in statuses: + return status.value + return RunStatus.NOT_RUN.value + + +def _strategy_line(strategy: Strategy, run: HarnessRun) -> str: + results: Final = tuple(run.results[case.key] for case in strategy.cases if case.key in run.results) + outcomes: dict[str, RunStatus] = {} + collected: set[str] = set() + for result in results: + outcomes.update(result.outcomes) + collected.update(result.collected) + values: Final = tuple(outcomes.values()) + statuses: Final = tuple(result.status for result in results) + state: Final = _strategy_state(statuses, values) + completed: Final = len(outcomes) + total: Final = len(collected) + progress: Final = f", {completed}/{total} checks" if total else "" + counts: Final = ( + f", {values.count(RunStatus.PASSED)} passed, " + f"{values.count(RunStatus.FAILED) + values.count(RunStatus.ERROR)} failed, " + f"{values.count(RunStatus.SKIPPED)} skipped" + if total + else "" ) + duration: Final = run.strategy_durations.get(strategy.id, 0.0) + return f"- {strategy.label}: {state}{progress}{counts}, {_format_duration(duration)}" + + +def _rendered_sections(run: HarnessRun, strategies: Sequence[Strategy]) -> tuple[ReportSection, ...]: + return tuple( + section + for strategy in strategies + if any(case.key in run.results for case in strategy.cases) + for section in strategy.definition.render( + tuple(run.results[case.key] for case in strategy.cases if case.key in run.results) + ) + ) + + +def _format_section(section: ReportSection) -> str: + return f"{section.title}\n" + ("\n\n".join(section.blocks) or "- No results") + + +def _run_result(run: HarnessRun, exit_code: int) -> str: + statuses: Final = tuple(result.status for result in run.results.values()) + if exit_code: + return "FAILED" + if not statuses or all(status is RunStatus.NOT_IMPLEMENTED for status in statuses): + return "NOT RUN" + if all(status is RunStatus.SKIPPED for status in statuses): + return "SKIPPED" + return "PASSED" + + +def final_report(run: HarnessRun, exit_code: int, strategies: Sequence[Strategy]) -> str: + passed, failed, errors, skipped = _summary(run) + run_result: Final = _run_result(run, exit_code) + statuses: Final = tuple(case_result.status for case_result in run.results.values()) + not_implemented: Final = statuses.count(RunStatus.NOT_IMPLEMENTED) + implemented: Final = len(statuses) - not_implemented + skipped_cells: Final = statuses.count(RunStatus.SKIPPED) + failure_lines: Final = tuple( + f"{index}. {nodeid}\n{indent(detail.strip(), ' ')}" + for index, (nodeid, detail) in enumerate(run.failures[:5], start=1) + ) + rendered: Final = tuple(_format_section(section) for section in _rendered_sections(run, strategies)) + summary: Final = ( + "Rust <-> Python parity report\n\n" + f"Result: {run_result}\n" + f"Harness support: {implemented}/{len(statuses)} cases implemented\n" + f"Cases: {len(statuses)} selected, {not_implemented} not implemented, {skipped_cells} skipped\n" + f"Checks: {run.completed_checks}/{run.unique_checks} completed, {passed} passed, " + f"{failed} failed, {errors} errors, {skipped} skipped\n" + f"Duration: {_format_duration(run.duration)}\n" + f"Exit code: {exit_code}" + ) + failures: Final = ( + (f"Failures (showing {len(failure_lines)} of {len(run.failures)})\n" + "\n\n".join(failure_lines)) + if failure_lines + else "" + ) + return "\n\n".join((summary, *rendered, *((failures,) if failures else ()))) class RichDashboard(AbstractContextManager["RichDashboard"]): - def __init__( - self, - strategies: Sequence[Strategy], - confidence_strategies: Sequence[Strategy], - ) -> None: + def __init__(self, strategies: Sequence[Strategy]) -> None: from rich.console import Console from rich.live import Live self.strategies = strategies - self.confidence_strategies = confidence_strategies self.console = Console() - self.live: Any = Live( - console=self.console, refresh_per_second=12, transient=False - ) + self.live: Live = Live(console=self.console, refresh_per_second=12, transient=True) + self._live_active = False - def _table(self, run: HarnessRun) -> Any: - from rich import box - from rich.table import Table - from rich.text import Text - - columns = tuple(dict.fromkeys((case.surface, case.sdk_function) for strategy in self.strategies for case in strategy.cases)) - narrow = self.console.width < 96 - if narrow: - table = Table(box=box.SIMPLE_HEAVY, expand=True, show_header=False) - table.add_column("Strategy", ratio=3) - table.add_column("Results", ratio=5) - for strategy in self.strategies: - values = [] - for surface, sdk_function in columns: - value, style = _cell_text(run, strategy.id, sdk_function, surface) - if value: - values.append( - Text.assemble((f"{surface}/{sdk_function} ", "dim"), (value, style)) - ) - table.add_row(strategy.label, Text(" ").join(values)) - return table - - table = Table(box=box.ROUNDED, expand=True, title="Strategy × API") - table.add_column("Strategy", ratio=3) - for surface, label in columns: - table.add_column(label if surface == "sdk" else f"gateway/{label}", justify="center", ratio=1) - for strategy in self.strategies: - cells = [] - for surface, sdk_function in columns: - value, style = _cell_text(run, strategy.id, sdk_function, surface) - cells.append(Text(value, style=style)) - table.add_row(strategy.label, *cells) - return table - - def __enter__(self) -> "RichDashboard": + def __enter__(self) -> RichDashboard: + _start_output_filtering() self.live.__enter__() + self._live_active = True return self - def __exit__(self, *args: object) -> None: - self.live.__exit__(*args) + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + _stop_output_filtering() + if self._live_active: + self.live.__exit__(exc_type, exc_value, traceback) + self._live_active = False def update(self, run: HarnessRun) -> None: from rich.markup import escape - from rich.panel import Panel - active = run.current_nodeid or "Waiting for test events…" - if len(active) > max(40, self.console.width - 16): - active = f"…{active[-(self.console.width - 17):]}" + active: Final = run.current_nodeid or "Waiting for test events..." + available_width: Final = max(40, self.console.width - 10) + visible_active: Final = active if len(active) <= available_width else f"...{active[-(available_width - 3) :]}" passed, failed, errors, skipped = _summary(run) - progress = ( - f"[bold]{run.completed_tests}/{run.unique_tests}[/bold] tests " - f"[green]{passed} passed[/green] [red]{failed + errors} failed[/red] " - f"[yellow]{skipped} skipped[/yellow] [dim]{_format_duration(run.duration)}[/dim]" - ) - legend = "✓ pass ✗ fail ! error ↷ skip\n? configured test missing — planned ◐ partial coverage" + total: Final = run.unique_checks + percentage: Final = round(100 * run.completed_checks / total) if total else 0 + strategy_lines: Final = "\n".join(escape(_strategy_line(strategy, run)) for strategy in self.strategies) self.live.update( - Panel( - self._table(run), - title="⚡ Rust ↔ Python parity lab", - subtitle=f"{progress}\n[dim]{escape(active)}[/dim]\n{legend}", - border_style="cyan", - ) + "[bold]Running Rust <-> Python parity[/bold]\n" + f"Progress: [bold]{run.completed_checks}/{total} ({percentage}%)[/bold] | " + f"[green]{passed} passed[/green] | [red]{failed + errors} failed[/red] | " + f"[yellow]{skipped} skipped[/yellow] | [dim]{_format_duration(run.duration)}[/dim]\n" + f"Strategies:\n{strategy_lines}\n" + f"Current: [dim]{escape(visible_active)}[/dim]" ) def finish(self, run: HarnessRun, exit_code: int) -> None: - self.update(run) - if run.failures: - from rich.markup import escape - from rich.panel import Panel - - for nodeid, detail in run.failures[:5]: - rerun = _rerun_command(nodeid) - self.console.print( - Panel( - f"{escape(detail)}\n\n[bold]Rerun just this test[/bold]\n" - f"[cyan]{escape(rerun)}[/cyan]", - title=f"✗ {escape(nodeid)}", - border_style="red", - ) - ) - durations: dict[str, float] = {} - for result in run.results.values(): - for nodeid, duration in result.durations.items(): - durations[nodeid] = max(duration, durations.get(nodeid, 0.0)) - if durations: - slow = sorted(durations.items(), key=lambda item: item[1], reverse=True)[:3] - self.console.print( - "[bold]Slowest tests[/bold] " - + " • ".join( - f"{Path(nodeid).name} [dim]{_format_duration(duration)}[/dim]" - for nodeid, duration in slow - ) - ) - from rich import box - from rich.table import Table - - confidence_table = Table( - title="Port confidence by API", box=box.ROUNDED, expand=True - ) - confidence_table.add_column("SDK section") - confidence_table.add_column("Score", justify="right") - confidence_table.add_column("Confidence") - confidence_table.add_column("Strategy evidence", ratio=4) - confidence_styles = {"HIGH": "green", "MEDIUM": "yellow", "LOW": "red"} - for score in section_confidence(run, self.confidence_strategies): - confidence_table.add_row( - score.sdk_function, - f"{score.verified_strategies}/{score.required_strategies} {score.percentage}%", - f"[{confidence_styles[score.level.value]}]{score.level.value}[/]", - " ".join(score.details), - ) - self.console.print(confidence_table) - self.console.print( - "[dim]Score = required strategies with passing evidence. " - "LOC coverage remains a separate report.[/dim]" - ) - style = "green" if exit_code == 0 else "red" - self.console.print( - f"[{style}]Harness finished in {_format_duration(run.duration)} " - f"(exit {exit_code})[/{style}]" - ) + if self._live_active: + self.live.stop() + self._live_active = False + print(final_report(run, exit_code, self.strategies), flush=True) # noqa: T201 # CLI output class PlainDashboard(AbstractContextManager["PlainDashboard"]): - def __init__( - self, - strategies: Sequence[Strategy], - confidence_strategies: Sequence[Strategy], - ) -> None: + def __init__(self, strategies: Sequence[Strategy]) -> None: self.strategies = strategies - self.confidence_strategies = confidence_strategies self._seen: dict[str, tuple[RunStatus, int]] = {} - def __enter__(self) -> "PlainDashboard": - print("Rust <-> Python SDK parity harness", flush=True) + def __enter__(self) -> PlainDashboard: + _start_output_filtering() + print("Running Rust <-> Python parity", flush=True) # noqa: T201 # CLI output return self - def __exit__(self, *args: object) -> None: - return None + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + del exc_type, exc_value, traceback + _stop_output_filtering() def update(self, run: HarnessRun) -> None: for key, result in run.results.items(): - state = (result.status, len(result.completed)) - if self._seen.get(key) != state: - self._seen[key] = state - progress = ( - f" {len(result.completed)}/{result.total}" if result.total else "" - ) - print( - f"{STATUS_GLYPHS[result.status]} {key}: {result.status.value}{progress}", - flush=True, - ) + self._update_result(key, result.case.display_name, result.status, len(result.completed), result.total) + + def _update_result(self, key: str, label: str, status: RunStatus, completed: int, total: int) -> None: + state: Final = (status, completed) + previous: Final = self._seen.get(key) + self._seen[key] = state + visible: Final = status not in { + RunStatus.NOT_RUN, + RunStatus.QUEUED, + RunStatus.NOT_IMPLEMENTED, + RunStatus.SKIPPED, + } + should_print: Final = visible and ( + previous is None or previous[0] is not status or (completed > 0 and completed % 25 == 0) + ) + if should_print: + progress: Final = f" {completed}/{total}" if total else "" + print( # noqa: T201 # CLI output + f"{label}: {status.value}{progress}", flush=True + ) def finish(self, run: HarnessRun, exit_code: int) -> None: - self.update(run) - passed, failed, errors, skipped = _summary(run) - print( - f"Summary: {passed} passed, {failed} failed, {errors} errors, " - f"{skipped} skipped in {_format_duration(run.duration)}", - flush=True, + print( # noqa: T201 # CLI output + f"\n{final_report(run, exit_code, self.strategies)}", flush=True ) - for nodeid, detail in run.failures[:5]: - print(f"{nodeid}: {detail}", flush=True) - print(f"Rerun: {_rerun_command(nodeid)}", flush=True) - print("Port confidence by API", flush=True) - for score in section_confidence(run, self.confidence_strategies): - print( - f" {score.sdk_function:12} " - f"{score.verified_strategies}/{score.required_strategies} " - f"{score.percentage:3}% {score.level.value:6} " - f"{' | '.join(score.details)}", - flush=True, - ) - print( - " Score = required strategies with passing evidence; LOC is reported separately.", - flush=True, - ) - print(f"Harness finished with exit code {exit_code}", flush=True) -def make_dashboard( - strategies: Sequence[Strategy], - plain: bool = False, - confidence_strategies: Sequence[Strategy] | None = None, -) -> RichDashboard | PlainDashboard: - confidence_strategies = confidence_strategies or strategies - interactive_terminal = ( - sys.stdout.isatty() - and not os.environ.get("CI") - and os.environ.get("TERM") != "dumb" - ) - if not plain and interactive_terminal: - try: - import rich # noqa: F401 - - return RichDashboard(strategies, confidence_strategies) - except ImportError: - pass - return PlainDashboard(strategies, confidence_strategies) +def make_dashboard(strategies: Sequence[Strategy]) -> PlainDashboard: + return PlainDashboard(strategies) diff --git a/tests/rust-python-harness/shared/test_native_build.py b/tests/rust-python-harness/shared/test_native_build.py new file mode 100644 index 00000000000..f3e5aead846 --- /dev/null +++ b/tests/rust-python-harness/shared/test_native_build.py @@ -0,0 +1,116 @@ +from __future__ import annotations + +import os +from types import SimpleNamespace +from typing import Final + +import pytest + +from . import native_build + + +def test_needs_rebuild_when_bridge_is_missing() -> None: + assert native_build.needs_rebuild(None, 1.0) + + +def test_needs_rebuild_when_sources_are_newer_than_bridge() -> None: + assert native_build.needs_rebuild(1.0, 2.0) + + +def test_fresh_bridge_with_older_sources_needs_no_rebuild() -> None: + assert not native_build.needs_rebuild(2.0, 1.0) + + +def test_bridge_without_rust_sources_needs_no_rebuild() -> None: + assert not native_build.needs_rebuild(2.0, None) + + +def test_newest_source_mtime_tracks_rust_sources_and_skips_target(tmp_path: Final) -> None: + source: Final = tmp_path / "litellm-rust" / "crates" / "bridge" / "src" + source.mkdir(parents=True) + (source / "lib.rs").write_text("fn main() {}\n") + os.utime(source / "lib.rs", (1_000, 1_000)) + manifest: Final = tmp_path / "litellm-rust" / "crates" / "bridge" / "Cargo.toml" + manifest.write_text("[package]\n") + os.utime(manifest, (2_000, 2_000)) + lockfile: Final = tmp_path / "litellm-rust" / "Cargo.lock" + lockfile.write_text("") + os.utime(lockfile, (1_500, 1_500)) + target: Final = tmp_path / "litellm-rust" / "target" / "debug" / "junk.rs" + target.parent.mkdir(parents=True) + target.write_text("fn main() {}\n") + os.utime(target, (9_999, 9_999)) + + assert native_build._newest_source_mtime(tmp_path) == 2_000.0 + + +def test_newest_source_mtime_is_none_without_rust_workspace(tmp_path: Final) -> None: + assert native_build._newest_source_mtime(tmp_path) is None + + +def test_ensure_trace_bridge_rebuilds_when_stale( + tmp_path: Final, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] +) -> None: + native: Final = tmp_path / "_native.abi3.so" + native.write_bytes(b"") + os.utime(native, (1_000, 1_000)) + source: Final = tmp_path / "litellm-rust" / "crates" / "bridge" / "src" / "lib.rs" + source.parent.mkdir(parents=True) + source.write_text("fn main() {}\n") + os.utime(source, (2_000, 2_000)) + state: Final = SimpleNamespace(rebuilt=False) + + def fake_rebuild(repo_root: object) -> tuple[bool, str]: + state.rebuilt = True + return True, "" + + monkeypatch.setattr(native_build, "_native_module_path", lambda: native) + monkeypatch.setattr(native_build, "_rebuild", fake_rebuild) + monkeypatch.setattr(native_build, "_drop_imported_bridge", lambda: None) + monkeypatch.setattr(native_build, "get_native_bridge", lambda: SimpleNamespace(_trace=object())) + + assert native_build.ensure_trace_bridge(tmp_path) is None + assert state.rebuilt is True + assert "Rebuilding native Rust bridge" in capsys.readouterr().out + + +def test_ensure_trace_bridge_reports_failed_rebuild(tmp_path: Final, monkeypatch: pytest.MonkeyPatch) -> None: + source: Final = tmp_path / "litellm-rust" / "crates" / "bridge" / "src" / "lib.rs" + source.parent.mkdir(parents=True) + source.write_text("fn main() {}\n") + + monkeypatch.setattr(native_build, "_native_module_path", lambda: None) + monkeypatch.setattr(native_build, "_rebuild", lambda repo_root: (False, "boom")) + + message: Final = native_build.ensure_trace_bridge(tmp_path) + + assert message is not None + assert "rebuild failed" in message + assert "boom" in message + + +def test_ensure_trace_bridge_flags_missing_trace_feature_without_rebuild( + tmp_path: Final, monkeypatch: pytest.MonkeyPatch +) -> None: + native: Final = tmp_path / "_native.abi3.so" + native.write_bytes(b"") + os.utime(native, (9_999, 9_999)) + source: Final = tmp_path / "litellm-rust" / "crates" / "bridge" / "src" / "lib.rs" + source.parent.mkdir(parents=True) + source.write_text("fn main() {}\n") + os.utime(source, (1_000, 1_000)) + state: Final = SimpleNamespace(rebuilt=False) + + def fake_rebuild(repo_root: object) -> tuple[bool, str]: + state.rebuilt = True + return True, "" + + monkeypatch.setattr(native_build, "_native_module_path", lambda: native) + monkeypatch.setattr(native_build, "_rebuild", fake_rebuild) + monkeypatch.setattr(native_build, "get_native_bridge", lambda: SimpleNamespace(_trace=None)) + + message: Final = native_build.ensure_trace_bridge(tmp_path) + + assert message is not None + assert "_trace" in message + assert state.rebuilt is False diff --git a/tests/rust-python-harness/shared/tracing/compare.py b/tests/rust-python-harness/shared/tracing/compare.py deleted file mode 100644 index 9c43bea6c0e..00000000000 --- a/tests/rust-python-harness/shared/tracing/compare.py +++ /dev/null @@ -1,57 +0,0 @@ -from __future__ import annotations - -from collections import Counter -from collections.abc import Mapping, Sequence -from dataclasses import dataclass -from typing import Final - - -@dataclass(frozen=True, slots=True) -class Operation: - name: str - started: int - finished: int - - -def compare_traces( - python: Sequence[Operation], - rust: Sequence[Operation], - mapping: Mapping[str, str], - required_order: Sequence[tuple[str, str]] = (), -) -> tuple[str, ...]: - python_names: Final = {operation.name for operation in python} - rust_names: Final = {operation.name for operation in rust} - problems: Final = ( - *(f"unmapped Python operation: {name}" for name in sorted(python_names - mapping.keys())), - *(f"unmapped Rust operation: {name}" for name in sorted(rust_names - set(mapping.values()))), - *(f"ambiguous Rust operation: {name}" for name, count in Counter(mapping.values()).items() if count > 1), - *( - f"invalid interval: {operation.name}" - for operation in (*python, *rust) - if operation.started > operation.finished - ), - ) - if problems: - return problems - python_counts: Final = Counter(operation.name for operation in python) - rust_counts: Final = Counter(operation.name for operation in rust) - counts: Final = tuple( - f"call count differs for {name}: Python={python_counts[name]}, Rust={rust_counts[target]}" - for name, target in mapping.items() - if python_counts[name] != rust_counts[target] - ) - ordering: Final = tuple( - f"{label}: required order {before} before {after} was not observed" - for before, after in required_order - for label, operations, first, second in ( - ("Python", python, before, after), - ("Rust", rust, mapping.get(before), mapping.get(after)), - ) - if not first - or not second - or not any(operation.name == first for operation in operations) - or not any(operation.name == second for operation in operations) - or max(operation.finished for operation in operations if operation.name == first) - > min(operation.started for operation in operations if operation.name == second) - ) - return (*counts, *ordering) diff --git a/tests/rust-python-harness/shared/tracing/native.py b/tests/rust-python-harness/shared/tracing/native.py new file mode 100644 index 00000000000..4f988f65294 --- /dev/null +++ b/tests/rust-python-harness/shared/tracing/native.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +from typing import Final + +from pydantic import BaseModel, ConfigDict + +from .profiler import FunctionTraceEvent + + +class _TraceEventPayload(BaseModel): + model_config = ConfigDict(strict=True, extra="forbid") + id: int + parent_id: int | None + function: str + module_path: str | None = None + file: str | None = None + line: int | None = None + + +class TraceResponsePayload(BaseModel): + model_config = ConfigDict(strict=True, extra="forbid") + response: object + trace: tuple[_TraceEventPayload, ...] | list[_TraceEventPayload] + + +def native_trace_events(payload: object) -> tuple[FunctionTraceEvent, ...]: + response: Final = TraceResponsePayload.model_validate(payload) + return tuple( + FunctionTraceEvent( + event.id, + event.parent_id, + event.function, + event.module_path, + event.file, + event.line, + ) + for event in response.trace + ) diff --git a/tests/rust-python-harness/shared/tracing/profiler.py b/tests/rust-python-harness/shared/tracing/profiler.py new file mode 100644 index 00000000000..55d9818f507 --- /dev/null +++ b/tests/rust-python-harness/shared/tracing/profiler.py @@ -0,0 +1,117 @@ +from __future__ import annotations + +import sys +import threading +from collections.abc import Generator +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from types import CodeType, FrameType +from typing import Final + + +@dataclass(frozen=True, slots=True) +class FunctionTraceEvent: + id: int + parent_id: int | None + function: str + module_path: str | None = None + file: str | None = None + line: int | None = None + + @property + def raw(self) -> str: + location: Final = f"{self.file}:{self.line}" if self.file is not None and self.line is not None else "" + qualified: Final = f"{self.module_path}::{self.function}" if self.module_path is not None else self.function + return f"{location} {qualified}" if location else qualified + + +class PythonProfiler: + def __init__(self, source_root: Path) -> None: + self._source_root: Final = str(source_root.resolve()) + "/" + self._seen_frames: Final[set[FrameType]] = set() + self._event_ids: Final[dict[FrameType, int]] = {} + self.events: Final[list[FunctionTraceEvent]] = [] + + def __call__(self, frame: FrameType, event: str, _arg: object) -> None: + if event != "call" or frame in self._seen_frames: + return + function_name: Final = self.function_name(frame.f_code) + if function_name is None: + return + event_id: Final = len(self.events) + parent_id: Final = next( + (self._event_ids[ancestor] for ancestor in _frame_ancestors(frame) if ancestor in self._event_ids), + None, + ) + self._seen_frames.add(frame) + self._event_ids[frame] = event_id + self.events.append(FunctionTraceEvent(id=event_id, parent_id=parent_id, function=function_name)) + + def function_name(self, code: CodeType) -> str | None: + if not code.co_filename.startswith(self._source_root): + return None + relative: Final = code.co_filename.removeprefix(self._source_root) + return f"{relative}:{code.co_firstlineno} {getattr(code, 'co_qualname', code.co_name)}" + + +class PythonFunctionUsageProfiler: + def __init__(self, source_root: Path, functions: frozenset[str]) -> None: + self._source_root: Final = str(source_root.resolve()) + "/" + self._functions: Final = functions + self.called: Final[set[str]] = set() + + def __call__(self, frame: FrameType, event: str, _arg: object) -> None: + if event != "call": + return + code: Final = frame.f_code + if not code.co_filename.startswith(self._source_root): + return + relative: Final = code.co_filename.removeprefix(self._source_root) + function: Final = f"{relative}:{code.co_firstlineno} {getattr(code, 'co_qualname', code.co_name)}" + if function in self._functions: + self.called.add(function) + + +def _frame_ancestors(frame: FrameType) -> Generator[FrameType]: + ancestor: Final = frame.f_back + if ancestor is not None: + yield ancestor + yield from _frame_ancestors(ancestor) + + +@contextmanager +def profile_python(source_root: Path, *, threads: bool = False) -> Generator[PythonProfiler]: + profiler: Final = PythonProfiler(source_root) + previous_thread: Final = threading.getprofile() + if threads: + threading.setprofile(profiler) + previous: Final = sys.getprofile() + sys.setprofile(profiler) + try: + yield profiler + finally: + sys.setprofile(previous) + if threads: + threading.setprofile(previous_thread) + + +@contextmanager +def profile_python_function_usage( + source_root: Path, + functions: frozenset[str], + *, + threads: bool = False, +) -> Generator[PythonFunctionUsageProfiler]: + profiler: Final = PythonFunctionUsageProfiler(source_root, functions) + previous_thread: Final = threading.getprofile() + if threads: + threading.setprofile(profiler) + previous: Final = sys.getprofile() + sys.setprofile(profiler) + try: + yield profiler + finally: + sys.setprofile(previous) + if threads: + threading.setprofile(previous_thread) diff --git a/tests/rust-python-harness/shared/tracing/pytest_usage.py b/tests/rust-python-harness/shared/tracing/pytest_usage.py new file mode 100644 index 00000000000..58af174df38 --- /dev/null +++ b/tests/rust-python-harness/shared/tracing/pytest_usage.py @@ -0,0 +1,335 @@ +from __future__ import annotations + +import argparse +import ast +import importlib +import inspect +import os +import subprocess +import sys +import tempfile +import warnings +from collections.abc import Generator, Sequence +from pathlib import Path +from typing import TYPE_CHECKING, Final + +from pluggy import HookimplMarker +from pydantic import BaseModel, ConfigDict + +from .profiler import profile_python_function_usage + +if TYPE_CHECKING: + import pytest + +hookimpl: Final = HookimplMarker("pytest") + + +class PythonFunctionIdentity(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + file: str + line: int + qualname: str + + @property + def raw(self) -> str: + return f"{self.file}:{self.line} {self.qualname}" + + @property + def key(self) -> str: + return f"{self.file}::{self.qualname}" + + @classmethod + def from_trace(cls, raw: str) -> PythonFunctionIdentity: + location, separator, qualname = raw.partition(" ") + file, line_separator, line = location.rpartition(":") + if not separator or not line_separator or not file or not line.isdigit() or not qualname: + raise ValueError(f"Unrecognized Python trace function: {raw}") + return cls(file=file, line=int(line), qualname=qualname) + + +class PythonFunctionReference(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + module: str + qualname: str + + @property + def owner(self) -> str: + return self.qualname.partition(".")[0] + + def resolve(self, source_root: Path) -> PythonFunctionIdentity: + value: object = importlib.import_module(self.module) + for component in self.qualname.split("."): + value = getattr(value, component) + function: Final = inspect.unwrap(value) + code: Final = getattr(function, "__code__", None) + if code is None: + raise ValueError(f"Python function has no code object: {self.module}:{self.qualname}") + source: Final = Path(code.co_filename).resolve() + try: + relative: Final = source.relative_to(source_root.resolve()) + except ValueError as error: + raise ValueError(f"Python function is outside {source_root}: {source}") from error + return PythonFunctionIdentity( + file=relative.as_posix(), + line=code.co_firstlineno, + qualname=code.co_qualname, + ) + + +class RustFunctionIdentity(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + file: str + line: int + module_path: str + function: str + + @property + def test_module(self) -> str: + _, separator, module = self.module_path.partition("::") + if not separator: + raise ValueError(f"Rust function has no crate-qualified module: {self.module_path}") + return f"{module}::tests" + + @classmethod + def from_trace(cls, raw: str) -> RustFunctionIdentity: + location, separator, qualified = raw.partition(" ") + file, line_separator, line = location.rpartition(":") + module_path, function_separator, function = qualified.rpartition("::") + if ( + not separator + or not line_separator + or not function_separator + or not file + or not line.isdigit() + or not module_path + or not function + ): + raise ValueError(f"Unrecognized Rust trace function: {raw}") + return cls(file=file, line=int(line), module_path=module_path, function=function) + + +class PythonFunctionUsage(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + function: PythonFunctionIdentity + tests: tuple[str, ...] + + +class PythonUsageReport(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + usages: tuple[PythonFunctionUsage, ...] + collected_tests: tuple[str, ...] + exit_code: int + problems: tuple[str, ...] = () + + +def candidate_test_files( + functions: Sequence[PythonFunctionReference | PythonFunctionIdentity], + search_roots: Sequence[str], + repo_root: Path, + *, + exclude_roots: Sequence[str] = (), +) -> tuple[str, ...]: + owners: Final = frozenset( + function.owner if isinstance(function, PythonFunctionReference) else function.qualname.partition(".")[0] + for function in functions + if "." in function.qualname + and ( + isinstance(function, PythonFunctionReference) + or function.file.startswith("ocr/") + or "/ocr/" in function.file + ) + ) + top_level_functions: Final = frozenset( + function.qualname for function in functions if "." not in function.qualname and function.qualname.isidentifier() + ) + candidates: Final = tuple( + path.relative_to(repo_root).as_posix() + for root in search_roots + for path in sorted((repo_root / root).rglob("test*.py")) + if not any( + path == repo_root / excluded or path.is_relative_to(repo_root / excluded) for excluded in exclude_roots + ) + if _references_function(path, owners, top_level_functions) + ) + return tuple(dict.fromkeys(candidates)) + + +def _references_function(path: Path, owners: frozenset[str], top_level_functions: frozenset[str]) -> bool: + contents: Final = path.read_text(errors="ignore") + if any(owner in contents for owner in owners): + return True + try: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", SyntaxWarning) + tree: Final = ast.parse(contents) + except SyntaxError: + return False + aliases: Final = frozenset( + alias.asname or alias.name + for node in ast.walk(tree) + if isinstance(node, ast.ImportFrom) + for alias in node.names + if alias.name in top_level_functions + ) + names: Final = top_level_functions | aliases + return any( + isinstance(node, ast.Call) + and ( + (isinstance(node.func, ast.Name) and node.func.id in names) + or (isinstance(node.func, ast.Attribute) and node.func.attr in top_level_functions) + ) + for node in ast.walk(tree) + ) + + +class _WorkerConfig(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + functions: tuple[PythonFunctionIdentity, ...] + source_root: Path + output: Path + pytest_args: tuple[str, ...] + + +class _FunctionUsagePlugin: + def __init__(self, functions: tuple[PythonFunctionIdentity, ...], source_root: Path) -> None: + self._functions: Final = functions + self._function_names: Final = frozenset(function.raw for function in functions) + self._source_root: Final = source_root + self._tests_by_function: Final[dict[str, set[str]]] = {function.raw: set() for function in functions} + self.collected_tests: tuple[str, ...] = () + self.problems: tuple[str, ...] = () + + def pytest_collection_finish(self, session: pytest.Session) -> None: + self.collected_tests = tuple(item.nodeid for item in session.items) + + def pytest_collectreport(self, report: pytest.CollectReport) -> None: + if report.failed: + self.problems = (*self.problems, str(report.longrepr)) + + @hookimpl(hookwrapper=True) + def pytest_runtest_protocol(self, item: pytest.Item, nextitem: pytest.Item | None) -> Generator[None, object, None]: + del nextitem + with profile_python_function_usage(self._source_root, self._function_names, threads=True) as profiler: + yield + for function in self._functions: + if function.raw in profiler.called: + self._tests_by_function[function.raw].add(item.nodeid) + + def usages(self) -> tuple[PythonFunctionUsage, ...]: + return tuple( + PythonFunctionUsage( + function=function, + tests=tuple(sorted(self._tests_by_function[function.raw])), + ) + for function in self._functions + ) + + +def collect_python_function_tests( + functions: Sequence[PythonFunctionIdentity], + selectors: Sequence[str], + repo_root: Path, + *, + source_root: Path | None = None, + exclusions: Sequence[str] = (), +) -> PythonUsageReport: + selected_functions: Final = tuple(dict.fromkeys(functions)) + if not selected_functions: + raise ValueError("Python function discovery needs at least one function") + if not selectors: + raise ValueError("Python function discovery needs at least one test selector") + with tempfile.TemporaryDirectory(prefix="litellm-function-tests-") as directory: + temporary: Final = Path(directory) + config_path: Final = temporary / "config.json" + output_path: Final = temporary / "report.json" + config: Final = _WorkerConfig( + functions=selected_functions, + source_root=source_root or repo_root / "litellm", + output=output_path, + pytest_args=tuple( + ( + "-o", + "consider_namespace_packages=true", + "-p", + "no:cacheprovider", + *selectors, + *(f"--deselect={nodeid}" for nodeid in exclusions), + ) + ), + ) + config_path.write_text(config.model_dump_json()) + import_roots: Final = tuple( + dict.fromkeys( + ( + str(repo_root), + str(source_root or repo_root / "litellm"), + *( + str(path.parent if path.suffix == ".py" else path) + for selector in selectors + if (path := repo_root / selector.partition("::")[0]).exists() + ), + os.environ.get("PYTHONPATH", ""), + ) + ) + ) + env: Final = { + **os.environ, + "PYTHONPATH": os.pathsep.join(import_roots), + } + try: + result: Final = subprocess.run( + (sys.executable, "-m", __name__, str(config_path)), + cwd=repo_root, + env=env, + capture_output=True, + text=True, + timeout=600, + check=False, + ) + except (OSError, subprocess.TimeoutExpired) as error: + return PythonUsageReport(usages=(), collected_tests=(), exit_code=1, problems=(str(error),)) + if not output_path.exists(): + return PythonUsageReport( + usages=(), + collected_tests=(), + exit_code=result.returncode or 1, + problems=((result.stdout + result.stderr).strip(),), + ) + report: Final = PythonUsageReport.model_validate_json(output_path.read_text()) + process_output: Final = (result.stdout + result.stderr).strip() + if result.returncode and not report.problems and process_output: + return report.model_copy(update={"problems": (process_output,)}) + return report + + +def _run_worker(config: _WorkerConfig) -> int: + import pytest + + plugin: Final = _FunctionUsagePlugin(config.functions, config.source_root) + exit_code: Final = int(pytest.main(list(config.pytest_args), plugins=[plugin])) + report: Final = PythonUsageReport( + usages=plugin.usages(), + collected_tests=plugin.collected_tests, + exit_code=exit_code, + problems=plugin.problems, + ) + config.output.write_text(report.model_dump_json()) + return exit_code + + +def main(argv: Sequence[str] | None = None) -> int: + parser: Final = argparse.ArgumentParser() + parser.add_argument("config", type=Path) + namespace: Final = parser.parse_args(argv) + config: Final = _WorkerConfig.model_validate_json(namespace.config.read_text()) + return _run_worker(config) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/rust-python-harness/shared/tracing/steps.py b/tests/rust-python-harness/shared/tracing/steps.py new file mode 100644 index 00000000000..2475f0fdcc5 --- /dev/null +++ b/tests/rust-python-harness/shared/tracing/steps.py @@ -0,0 +1,271 @@ +from __future__ import annotations + +import re +from collections import Counter +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Final, Literal + +from .profiler import FunctionTraceEvent + +Engine = Literal["python", "rust"] + + +@dataclass(frozen=True, slots=True) +class TraceMapping: + span: str + python: re.Pattern[str] | None + rust: str | None + + +def mapping( + *, + python_frame: str | None = None, + rust_span: str | None = None, + span: str | None = None, +) -> TraceMapping: + if rust_span is None: + if python_frame is None: + raise ValueError("mapping needs a python_frame pattern, a rust_span name, or both") + if span is None: + raise ValueError("a python-only mapping needs an explicit span to compare under") + return TraceMapping(span, re.compile(python_frame), None) + if python_frame is None: + return TraceMapping(rust_span, None, rust_span) + if span is not None and span != rust_span: + raise ValueError(f"span {span!r} disagrees with rust_span {rust_span!r}") + return TraceMapping(rust_span, re.compile(python_frame), rust_span) + + +@dataclass(frozen=True, slots=True) +class TraceContract: + unordered_children_of: frozenset[str] = frozenset() + + +@dataclass(frozen=True, slots=True) +class PipelineStep: + id: int + parent_id: int | None + span: str + raw: str + + +@dataclass(frozen=True, slots=True) +class PipelineProjection: + steps: tuple[PipelineStep, ...] = () + unmatched: int = 0 + + +def _span_for(engine: Engine, function: str, mappings: Sequence[TraceMapping]) -> str | None: + matches: Final = tuple( + item.span + for item in mappings + if ( + engine == "python" + and item.python is not None + and item.python.search(function) + or engine == "rust" + and item.rust == function + ) + ) + if len(matches) > 1: + raise ValueError(f"{engine} event {function!r} matches multiple trace mappings: {matches}") + if matches: + return matches[0] + return function if engine == "rust" else None + + +def pipeline_projection( + engine: Engine, events: Sequence[FunctionTraceEvent], mappings: Sequence[TraceMapping] +) -> PipelineProjection: + raw_parents: dict[int, int | None] = {} + projected_ids: set[int] = set() + shown: list[PipelineStep] = [] + unmatched: int = 0 + for event in events: + if event.id in raw_parents: + raise ValueError(f"duplicate trace event id {event.id}") + if event.parent_id is not None and event.parent_id not in raw_parents: + raise ValueError(f"trace event {event.id} references unknown or later parent {event.parent_id}") + raw_parents[event.id] = event.parent_id + span = _span_for(engine, event.function, mappings) + if span is None: + unmatched += 1 + continue + parent_id: int | None = event.parent_id + while parent_id is not None and parent_id not in projected_ids: + parent_id = raw_parents[parent_id] + shown.append(PipelineStep(event.id, parent_id, span, event.raw)) + projected_ids.add(event.id) + return PipelineProjection(tuple(shown), unmatched) + + +@dataclass(frozen=True, slots=True) +class TraceNode: + id: int + span: str + children: tuple[TraceNode, ...] + + +def trace_depths(steps: Sequence[PipelineStep]) -> dict[int, int]: + depths: dict[int, int] = {} + for step in steps: + depths[step.id] = 0 if step.parent_id is None else depths[step.parent_id] + 1 + return depths + + +def _forest(steps: Sequence[PipelineStep]) -> tuple[TraceNode, ...]: + children: dict[int | None, list[PipelineStep]] = {} + known: set[int] = set() + for step in steps: + if step.id in known: + raise ValueError(f"duplicate projected event id {step.id}") + if step.parent_id is not None and step.parent_id not in known: + raise ValueError(f"projected event {step.id} references unknown or later parent {step.parent_id}") + known.add(step.id) + children.setdefault(step.parent_id, []).append(step) + + def node(step: PipelineStep) -> TraceNode: + return TraceNode(step.id, step.span, tuple(node(child) for child in children.get(step.id, ()))) + + return tuple(node(step) for step in children.get(None, ())) + + +def _exclusive_spans(engine: Engine, mappings: Sequence[TraceMapping]) -> frozenset[str]: + return frozenset( + item.span + for item in mappings + if (engine == "python" and item.rust is None) or (engine == "rust" and item.python is None) + ) + + +def _comparable_steps( + engine: Engine, steps: Sequence[PipelineStep], mappings: Sequence[TraceMapping] +) -> tuple[PipelineStep, ...]: + exclusive: Final = _exclusive_spans(engine, mappings) + raw_parents: Final = {step.id: step.parent_id for step in steps} + included: Final = {step.id for step in steps if step.span not in exclusive} + comparable: list[PipelineStep] = [] + for step in steps: + if step.id not in included: + continue + parent_id: int | None = step.parent_id + while parent_id is not None and parent_id not in included: + parent_id = raw_parents[parent_id] + comparable.append(PipelineStep(step.id, parent_id, step.span, step.raw)) + return tuple(comparable) + + +def _signature(node: TraceNode, contract: TraceContract) -> tuple[object, ...]: + children: tuple[tuple[object, ...], ...] = tuple(_signature(child, contract) for child in node.children) + normalized: Final = tuple(sorted(children, key=repr)) if node.span in contract.unordered_children_of else children + return (node.span, normalized) + + +def trace_signature( + engine: Engine, + steps: Sequence[PipelineStep], + mappings: Sequence[TraceMapping], + contract: TraceContract, +) -> tuple[tuple[object, ...], ...]: + return tuple(_signature(root, contract) for root in _forest(_comparable_steps(engine, steps, mappings))) + + +@dataclass(frozen=True, slots=True) +class TraceDiff: + python_only: tuple[str, ...] + rust_only: tuple[str, ...] + shared_order_matches: bool + missing_mappings: tuple[str, ...] = () + first_difference: str | None = None + + @property + def matches(self) -> bool: + return ( + not self.python_only + and not self.rust_only + and not self.missing_mappings + and self.shared_order_matches + ) + + +def _missing_mappings( + python: Sequence[PipelineStep], rust: Sequence[PipelineStep], mappings: Sequence[TraceMapping] +) -> tuple[str, ...]: + python_seen: Final = frozenset(step.span for step in python) + rust_seen: Final = frozenset(step.span for step in rust) + return tuple( + item.span + for item in mappings + if (item.python is not None and item.span not in python_seen) + or (item.rust is not None and item.span not in rust_seen) + ) + + +def _first_difference( + python: Sequence[PipelineStep], + rust: Sequence[PipelineStep], + mappings: Sequence[TraceMapping], + contract: TraceContract, +) -> str | None: + python_forest: Final = _forest(_comparable_steps("python", python, mappings)) + rust_forest: Final = _forest(_comparable_steps("rust", rust, mappings)) + + def compare_children( + python_nodes: Sequence[TraceNode], rust_nodes: Sequence[TraceNode], path: str, *, unordered: bool + ) -> str | None: + if unordered: + python_signatures: Final = Counter(_signature(node, contract) for node in python_nodes) + rust_signatures: Final = Counter(_signature(node, contract) for node in rust_nodes) + if python_signatures != rust_signatures: + return f"{path}: unordered child subtree multiset differs" + return None + for index in range(max(len(python_nodes), len(rust_nodes))): + child_path = f"{path}/child[{index + 1}]" + if index >= len(python_nodes): + return f"{child_path}: Rust has extra {rust_nodes[index].span!r}" + if index >= len(rust_nodes): + return f"{child_path}: Python has extra {python_nodes[index].span!r}" + python_node = python_nodes[index] + rust_node = rust_nodes[index] + if python_node.span != rust_node.span: + return f"{child_path}: Python={python_node.span!r}, Rust={rust_node.span!r}" + difference = compare_children( + python_node.children, + rust_node.children, + f"{child_path}/{python_node.span}", + unordered=python_node.span in contract.unordered_children_of, + ) + if difference is not None: + return difference + return None + + return compare_children(python_forest, rust_forest, "root", unordered=False) + + +def trace_diff( + python: Sequence[PipelineStep], + rust: Sequence[PipelineStep], + mappings: Sequence[TraceMapping] = (), + contract: TraceContract = TraceContract(), +) -> TraceDiff: + python_comparable: Final = _comparable_steps("python", python, mappings) + rust_comparable: Final = _comparable_steps("rust", rust, mappings) + python_spans: Final = tuple(step.span for step in python_comparable) + rust_spans: Final = tuple(step.span for step in rust_comparable) + python_counts: Final = Counter(python_spans) + rust_counts: Final = Counter(rust_spans) + python_only_counts: Final = python_counts - rust_counts + rust_only_counts: Final = rust_counts - python_counts + python_only: Final = tuple( + span for span, count in python_only_counts.items() for _ in range(count) + ) + rust_only: Final = tuple(span for span, count in rust_only_counts.items() for _ in range(count)) + first_difference: Final = _first_difference(python, rust, mappings, contract) + return TraceDiff( + python_only=python_only, + rust_only=rust_only, + shared_order_matches=bool(python_comparable or rust_comparable) and first_difference is None, + missing_mappings=_missing_mappings(python, rust, mappings), + first_difference=first_difference, + ) diff --git a/tests/rust-python-harness/shared/tracing/test_compare.py b/tests/rust-python-harness/shared/tracing/test_compare.py deleted file mode 100644 index 2dfad24846b..00000000000 --- a/tests/rust-python-harness/shared/tracing/test_compare.py +++ /dev/null @@ -1,33 +0,0 @@ -from __future__ import annotations - -import pytest - -from .compare import Operation, compare_traces - - -@pytest.mark.parametrize( - ("rust", "message"), - ( - ((Operation("decode", 0, 1), Operation("send", 2, 3)), None), - ((Operation("decode", 0, 1), Operation("send", 2, 3), Operation("send", 4, 5)), "call count differs"), - ((Operation("send", 0, 1), Operation("decode", 2, 3)), "required order"), - ((Operation("decode", 0, 4), Operation("send", 2, 3)), "required order"), - ((Operation("decode", 0, 1), Operation("unknown", 2, 3)), "unmapped Rust"), - ), -) -def test_compare_mapped_calls_and_required_completion_order(rust: tuple[Operation, ...], message: str | None) -> None: - problems = compare_traces( - (Operation("parse", 0, 1), Operation("request", 2, 3)), - rust, - {"parse": "decode", "request": "send"}, - (("parse", "request"),), - ) - if message is None: - assert problems == () - else: - assert any(message in problem for problem in problems) - - -def test_missing_required_operations_and_ambiguous_mappings_fail() -> None: - assert compare_traces((), (), {"parse": "decode"}, (("parse", "request"),)) - assert compare_traces((), (), {"parse": "decode", "request": "decode"}) == ("ambiguous Rust operation: decode",) diff --git a/tests/rust-python-harness/shared/tracing/test_profiler.py b/tests/rust-python-harness/shared/tracing/test_profiler.py new file mode 100644 index 00000000000..ba85ffe63cd --- /dev/null +++ b/tests/rust-python-harness/shared/tracing/test_profiler.py @@ -0,0 +1,108 @@ +from __future__ import annotations + +import asyncio +import sys +import threading +from pathlib import Path +from typing import Final + +import pytest + +from .profiler import FunctionTraceEvent, PythonProfiler, profile_python, profile_python_function_usage + + +def _events_named(profiler: PythonProfiler, name: str) -> tuple[FunctionTraceEvent, ...]: + return tuple(event for event in profiler.events if event.function.endswith(name)) + + +def test_profiler_keeps_repeated_calls() -> None: + def called() -> None: + return None + + with profile_python(Path(__file__).parent) as profiler: + called() + called() + + assert len(_events_named(profiler, "called")) == 2 + + +def test_profiler_records_real_frame_ancestry() -> None: + def called() -> None: + return None + + def outer() -> None: + called() + + with profile_python(Path(__file__).parent) as profiler: + outer() + + outer_event, called_event = (event for event in profiler.events if event.function.endswith(("outer", "called"))) + assert called_event.parent_id == outer_event.id + + +def test_profiler_restores_previous_profiler_after_failure() -> None: + previous: Final = sys.getprofile() + + with pytest.raises(RuntimeError, match="stop"): + with profile_python(Path(__file__).parent): + raise RuntimeError("stop") + + assert sys.getprofile() is previous + + +def test_profiler_does_not_count_coroutine_resumption_as_another_call() -> None: + async def suspended() -> None: + await asyncio.sleep(0) + await asyncio.sleep(0) + + with profile_python(Path(__file__).parent) as profiler: + asyncio.run(suspended()) + + assert len(_events_named(profiler, "suspended")) == 1 + + +def test_profiler_preserves_parent_across_coroutine_suspension() -> None: + def called() -> None: + return None + + async def suspended() -> None: + await asyncio.sleep(0) + called() + + with profile_python(Path(__file__).parent) as profiler: + asyncio.run(suspended()) + + suspended_event: Final = _events_named(profiler, "suspended")[0] + called_event: Final = _events_named(profiler, "called")[0] + assert called_event.parent_id == suspended_event.id + + +def test_profiler_captures_worker_threads_when_enabled() -> None: + def called() -> None: + return None + + with profile_python(Path(__file__).parent, threads=True) as profiler: + thread: Final = threading.Thread(target=called) + thread.start() + thread.join() + + called_event: Final = _events_named(profiler, "called")[0] + assert called_event.parent_id is None + + +def test_function_usage_profiler_records_only_selected_functions() -> None: + def selected() -> None: + return None + + def ignored() -> None: + return None + + source_root: Final = Path(__file__).parent + function: Final = PythonProfiler(source_root).function_name(selected.__code__) + assert function is not None + + with profile_python_function_usage(source_root, frozenset((function,))) as profiler: + selected() + ignored() + + assert profiler.called == {function} diff --git a/tests/rust-python-harness/shared/tracing/test_pytest_usage.py b/tests/rust-python-harness/shared/tracing/test_pytest_usage.py new file mode 100644 index 00000000000..9f40da19010 --- /dev/null +++ b/tests/rust-python-harness/shared/tracing/test_pytest_usage.py @@ -0,0 +1,168 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Final + +import pytest + +from .pytest_usage import ( + PythonFunctionIdentity, + PythonFunctionReference, + RustFunctionIdentity, + candidate_test_files, + collect_python_function_tests, +) + + +def test_collects_parameterized_tests_that_execute_function(tmp_path: Path) -> None: + (tmp_path / "pytest.ini").write_text("[pytest]\n") + (tmp_path / "source.py").write_text("def target():\n return 1\n\ndef other():\n return 2\n") + (tmp_path / "test_source.py").write_text( + "import pytest\n" + "from source import other, target\n" + "@pytest.mark.parametrize('value', [1, 2])\n" + "def test_target(value): assert target() + value > 0\n" + "def test_other(): assert other() == 2\n" + ) + target: Final = PythonFunctionIdentity(file="source.py", line=1, qualname="target") + + report: Final = collect_python_function_tests( + (target,), + ("test_source.py",), + tmp_path, + source_root=tmp_path, + ) + + assert report.exit_code == 0, report.problems + assert report.usages[0].tests == ( + "test_source.py::test_target[1]", + "test_source.py::test_target[2]", + ) + + +def test_collects_async_and_threaded_function_calls(tmp_path: Path) -> None: + (tmp_path / "pytest.ini").write_text("[pytest]\n") + (tmp_path / "source.py").write_text( + "async def async_target():\n return 1\n\ndef threaded_target():\n return 2\n" + ) + (tmp_path / "test_source.py").write_text( + "import asyncio\n" + "from threading import Thread\n" + "from source import async_target, threaded_target\n" + "def test_async(): assert asyncio.run(async_target()) == 1\n" + "def test_thread():\n" + " thread = Thread(target=threaded_target)\n" + " thread.start()\n" + " thread.join()\n" + ) + functions: Final = ( + PythonFunctionIdentity(file="source.py", line=1, qualname="async_target"), + PythonFunctionIdentity(file="source.py", line=4, qualname="threaded_target"), + ) + + report: Final = collect_python_function_tests( + functions, + ("test_source.py",), + tmp_path, + source_root=tmp_path, + ) + + assert report.exit_code == 0, report.problems + assert report.usages[0].tests == ("test_source.py::test_async",) + assert report.usages[1].tests == ("test_source.py::test_thread",) + + +def test_adds_candidate_directory_to_worker_import_path(tmp_path: Path) -> None: + (tmp_path / "pytest.ini").write_text("[pytest]\n") + source: Final = tmp_path / "source" + tests: Final = tmp_path / "tests" + source.mkdir() + tests.mkdir() + (source / "implementation.py").write_text("def target():\n return 1\n") + (tests / "helper.py").write_text("VALUE = 1\n") + (tests / "test_source.py").write_text( + "from helper import VALUE\nfrom implementation import target\ndef test_target(): assert target() == VALUE\n" + ) + target: Final = PythonFunctionIdentity(file="implementation.py", line=1, qualname="target") + + report: Final = collect_python_function_tests( + (target,), + ("tests/test_source.py",), + tmp_path, + source_root=source, + ) + + assert report.exit_code == 0, report.problems + assert report.usages[0].tests == ("tests/test_source.py::test_target",) + + +def test_parses_function_identity_from_trace() -> None: + function: Final = PythonFunctionIdentity.from_trace("llms/mistral/ocr/transformation.py:72 Config.map") + + assert function.file == "llms/mistral/ocr/transformation.py" + assert function.line == 72 + assert function.qualname == "Config.map" + + +def test_resolves_function_and_finds_candidate_test_files(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + package: Final = tmp_path / "package" + tests: Final = tmp_path / "tests" + package.mkdir() + tests.mkdir() + (package / "__init__.py").write_text("") + (package / "implementation.py").write_text("class Config:\n def transform(self):\n return 1\n") + (tests / "test_implementation.py").write_text("from package.implementation import Config\n") + (tests / "test_unrelated.py").write_text("def test_other(): pass\n") + monkeypatch.syspath_prepend(tmp_path) + reference: Final = PythonFunctionReference(module="package.implementation", qualname="Config.transform") + + function: Final = reference.resolve(tmp_path) + candidates: Final = candidate_test_files((reference,), ("tests",), tmp_path) + + assert function.file == "package/implementation.py" + assert function.qualname == "Config.transform" + assert candidates == ("tests/test_implementation.py",) + + +def test_candidate_test_files_excludes_harness_roots(tmp_path: Path) -> None: + tests: Final = tmp_path / "tests" + harness: Final = tests / "harness" + harness.mkdir(parents=True) + (tests / "test_implementation.py").write_text("from package.implementation import Config\n") + (harness / "test_fixture.py").write_text("from package.implementation import Config\n") + function: Final = PythonFunctionReference(module="package.implementation", qualname="Config.transform") + + candidates: Final = candidate_test_files( + (function,), + ("tests",), + tmp_path, + exclude_roots=("tests/harness",), + ) + + assert candidates == ("tests/test_implementation.py",) + + +def test_candidate_test_files_finds_top_level_calls_and_import_aliases(tmp_path: Path) -> None: + tests: Final = tmp_path / "tests" + tests.mkdir() + (tests / "test_attribute.py").write_text("import package\ndef test_call(): package.ocr()\n") + (tests / "test_alias.py").write_text("from package import ocr as run_ocr\ndef test_call(): run_ocr()\n") + (tests / "test_unrelated.py").write_text("def test_call(): return 'ocr'\n") + function: Final = PythonFunctionIdentity(file="ocr/main.py", line=1, qualname="ocr") + + candidates: Final = candidate_test_files((function,), ("tests",), tmp_path) + + assert candidates == ( + "tests/test_alias.py", + "tests/test_attribute.py", + ) + + +def test_parses_rust_function_identity_and_derives_test_module() -> None: + function: Final = RustFunctionIdentity.from_trace( + "crates/core/src/providers/mistral/ocr/transformation.rs:73 " + "litellm_core::providers::mistral::ocr::transformation::supported_ocr_params" + ) + + assert function.file == "crates/core/src/providers/mistral/ocr/transformation.rs" + assert function.test_module == "providers::mistral::ocr::transformation::tests" diff --git a/tests/rust-python-harness/shared/tracing/test_steps.py b/tests/rust-python-harness/shared/tracing/test_steps.py new file mode 100644 index 00000000000..2efc6a3c579 --- /dev/null +++ b/tests/rust-python-harness/shared/tracing/test_steps.py @@ -0,0 +1,157 @@ +from __future__ import annotations + +from typing import Final + +import pytest + +from .profiler import FunctionTraceEvent +from .steps import Engine, TraceContract, mapping, pipeline_projection, trace_depths, trace_diff + +MAPPINGS: Final = ( + mapping(rust_span="route", python_frame=r"entry$"), + mapping(rust_span="provider", python_frame=r"provider$"), + mapping(rust_span="request", python_frame=r"request$"), + mapping(rust_span="http", python_frame=r"post$"), + mapping(rust_span="response", python_frame=r"response$"), +) + + +def event(event_id: int, function: str, parent_id: int | None = None) -> FunctionTraceEvent: + return FunctionTraceEvent(event_id, parent_id, function) + + +def test_python_projection_collapses_unmapped_parents_and_counts_noise() -> None: + events: Final = ( + event(0, "module.py:1 entry"), + event(1, "noise", 0), + event(2, "module.py:2 provider", 1), + event(3, "module.py:3 request", 0), + event(4, "client.py:4 post", 3), + event(5, "module.py:5 response", 0), + ) + projection: Final = pipeline_projection("python", events, MAPPINGS) + assert projection.unmatched == 1 + assert [(step.id, step.parent_id, step.span, step.raw) for step in projection.steps] == [ + (0, None, "route", "module.py:1 entry"), + (2, 0, "provider", "module.py:2 provider"), + (3, 0, "request", "module.py:3 request"), + (4, 3, "http", "client.py:4 post"), + (5, 0, "response", "module.py:5 response"), + ] + + +def test_rust_projection_keeps_unknown_spans() -> None: + projection: Final = pipeline_projection("rust", (event(0, "route"), event(1, "new_span", 0)), MAPPINGS) + assert [(step.span, step.parent_id) for step in projection.steps] == [("route", None), ("new_span", 0)] + + +def test_projection_preserves_repeated_occurrences() -> None: + projection: Final = pipeline_projection( + "rust", + (event(0, "route"), event(1, "http", 0), event(2, "http", 0)), + MAPPINGS, + ) + assert [step.span for step in projection.steps] == ["route", "http", "http"] + + +def test_projection_preserves_multiple_roots() -> None: + projection: Final = pipeline_projection("rust", (event(0, "route"), event(1, "request")), MAPPINGS) + assert trace_depths(projection.steps) == {0: 0, 1: 0} + + +def test_projection_rejects_duplicate_and_unknown_parent_ids() -> None: + with pytest.raises(ValueError, match="duplicate trace event id"): + pipeline_projection("rust", (event(0, "route"), event(0, "request")), MAPPINGS) + with pytest.raises(ValueError, match="unknown or later parent"): + pipeline_projection("rust", (event(1, "request", 0),), MAPPINGS) + + +@pytest.mark.parametrize("engine", ("python", "rust")) +def test_rust_only_mappings_do_not_swallow_python_frames(engine: Engine) -> None: + projection: Final = pipeline_projection( + engine, + (event(0, "anything"),), + (mapping(rust_span="rust_only_span"),), + ) + if engine == "python": + assert projection.unmatched == 1 + assert projection.steps == () + else: + assert projection.unmatched == 0 + assert projection.steps[0].span == "anything" + + +def test_mapping_builder_rejects_empty_and_ambiguous_declarations() -> None: + with pytest.raises(ValueError, match="mapping needs"): + mapping() + with pytest.raises(ValueError, match="python-only mapping needs"): + mapping(python_frame=r"frame$") + with pytest.raises(ValueError, match="disagrees with"): + mapping(rust_span="span_a", python_frame=r"frame$", span="span_b") + + +def test_projection_rejects_ambiguous_python_mapping() -> None: + mappings: Final = ( + mapping(rust_span="first", python_frame=r"same$"), + mapping(rust_span="second", python_frame=r"same$"), + ) + with pytest.raises(ValueError, match="multiple trace mappings"): + pipeline_projection("python", (event(0, "module.py:1 same"),), mappings) + + +def test_trace_diff_matches_identical_occurrence_trees() -> None: + mappings: Final = (MAPPINGS[0], MAPPINGS[2]) + steps: Final = pipeline_projection( + "rust", (event(0, "route"), event(1, "request", 0), event(2, "request", 0)), mappings + ).steps + assert trace_diff(steps, steps, mappings).matches + + +def test_trace_diff_rejects_missing_occurrence_and_parent_drift() -> None: + python: Final = pipeline_projection( + "rust", (event(0, "route"), event(1, "request", 0), event(2, "request", 0)), MAPPINGS + ).steps + missing: Final = pipeline_projection("rust", (event(0, "route"), event(1, "request", 0)), MAPPINGS).steps + reparented: Final = pipeline_projection( + "rust", (event(0, "route"), event(1, "request", 0), event(2, "request", 1)), MAPPINGS + ).steps + assert trace_diff(python, missing, MAPPINGS).python_only == ("request",) + assert not trace_diff(python, reparented, MAPPINGS).matches + + +def test_trace_diff_rejects_sequential_reorder() -> None: + first: Final = pipeline_projection( + "rust", (event(0, "route"), event(1, "request", 0), event(2, "response", 0)), MAPPINGS + ).steps + second: Final = pipeline_projection( + "rust", (event(0, "route"), event(1, "response", 0), event(2, "request", 0)), MAPPINGS + ).steps + diff: Final = trace_diff(first, second, MAPPINGS) + assert not diff.matches + assert diff.first_difference == "root/child[1]/route/child[1]: Python='request', Rust='response'" + + +def test_trace_diff_allows_reordered_concurrent_children() -> None: + mappings: Final = (MAPPINGS[0], MAPPINGS[2], MAPPINGS[4]) + first: Final = pipeline_projection( + "rust", (event(0, "route"), event(1, "request", 0), event(2, "response", 0)), mappings + ).steps + second: Final = pipeline_projection( + "rust", (event(0, "route"), event(1, "response", 0), event(2, "request", 0)), mappings + ).steps + contract: Final = TraceContract(frozenset({"route"})) + assert trace_diff(first, second, mappings, contract).matches + + +def test_trace_diff_prunes_declared_engine_only_nodes_but_requires_them() -> None: + mappings: Final = (MAPPINGS[0], mapping(rust_span="rust_prepare")) + python: Final = pipeline_projection("python", (event(0, "module.py:1 entry"),), mappings).steps + rust: Final = pipeline_projection( + "rust", (event(0, "route"), event(1, "rust_prepare", 0)), mappings + ).steps + assert trace_diff(python, rust, mappings).matches + assert trace_diff(python, rust[:1], mappings).missing_mappings == ("rust_prepare",) + + +def test_trace_diff_does_not_claim_empty_traces_match() -> None: + assert not trace_diff((), ()).matches diff --git a/tests/rust-python-harness/strategies/e2e_parity/gateway/__init__.py b/tests/rust-python-harness/shared/unit_runners/__init__.py similarity index 100% rename from tests/rust-python-harness/strategies/e2e_parity/gateway/__init__.py rename to tests/rust-python-harness/shared/unit_runners/__init__.py diff --git a/tests/rust-python-harness/strategies/unit_tests/python_runner.py b/tests/rust-python-harness/shared/unit_runners/python_runner.py similarity index 63% rename from tests/rust-python-harness/strategies/unit_tests/python_runner.py rename to tests/rust-python-harness/shared/unit_runners/python_runner.py index 900b30180dc..86ea0dbf1b9 100644 --- a/tests/rust-python-harness/strategies/unit_tests/python_runner.py +++ b/tests/rust-python-harness/shared/unit_runners/python_runner.py @@ -1,7 +1,6 @@ from __future__ import annotations import argparse -import ast import importlib import os import subprocess @@ -9,11 +8,16 @@ import sys import tempfile from collections.abc import Callable, Sequence from pathlib import Path -from typing import Final, Literal, cast +from typing import TYPE_CHECKING, Final, Literal, cast -import pytest +from pluggy import HookimplMarker from pydantic import BaseModel, ConfigDict +if TYPE_CHECKING: + import pytest + +hookimpl: Final = HookimplMarker("pytest") + Backend = Literal["python", "rust"] @@ -32,22 +36,20 @@ class BackendSpec(BaseModel): model_config = ConfigDict(frozen=True, extra="forbid") environment_variable: str + probe: str = "" + + +class WorkerArgs(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + backend: Backend probe: str - - -def ocr_backend() -> Backend: - from litellm.rust_bridge import native_bridge_available - from litellm.rust_bridge.configuration import rust_ocr_enabled - - if not rust_ocr_enabled(): - return "python" - if not native_bridge_available(): - raise RuntimeError("Rust OCR was enabled but the native extension is unavailable") - return "rust" + output: Path + pytest_args: tuple[str, ...] class ResultPlugin: - def __init__(self, backend: Backend, probe: Callable[[], object]) -> None: + def __init__(self, backend: Backend, probe: Callable[[], object] | None) -> None: self.backend: Final = backend self.probe: Final = probe self.tests: tuple[str, ...] = () @@ -55,13 +57,13 @@ class ResultPlugin: self.problems: tuple[str, ...] = () def verify(self) -> None: - if self.probe() != self.backend: + if self.probe is not None and self.probe() != self.backend: raise RuntimeError(f"backend probe did not select {self.backend}") def pytest_collection_finish(self, session: pytest.Session) -> None: self.tests = tuple(item.nodeid for item in session.items) - @pytest.hookimpl(tryfirst=True) + @hookimpl(tryfirst=True) def pytest_runtest_call(self, item: pytest.Item) -> None: del item self.verify() @@ -91,8 +93,7 @@ def run_python_tests( __name__, "--backend", backend, - "--probe", - spec.probe, + *(("--probe", spec.probe) if spec.probe else ()), "--output", str(output), "--", @@ -118,6 +119,9 @@ def run_python_tests( problems=(result.stdout + result.stderr,), ) report: Final = PythonReport.model_validate_json(output.read_text()) + process_output: Final = (result.stdout + result.stderr).strip() + if result.returncode and not report.problems and process_output: + return report.model_copy(update={"problems": (process_output,)}) if report.exit_code != result.returncode: return report.model_copy( update={ @@ -129,43 +133,42 @@ def run_python_tests( def compare_python_runs(python: PythonReport, rust: PythonReport) -> tuple[str, ...]: + python_only: Final = tuple(sorted(set(python.outcomes) - set(rust.outcomes))) + rust_only: Final = tuple(sorted(set(rust.outcomes) - set(python.outcomes))) return ( *(("backend selection was not verified",) if not python.verified or not rust.verified else ()), *(("Python run used the wrong backend",) if python.backend != "python" else ()), *(("Rust run used the wrong backend",) if rust.backend != "rust" else ()), *(("Python/Rust test inventories differ",) if python.tests != rust.tests else ()), - *(("Python/Rust test outcomes differ",) if sorted(python.outcomes) != sorted(rust.outcomes) else ()), + *(("Python/Rust test outcomes differ",) if python_only or rust_only else ()), + *(f"Python only: {nodeid} [{stage}] {outcome}" for nodeid, stage, outcome in python_only), + *(f"Rust only: {nodeid} [{stage}] {outcome}" for nodeid, stage, outcome in rust_only), + *(f"Python run: {problem}" for problem in python.problems if not python.verified or not python.tests), + *(f"Rust run: {problem}" for problem in rust.problems if not rust.verified or not rust.tests), *(("no Python tests collected",) if not python.tests else ()), - *( - ("Python tests did not all pass",) - if set(python.tests) - != {node for node, phase, status in python.outcomes if phase == "call" and status == "passed"} - else () - ), - *( - ("Rust-enabled Python tests did not all pass",) - if set(rust.tests) - != {node for node, phase, status in rust.outcomes if phase == "call" and status == "passed"} - else () - ), - *(("Python test run failed",) if python.exit_code else ()), - *(("Rust-enabled Python test run failed",) if rust.exit_code else ()), - *python.problems, - *rust.problems, + *(("Python/Rust exit codes differ",) if python.exit_code != rust.exit_code else ()), ) +def _load_probe(reference: str) -> Callable[[], object] | None: + if not reference: + return None + module, name = reference.rsplit(":", 1) + return cast(Callable[[], object], getattr(importlib.import_module(module), name)) + + def main(argv: Sequence[str] | None = None) -> int: + import pytest + parser: Final = argparse.ArgumentParser() parser.add_argument("--backend", required=True, choices=("python", "rust")) - parser.add_argument("--probe", required=True) + parser.add_argument("--probe", default="") parser.add_argument("--output", required=True, type=Path) parser.add_argument("pytest_args", nargs=argparse.REMAINDER) - args: Final = parser.parse_args(argv) + namespace: Final = parser.parse_args(argv) + args: Final = WorkerArgs.model_validate(vars(namespace)) try: - module, name = args.probe.rsplit(":", 1) - probe: Final = cast(Callable[[], object], getattr(importlib.import_module(module), name)) - plugin: Final = ResultPlugin(args.backend, probe) + plugin: Final = ResultPlugin(args.backend, _load_probe(args.probe)) plugin.verify() code: Final = int( pytest.main(["-o", "consider_namespace_packages=true", *args.pytest_args[1:]], plugins=[plugin]) @@ -186,22 +189,29 @@ def main(argv: Sequence[str] | None = None) -> int: return report.exit_code -def enumerate_python_tests(repo_root: Path, relative_path: str) -> frozenset[str]: - source = (repo_root / relative_path).read_text(encoding="utf-8") - tree = ast.parse(source, filename=relative_path) +def contract_nodeid(nodeid: str) -> str: + owner, separator, test = nodeid.rpartition("::") + function: Final = test.partition("[")[0] + if not separator or not function.startswith("test_"): + raise ValueError(f"Unrecognized pytest node id: {nodeid}") + return f"{owner}::{function}" - module_level: list[str] = [] - for node in ast.iter_child_nodes(tree): - if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name.startswith("test_"): - module_level.append(node.name) - elif isinstance(node, ast.ClassDef): - for child in ast.iter_child_nodes(node): - if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)) and child.name.startswith( - "test_" - ): - module_level.append(f"{node.name}::{child.name}") - return frozenset(module_level) +def collect_python_tests(selectors: Sequence[str], repo_root: Path) -> frozenset[str]: + report: Final = run_python_tests( + selectors, + repo_root, + "python", + BackendSpec(environment_variable="LITELLM_RUST"), + ("--collect-only", "-p", "no:cacheprovider"), + ) + if report.exit_code or report.problems: + details: Final = "\n".join(report.problems) or f"pytest exited with code {report.exit_code}" + raise ValueError(f"Python test collection failed:\n{details}") + tests: Final = frozenset(contract_nodeid(nodeid) for nodeid in report.tests) + if not tests: + raise ValueError(f"pytest collected no tests for: {', '.join(selectors)}") + return tests if __name__ == "__main__": diff --git a/tests/rust-python-harness/shared/unit_runners/rust_runner.py b/tests/rust-python-harness/shared/unit_runners/rust_runner.py new file mode 100644 index 00000000000..9df77895519 --- /dev/null +++ b/tests/rust-python-harness/shared/unit_runners/rust_runner.py @@ -0,0 +1,253 @@ +from __future__ import annotations + +import subprocess +from collections.abc import Callable +from dataclasses import dataclass +from itertools import groupby +from pathlib import Path +from typing import Annotated, Final, Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict, Field, model_validator +from typing_extensions import Self + +CommandRunner: TypeAlias = Callable[[tuple[str, ...], Path], str] +_MODEL_CONFIG: Final = ConfigDict(extra="forbid", frozen=True, strict=True) + + +class RustTarget(BaseModel): + model_config = _MODEL_CONFIG + + package: str + name: str + kind: Literal["lib", "bin", "test"] + + @property + def key(self) -> str: + return f"{self.package}/{self.kind}/{self.name}" + + +class RustTestIdentity(BaseModel): + model_config = _MODEL_CONFIG + + target: RustTarget + name: str + + @property + def key(self) -> str: + return f"{self.target.key}::{self.name}" + + +class RustTestScope(BaseModel): + model_config = _MODEL_CONFIG + + target: RustTarget + modules: Annotated[tuple[str, ...], Field(min_length=1)] + features: tuple[str, ...] = () + default_features: bool = True + + @model_validator(mode="after") + def validate_scope(self) -> Self: + duplicate_features: Final = tuple( + feature for feature, values in groupby(sorted(self.features)) if sum(1 for _ in values) > 1 + ) + duplicate_modules: Final = tuple( + module for module, values in groupby(sorted(self.modules)) if sum(1 for _ in values) > 1 + ) + if duplicate_features: + raise ValueError(f"Rust features contain duplicates: {', '.join(duplicate_features)}") + if duplicate_modules: + raise ValueError(f"Rust modules contain duplicates: {', '.join(duplicate_modules)}") + if any(not module or module.endswith("::") for module in self.modules): + raise ValueError("Rust modules must be non-empty and omit the trailing :: separator") + overlaps: Final = tuple( + f"{outer} includes {inner}" + for outer in self.modules + for inner in self.modules + if inner.startswith(f"{outer}::") + ) + if overlaps: + raise ValueError(f"Rust modules overlap: {', '.join(overlaps)}") + return self + + def contains(self, identity: RustTestIdentity) -> bool: + return identity.target == self.target and any( + identity.name.startswith(f"{module}::") for module in self.modules + ) + + +class _CargoPackage(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True, strict=True) + + id: str + name: str + + +class _CargoMetadata(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True, strict=True) + + packages: tuple[_CargoPackage, ...] + + +class _CargoMessage(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True, strict=True) + + reason: str + + +class _CargoTarget(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True, strict=True) + + name: str + kind: tuple[str, ...] + + +class _CargoProfile(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True, strict=True) + + test: bool + + +class _CargoArtifact(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True, strict=True) + + reason: Literal["compiler-artifact"] + package_id: str + target: _CargoTarget + profile: _CargoProfile + executable: str | None + + +@dataclass(frozen=True, slots=True) +class RustReport: + tests: tuple[str, ...] + exit_code: int + output: str + + +def run_command(command: tuple[str, ...], cwd: Path) -> str: + try: + result: Final = subprocess.run(command, cwd=cwd, capture_output=True, text=True, check=False, timeout=600) + except (OSError, subprocess.TimeoutExpired) as error: + raise ValueError(f"Rust inventory command failed: {error}") from error + if result.returncode != 0: + raise ValueError( + f"Rust inventory command failed ({result.returncode}): {' '.join(command)}\n" + f"{result.stderr}\n{result.stdout}" + ) + return result.stdout + + +def run_rust_tests(manifest: Path, package: str | None, test_filter: str, *, collect_only: bool = False) -> RustReport: + command: Final = ( + "cargo", + "test", + "--manifest-path", + str(manifest), + *(("--package", package) if package else ()), + "--lib", + test_filter, + "--", + *(("--list",) if collect_only else ("--format=pretty",)), + ) + try: + result: Final = subprocess.run(command, capture_output=True, text=True, check=False, timeout=600) + except (OSError, subprocess.TimeoutExpired) as error: + return RustReport((), 1, str(error)) + tests: Final = ( + tuple(line.removesuffix(": test") for line in result.stdout.splitlines() if line.endswith(": test")) + if collect_only + else tuple( + line.removeprefix("test ").removesuffix(" ... ok") + for line in result.stdout.splitlines() + if line.startswith("test ") and line.endswith(" ... ok") + ) + ) + return RustReport(tests, result.returncode, result.stdout + result.stderr) + + +def _build_command(scope: RustTestScope) -> tuple[str, ...]: + selector: Final = ("--lib",) if scope.target.kind == "lib" else (f"--{scope.target.kind}", scope.target.name) + features: Final = ("--features", ",".join(scope.features)) if scope.features else () + defaults: Final = () if scope.default_features else ("--no-default-features",) + return ( + "cargo", + "test", + "--package", + scope.target.package, + *selector, + *features, + *defaults, + "--locked", + "--no-run", + "--message-format=json", + "--color", + "never", + ) + + +def _test_names(output: str) -> frozenset[str]: + lines: Final = tuple(line for line in output.splitlines() if line) + invalid: Final = tuple(line for line in lines if not line.endswith((": test", ": benchmark"))) + if invalid: + raise ValueError(f"Unrecognized libtest inventory output: {invalid!r}") + names: Final = tuple(line.removesuffix(": test") for line in lines if line.endswith(": test")) + if len(names) != len(frozenset(names)): + raise ValueError("Duplicate test names in libtest inventory") + return frozenset(names) + + +def _scope_tests( + scope: RustTestScope, + metadata: _CargoMetadata, + cwd: Path, + command_runner: CommandRunner, +) -> frozenset[RustTestIdentity]: + package_ids: Final = tuple(package.id for package in metadata.packages if package.name == scope.target.package) + if len(package_ids) != 1: + raise ValueError(f"Expected one Cargo package for {scope.target.package}, found {len(package_ids)}") + output: Final = command_runner(_build_command(scope), cwd) + artifacts: Final = tuple( + _CargoArtifact.model_validate_json(line) + for line in output.splitlines() + if _CargoMessage.model_validate_json(line).reason == "compiler-artifact" + ) + executables: Final = frozenset( + artifact.executable + for artifact in artifacts + if artifact.package_id == package_ids[0] + and artifact.target.name == scope.target.name + and scope.target.kind in artifact.target.kind + and artifact.profile.test + and artifact.executable is not None + ) + if len(executables) != 1: + raise ValueError(f"Expected one test executable for {scope.target.key}, found {len(executables)}") + executable: Final = next(iter(executables)) + names: Final = _test_names(command_runner((executable, "--list", "--format", "terse"), cwd)) + ignored: Final = _test_names(command_runner((executable, "--list", "--ignored", "--format", "terse"), cwd)) + identities: Final = frozenset(RustTestIdentity(target=scope.target, name=name) for name in names) + scoped: Final = frozenset(identity for identity in identities if scope.contains(identity)) + ignored_scoped: Final = tuple(sorted(identity.key for identity in scoped if identity.name in ignored)) + if ignored_scoped: + raise ValueError(f"Ignored Rust tests cannot satisfy the mapping: {', '.join(ignored_scoped)}") + empty_modules: Final = tuple( + module for module in scope.modules if not any(identity.name.startswith(f"{module}::") for identity in scoped) + ) + if empty_modules: + raise ValueError(f"No compiled tests in {scope.target.key} modules: {', '.join(empty_modules)}") + return scoped + + +def enumerate_rust_tests( + repo_root: Path, + scopes: tuple[RustTestScope, ...], + *, + command_runner: CommandRunner = run_command, +) -> frozenset[RustTestIdentity]: + if not scopes: + return frozenset() + cwd: Final = repo_root / "litellm-rust" + metadata: Final = _CargoMetadata.model_validate_json( + command_runner(("cargo", "metadata", "--format-version", "1", "--no-deps", "--locked"), cwd) + ) + return frozenset(identity for scope in scopes for identity in _scope_tests(scope, metadata, cwd, command_runner)) diff --git a/tests/rust-python-harness/shared/unit_runners/suite_runner.py b/tests/rust-python-harness/shared/unit_runners/suite_runner.py new file mode 100644 index 00000000000..b0374bdf833 --- /dev/null +++ b/tests/rust-python-harness/shared/unit_runners/suite_runner.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from time import monotonic +from typing import Final, TypeVar + +from pydantic import BaseModel + +from ..reporting.models import HarnessCase, HarnessRun, ResultArtifact, RunStatus, SdkFunction +from ..reporting.strategy import SuiteCaseSpec, UpdateCallback + +S = TypeVar("S", bound=BaseModel) + + +@dataclass(frozen=True, slots=True) +class SuiteExecution: + problems: tuple[str, ...] = () + artifacts: tuple[ResultArtifact, ...] = () + + +SuiteExecutor = Callable[[S, Path, Sequence[str]], SuiteExecution] + + +def suite_nodeid(case: HarnessCase) -> str: + spec = case.spec + suite = spec.suite if isinstance(spec, SuiteCaseSpec) else "invalid" + return f"suite:{case.strategy_id}:{case.sdk_function}:{suite}" + + +def run_suites( + cases: Sequence[HarnessCase], + repo_root: Path, + on_update: UpdateCallback, + runner_args: Sequence[str] = (), + *, + suites: Mapping[SdkFunction, S], + execute: SuiteExecutor[S], +) -> tuple[int, HarnessRun]: + report = HarnessRun.from_cases(cases) + for case in cases: + result = report.results[case.key] + spec = case.spec + if not isinstance(spec, SuiteCaseSpec): + continue + nodeid = suite_nodeid(case) + result.collected.add(nodeid) + result.status = RunStatus.RUNNING + on_update(report) + suite = suites.get(case.sdk_function) + if suite is None: + result.record(nodeid, RunStatus.ERROR) + report.failures.append((nodeid, f"no suite registered for {case.sdk_function}")) + continue + try: + execution: Final = execute(suite, repo_root, runner_args) + except (OSError, ValueError) as error: + result.record(nodeid, RunStatus.ERROR) + report.failures.append((nodeid, str(error))) + continue + result.record( + nodeid, + RunStatus.FAILED if execution.problems else RunStatus.PASSED, + artifacts=execution.artifacts, + ) + report.failures.extend((nodeid, problem) for problem in execution.problems) + on_update(report) + report.finished_at = monotonic() + on_update(report) + return int( + any( + result.status in {RunStatus.ERROR, RunStatus.FAILED, RunStatus.MISSING} + for result in report.results.values() + ) + ), report diff --git a/tests/rust-python-harness/strategies/unit_tests/test_python_runner.py b/tests/rust-python-harness/shared/unit_runners/test_python_runner.py similarity index 54% rename from tests/rust-python-harness/strategies/unit_tests/test_python_runner.py rename to tests/rust-python-harness/shared/unit_runners/test_python_runner.py index ed4920b7525..abd4316ca98 100644 --- a/tests/rust-python-harness/strategies/unit_tests/test_python_runner.py +++ b/tests/rust-python-harness/shared/unit_runners/test_python_runner.py @@ -4,9 +4,7 @@ import os from pathlib import Path from typing import Final -from .python_runner import BackendSpec, compare_python_runs, run_python_tests - -HARNESS_ROOT: Final = Path(__file__).resolve().parents[4] +from .python_runner import BackendSpec, collect_python_tests, compare_python_runs, run_python_tests def _suite(root: Path, *, mismatch: bool = False) -> BackendSpec: @@ -23,9 +21,7 @@ def _suite(root: Path, *, mismatch: bool = False) -> BackendSpec: return BackendSpec(environment_variable="TEST_USE_RUST", probe="backend_probe:selected") -def test_runs_existing_python_tests_in_separate_verified_backends(tmp_path: Path, monkeypatch) -> None: - monkeypatch.setenv("PYTHONPATH", str(HARNESS_ROOT)) - monkeypatch.setenv("PYTEST_DISABLE_PLUGIN_AUTOLOAD", "1") +def test_runs_existing_python_tests_in_separate_verified_backends(tmp_path: Path) -> None: spec: Final = _suite(tmp_path) python: Final = run_python_tests(("test_backend.py",), tmp_path, "python", spec) rust: Final = run_python_tests(("test_backend.py",), tmp_path, "rust", spec) @@ -35,9 +31,7 @@ def test_runs_existing_python_tests_in_separate_verified_backends(tmp_path: Path assert (tmp_path / "python.pid").read_text() != str(os.getpid()) -def test_rejects_wrong_backend_and_different_test_results(tmp_path: Path, monkeypatch) -> None: - monkeypatch.setenv("PYTHONPATH", str(HARNESS_ROOT)) - monkeypatch.setenv("PYTEST_DISABLE_PLUGIN_AUTOLOAD", "1") +def test_rejects_wrong_backend_and_different_test_results(tmp_path: Path) -> None: spec: Final = _suite(tmp_path, mismatch=True) python: Final = run_python_tests(("test_backend.py",), tmp_path, "python", spec) rust: Final = run_python_tests(("test_backend.py",), tmp_path, "rust", spec) @@ -51,3 +45,44 @@ def test_rejects_wrong_backend_and_different_test_results(tmp_path: Path, monkey assert wrong.exit_code == 1 assert not wrong.verified assert "backend probe did not select rust" in wrong.problems[0] + + +def test_matches_outcomes_without_a_probe_when_both_backends_fail_identically(tmp_path: Path) -> None: + (tmp_path / "pytest.ini").write_text("[pytest]\n") + (tmp_path / "test_backend.py").write_text("def test_fails():\n assert False\n") + spec: Final = BackendSpec(environment_variable="TEST_USE_RUST") + + python: Final = run_python_tests(("test_backend.py",), tmp_path, "python", spec) + rust: Final = run_python_tests(("test_backend.py",), tmp_path, "rust", spec) + + assert compare_python_runs(python, rust) == () + + +def test_reports_worker_output_when_pytest_exits_before_collection(tmp_path: Path) -> None: + report: Final = run_python_tests( + ("missing.py",), + tmp_path, + "python", + BackendSpec(environment_variable="TEST_USE_RUST"), + ) + + assert report.exit_code != 0 + assert report.problems + assert "missing.py" in report.problems[0] + + +def test_collects_tests_with_pytest_semantics_and_collapses_parameters(tmp_path: Path) -> None: + (tmp_path / "pytest.ini").write_text("[pytest]\n") + (tmp_path / "test_inventory.py").write_text( + "import pytest\n" + "class Helper:\n" + " def test_not_collected(self): pass\n" + "class TestCollected:\n" + " @pytest.mark.parametrize('value', [1, 2])\n" + " def test_parameterized(self, value): pass\n", + encoding="utf-8", + ) + + tests: Final = collect_python_tests(("test_inventory.py",), tmp_path) + + assert tests == frozenset(("test_inventory.py::TestCollected::test_parameterized",)) diff --git a/tests/rust-python-harness/shared/unit_runners/test_rust_runner.py b/tests/rust-python-harness/shared/unit_runners/test_rust_runner.py new file mode 100644 index 00000000000..a12d00a7ed6 --- /dev/null +++ b/tests/rust-python-harness/shared/unit_runners/test_rust_runner.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +import shutil +from collections.abc import Callable +from pathlib import Path +from typing import Final + +import pytest + +from .rust_runner import RustTarget, RustTestScope, enumerate_rust_tests, run_command, run_rust_tests + + +@pytest.mark.skipif(shutil.which("cargo") is None, reason="Cargo is required for native runner integration") +def test_collects_and_runs_native_tests_and_propagates_failure( + tmp_path: Path, + cargo_project: Callable[[str, str], Path], +) -> None: + manifest: Final = cargo_project("harness-runner-check", "#[test] fn test_parity() { assert_eq!(2 + 2, 4); }\n") + source: Final = tmp_path / "src/lib.rs" + inventory: Final = run_rust_tests(manifest, "harness-runner-check", "test_parity", collect_only=True) + assert inventory.exit_code == 0, inventory.output + assert inventory.tests == ("test_parity",) + passing: Final = run_rust_tests(manifest, "harness-runner-check", "test_parity") + assert passing.exit_code == 0, passing.output + source.write_text("#[test] fn test_parity() { assert_eq!(2 + 2, 5); }\n") + failed: Final = run_rust_tests(manifest, "harness-runner-check", "test_parity") + assert failed.exit_code != 0 + assert "test_parity" in failed.output + + +@pytest.mark.skipif(shutil.which("cargo") is None, reason="Cargo is required for compiled inventory tests") +def test_discovers_compiled_fully_qualified_tests(tmp_path: Path) -> None: + workspace: Final = tmp_path / "litellm-rust" + source: Final = workspace / "src" + external: Final = source / "ocr" / "external.rs" + external.parent.mkdir(parents=True) + (workspace / "Cargo.toml").write_text( + '[package]\nname = "inventory-fixture"\nversion = "0.1.0"\nedition = "2021"\n[workspace]\n', + encoding="utf-8", + ) + (source / "lib.rs").write_text( + "#[cfg(test)]\n" + "mod ocr {\n" + " mod external;\n" + " #[test] fn same_name() {}\n" + " #[cfg(any())] #[test] fn compiled_out() {}\n" + " macro_rules! generate_test { ($name:ident) => { #[test] fn $name() {} }; }\n" + " generate_test!(generated_case);\n" + "}\n", + encoding="utf-8", + ) + external.write_text("#[test] fn same_name() {}\n", encoding="utf-8") + run_command(("cargo", "generate-lockfile", "--offline"), workspace) + target: Final = RustTarget(package="inventory-fixture", name="inventory_fixture", kind="lib") + scope: Final = RustTestScope(target=target, modules=("ocr",)) + + inventory: Final = enumerate_rust_tests(tmp_path, (scope,)) + + assert frozenset(identity.name for identity in inventory) == frozenset( + ("ocr::same_name", "ocr::external::same_name", "ocr::generated_case") + ) diff --git a/tests/rust-python-harness/shared/unit_runners/test_suite_runner.py b/tests/rust-python-harness/shared/unit_runners/test_suite_runner.py new file mode 100644 index 00000000000..92670da666e --- /dev/null +++ b/tests/rust-python-harness/shared/unit_runners/test_suite_runner.py @@ -0,0 +1,88 @@ +from __future__ import annotations + +from collections.abc import Sequence +from pathlib import Path +from typing import Final + +from pydantic import BaseModel + +from ..reporting.models import Coverage, HarnessCase, ResultArtifact, RunStatus +from ..reporting.strategy import CaseSpec, NotImplementedCaseSpec, SuiteCaseSpec +from .suite_runner import SuiteExecution, run_suites + + +class _Suite(BaseModel): + problems: tuple[str, ...] = () + + +def _execute(suite: _Suite, repo_root: Path, pytest_args: Sequence[str]) -> SuiteExecution: + del repo_root, pytest_args + return SuiteExecution(problems=suite.problems) + + +def _case(spec: CaseSpec) -> HarnessCase: + return HarnessCase( + strategy_id="example", + strategy_label="Example", + sdk_function="ocr", + spec=spec, + ) + + +def test_not_implemented_cell_finalizes_without_running(tmp_path: Path) -> None: + case = _case(NotImplementedCaseSpec(reason="No suite is registered.")) + + code, report = run_suites((case,), tmp_path, lambda _: None, (), suites={}, execute=_execute) + + assert code == 0 + assert report.results[case.key].status is RunStatus.NOT_IMPLEMENTED + assert not report.failures + + +def test_missing_registered_suite_marks_the_cell_as_error(tmp_path: Path) -> None: + case = _case(SuiteCaseSpec(coverage=Coverage.COMPLETE, suite="ocr")) + + code, report = run_suites((case,), tmp_path, lambda _: None, (), suites={}, execute=_execute) + + assert code == 1 + assert report.results[case.key].status is RunStatus.ERROR + assert report.failures + + +def test_suite_problems_mark_the_cell_as_failed(tmp_path: Path) -> None: + case = _case(SuiteCaseSpec(coverage=Coverage.COMPLETE, suite="ocr")) + + code, report = run_suites( + (case,), tmp_path, lambda _: None, (), suites={"ocr": _Suite(problems=("boom",))}, execute=_execute + ) + + assert code == 1 + assert report.results[case.key].status is RunStatus.FAILED + assert ("suite:example:ocr:ocr", "boom") in report.failures + + +def test_suite_without_problems_passes(tmp_path: Path) -> None: + case = _case(SuiteCaseSpec(coverage=Coverage.COMPLETE, suite="ocr")) + + code, report = run_suites((case,), tmp_path, lambda _: None, (), suites={"ocr": _Suite()}, execute=_execute) + + assert code == 0 + assert report.results[case.key].status is RunStatus.PASSED + assert not report.failures + + +def test_suite_attaches_artifacts_to_passing_and_failing_results(tmp_path: Path) -> None: + case: Final = _case(SuiteCaseSpec(coverage=Coverage.COMPLETE, suite="ocr")) + artifact: Final = ResultArtifact("example", "body") + + def execute(suite: _Suite, repo_root: Path, pytest_args: Sequence[str]) -> SuiteExecution: + del repo_root, pytest_args + return SuiteExecution(problems=suite.problems, artifacts=(artifact,)) + + _, passing = run_suites((case,), tmp_path, lambda _: None, suites={"ocr": _Suite()}, execute=execute) + _, failing = run_suites( + (case,), tmp_path, lambda _: None, suites={"ocr": _Suite(problems=("boom",))}, execute=execute + ) + + assert passing.results[case.key].artifacts == {"suite:example:ocr:ocr": (artifact,)} + assert failing.results[case.key].artifacts == {"suite:example:ocr:ocr": (artifact,)} diff --git a/tests/rust-python-harness/strategies/e2e_parity/AGENTS.md b/tests/rust-python-harness/strategies/e2e_parity/AGENTS.md new file mode 100644 index 00000000000..27c87d0f100 --- /dev/null +++ b/tests/rust-python-harness/strategies/e2e_parity/AGENTS.md @@ -0,0 +1 @@ +Switches between the Rust implementation and the existing Python core, then compares their observable behavior for parity across SDK objects, exceptions, callbacks, streams, and gateway HTTP responses using generated and recorded inputs. diff --git a/tests/rust-python-harness/strategies/e2e_parity/README.md b/tests/rust-python-harness/strategies/e2e_parity/README.md deleted file mode 100644 index 17643e69676..00000000000 --- a/tests/rust-python-harness/strategies/e2e_parity/README.md +++ /dev/null @@ -1,5 +0,0 @@ -# E2E Parity - -Run independently with `uv run python -m tests.rust-python-harness.strategies.e2e_parity.runner --plain`. Configure SDK and gateway selectors in `strategy.json`; keep API-specific execution and fixtures in their owning surface folder - -See [the harness guide](../../README.md) for coverage status and shared comparison tools diff --git a/tests/rust-python-harness/strategies/e2e_parity/__init__.py b/tests/rust-python-harness/strategies/e2e_parity/__init__.py index e69de29bb2d..f668e178eef 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/__init__.py +++ b/tests/rust-python-harness/strategies/e2e_parity/__init__.py @@ -0,0 +1,103 @@ +from pathlib import Path +from typing import Final + +from ...shared.reporting.models import SURFACES, Coverage +from ...shared.reporting.strategy import ( + CaseDefinition, + ModuleCaseSpec, + NotImplementedCaseSpec, + StrategyDefinition, +) +from .reporting import render_e2e_results +from .runner import run_e2e_cases + +CASES: Final[tuple[CaseDefinition, ...]] = ( + CaseDefinition( + "ocr", + ModuleCaseSpec( + coverage=Coverage.PARTIAL, + module="tests.rust-python-harness.strategies.e2e_parity.sdk.ocr.test_sdk_parity", + note=( + "Recorded sync/async SDK parity; invalid-model provider errors differ, " + "and Reducto lacks a Rust contract." + ), + ), + surface="sdk", + ), + CaseDefinition( + "messages", + NotImplementedCaseSpec( + reason="Bridge unit tests exist, but no standalone end-to-end parity case is registered." + ), + surface="sdk", + ), + CaseDefinition( + "responses", + NotImplementedCaseSpec( + reason="Bridge unit tests exist, but no standalone end-to-end parity case is registered." + ), + surface="sdk", + ), + CaseDefinition( + "count_tokens", + NotImplementedCaseSpec(reason="No Rust count_tokens parity test is present yet."), + surface="sdk", + ), + CaseDefinition( + "chat_completions", + NotImplementedCaseSpec( + reason="Bridge unit tests exist, but no standalone end-to-end parity case is registered." + ), + surface="sdk", + ), + CaseDefinition( + "transcription", + NotImplementedCaseSpec( + reason="Bridge unit tests exist, but no standalone end-to-end parity case is registered." + ), + surface="sdk", + ), + CaseDefinition( + "ocr", + NotImplementedCaseSpec(reason="No gateway end-to-end OCR parity case is registered."), + surface="gateway", + ), + CaseDefinition( + "messages", + NotImplementedCaseSpec(reason="No gateway end-to-end Messages parity case is registered."), + surface="gateway", + ), + CaseDefinition( + "responses", + NotImplementedCaseSpec(reason="No gateway end-to-end Responses parity case is registered."), + surface="gateway", + ), + CaseDefinition( + "count_tokens", + NotImplementedCaseSpec(reason="No gateway end-to-end token-count parity case is registered."), + surface="gateway", + ), + CaseDefinition( + "chat_completions", + NotImplementedCaseSpec(reason="No gateway end-to-end chat parity case is registered."), + surface="gateway", + ), + CaseDefinition( + "transcription", + NotImplementedCaseSpec(reason="No gateway end-to-end transcription parity case is registered."), + surface="gateway", + ), +) + +STRATEGY: Final = StrategyDefinition( + id="e2e_parity", + order=10, + label="End-to-end parity", + description="Compare observable Python and Rust behavior over generated and recorded inputs.", + directory=Path(__file__).parent, + runnable_spec=ModuleCaseSpec, + cases=CASES, + run=run_e2e_cases, + render=render_e2e_results, + surfaces=SURFACES, +) diff --git a/tests/rust-python-harness/strategies/e2e_parity/reporting.py b/tests/rust-python-harness/strategies/e2e_parity/reporting.py new file mode 100644 index 00000000000..d1a5c389bdf --- /dev/null +++ b/tests/rust-python-harness/strategies/e2e_parity/reporting.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +from collections.abc import Sequence +from typing import Final + +from ...shared.reporting.models import CaseResult +from ...shared.reporting.rendering import ReportSection, render_case_outcome + + +def render_e2e_results(results: Sequence[CaseResult]) -> tuple[ReportSection, ...]: + blocks: Final = tuple(render_case_outcome(result) for result in results) + return (ReportSection("End-to-end parity outcomes", blocks or ("No end-to-end cases selected",)),) diff --git a/tests/rust-python-harness/strategies/e2e_parity/runner.py b/tests/rust-python-harness/strategies/e2e_parity/runner.py index 2886c823370..109df68e2ec 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/runner.py +++ b/tests/rust-python-harness/strategies/e2e_parity/runner.py @@ -1,26 +1,125 @@ from __future__ import annotations -from collections.abc import Sequence +import importlib +from collections.abc import Callable, Sequence +from contextlib import AbstractContextManager, nullcontext +from dataclasses import dataclass from pathlib import Path +from time import monotonic +from typing import Final, cast -from ...shared.reporting.models import HarnessCase, HarnessRun -from ...shared.reporting.pytest_runner import UpdateCallback, run_pytest +from ...shared.reporting.models import CaseResult, HarnessCase, HarnessRun, RunStatus +from ...shared.reporting.strategy import ModuleCaseSpec, UpdateCallback -def run( +@dataclass(frozen=True, slots=True) +class E2ECheck: + name: str + execute: Callable[[], None] + + +@dataclass(frozen=True, slots=True) +class E2ELoadFailure: + message: str + + +def _load_checks(reference: str) -> AbstractContextManager[object] | E2ELoadFailure: + try: + module: Final = importlib.import_module(reference) + factory_value: Final[object] = getattr(module, "parity_checks", None) + if not callable(factory_value): + return E2ELoadFailure(f"{reference} must export parity_checks()") + factory: Final = cast(Callable[[], object], factory_value) + checks_value: Final = factory() + except Exception as error: + return E2ELoadFailure(f"cannot load {reference}: {type(error).__name__}: {error}") + if isinstance(checks_value, tuple): + return nullcontext(cast(object, checks_value)) + if isinstance(checks_value, AbstractContextManager): + return cast(AbstractContextManager[object], checks_value) + return E2ELoadFailure( + f"{reference}.parity_checks() must return tuple[E2ECheck, ...] or a context manager yielding one" + ) + + +def _validate_checks(reference: str, checks_value: object) -> tuple[E2ECheck, ...] | E2ELoadFailure: + if not isinstance(checks_value, tuple): + return E2ELoadFailure(f"{reference}.parity_checks() context manager must yield tuple[E2ECheck, ...]") + untyped_checks: Final = cast(tuple[object, ...], checks_value) + if not all(isinstance(check, E2ECheck) for check in untyped_checks): + return E2ELoadFailure(f"{reference}.parity_checks() must return tuple[E2ECheck, ...]") + return cast(tuple[E2ECheck, ...], untyped_checks) + + +def _run_check( + run: HarnessRun, + result: CaseResult, + check: E2ECheck, + nodeid: str, + on_update: UpdateCallback, +) -> None: + started_at: Final = monotonic() + try: + check.execute() + except Exception as error: + result.record(nodeid, RunStatus.FAILED, monotonic() - started_at) + run.failures.append((nodeid, f"{type(error).__name__}: {error}")) + else: + result.record(nodeid, RunStatus.PASSED, monotonic() - started_at) + on_update(run) + + +def _run_case(run: HarnessRun, harness_case: HarnessCase, on_update: UpdateCallback) -> None: + result: Final = run.results[harness_case.key] + spec: Final = harness_case.spec + if not isinstance(spec, ModuleCaseSpec): + return + loaded: Final = _load_checks(spec.module) + if isinstance(loaded, E2ELoadFailure): + load_nodeid: Final = f"e2e:{harness_case.surface}:{harness_case.sdk_function}:load" + result.collected.add(load_nodeid) + result.record(load_nodeid, RunStatus.ERROR) + run.failures.append((load_nodeid, loaded.message)) + on_update(run) + return + try: + with loaded as checks_value: + checks: Final = _validate_checks(spec.module, checks_value) + if isinstance(checks, E2ELoadFailure): + raise TypeError(checks.message) + nodeids: Final = tuple( + (check, f"e2e:{harness_case.surface}:{harness_case.sdk_function}:{check.name}") for check in checks + ) + result.collected.update(nodeid for _, nodeid in nodeids) + if not nodeids: + result.status = RunStatus.SKIPPED + on_update(run) + return + result.status = RunStatus.RUNNING + on_update(run) + for check, nodeid in nodeids: + _run_check(run, result, check, nodeid, on_update) + except Exception as error: + session_nodeid: Final = f"e2e:{harness_case.surface}:{harness_case.sdk_function}:session" + result.collected.add(session_nodeid) + result.record(session_nodeid, RunStatus.ERROR) + run.failures.append((session_nodeid, f"{type(error).__name__}: {error}")) + on_update(run) + + +def run_e2e_cases( cases: Sequence[HarnessCase], repo_root: Path, on_update: UpdateCallback, - pytest_args: Sequence[str] = (), + runner_args: Sequence[str] = (), ) -> tuple[int, HarnessRun]: - return run_pytest(cases, repo_root, on_update, pytest_args) - - -def main(argv: Sequence[str] | None = None) -> int: - from ...cli import main as harness_main - - return harness_main(argv, strategy_id="e2e_parity") - - -if __name__ == "__main__": - raise SystemExit(main()) + del repo_root, runner_args + run: Final = HarnessRun.from_cases(cases) + for harness_case in cases: + _run_case(run, harness_case, on_update) + run.finished_at = monotonic() + on_update(run) + failed: Final = any( + result.status in {RunStatus.ERROR, RunStatus.FAILED, RunStatus.MISSING} for result in run.results.values() + ) + return int(failed), run diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/__init__.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/chat_completions/__init__.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/chat_completions/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/__init__.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/__init__.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/config.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/config.py index 2d65189790a..ffd8d903551 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/config.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/config.py @@ -6,6 +6,8 @@ from collections.abc import Callable, Mapping from pathlib import Path from typing import Final +from ......shared.parity.fixtures.store import fixture_directory + FIXTURE_DIR_ENV: Final = "LITELLM_OCR_FIXTURE_DIR" DEFAULT_FIXTURE_DIRECTORY: Final = Path(__file__).with_name("data") @@ -40,5 +42,4 @@ def recording_environment( def configured_fixture_directory() -> Path: - configured: Final = os.environ.get(FIXTURE_DIR_ENV) - return Path(configured).expanduser() if configured is not None else DEFAULT_FIXTURE_DIRECTORY + return fixture_directory(None, os.environ.get(FIXTURE_DIR_ENV), DEFAULT_FIXTURE_DIRECTORY) diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/reducto.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/reducto.py index fe8fab50518..bcca8ac6d42 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/reducto.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/reducto.py @@ -89,13 +89,19 @@ ReductoBlockType = Literal[ "Comment", "Signature", ] -_REDUCTO_FILTER_BLOCK_GROUPS: Final[tuple[tuple[ReductoBlockType, ...], ...]] = ( +REDUCTO_FORMATTING_INCLUDE_GROUPS: Final[tuple[tuple[ReductoFormattingInclude, ...], ...]] = ( + (), + ("hyperlinks",), + ("change_tracking", "highlight", "comments"), + ("signatures", "ignore_watermarks"), +) +REDUCTO_FILTER_BLOCK_GROUPS: Final[tuple[tuple[ReductoBlockType, ...], ...]] = ( (), ("Header",), ("Header", "Footer", "Page Number"), ("Figure", "Table", "Key Value"), ) -_REDUCTO_RETURN_IMAGE_GROUPS: Final[tuple[tuple[ReductoReturnImage, ...], ...]] = ( +REDUCTO_RETURN_IMAGE_GROUPS: Final[tuple[tuple[ReductoReturnImage, ...], ...]] = ( (), ("figure",), ("table",), @@ -258,14 +264,7 @@ def _formatting_strategy() -> SearchStrategy[ReductoFormatting]: ), st.sampled_from((False, True)).map(lambda value: {"add_page_markers": value}), st.sampled_from((False, True)).map(lambda value: {"merge_tables": value}), - st.sampled_from( - ( - (), - ("hyperlinks",), - ("change_tracking", "highlight", "comments"), - ("signatures", "ignore_watermarks"), - ) - ) + st.sampled_from(REDUCTO_FORMATTING_INCLUDE_GROUPS) .map(list) .map(lambda value: {"include": value}), ) @@ -288,7 +287,7 @@ def _chunking_strategy() -> SearchStrategy[ReductoChunking]: def _retrieval_strategy() -> SearchStrategy[ReductoRetrieval]: filter_blocks: Final = cast( SearchStrategy[list[ReductoBlockType]], - st.sampled_from(_REDUCTO_FILTER_BLOCK_GROUPS).map(list), + st.sampled_from(REDUCTO_FILTER_BLOCK_GROUPS).map(list), ) return st.one_of( _chunking_strategy().map(lambda chunking: ReductoRetrieval(chunking=chunking)), @@ -305,7 +304,7 @@ def _retrieval_strategy() -> SearchStrategy[ReductoRetrieval]: def _settings_strategy() -> SearchStrategy[ReductoSettings]: # force_url_result stays model-compatible but is not recorded until the # response transform follows and downloads result.url. - return_images: Final[SearchStrategy[list[ReductoReturnImage]]] = st.sampled_from(_REDUCTO_RETURN_IMAGE_GROUPS).map( + return_images: Final[SearchStrategy[list[ReductoReturnImage]]] = st.sampled_from(REDUCTO_RETURN_IMAGE_GROUPS).map( list ) page_ranges: Final = st.one_of( diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_fixture_models.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_fixture_models.py index bdccddd9cfa..80b830369e6 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_fixture_models.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_fixture_models.py @@ -4,16 +4,16 @@ import base64 from collections.abc import Callable from datetime import date from pathlib import Path -from typing import Final, TypeVar, cast +from typing import Final, cast from unittest.mock import patch from urllib.parse import parse_qs, urlparse import httpx import pytest import respx -from hypothesis import find, given, settings +from hypothesis import given, settings from hypothesis import strategies as st -from hypothesis.strategies import DataObject, SearchStrategy +from hypothesis.strategies import DataObject from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError from litellm.llms.azure_ai.ocr.document_intelligence.transformation import AzureDocumentIntelligenceOCRConfig @@ -27,6 +27,7 @@ from litellm.llms.mistral.ocr.transformation import MistralOCRConfig from litellm.llms.reducto.ocr.transformation import ReductoParseLegacyConfig, ReductoParseV3Config from litellm.llms.vertex_ai.ocr.deepseek_transformation import VertexAIDeepSeekOCRConfig from litellm.llms.vertex_ai.ocr.transformation import VertexAIOCRConfig + from .....shared.parity.fixtures.media import structured_pdf_data_uri from .conftest import ocr_fixture_marks from .fixtures.azure import ( @@ -49,7 +50,10 @@ from .fixtures.base import ( from .fixtures.mistral import MISTRAL_MODELS, MistralOcrSdkInput, mistral_input_strategy from .fixtures.models import OcrParityCase, OcrSdkInput from .fixtures.reducto import ( + REDUCTO_FILTER_BLOCK_GROUPS, + REDUCTO_FORMATTING_INCLUDE_GROUPS, REDUCTO_LEGACY_MODELS, + REDUCTO_RETURN_IMAGE_GROUPS, REDUCTO_V3_MODELS, ReductoChunking, ReductoDocumentUrlDocument, @@ -71,6 +75,7 @@ from .fixtures.vertex import ( vertex_deepseek_input_strategy, vertex_mistral_input_strategy, ) +from .test_support import find_fixture as _find_fixture COMMON_FIELDS: Final = frozenset( {"contract", "model", "document", "custom_llm_provider", "vertex_project", "vertex_location"} @@ -150,27 +155,6 @@ _MISTRAL_2505_OPTION_GROUPS: Final = frozenset( _AZURE_MISTRAL_OPTION_GROUPS: Final = _MISTRAL_2505_OPTION_GROUPS - { frozenset({"document_annotation_format", "document_annotation_prompt"}) } -_REDUCTO_FORMATTING_INCLUDE_GROUPS: Final = ( - (), - ("hyperlinks",), - ("change_tracking", "highlight", "comments"), - ("signatures", "ignore_watermarks"), -) -_REDUCTO_FILTER_BLOCK_GROUPS: Final = ( - (), - ("Header",), - ("Header", "Footer", "Page Number"), - ("Figure", "Table", "Key Value"), -) -_REDUCTO_RETURN_IMAGE_GROUPS: Final = ( - (), - ("figure",), - ("table",), - ("page",), - ("figure", "table"), -) -_FIND_SETTINGS: Final = settings(max_examples=2_000, deadline=None, derandomize=True, database=None) -_FixtureInputT = TypeVar("_FixtureInputT") INLINE_IMAGE_DATA_URI: Final = "data:image/png;base64,dGVzdA==" _MapOcrParams = Callable[[dict[str, object], dict[str, object], str], dict[str, object]] _TransformOcrRequest = Callable[ @@ -198,13 +182,6 @@ def _transform_with_stubbed_download( return transform_request(model, document, mapped, {}) -def _find_fixture( - strategy: SearchStrategy[_FixtureInputT], - predicate: Callable[[_FixtureInputT], bool], -) -> _FixtureInputT: - return find(strategy, predicate, settings=_FIND_SETTINGS) - - def _document_transport(document: ImageUrlDocument | DocumentUrlDocument) -> tuple[str, str]: if isinstance(document, ImageUrlDocument): source: Final = document.image_url.url if isinstance(document.image_url, ImageUrlValue) else document.image_url @@ -701,7 +678,7 @@ def test_reducto_v3_strategy_only_generates_bounded_valid_sdk_inputs(sdk_input: if "merge_tables" in formatting_fields: assert sdk_input.formatting.merge_tables in {False, True} if "include" in formatting_fields: - assert tuple(sdk_input.formatting.include) in _REDUCTO_FORMATTING_INCLUDE_GROUPS + assert tuple(sdk_input.formatting.include) in REDUCTO_FORMATTING_INCLUDE_GROUPS if "retrieval" in option_groups: retrieval_fields: Final = frozenset(sdk_input.retrieval.model_fields_set) assert retrieval_fields in { @@ -719,7 +696,7 @@ def test_reducto_v3_strategy_only_generates_bounded_valid_sdk_inputs(sdk_input: if chunking.chunk_overlap: assert chunking.chunk_size == 1000 if "filter_blocks" in retrieval_fields: - assert tuple(sdk_input.retrieval.filter_blocks) in _REDUCTO_FILTER_BLOCK_GROUPS + assert tuple(sdk_input.retrieval.filter_blocks) in REDUCTO_FILTER_BLOCK_GROUPS if "embedding_optimized" in retrieval_fields: assert chunking.chunk_mode == "variable" assert chunking.chunk_size is None @@ -757,7 +734,7 @@ def test_reducto_v3_strategy_only_generates_bounded_valid_sdk_inputs(sdk_input: if "return_ocr_data" in settings_fields: assert sdk_input.settings.return_ocr_data is True if "return_images" in settings_fields: - assert tuple(sdk_input.settings.return_images) in _REDUCTO_RETURN_IMAGE_GROUPS + assert tuple(sdk_input.settings.return_images) in REDUCTO_RETURN_IMAGE_GROUPS if "embed_pdf_metadata_dpi" in settings_fields: assert sdk_input.settings.embed_pdf_metadata is True assert sdk_input.settings.embed_pdf_metadata_dpi in {50, 100, 250} @@ -859,7 +836,7 @@ def test_reducto_v3_strategy_reaches_every_formatting_boolean(field: str, value: assert getattr(sdk_input.formatting, field) is value -@pytest.mark.parametrize("include", _REDUCTO_FORMATTING_INCLUDE_GROUPS) +@pytest.mark.parametrize("include", REDUCTO_FORMATTING_INCLUDE_GROUPS) def test_reducto_v3_strategy_reaches_every_formatting_include(include: tuple[str, ...]) -> None: sdk_input: Final = _find_fixture( reducto_v3_input_strategy(INLINE_IMAGE_DATA_URI), @@ -910,7 +887,7 @@ def test_reducto_v3_strategy_reaches_every_chunk_overlap(chunk_overlap: int) -> assert sdk_input.retrieval.chunking.chunk_overlap == chunk_overlap -@pytest.mark.parametrize("filter_blocks", _REDUCTO_FILTER_BLOCK_GROUPS) +@pytest.mark.parametrize("filter_blocks", REDUCTO_FILTER_BLOCK_GROUPS) def test_reducto_v3_strategy_reaches_every_filter_block_group(filter_blocks: tuple[str, ...]) -> None: sdk_input: Final = _find_fixture( reducto_v3_input_strategy(INLINE_IMAGE_DATA_URI), @@ -992,7 +969,7 @@ def test_reducto_v3_strategy_reaches_every_scalar_setting(field: str, value: obj assert getattr(sdk_input.settings, field) == value -@pytest.mark.parametrize("return_images", _REDUCTO_RETURN_IMAGE_GROUPS) +@pytest.mark.parametrize("return_images", REDUCTO_RETURN_IMAGE_GROUPS) def test_reducto_v3_strategy_reaches_every_return_image_group(return_images: tuple[str, ...]) -> None: sdk_input: Final = _find_fixture( reducto_v3_input_strategy(INLINE_IMAGE_DATA_URI), diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_record_fixtures.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_record_fixtures.py index 6c6dcea17d1..bc770131d85 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_record_fixtures.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_record_fixtures.py @@ -7,7 +7,6 @@ from pathlib import Path from typing import Final, cast import pytest -from hypothesis import find, settings from hypothesis.strategies import SearchStrategy from .....shared.parity.fixtures.cli import parse_recording_args @@ -41,6 +40,7 @@ from .fixtures.vertex import ( vertex_deepseek_provider_rejected_inputs, vertex_mistral_provider_rejected_inputs, ) +from .test_support import find_fixture class _UnusedOcrClient: @@ -76,7 +76,6 @@ _MISTRAL_PARAMS: Final = frozenset( _MISTRAL_2512_PARAMS: Final = _MISTRAL_PARAMS - {"include_blocks"} _MISTRAL_2505_PARAMS: Final = _MISTRAL_2512_PARAMS - {"extract_header", "extract_footer", "table_format"} _AZURE_MISTRAL_PARAMS: Final = _MISTRAL_2505_PARAMS - {"document_annotation_prompt"} -_FIND_SETTINGS: Final = settings(max_examples=2_000, deadline=None, derandomize=True, database=None) _INLINE_IMAGE_DATA_URI: Final = "data:image/png;base64,dGVzdA==" @@ -94,7 +93,7 @@ def _find_input( strategy: SearchStrategy[OcrSdkInputBase], predicate: Callable[[OcrSdkInputBase], bool], ) -> OcrSdkInputBase: - return find(strategy, predicate, settings=_FIND_SETTINGS) + return find_fixture(strategy, predicate) def _document_transport(case_input: OcrSdkInputBase) -> tuple[str, str]: diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py index bedbdeb6a13..e72980752f2 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py @@ -2,23 +2,21 @@ from __future__ import annotations import asyncio import sys +import tempfile import traceback -from collections.abc import Awaitable, Callable, Coroutine, Generator +from collections.abc import Callable, Coroutine, Generator from contextlib import contextmanager -from dataclasses import dataclass from enum import Enum +from functools import partial from pathlib import Path -from typing import Final, cast +from typing import Annotated, Final, Literal, cast -import pytest +from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter from litellm.llms.base_llm.ocr.transformation import OCRResponse -from litellm.rust_bridge import get_native_bridge -from litellm.rust_bridge import ocr as rust_ocr_bridge -from litellm.rust_bridge.ocr import RustAocr, RustOcr -from .....shared.parity.compare import assert_model_parity, assert_parity, assert_request_parity -from .....shared.parity.fixtures.store import recorded_fixtures -from .....shared.parity.inprocess import run_in_process + +from .....shared.parity.compare import assert_parity +from .....shared.parity.fixtures.store import fixture_id, recorded_fixtures from .....shared.parity.models import ( SDKCommand, SDKError, @@ -29,22 +27,21 @@ from .....shared.parity.models import ( WorkerSuccess, sdk_error_report, ) -from .....shared.parity.replay import replay_server from .....shared.parity.runner import ( ExecutionVariant, SubprocessRunner, SubprocessWorker, execution_worker_pair, parity_worker_main, - run_execution, ) +from ...runner import E2ECheck from .fixtures.config import configured_fixture_directory from .fixtures.models import OcrParityCase, OcrSdkInput API_KEY: Final = "test-key" PYTHON_HTTP_SENTINEL: Final = "python-ocr-parity-fallback" -PYTHON_VARIANT: Final = ExecutionVariant(name="Python", environment=(("LITELLM_USE_RUST_OCR", "0"),)) -RUST_VARIANT: Final = ExecutionVariant(name="Rust", environment=(("LITELLM_USE_RUST_OCR", "1"),)) +PYTHON_VARIANT: Final = ExecutionVariant(name="Python", environment=(("LITELLM_RUST", "0"),)) +RUST_VARIANT: Final = ExecutionVariant(name="Rust", environment=(("LITELLM_RUST", "1"),)) class SDKRoute(str, Enum): @@ -52,16 +49,34 @@ class SDKRoute(str, Enum): AOCR = "aocr" -@dataclass(frozen=True, slots=True) -class InvalidOcrCase: +class InvalidOcrCase(BaseModel): + model_config = ConfigDict(frozen=True) + name: str model: str - document: object + document: JsonValue expected_exception_type: str expected_status_code: int expected_message: str - extra_kwargs: tuple[tuple[str, object], ...] = () - expected_rust_calls: int = 0 + extra_kwargs: tuple[tuple[str, JsonValue], ...] = () + + +class RecordedOcrWorkerCase(BaseModel): + model_config = ConfigDict(frozen=True) + + kind: Literal["recorded"] = "recorded" + case: OcrParityCase + + +class InvalidOcrWorkerCase(BaseModel): + model_config = ConfigDict(frozen=True) + + kind: Literal["invalid"] = "invalid" + case: InvalidOcrCase + + +OcrWorkerCase = Annotated[RecordedOcrWorkerCase | InvalidOcrWorkerCase, Field(discriminator="kind")] +OCR_WORKER_CASE_ADAPTER: Final[TypeAdapter[OcrWorkerCase]] = TypeAdapter(OcrWorkerCase) INVALID_OCR_CASES: Final = ( @@ -120,7 +135,6 @@ INVALID_OCR_CASES: Final = ( expected_exception_type="litellm.exceptions.APIConnectionError", expected_status_code=500, expected_message="Document URL is required", - expected_rust_calls=1, ), InvalidOcrCase( name="missing_image_url", @@ -129,7 +143,6 @@ INVALID_OCR_CASES: Final = ( expected_exception_type="litellm.exceptions.APIConnectionError", expected_status_code=500, expected_message="Document URL is required", - expected_rust_calls=1, ), InvalidOcrCase( name="invalid_request_format", @@ -199,32 +212,13 @@ def _execute_sdk_case( return _execute_sdk_call(call_kwargs, route, event_loop) -def _execute_recorded_sdk_case( - sdk_input: OcrSdkInput, - route: SDKRoute, - mock_url: str, - event_loop: asyncio.AbstractEventLoop, -) -> OCRResponse | SDKError: - import litellm - - call_kwargs: Final = _call_kwargs(sdk_input, mock_url, route) - try: - if route is SDKRoute.OCR: - sync_route: Final = cast(Callable[..., OCRResponse], litellm.ocr) - return sync_route(**call_kwargs) - async_route: Final = cast(Callable[..., Coroutine[object, object, OCRResponse]], litellm.aocr) - return event_loop.run_until_complete(async_route(**call_kwargs)) - except Exception as error: - return sdk_error_report(error) - - def _execute_invalid_sdk_case( case: InvalidOcrCase, route: SDKRoute, mock_url: str, event_loop: asyncio.AbstractEventLoop, ) -> SDKReport: - call_kwargs: Final = { + call_kwargs: Final[dict[str, object]] = { "model": case.model, "document": case.document, "api_base": mock_url, @@ -235,204 +229,93 @@ def _execute_invalid_sdk_case( return _execute_sdk_call(call_kwargs, route, event_loop) -class _RustOcrSpy: - def __init__(self, delegate: RustOcr) -> None: - self.delegate: Final = delegate - self.calls = 0 +def _check_recorded_ocr_sdk_parity( + ocr_fixture: OcrParityCase, + route: SDKRoute, + case_file: Path, + sdk_workers: tuple[SubprocessWorker, SubprocessWorker], +) -> None: + python_worker, rust_worker = sdk_workers + python: Final = python_worker.execute(case_file, route.value, ocr_fixture.provider_responses) + rust: Final = rust_worker.execute(case_file, route.value, ocr_fixture.provider_responses) - def __call__( - self, - model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - timeout_seconds: float | None, - ) -> dict[str, object]: - self.calls += 1 - return self.delegate( - model=model, - document=document, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - optional_params=optional_params, - timeout_seconds=timeout_seconds, - ) + assert_parity(python, rust, PYTHON_HTTP_SENTINEL) + if any(response.status_code >= 400 for response in ocr_fixture.provider_responses): + assert isinstance(python.report, SDKError) -class _RustAocrSpy: - def __init__(self, delegate: RustAocr) -> None: - self.delegate: Final = delegate - self.calls = 0 +def _check_invalid_ocr_sdk_parity( + case: InvalidOcrCase, + route: SDKRoute, + case_file: Path, + sdk_workers: tuple[SubprocessWorker, SubprocessWorker], +) -> None: + python_worker, rust_worker = sdk_workers + python: Final = python_worker.execute(case_file, route.value, ()) + rust: Final = rust_worker.execute(case_file, route.value, ()) - async def __call__( - self, - model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - timeout_seconds: float | None, - ) -> dict[str, object]: - self.calls += 1 - result: Final[Awaitable[dict[str, object]]] = self.delegate( - model=model, - document=document, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - optional_params=optional_params, - timeout_seconds=timeout_seconds, - ) - return await result + assert_parity(python, rust, PYTHON_HTTP_SENTINEL) + assert python.requests == () + assert rust.requests == () + assert isinstance(python.report, SDKError) + assert python.report.exception_type == case.expected_exception_type + assert python.report.status_code == case.expected_status_code + assert case.expected_message in python.report.message + + +def _recorded_check_name(fixture: OcrParityCase, route: SDKRoute) -> str: + case_input: Final = fixture.litellm_input + provider: Final = case_input.custom_llm_provider + prefix: Final = f"{provider}/{case_input.model}" if provider else case_input.model + return f"recorded:{route.value}:{fixture_id(case_input, prefix)}" + + +def _write_worker_case(directory: Path, index: int, case: OcrWorkerCase) -> Path: + case_file: Final = directory / f"case-{index}.json" + case_file.write_text(OCR_WORKER_CASE_ADAPTER.dump_json(case).decode("utf-8"), encoding="utf-8") + return case_file @contextmanager -def _restore_rust_ocr_state() -> Generator[None]: - enabled: Final = rust_ocr_bridge.rust_ocr_enabled() - ocr_impl: Final = rust_ocr_bridge._rust_ocr_impl # pyright: ignore[reportPrivateUsage] # preserve injected test binding - aocr_impl: Final = rust_ocr_bridge._rust_aocr_impl # pyright: ignore[reportPrivateUsage] # preserve injected test binding - try: - yield - finally: - rust_ocr_bridge.use_litellm_rust(enabled, ocr=ocr_impl, aocr=aocr_impl) - - -def _native_spies() -> tuple[_RustOcrSpy, _RustAocrSpy]: - native_bridge: Final = get_native_bridge() - if native_bridge is None: - pytest.fail("native Rust bridge is required for OCR parity testing") - sync_spy: Final = _RustOcrSpy(cast(RustOcr, getattr(native_bridge, "ocr"))) - async_spy: Final = _RustAocrSpy(cast(RustAocr, getattr(native_bridge, "aocr"))) - return sync_spy, async_spy - - -@pytest.fixture(scope="module") -def sdk_workers() -> Generator[tuple[SubprocessWorker, SubprocessWorker]]: +def parity_checks() -> Generator[tuple[E2ECheck, ...]]: + fixtures: Final = tuple( + fixture + for fixture in recorded_fixtures(configured_fixture_directory(), OcrParityCase) + if fixture.litellm_input.contract not in {"reducto_v3", "reducto_legacy"} + ) runner: Final = SubprocessRunner( entrypoint=Path(__file__), baseline_user_agent=PYTHON_HTTP_SENTINEL, route_label="OCR", ) - with execution_worker_pair(runner, PYTHON_VARIANT, RUST_VARIANT) as workers: - yield workers - - -@pytest.fixture(scope="module") -def startup_ocr_fixture() -> OcrParityCase: - directory: Final = configured_fixture_directory() - fixtures: Final = recorded_fixtures(directory, OcrParityCase) - if not fixtures: - pytest.skip(f"no recorded fixtures in {directory}") - return fixtures[0] - - -@pytest.mark.parametrize("route", tuple(SDKRoute), ids=tuple(route.value for route in SDKRoute)) -def test_recorded_ocr_sdk_parity( - ocr_fixture: OcrParityCase, - route: SDKRoute, -) -> None: - sync_spy, async_spy = _native_spies() - event_loop: Final = asyncio.new_event_loop() - try: - with _restore_rust_ocr_state(), replay_server() as provider: - rust_ocr_bridge.use_litellm_rust(False, ocr=sync_spy, aocr=async_spy) - rust_ocr_bridge.use_litellm_rust(False) - python: Final = run_in_process( - provider, - ocr_fixture.provider_responses, - lambda mock_url: _execute_recorded_sdk_case(ocr_fixture.litellm_input, route, mock_url, event_loop), + with tempfile.TemporaryDirectory(prefix="litellm-ocr-parity-") as raw_directory: + directory: Final = Path(raw_directory) + recorded_files: Final = tuple( + _write_worker_case(directory, index, RecordedOcrWorkerCase(case=fixture)) + for index, fixture in enumerate(fixtures) + ) + invalid_files: Final = tuple( + _write_worker_case(directory, len(recorded_files) + index, InvalidOcrWorkerCase(case=case)) + for index, case in enumerate(INVALID_OCR_CASES) + ) + with execution_worker_pair(runner, PYTHON_VARIANT, RUST_VARIANT) as workers: + recorded: Final = tuple( + E2ECheck( + _recorded_check_name(fixture, route), + partial(_check_recorded_ocr_sdk_parity, fixture, route, case_file, workers), + ) + for fixture, case_file in zip(fixtures, recorded_files, strict=True) + for route in SDKRoute ) - assert sync_spy.calls == 0 - assert async_spy.calls == 0 - - rust_ocr_bridge.use_litellm_rust(True) - rust: Final = run_in_process( - provider, - ocr_fixture.provider_responses, - lambda mock_url: _execute_recorded_sdk_case(ocr_fixture.litellm_input, route, mock_url, event_loop), + invalid: Final = tuple( + E2ECheck( + f"invalid:{route.value}:{case.name}", + partial(_check_invalid_ocr_sdk_parity, case, route, case_file, workers), + ) + for case, case_file in zip(INVALID_OCR_CASES, invalid_files, strict=True) + for route in SDKRoute ) - finally: - event_loop.close() - - assert sync_spy.calls == (1 if route is SDKRoute.OCR else 0) - assert async_spy.calls == (1 if route is SDKRoute.AOCR else 0) - assert_request_parity(python.requests, rust.requests) - if any(response.status_code >= 400 for response in ocr_fixture.provider_responses): - assert isinstance(python.response, SDKError) - if isinstance(python.response, SDKError): - assert python.response == rust.response - else: - assert isinstance(rust.response, OCRResponse) - assert_model_parity(python.response, rust.response) - - -@pytest.mark.parametrize("case", INVALID_OCR_CASES, ids=tuple(case.name for case in INVALID_OCR_CASES)) -@pytest.mark.parametrize("route", tuple(SDKRoute), ids=tuple(route.value for route in SDKRoute)) -def test_invalid_ocr_sdk_parity(case: InvalidOcrCase, route: SDKRoute) -> None: - sync_spy, async_spy = _native_spies() - event_loop: Final = asyncio.new_event_loop() - try: - with _restore_rust_ocr_state(), replay_server() as provider: - rust_ocr_bridge.use_litellm_rust(False, ocr=sync_spy, aocr=async_spy) - rust_ocr_bridge.use_litellm_rust(False) - python: Final = run_in_process( - provider, - (), - lambda mock_url: _execute_invalid_sdk_case(case, route, mock_url, event_loop), - ) - assert sync_spy.calls == 0 - assert async_spy.calls == 0 - - rust_ocr_bridge.use_litellm_rust(True) - rust: Final = run_in_process( - provider, - (), - lambda mock_url: _execute_invalid_sdk_case(case, route, mock_url, event_loop), - ) - finally: - event_loop.close() - - assert sync_spy.calls == (case.expected_rust_calls if route is SDKRoute.OCR else 0) - assert async_spy.calls == (case.expected_rust_calls if route is SDKRoute.AOCR else 0) - assert python.requests == () - assert rust.requests == () - assert python.response == rust.response - assert isinstance(python.response, SDKError) - assert python.response.exception_type == case.expected_exception_type - assert python.response.status_code == case.expected_status_code - assert case.expected_message in python.response.message - - -def test_ocr_subprocess_startup_smoke( - startup_ocr_fixture: OcrParityCase, - tmp_path: Path, - sdk_workers: tuple[SubprocessWorker, SubprocessWorker], -) -> None: - case_file: Final = tmp_path / "ocr-startup-smoke.json" - case_file.write_text(startup_ocr_fixture.model_dump_json(indent=2, exclude_unset=True), encoding="utf-8") - python_worker, rust_worker = sdk_workers - python: Final = run_execution( - python_worker, - case_file, - SDKRoute.OCR.value, - startup_ocr_fixture.provider_responses, - ) - rust: Final = run_execution( - rust_worker, - case_file, - SDKRoute.OCR.value, - startup_ocr_fixture.provider_responses, - ) - - assert_parity(python, rust, PYTHON_HTTP_SENTINEL) + yield (*recorded, *invalid) def _execute_worker_command( @@ -444,8 +327,12 @@ def _execute_worker_command( command: Final = SDKCommand.model_validate_json(command_json) case_file: Final = Path(command.case_file) route: Final = SDKRoute(command.route) - case: Final = OcrParityCase.model_validate_json(case_file.read_text(encoding="utf-8")) - return WorkerSuccess(report=_execute_sdk_case(case.litellm_input, route, mock_url, event_loop)) + worker_case: Final = OCR_WORKER_CASE_ADAPTER.validate_json(case_file.read_bytes()) + match worker_case: + case RecordedOcrWorkerCase(case=recorded): + return WorkerSuccess(report=_execute_sdk_case(recorded.litellm_input, route, mock_url, event_loop)) + case InvalidOcrWorkerCase(case=invalid): + return WorkerSuccess(report=_execute_invalid_sdk_case(invalid, route, mock_url, event_loop)) except Exception: return WorkerFailure(error=traceback.format_exc()) diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_support.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_support.py new file mode 100644 index 00000000000..b6526fb12f5 --- /dev/null +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_support.py @@ -0,0 +1,17 @@ +from __future__ import annotations + +from collections.abc import Callable +from typing import TypeVar + +from hypothesis import find, settings +from hypothesis.strategies import SearchStrategy + +FixtureT = TypeVar("FixtureT") +FIND_SETTINGS = settings(max_examples=2_000, deadline=None, derandomize=True, database=None) + + +def find_fixture( + strategy: SearchStrategy[FixtureT], + predicate: Callable[[FixtureT], bool], +) -> FixtureT: + return find(strategy, predicate, settings=FIND_SETTINGS) diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/responses/__init__.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/responses/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/rust-python-harness/strategies/e2e_parity/strategy.json b/tests/rust-python-harness/strategies/e2e_parity/strategy.json deleted file mode 100644 index d791b9373aa..00000000000 --- a/tests/rust-python-harness/strategies/e2e_parity/strategy.json +++ /dev/null @@ -1,50 +0,0 @@ -{ - "order": 10, - "id": "e2e_parity", - "label": "End-to-end parity", - "description": "Compare observable Python and Rust SDK behavior over generated and recorded inputs.", - "functions": { - "ocr": { - "coverage": "partial", - "selectors": [ - "tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py", - "tests/test_litellm/ocr/test_rust_bridge.py" - ], - "note": "Recorded sync/async SDK parity; invalid-model provider errors differ, and Reducto lacks a Rust contract." - }, - "messages": { - "coverage": "partial", - "selectors": [ - "tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py" - ], - "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added." - }, - "responses": { - "coverage": "partial", - "selectors": [ - "tests/test_litellm/responses/test_rust_bridge_websocket.py" - ], - "note": "Covers the websocket bridge; full responses parity is still being added." - }, - "count_tokens": { - "coverage": "planned", - "selectors": [], - "note": "No Rust count_tokens parity test is present yet." - }, - "chat_completions": { - "coverage": "partial", - "selectors": [ - "tests/test_litellm/rust_bridge/test_chat_completions.py" - ], - "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added." - }, - "transcription": { - "coverage": "partial", - "selectors": [ - "tests/test_litellm/test_audio_transcription_rust_bridge.py" - ], - "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added." - } - }, - "gateway": {} -} diff --git a/tests/rust-python-harness/strategies/e2e_parity/test_runner.py b/tests/rust-python-harness/strategies/e2e_parity/test_runner.py new file mode 100644 index 00000000000..31ef907465a --- /dev/null +++ b/tests/rust-python-harness/strategies/e2e_parity/test_runner.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +from collections.abc import Generator +from contextlib import contextmanager +from functools import partial +from pathlib import Path +from types import SimpleNamespace +from typing import Final +from unittest.mock import Mock, call + +from pytest import MonkeyPatch + +from ...shared.reporting.models import Coverage, HarnessCase, RunStatus +from ...shared.reporting.strategy import ModuleCaseSpec +from . import runner as e2e_runner +from .runner import E2ECheck, run_e2e_cases + + +def test_runs_checks_inside_suite_context(monkeypatch: MonkeyPatch, tmp_path: Path) -> None: + lifecycle: Final = Mock() + + @contextmanager + def parity_checks() -> Generator[tuple[E2ECheck, ...]]: + lifecycle("entered") + try: + yield (E2ECheck("check", partial(lifecycle, "checked")),) + finally: + lifecycle("exited") + + module: Final = SimpleNamespace(parity_checks=parity_checks) + + def import_module(_name: str, _package: str | None = None) -> SimpleNamespace: + return module + + monkeypatch.setattr(e2e_runner.importlib, "import_module", import_module) + case: Final = HarnessCase( + strategy_id="e2e_parity", + strategy_label="End-to-end parity", + sdk_function="ocr", + spec=ModuleCaseSpec(coverage=Coverage.PARTIAL, module="example"), + surface="sdk", + ) + + code, run = run_e2e_cases((case,), tmp_path, lambda _: None) + + assert code == 0, run.failures + assert run.results[case.key].status is RunStatus.PASSED + assert lifecycle.call_args_list == [call("entered"), call("checked"), call("exited")] diff --git a/tests/rust-python-harness/strategies/existing_e2e_test_sdk/README.md b/tests/rust-python-harness/strategies/existing_e2e_test_sdk/README.md deleted file mode 100644 index fb84f170703..00000000000 --- a/tests/rust-python-harness/strategies/existing_e2e_test_sdk/README.md +++ /dev/null @@ -1,3 +0,0 @@ -# Existing e2e SDK tests - -Wires already-existing live-API SDK tests into the matrix instead of writing new parity tests. Selectors point at real test files and folders, such as `tests/ocr_tests/`, rather than individual node IDs, so future tests added to those folders are picked up automatically. diff --git a/tests/rust-python-harness/strategies/existing_e2e_test_sdk/__init__.py b/tests/rust-python-harness/strategies/existing_e2e_test_sdk/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/rust-python-harness/strategies/existing_e2e_test_sdk/runner.py b/tests/rust-python-harness/strategies/existing_e2e_test_sdk/runner.py deleted file mode 100644 index f5ea17735fc..00000000000 --- a/tests/rust-python-harness/strategies/existing_e2e_test_sdk/runner.py +++ /dev/null @@ -1,26 +0,0 @@ -from __future__ import annotations - -from collections.abc import Sequence -from pathlib import Path - -from ...shared.reporting.models import HarnessCase, HarnessRun -from ...shared.reporting.pytest_runner import UpdateCallback, run_pytest - - -def run( - cases: Sequence[HarnessCase], - repo_root: Path, - on_update: UpdateCallback, - pytest_args: Sequence[str] = (), -) -> tuple[int, HarnessRun]: - return run_pytest(cases, repo_root, on_update, pytest_args) - - -def main(argv: Sequence[str] | None = None) -> int: - from ...cli import main as harness_main - - return harness_main(argv, strategy_id="existing_e2e_test_sdk") - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tests/rust-python-harness/strategies/existing_e2e_test_sdk/strategy.json b/tests/rust-python-harness/strategies/existing_e2e_test_sdk/strategy.json deleted file mode 100644 index eefceea1a75..00000000000 --- a/tests/rust-python-harness/strategies/existing_e2e_test_sdk/strategy.json +++ /dev/null @@ -1,14 +0,0 @@ -{ - "order": 40, - "id": "existing_e2e_test_sdk", - "label": "Existing e2e SDK tests", - "description": "Wire already-existing live-API SDK tests into the matrix instead of writing new parity tests.", - "functions": { - "ocr": {"coverage": "partial", "selectors": ["tests/ocr_tests/"], "note": "Existing live OCR provider tests; not yet a frozen Rust/Python oracle comparison."}, - "messages": {"coverage": "planned", "selectors": []}, - "responses": {"coverage": "planned", "selectors": []}, - "count_tokens": {"coverage": "planned", "selectors": []}, - "chat_completions": {"coverage": "partial", "selectors": ["tests/llm_translation/test_anthropic_completion.py", "tests/llm_translation/test_bedrock_completion.py"], "note": "Existing live chat completion tests for providers with confirmed Rust bridge regressions."}, - "transcription": {"coverage": "partial", "selectors": ["tests/audio_tests/test_whisper.py"], "note": "Existing live Whisper transcription test."} - } -} diff --git a/tests/rust-python-harness/strategies/trace_parity/AGENTS.md b/tests/rust-python-harness/strategies/trace_parity/AGENTS.md new file mode 100644 index 00000000000..bb7cb8c91d8 --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/AGENTS.md @@ -0,0 +1 @@ +Maps Python profiler frames onto feature-gated Rust span names via an explicit per-case mapping (Rust span name is the identity) and compares steps, order, and nesting of both live traces against a replayed provider response. diff --git a/tests/rust-python-harness/strategies/trace_parity/README.md b/tests/rust-python-harness/strategies/trace_parity/README.md deleted file mode 100644 index 6520a510112..00000000000 --- a/tests/rust-python-harness/strategies/trace_parity/README.md +++ /dev/null @@ -1,5 +0,0 @@ -# Trace Parity - -Run independently with `uv run python -m tests.rust-python-harness.strategies.trace_parity.runner --plain`. Configure SDK and gateway selectors in `strategy.json`; keep API-specific execution and fixtures in their owning surface folder - -See [the harness guide](../../README.md) for coverage status and shared comparison tools diff --git a/tests/rust-python-harness/strategies/trace_parity/__init__.py b/tests/rust-python-harness/strategies/trace_parity/__init__.py index e69de29bb2d..ec88b0169fa 100644 --- a/tests/rust-python-harness/strategies/trace_parity/__init__.py +++ b/tests/rust-python-harness/strategies/trace_parity/__init__.py @@ -0,0 +1,115 @@ +from pathlib import Path +from typing import Final + +from ...shared.reporting.models import SURFACES, Coverage +from ...shared.reporting.strategy import ( + CaseDefinition, + ModuleCaseSpec, + NotImplementedCaseSpec, + RunnerArgumentDefinition, + StrategyDefinition, +) +from .reporting import render_trace_results +from .runner import run_trace_cases + +CASES: Final[tuple[CaseDefinition, ...]] = ( + CaseDefinition( + "ocr", + ModuleCaseSpec( + coverage=Coverage.PARTIAL, + module="tests.rust-python-harness.strategies.trace_parity.sdk.ocr.case", + ), + surface="sdk", + ), + CaseDefinition( + "messages", + ModuleCaseSpec( + coverage=Coverage.PARTIAL, + module="tests.rust-python-harness.strategies.trace_parity.sdk.messages.case", + note="Async only until anthropic_messages_handler supports sync calls.", + ), + surface="sdk", + ), + CaseDefinition( + "responses", + NotImplementedCaseSpec(reason="No Responses trace-parity case is registered."), + surface="sdk", + ), + CaseDefinition( + "count_tokens", + NotImplementedCaseSpec(reason="No token-count trace-parity case is registered."), + surface="sdk", + ), + CaseDefinition( + "chat_completions", + ModuleCaseSpec( + coverage=Coverage.PARTIAL, + module="tests.rust-python-harness.strategies.trace_parity.sdk.chat_completions.case", + ), + surface="sdk", + ), + CaseDefinition( + "transcription", + ModuleCaseSpec( + coverage=Coverage.PARTIAL, + module="tests.rust-python-harness.strategies.trace_parity.sdk.transcription.case", + note=( + "The Python SDK delegates this provider to the Rust pipeline, so only dispatch is visible " + "to the Python profiler." + ), + ), + surface="sdk", + ), + CaseDefinition( + "ocr", + NotImplementedCaseSpec(reason="No gateway OCR trace-parity case is registered."), + surface="gateway", + ), + CaseDefinition( + "messages", + ModuleCaseSpec( + coverage=Coverage.PARTIAL, + module="tests.rust-python-harness.strategies.trace_parity.gateway.messages.case", + note="Non-streaming success paths only.", + ), + surface="gateway", + ), + CaseDefinition( + "responses", + NotImplementedCaseSpec(reason="No gateway Responses trace-parity case is registered."), + surface="gateway", + ), + CaseDefinition( + "count_tokens", + NotImplementedCaseSpec(reason="No gateway token-count trace-parity case is registered."), + surface="gateway", + ), + CaseDefinition( + "chat_completions", + NotImplementedCaseSpec(reason="No gateway chat trace-parity case is registered."), + surface="gateway", + ), + CaseDefinition( + "transcription", + NotImplementedCaseSpec(reason="No gateway transcription trace-parity case is registered."), + surface="gateway", + ), +) + +STRATEGY: Final = StrategyDefinition( + id="trace_parity", + order=20, + label="Trace parity", + description="Compare pipeline steps, order, and nesting between Python profiler frames and Rust spans via an explicit mapping.", + directory=Path(__file__).parent, + runnable_spec=ModuleCaseSpec, + cases=CASES, + run=run_trace_cases, + render=render_trace_results, + surfaces=SURFACES, + runner_argument=RunnerArgumentDefinition( + option="--scenario", + metavar="NAME", + help="run only this named trace scenario; repeat to select more than one", + ), +) diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/__init__.py b/tests/rust-python-harness/strategies/trace_parity/gateway/__init__.py index e69de29bb2d..f999dfecfc6 100644 --- a/tests/rust-python-harness/strategies/trace_parity/gateway/__init__.py +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/__init__.py @@ -0,0 +1 @@ +"""In-process gateway trace adapters.""" diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py b/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py new file mode 100644 index 00000000000..2bd3a50f39f --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py @@ -0,0 +1,160 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable +from pathlib import Path +from typing import Final, Protocol, cast + +import httpx +from pydantic import BaseModel, ConfigDict + +from ....shared.parity.replay import replay_server +from ....shared.tracing.native import TraceResponsePayload, native_trace_events +from ....shared.tracing.profiler import FunctionTraceEvent, profile_python +from ....shared.tracing.steps import Engine, PipelineProjection, pipeline_projection +from ..models import GatewayRouteSpec, RouteFixture, TraceExecutionFailure, TraceMode, TraceScenario +from ..reporting import TraceComparisonArtifact + + +class _GatewayResponsePayload(BaseModel): + model_config = ConfigDict(strict=True, extra="forbid") + + status: int + body: object + + +class _GatewayClient(Protocol): + def post(self, url: str, *, json: object, headers: dict[str, str]) -> httpx.Response: ... + + +def _collect_python(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]: + import litellm + from fastapi.testclient import TestClient + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.anthropic_endpoints.endpoints import user_api_key_auth + from litellm.proxy import proxy_server + + provider_model: Final = cast(str, fixture.kwargs["provider_model"]) + model_alias: Final = cast(str, fixture.kwargs["model_alias"]) + old_router: Final = proxy_server.llm_router + old_override: Final = proxy_server.app.dependency_overrides.get(user_api_key_auth) + + async def authorize() -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="trace-key") + + proxy_server.llm_router = litellm.Router( + model_list=[ + { + "model_name": model_alias, + "litellm_params": { + "model": provider_model, + "api_key": "trace-provider-key", + "api_base": fixture.kwargs["api_base"], + }, + } + ] + ) + proxy_server.app.dependency_overrides[user_api_key_auth] = authorize + try: + with profile_python(Path(litellm.__file__).parent, threads=True) as profiler: + client: Final = cast(_GatewayClient, TestClient(proxy_server.app)) + response: Final = client.post( + "/v1/messages", + json=fixture.kwargs["body"], + headers={"authorization": "Bearer trace-key"}, + ) + if response.status_code != 200: + raise RuntimeError(f"Python gateway returned {response.status_code}: {response.text}") + return tuple(profiler.events) + finally: + proxy_server.llm_router = old_router + if old_override is None: + proxy_server.app.dependency_overrides.pop(user_api_key_auth, None) + else: + proxy_server.app.dependency_overrides[user_api_key_auth] = old_override + + +def _collect_rust(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]: + from litellm.rust_bridge import get_native_bridge + + bridge: Final[object | None] = get_native_bridge() + trace: Final[object | None] = getattr(bridge, "_trace", None) if bridge is not None else None + gateway_messages: Final[object | None] = getattr(trace, "gateway_messages", None) + if gateway_messages is None or not callable(gateway_messages): + raise RuntimeError("native Rust trace bridge does not expose gateway_messages") + invoke_gateway: Final = cast(Callable[[str, str, str, object], Awaitable[object]], gateway_messages) + + async def invoke() -> object: + return await invoke_gateway( + cast(str, fixture.kwargs["model_alias"]), + cast(str, fixture.kwargs["provider_model"]), + cast(str, fixture.kwargs["api_base"]), + fixture.kwargs["body"], + ) + + result: Final = asyncio.run(invoke()) + payload: Final = TraceResponsePayload.model_validate(result) + response: Final = _GatewayResponsePayload.model_validate(payload.response) + if response.status != 200: + raise RuntimeError(f"Rust gateway returned {response.status}: {response.body}") + return native_trace_events(payload) + + +def _collect(scenario: TraceScenario, engine: Engine) -> tuple[FunctionTraceEvent, ...] | TraceExecutionFailure: + try: + with replay_server() as provider: + base_fixture: Final = scenario.fixture(engine, provider.url) + fixture: Final = RouteFixture( + kwargs={**base_fixture.kwargs, "api_base": provider.url}, + provider_responses=base_fixture.provider_responses, + ) + for response in fixture.provider_responses: + provider.enqueue_response(response) + events: Final = _collect_python(fixture) if engine == "python" else _collect_rust(fixture) + provider.take_requests(len(fixture.provider_responses)) + return events + except Exception as error: + return TraceExecutionFailure(engine, f"{type(error).__name__}: {error}") + + +def _projections( + python_events: tuple[FunctionTraceEvent, ...], + rust_events: tuple[FunctionTraceEvent, ...], + scenario: TraceScenario, + mode: TraceMode, +) -> tuple[PipelineProjection, PipelineProjection, str | None]: + mappings: Final = scenario.mappings_for(mode) + try: + return ( + pipeline_projection("python", python_events, mappings), + pipeline_projection("rust", rust_events, mappings), + None, + ) + except ValueError as error: + return PipelineProjection(), PipelineProjection(), f"harness: {error}" + + +def execute_gateway_trace(route: GatewayRouteSpec, scenario: TraceScenario, mode: TraceMode) -> TraceComparisonArtifact: + mappings: Final = scenario.mappings_for(mode) + python_trace: Final = _collect(scenario, "python") + rust_trace: Final = _collect(scenario, "rust") + collection_python_error: Final = None if isinstance(python_trace, tuple) else f"python: {python_trace.message}" + rust_error: Final = None if isinstance(rust_trace, tuple) else f"rust: {rust_trace.message}" + python_events: Final = python_trace if isinstance(python_trace, tuple) else () + rust_events: Final = rust_trace if isinstance(rust_trace, tuple) else () + python, rust, projection_error = _projections(python_events, rust_events, scenario, mode) + python_error: Final = projection_error or collection_python_error + return TraceComparisonArtifact.from_traces( + surface="gateway", + sdk_function=route.route, + scenario=scenario.name, + mode=mode, + mappings=mappings, + contract=scenario.contract, + python=python.steps, + rust=rust.steps, + python_unmatched=python.unmatched, + python_error=python_error, + rust_error=rust_error, + ) diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/messages/__init__.py b/tests/rust-python-harness/strategies/trace_parity/gateway/messages/__init__.py new file mode 100644 index 00000000000..bd9195b7c22 --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/messages/__init__.py @@ -0,0 +1 @@ +"""Messages gateway trace cases.""" diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/messages/case.py b/tests/rust-python-harness/strategies/trace_parity/gateway/messages/case.py new file mode 100644 index 00000000000..30f51cee353 --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/messages/case.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +import json +from typing import Final + +from .....shared.parity.recorded_http import HttpHeader, RecordedHttpResponse +from .....shared.tracing.steps import Engine, mapping +from ...models import GatewayRouteSpec, RouteFixture, TraceScenario, TraceSuite + + +GATEWAY_MAPPINGS: Final = ( + mapping( + span="python_messages_gateway_route", + python_frame=r"anthropic_endpoints/endpoints\.py:\d+ anthropic_response$", + ), + mapping(rust_span="messages_gateway_route"), + mapping( + span="python_messages_gateway_service", + python_frame=r"ProxyBaseLLMRequestProcessing\.base_process_llm_request$", + ), + mapping(rust_span="messages_gateway_service"), + mapping(rust_span="messages"), + mapping( + span="python_messages_provider_config", + python_frame=r"ProviderConfigManager\.get_provider_anthropic_messages_config$", + ), + mapping(rust_span="messages_provider_config"), + mapping(rust_span="validate_environment", python_frame=r"validate_anthropic_messages_environment$"), + mapping(rust_span="complete_url", python_frame=r"get_complete_url$"), + mapping(span="python_messages_entry_handler", python_frame=r"messages/handler\.py:\d+ anthropic_messages_handler$"), + mapping(span="python_messages_handler_wrapper", python_frame=r"BaseLLMHTTPHandler\.anthropic_messages_handler$"), + mapping( + rust_span="execute_messages_provider_call", + python_frame=r"BaseLLMHTTPHandler\.async_anthropic_messages_handler$", + ), + mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"), + mapping(rust_span="transform_response", python_frame=r"(? RouteFixture: + return RouteFixture( + kwargs={ + "model_alias": "trace-model", + "provider_model": f"{provider}/claude-sonnet-5", + "body": { + "model": "trace-model", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 16, + }, + }, + provider_responses=( + RecordedHttpResponse.from_bytes( + 200, + (HttpHeader(name="content-type", value="application/json"),), + json.dumps( + { + "id": "msg_trace", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 2, "output_tokens": 3}, + } + ).encode(), + ), + ), + ) + + +def _anthropic_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return _fixture(engine, "anthropic") + + +def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return _fixture(engine, "azure_ai") + + +ANTHROPIC_MAPPINGS: Final = ( + *GATEWAY_MAPPINGS, + mapping( + rust_span="transform_request", + python_frame=r"(? tuple[TraceMapping, ...]: + selected: Final = self.async_mappings if mode == "async" else self.sync_mappings + return self.mappings if selected is None else selected + + +@dataclass(frozen=True, slots=True) +class TraceSuite: + route: TraceRouteSpec + scenarios: tuple[TraceScenario, ...] + + +@dataclass(frozen=True, slots=True) +class TraceExecutionFailure: + engine: TraceFailureSource + message: str diff --git a/tests/rust-python-harness/strategies/trace_parity/reporting.py b/tests/rust-python-harness/strategies/trace_parity/reporting.py new file mode 100644 index 00000000000..9c5bf9e88cd --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/reporting.py @@ -0,0 +1,367 @@ +from __future__ import annotations + +import os +import re +import sys +from collections.abc import Sequence +from typing import Final, Literal + +from pydantic import BaseModel, ConfigDict, ValidationError + +from ...shared.reporting.models import SURFACES, CaseResult, RunStatus, SdkFunction, Surface +from ...shared.reporting.rendering import ReportSection +from ...shared.reporting.strategy import NotImplementedCaseSpec, SkippedCaseSpec +from ...shared.tracing.steps import ( + PipelineStep, + TraceContract, + TraceDiff, + TraceMapping, + trace_depths, + trace_diff, +) + +TRACE_COMPARISON_ARTIFACT: Final = "trace_comparison" +TRACE_PARITY_HINT: Final = ( + "rebuild the native bridge with the trace-parity feature, e.g. `uvx maturin develop --features trace-parity`" +) + +_COLORS: Final[dict[str, str]] = {"green": "32", "yellow": "33", "red": "31", "cyan": "36"} +_RESET: Final = "\033[0m" + + +def _paint(text: str, color: str) -> str: + if not sys.stdout.isatty() or os.environ.get("NO_COLOR"): + return text + return f"\033[{_COLORS[color]}m{text}{_RESET}" + + +class TraceEventArtifact(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + id: int + parent_id: int | None + span: str + raw: str + + def step(self) -> PipelineStep: + return PipelineStep(self.id, self.parent_id, self.span, self.raw) + + +class TraceMappingArtifact(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + span: str + python: str | None + rust: str | None + + +class TraceComparisonArtifact(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + surface: Surface + sdk_function: SdkFunction + scenario: str + mode: Literal["sync", "async"] + mappings: tuple[TraceMappingArtifact, ...] + python: tuple[TraceEventArtifact, ...] + rust: tuple[TraceEventArtifact, ...] + python_unmatched: int + unordered_children_of: frozenset[str] + python_error: str | None = None + rust_error: str | None = None + + @classmethod + def from_traces( + cls, + *, + surface: Surface, + sdk_function: SdkFunction, + scenario: str, + mode: Literal["sync", "async"], + mappings: Sequence[TraceMapping], + contract: TraceContract, + python: Sequence[PipelineStep], + rust: Sequence[PipelineStep], + python_unmatched: int, + python_error: str | None = None, + rust_error: str | None = None, + ) -> TraceComparisonArtifact: + return cls( + surface=surface, + sdk_function=sdk_function, + scenario=scenario, + mode=mode, + mappings=tuple( + TraceMappingArtifact( + span=item.span, + python=item.python.pattern if item.python else None, + rust=item.rust, + ) + for item in mappings + ), + python=tuple( + TraceEventArtifact(id=step.id, parent_id=step.parent_id, span=step.span, raw=step.raw) + for step in python + ), + rust=tuple( + TraceEventArtifact(id=step.id, parent_id=step.parent_id, span=step.span, raw=step.raw) + for step in rust + ), + python_unmatched=python_unmatched, + unordered_children_of=contract.unordered_children_of, + python_error=python_error, + rust_error=rust_error, + ) + + def python_steps(self) -> tuple[PipelineStep, ...]: + return tuple(event.step() for event in self.python) + + def rust_steps(self) -> tuple[PipelineStep, ...]: + return tuple(event.step() for event in self.rust) + + def diff(self) -> TraceDiff: + return trace_diff( + self.python_steps(), + self.rust_steps(), + tuple( + TraceMapping( + item.span, + re.compile(item.python) if item.python is not None else None, + item.rust, + ) + for item in self.mappings + ), + TraceContract(self.unordered_children_of), + ) + + def exact_match(self) -> bool: + return self.diff().matches + + def has_errors(self) -> bool: + return self.python_error is not None or self.rust_error is not None + + def contract_matches(self) -> bool: + if self.has_errors(): + return False + return self.diff().matches + + +def _split_raw(raw: str) -> tuple[str, str]: + location, separator, name = raw.partition(" ") + if separator: + return name, location + return raw, "" + + +def _python_line(index: int, step: PipelineStep, depth: int, exclusive: frozenset[str]) -> str: + name: Final = _split_raw(step.raw)[0] + location: Final = _split_raw(step.raw)[1] + suffix: Final = f" ({location})" if location else "" + marker: Final = " [python only]" if step.span in exclusive else "" + return _paint(f"{index} {' ' * depth}{name}{suffix}{marker}", "cyan") + + +def _python_lines(steps: tuple[PipelineStep, ...], exclusive: frozenset[str]) -> str: + depths: Final = trace_depths(steps) + lines: Final = tuple( + _python_line(index, step, depths[step.id], exclusive) for index, step in enumerate(steps, start=1) + ) + return f"{_paint('PYTHON', 'cyan')} ({len(steps)} steps)\n" + ("\n".join(lines) if lines else "(empty)") + + +def _python_references(steps: tuple[PipelineStep, ...]) -> dict[tuple[str, int], str]: + references: dict[tuple[str, int], str] = {} + occurrences: dict[str, int] = {} + for index, step in enumerate(steps, start=1): + name = _split_raw(step.raw)[0] + occurrence = occurrences.get(step.span, 0) + 1 + occurrences[step.span] = occurrence + references[(step.span, occurrence)] = f"{index} {name}" + return references + + +def _rust_line( + step: PipelineStep, + depth: int, + occurrence: int, + references: dict[tuple[str, int], str], +) -> str: + span: Final = _paint(step.span, "yellow") + key: Final = (step.span, occurrence) + reference: Final = ( + _paint(references[key], "cyan") if key in references else _paint("[rust only]", "yellow") + ) + suffix: Final = f"#{occurrence}" if occurrence > 1 else "" + return f"{' ' * depth}{span}{suffix} -> {reference}" + + +def _rust_lines(steps: tuple[PipelineStep, ...], references: dict[tuple[str, int], str]) -> str: + depths: Final = trace_depths(steps) + occurrences: dict[str, int] = {} + lines: list[str] = [] + for step in steps: + occurrence = occurrences.get(step.span, 0) + 1 + occurrences[step.span] = occurrence + lines.append(_rust_line(step, depths[step.id], occurrence, references)) + return f"{_paint('RUST', 'yellow')} ({len(steps)} steps)\n" + ("\n".join(lines) if lines else "(empty)") + + +def _state_text(state: str, *, good: bool) -> str: + return _paint(state, "green" if good else "red") + + +def _contract_line(artifact: TraceComparisonArtifact) -> str: + matches: Final = artifact.contract_matches() + status: Final = _state_text("PASS" if matches else "FAIL", good=matches) + if artifact.python_error or artifact.rust_error: + return f"Contract: {status}" + return f"Contract: {status}" + + +def _error_lines(artifact: TraceComparisonArtifact) -> tuple[str, ...]: + lines: list[str] = [] + for engine, error in (("Python", artifact.python_error), ("Rust", artifact.rust_error)): + if error is None: + continue + lines.append(_paint(f"{engine} error: {error}", "red")) + if "trace-parity feature" in error: + lines.append(f"hint: {TRACE_PARITY_HINT}") + return tuple(lines) + + +def _unseen_mappings( + artifact: TraceComparisonArtifact, + python: tuple[PipelineStep, ...], + rust: tuple[PipelineStep, ...], +) -> tuple[str, ...]: + return artifact.diff().missing_mappings + + +def _comparison_status_lines( + artifact: TraceComparisonArtifact, + python: tuple[PipelineStep, ...], + rust: tuple[PipelineStep, ...], +) -> tuple[str, ...]: + diff: Final = artifact.diff() + exact_match: Final = artifact.exact_match() + if artifact.has_errors(): + return (*_error_lines(artifact), _contract_line(artifact)) + unseen: Final = _unseen_mappings(artifact, python, rust) + unseen_line: Final[tuple[str, ...]] = (f"Unseen mappings: {', '.join(unseen)}",) if unseen else () + drift_lines: Final[tuple[str, ...]] = ( + (_state_text("Same steps, order, and nesting", good=True),) + if exact_match + else ( + _paint(f"Python only: {', '.join(diff.python_only) or 'none'}", "cyan"), + _paint(f"Rust only: {', '.join(diff.rust_only) or 'none'}", "yellow"), + f"First difference: {diff.first_difference or 'none'}", + f"Python frames outside mapping: {artifact.python_unmatched}", + ) + ) + return ( + f"Trace: {_state_text('MATCH' if exact_match else 'DRIFT', good=exact_match)}", + *drift_lines, + *unseen_line, + _contract_line(artifact), + ) + + +def _render_comparison(artifact: TraceComparisonArtifact) -> str: + python: Final = artifact.python_steps() + rust: Final = artifact.rust_steps() + diff: Final = artifact.diff() + python_exclusive: Final = frozenset(item.span for item in artifact.mappings if item.rust is None) + status_lines: Final = _comparison_status_lines(artifact, python, rust) + return "\n\n".join( + ( + _python_lines(python, python_exclusive | frozenset(diff.python_only)), + _rust_lines(rust, _python_references(python)), + "\n".join(status_lines), + ) + ) + + +def _mode(nodeid: str) -> str: + if "[" in nodeid: + return nodeid.rsplit("[", 1)[-1].removesuffix("]") + head, _, tail = nodeid.rpartition(":") + return tail if head else "unknown mode" + + +def _scenario(nodeid: str) -> str: + parts: Final = nodeid.split(":") + return parts[-2] if len(parts) >= 5 else "default" + + +def _unavailable(status: RunStatus) -> str: + return f"Trace: NOT AVAILABLE\nTest outcome: {status.value}" + + +def _render_artifact(body: str) -> str: + try: + artifact: Final = TraceComparisonArtifact.model_validate_json(body) + except ValidationError as error: + return f"Trace comparison artifact is invalid: {error}" + return _render_comparison(artifact) + + +def _mode_section(result: CaseResult, nodeid: str, status: RunStatus) -> str: + artifacts: Final = tuple( + artifact for artifact in result.artifacts.get(nodeid, ()) if artifact.kind == TRACE_COMPARISON_ARTIFACT + ) + body: Final = ( + "\n\n".join(_render_artifact(artifact.body) for artifact in artifacts) if artifacts else _unavailable(status) + ) + label: Final = f"Scenario: {_scenario(nodeid)} / Mode: {_mode(nodeid)}" + return f"{label}\n{'-' * len(label)}\n\n{body}" + + +def _case_block(result: CaseResult) -> str: + header: Final = f"Case: {result.case.sdk_function}" + outcomes: Final = tuple(result.outcomes.items()) or ( + (nodeid, RunStatus.NOT_RUN) for nodeid in sorted(result.collected) + ) + sections: Final = tuple(_mode_section(result, nodeid, status) for nodeid, status in outcomes) + return "\n\n".join((f"{header}\n{'=' * len(header)}", *sections)) + + +def _unavailable_block(title: str, lines: tuple[str, ...]) -> str | None: + if not lines: + return None + return f"{title}\n{'-' * len(title)}\n" + "\n".join(lines) + + +def _surface_section(surface: Surface, results: Sequence[CaseResult]) -> ReportSection | None: + selected: Final = tuple(result for result in results if result.case.surface == surface) + if not selected: + return None + outcome_blocks: Final = tuple(_case_block(result) for result in selected if result.outcomes) + not_implemented: Final = _unavailable_block( + "Not implemented", + tuple( + f"- {result.case.sdk_function}: {spec.reason}" + for result in selected + if isinstance((spec := result.case.spec), NotImplementedCaseSpec) + ), + ) + skipped: Final = _unavailable_block( + "Skipped", + tuple( + f"- {result.case.sdk_function}: {spec.reason}" + for result in selected + if isinstance((spec := result.case.spec), SkippedCaseSpec) + ), + ) + blocks: Final = ( + *outcome_blocks, + *((not_implemented,) if not_implemented else ()), + *((skipped,) if skipped else ()), + ) + return ReportSection(f"{surface.upper()} trace comparisons", blocks or ("No runnable trace comparisons",)) + + +def render_trace_results(results: Sequence[CaseResult]) -> tuple[ReportSection, ...]: + sections: Final = tuple( + section for surface in SURFACES if (section := _surface_section(surface, results)) is not None + ) + return sections or (ReportSection("Trace comparisons", ("No trace comparisons selected",)),) diff --git a/tests/rust-python-harness/strategies/trace_parity/runner.py b/tests/rust-python-harness/strategies/trace_parity/runner.py index 127bf5dce40..ef373a6f2f6 100644 --- a/tests/rust-python-harness/strategies/trace_parity/runner.py +++ b/tests/rust-python-harness/strategies/trace_parity/runner.py @@ -1,26 +1,182 @@ from __future__ import annotations +import importlib from collections.abc import Sequence from pathlib import Path +from time import monotonic +from typing import Final -from ...shared.reporting.models import HarnessCase, HarnessRun -from ...shared.reporting.pytest_runner import UpdateCallback, run_pytest +from ...shared.reporting.models import CaseResult, HarnessCase, HarnessRun, ResultArtifact, RunStatus, Surface +from ...shared.reporting.strategy import ModuleCaseSpec, UpdateCallback +from ...shared.native_build import ensure_trace_bridge +from .models import GatewayRouteSpec, RouteSpec, TraceExecutionFailure, TraceMode, TraceScenario, TraceSuite +from .reporting import TRACE_COMPARISON_ARTIFACT, TraceComparisonArtifact +from .sdk.execution import execute_trace -def run( +def _load_case(reference: str, harness_case: HarnessCase) -> TraceSuite | TraceExecutionFailure: + try: + module: Final = importlib.import_module(reference) + except Exception as error: + return TraceExecutionFailure("harness", f"cannot import {reference}: {type(error).__name__}: {error}") + suite: Final = getattr(module, "TRACE_SUITE", None) + if not isinstance(suite, TraceSuite): + return TraceExecutionFailure("harness", f"{reference} must export TRACE_SUITE: TraceSuite") + validation_error: Final = validate_trace_suite(suite, harness_case) + if validation_error is not None: + return TraceExecutionFailure("harness", f"{reference} {validation_error}") + return suite + + +def validate_trace_suite(suite: TraceSuite, harness_case: HarnessCase) -> str | None: + names: Final = tuple(scenario.name for scenario in suite.scenarios) + if not names or len(names) != len(set(names)) or any(not name or ":" in name for name in names): + return "scenario names must be non-empty, unique, and colon-free" + invalid_modes: Final = tuple( + scenario.name + for scenario in suite.scenarios + if not scenario.modes + or len(scenario.modes) != len(set(scenario.modes)) + or any(mode not in {"sync", "async"} for mode in scenario.modes) + ) + if invalid_modes: + return f"scenarios must use non-empty, unique sync/async modes: {', '.join(invalid_modes)}" + surface: Final = harness_case.surface + if surface == "sdk" and not isinstance(suite.route, RouteSpec): + return "must use RouteSpec for the sdk surface" + if surface == "gateway" and not isinstance(suite.route, GatewayRouteSpec): + return "must use GatewayRouteSpec for the gateway surface" + if surface is None: + return "requires an sdk or gateway surface" + if suite.route.route != harness_case.sdk_function: + return f"route {suite.route.route} does not match case function {harness_case.sdk_function}" + return None + + +def scenario_nodeids( + trace_suite: TraceSuite, + harness_case: HarnessCase, + selected_scenarios: frozenset[str] = frozenset(), +) -> tuple[tuple[TraceScenario, TraceMode, str], ...]: + surface: Final = harness_case.surface + if surface is None: + return () + return tuple( + (scenario, mode, f"trace:{surface}:{harness_case.sdk_function}:{scenario.name}:{mode}") + for scenario in trace_suite.scenarios + if not selected_scenarios or scenario.name in selected_scenarios + for mode in scenario.modes + ) + + +def _record_setup_failure(run: HarnessRun, case: HarnessCase, message: str, stage: str) -> None: + result: Final = run.results[case.key] + nodeid: Final = f"trace:{case.surface}:{case.sdk_function}:{stage}" + result.collected.add(nodeid) + result.record(nodeid, RunStatus.ERROR) + run.failures.append((nodeid, message)) + + +def run_trace_mode( + run: HarnessRun, + result: CaseResult, + trace_suite: TraceSuite, + scenario: TraceScenario, + mode: TraceMode, + surface: Surface, + nodeid: str, + on_update: UpdateCallback, +) -> None: + started_at: Final = monotonic() + comparison: Final = _execute_mode(trace_suite, scenario, mode, surface) + duration: Final = monotonic() - started_at + if isinstance(comparison, TraceExecutionFailure): + result.record(nodeid, RunStatus.ERROR, duration) + run.failures.append((nodeid, comparison.message)) + on_update(run) + return + artifact: Final = ResultArtifact(TRACE_COMPARISON_ARTIFACT, comparison.model_dump_json()) + if comparison.has_errors(): + result.record(nodeid, RunStatus.ERROR, duration, (artifact,)) + run.failures.append( + (nodeid, "\n".join(error for error in (comparison.python_error, comparison.rust_error) if error)) + ) + else: + status: Final = RunStatus.PASSED if comparison.contract_matches() else RunStatus.FAILED + result.record(nodeid, status, duration, (artifact,)) + if status is RunStatus.FAILED: + run.failures.append((nodeid, "trace contract mismatch; see the rendered comparison")) + on_update(run) + + +def _execute_mode( + trace_suite: TraceSuite, + scenario: TraceScenario, + mode: TraceMode, + surface: Surface, +) -> TraceComparisonArtifact | TraceExecutionFailure: + route: Final = trace_suite.route + if isinstance(route, GatewayRouteSpec): + if surface != "gateway": + return TraceExecutionFailure("harness", "gateway route cannot run on the sdk surface") + from .gateway.execution import execute_gateway_trace + + return execute_gateway_trace(route, scenario, mode) + if surface != "sdk": + return TraceExecutionFailure("harness", "sdk route cannot run on the gateway surface") + return execute_trace(route, scenario, mode, surface) + + +def _run_case( + run: HarnessRun, + harness_case: HarnessCase, + selected_scenarios: frozenset[str], + on_update: UpdateCallback, +) -> None: + result: Final = run.results[harness_case.key] + spec: Final = harness_case.spec + if not isinstance(spec, ModuleCaseSpec): + return + surface: Final = harness_case.surface + if surface is None: + return + trace_suite: Final = _load_case(spec.module, harness_case) + if isinstance(trace_suite, TraceExecutionFailure): + _record_setup_failure(run, harness_case, trace_suite.message, "load") + on_update(run) + return + nodeids: Final = scenario_nodeids(trace_suite, harness_case, selected_scenarios) + result.collected.update(nodeid for _, _, nodeid in nodeids) + if not nodeids: + result.status = RunStatus.SKIPPED + on_update(run) + return + result.status = RunStatus.RUNNING + on_update(run) + for scenario, mode, nodeid in nodeids: + run_trace_mode(run, result, trace_suite, scenario, mode, surface, nodeid, on_update) + + +def run_trace_cases( cases: Sequence[HarnessCase], repo_root: Path, on_update: UpdateCallback, - pytest_args: Sequence[str] = (), + runner_args: Sequence[str] = (), ) -> tuple[int, HarnessRun]: - return run_pytest(cases, repo_root, on_update, pytest_args) - - -def main(argv: Sequence[str] | None = None) -> int: - from ...cli import main as harness_main - - return harness_main(argv, strategy_id="trace_parity") - - -if __name__ == "__main__": - raise SystemExit(main()) + selected_scenarios: Final = frozenset(runner_args) + run: Final = HarnessRun.from_cases(cases) + bridge_error: Final = ensure_trace_bridge(repo_root) + if bridge_error is not None: + for harness_case in cases: + _record_setup_failure(run, harness_case, bridge_error, "bridge") + run.finished_at = monotonic() + on_update(run) + return 1, run + for harness_case in cases: + _run_case(run, harness_case, selected_scenarios, on_update) + run.finished_at = monotonic() + on_update(run) + failed: Final = any( + result.status in {RunStatus.ERROR, RunStatus.FAILED, RunStatus.MISSING} for result in run.results.values() + ) + return int(failed), run diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/__init__.py b/tests/rust-python-harness/strategies/trace_parity/sdk/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/chat_completions/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/chat_completions/case.py new file mode 100644 index 00000000000..6be5afd60d6 --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/chat_completions/case.py @@ -0,0 +1,149 @@ +from __future__ import annotations + +import json +from typing import Final + +from .....shared.parity.recorded_http import HttpHeader, RecordedHttpResponse +from .....shared.tracing.steps import Engine, mapping +from ...models import RouteFixture, RouteSpec, TraceScenario, TraceSuite + +COMMON_MAPPINGS: Final = ( + mapping(span="python_provider_config", python_frame=r"ProviderConfigManager\.get_provider_chat_config$"), + mapping(rust_span="chat_completions_provider_config"), + mapping( + span="python_supported_openai_params", + python_frame=r"litellm_core_utils/get_supported_openai_params\.py:\d+ get_supported_openai_params$", + ), + mapping( + span="python_provider_supported_openai_params", + python_frame=r"AnthropicConfig\.get_supported_openai_params$", + ), + mapping(rust_span="supported_openai_params"), + mapping(rust_span="validate_environment", python_frame=r"(? RouteFixture: + response: Final = json.dumps( + { + "id": "msg_trace", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 2, "output_tokens": 3}, + } + ).encode() + return RouteFixture( + kwargs={ + "model": "anthropic/claude-sonnet-5", + "messages": [{"role": "user", "content": "hello"}], + **({"optional_params": {"max_tokens": 16}} if engine == "rust" else {"max_tokens": 16}), + }, + provider_responses=( + RecordedHttpResponse.from_bytes( + 200, (HttpHeader(name="content-type", value="application/json"),), response + ), + ), + ) + + +def _bedrock_fixture(engine: Engine, _base_url: str) -> RouteFixture: + response: Final = json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": "hello"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 2, "outputTokens": 3, "totalTokens": 5}, + "metrics": {"latencyMs": 1}, + } + ).encode() + credentials: Final = { + "aws_access_key_id": "test-access", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-east-1", + } + return RouteFixture( + kwargs={ + "model": "bedrock/us-east-1/anthropic.claude-v2", + "messages": [{"role": "user", "content": "hello"}], + **( + {"optional_params": {**credentials, "maxTokens": 16}} + if engine == "rust" + else {**credentials, "max_tokens": 16} + ), + }, + provider_responses=( + RecordedHttpResponse.from_bytes( + 200, (HttpHeader(name="content-type", value="application/json"),), response + ), + ), + ) + + +SPEC: Final = RouteSpec( + "chat_completions", + ("completion", "acompletion"), + ("chat_completions", "achat_completions"), + _anthropic_fixture, +) +BEDROCK_COMMON_MAPPINGS: Final = ( + mapping(rust_span="chat_completions_provider_config"), + mapping(rust_span="supported_openai_params"), + mapping(rust_span="execute_chat_completions_provider_call"), + mapping(rust_span="validate_environment"), + mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"), + mapping(span="python_transform_response", python_frame=r"AmazonConverseConfig\._transform_response$"), +) +BEDROCK_SYNC_MAPPINGS: Final = ( + mapping(span="python_chat_completions", python_frame=r"main\.py:\d+ completion$"), + mapping(rust_span="chat_completions"), + mapping(span="python_transform_request", python_frame=r"AmazonConverseConfig\._transform_request$"), + *BEDROCK_COMMON_MAPPINGS, +) +BEDROCK_ASYNC_MAPPINGS: Final = ( + mapping(span="python_chat_completions", python_frame=r"main\.py:\d+ acompletion$"), + mapping(span="python_completion_wrapper", python_frame=r"main\.py:\d+ completion$"), + mapping(rust_span="chat_completions"), + *BEDROCK_COMMON_MAPPINGS, +) +TRACE_SUITE: Final = TraceSuite( + route=SPEC, + scenarios=( + TraceScenario( + name="anthropic", + fixture=_anthropic_fixture, + mappings=COMMON_MAPPINGS, + sync_mappings=SYNC_MAPPINGS, + async_mappings=ASYNC_MAPPINGS, + ), + TraceScenario( + name="bedrock", + fixture=_bedrock_fixture, + mappings=BEDROCK_COMMON_MAPPINGS, + sync_mappings=BEDROCK_SYNC_MAPPINGS, + async_mappings=BEDROCK_ASYNC_MAPPINGS, + ), + ), +) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py new file mode 100644 index 00000000000..f8d7c55d4e2 --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py @@ -0,0 +1,140 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable +from pathlib import Path +from typing import Final, Protocol, cast + +from ....shared.parity.replay import replay_server +from ....shared.reporting.models import Surface +from ....shared.tracing.native import native_trace_events +from ....shared.tracing.profiler import FunctionTraceEvent, profile_python +from ....shared.tracing.steps import Engine, pipeline_projection +from ..models import RouteSpec, TraceExecutionFailure, TraceMode, TraceScenario +from ..reporting import TraceComparisonArtifact + + +class SdkCall(Protocol): + def __call__(self, **kwargs: object) -> object: ... + + +def _invoke(function: SdkCall, kwargs: dict[str, object], *, asynchronous: bool) -> object: + async def invoke_async() -> object: + return await cast(Awaitable[object], function(**kwargs)) + + if asynchronous: + return asyncio.run(invoke_async()) + return function(**kwargs) + + +def _entrypoint(spec: RouteSpec, engine: Engine, *, asynchronous: bool) -> SdkCall | TraceExecutionFailure: + import litellm + from litellm.anthropic_interface import messages as sdk_messages + from litellm.rust_bridge import get_native_bridge + + if engine == "rust": + bridge: Final = cast(object | None, get_native_bridge()) + if bridge is None: + return TraceExecutionFailure("rust", "native Rust bridge is required for trace parity") + trace_bridge: Final[object | None] = getattr(bridge, "_trace", None) + if trace_bridge is None: + return TraceExecutionFailure("rust", "native Rust bridge must include the trace-parity feature") + entrypoint: Final = spec.rust_entrypoints[int(asynchronous)] + function: Final[object | None] = getattr(trace_bridge, entrypoint, None) + if function is None: + return TraceExecutionFailure("rust", f"native Rust trace bridge does not expose {entrypoint}") + return cast(SdkCall, function) + owner: Final = sdk_messages if spec.route == "messages" else litellm + return cast(SdkCall, getattr(owner, spec.python_entrypoints[int(asynchronous)])) + + +def _collect( + function: SdkCall, kwargs: dict[str, object], engine: Engine, *, asynchronous: bool +) -> tuple[FunctionTraceEvent, ...]: + if engine == "rust": + return native_trace_events(_invoke(function, kwargs, asynchronous=asynchronous)) + import litellm + + with profile_python(Path(litellm.__file__).parent, threads=True) as profiler: + _invoke(function, kwargs, asynchronous=asynchronous) + return tuple(profiler.events) + + +def collect_trace( + spec: RouteSpec, engine: Engine, *, asynchronous: bool +) -> tuple[FunctionTraceEvent, ...] | TraceExecutionFailure: + function: Final = _entrypoint(spec, engine, asynchronous=asynchronous) + if isinstance(function, TraceExecutionFailure): + return function + try: + with replay_server() as provider: + fixture: Final = spec.fixture(engine, provider.url) + for response in fixture.provider_responses: + provider.enqueue_response(response) + kwargs: Final = { + **fixture.kwargs, + "api_key": "test-key", + "api_base": provider.url, + **({"timeout_seconds": 5} if engine == "rust" else {"timeout": 5}), + } + events: Final = _collect(function, kwargs, engine, asynchronous=asynchronous) + provider.take_requests(len(fixture.provider_responses)) + except Exception as error: + return TraceExecutionFailure(engine, f"{type(error).__name__}: {error}") + if not events: + return TraceExecutionFailure(engine, "trace is empty") + return events + + +def _failure_message(result: tuple[FunctionTraceEvent, ...] | TraceExecutionFailure) -> str | None: + if isinstance(result, tuple): + return None + return f"{result.engine}: {result.message}" + + +def execute_trace( + route: RouteSpec, scenario: TraceScenario, mode: TraceMode, surface: Surface +) -> TraceComparisonArtifact: + asynchronous: Final = mode == "async" + mappings: Final = scenario.mappings_for(mode) + scenario_route: Final = RouteSpec( + route=route.route, + python_entrypoints=route.python_entrypoints, + rust_entrypoints=route.rust_entrypoints, + fixture=scenario.fixture, + ) + python_trace: Final = collect_trace(scenario_route, "python", asynchronous=asynchronous) + rust_trace: Final = collect_trace(scenario_route, "rust", asynchronous=asynchronous) + python_error: Final = _failure_message(python_trace) + rust_error: Final = _failure_message(rust_trace) + python_events: Final = python_trace if isinstance(python_trace, tuple) else () + rust_events: Final = rust_trace if isinstance(rust_trace, tuple) else () + try: + python: Final = pipeline_projection("python", python_events, mappings) + rust: Final = pipeline_projection("rust", rust_events, mappings) + except ValueError as error: + return TraceComparisonArtifact.from_traces( + surface=surface, + sdk_function=route.route, + scenario=scenario.name, + mode=mode, + mappings=mappings, + contract=scenario.contract, + python=(), + rust=(), + python_unmatched=0, + python_error=f"harness: {error}", + ) + return TraceComparisonArtifact.from_traces( + surface=surface, + sdk_function=route.route, + scenario=scenario.name, + mode=mode, + mappings=mappings, + contract=scenario.contract, + python=python.steps, + rust=rust.steps, + python_unmatched=python.unmatched, + python_error=python_error, + rust_error=rust_error, + ) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/messages/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/messages/case.py new file mode 100644 index 00000000000..27079c28cd8 --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/messages/case.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +import json +from typing import Final + +from .....shared.parity.recorded_http import HttpHeader, RecordedHttpResponse +from .....shared.tracing.steps import Engine, mapping +from ...models import RouteFixture, RouteSpec, TraceScenario, TraceSuite + +COMMON_MAPPINGS: Final = ( + mapping(rust_span="messages", python_frame=r"anthropic_interface/messages/__init__\.py:\d+ a?create$"), + mapping( + span="python_messages_provider_config", + python_frame=r"ProviderConfigManager\.get_provider_anthropic_messages_config$", + ), + mapping(rust_span="messages_provider_config"), + mapping(rust_span="validate_environment", python_frame=r"validate_anthropic_messages_environment$"), + mapping(rust_span="complete_url", python_frame=r"get_complete_url$"), + mapping( + span="python_messages_entry_handler", + python_frame=r"messages/handler\.py:\d+ anthropic_messages_handler$", + ), + mapping( + span="python_messages_handler_wrapper", + python_frame=r"BaseLLMHTTPHandler\.anthropic_messages_handler$", + ), + mapping( + rust_span="execute_messages_provider_call", + python_frame=r"BaseLLMHTTPHandler\.async_anthropic_messages_handler$", + ), + mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"), + mapping(rust_span="transform_response", python_frame=r"(? RouteFixture: + conversation: Final = {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 16} + response: Final = json.dumps( + { + "id": "msg_trace", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 2, "output_tokens": 3}, + } + ).encode() + return RouteFixture( + kwargs={ + "model": f"{provider}/claude-sonnet-5", + **({"body": {**conversation, "model": "claude-sonnet-5"}} if engine == "rust" else conversation), + }, + provider_responses=( + RecordedHttpResponse.from_bytes( + 200, (HttpHeader(name="content-type", value="application/json"),), response + ), + ), + ) + + +def _anthropic_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return _fixture(engine, "anthropic") + + +def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return _fixture(engine, "azure_ai") + + +SPEC: Final = RouteSpec("messages", ("create", "acreate"), ("messages", "amessages"), _anthropic_fixture) +TRACE_SUITE: Final = TraceSuite( + route=SPEC, + scenarios=( + TraceScenario(name="anthropic", fixture=_anthropic_fixture, mappings=ANTHROPIC_MAPPINGS, modes=("async",)), + TraceScenario(name="azure-ai", fixture=_azure_fixture, mappings=AZURE_MAPPINGS, modes=("async",)), + ), +) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py new file mode 100644 index 00000000000..fe214f45339 --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py @@ -0,0 +1,339 @@ +from __future__ import annotations + +import json +from typing import Final, cast + +from .....shared.parity.recorded_http import HttpHeader, RecordedHttpResponse +from .....shared.tracing.steps import Engine, mapping +from ...models import RouteFixture, RouteSpec, TraceScenario, TraceSuite + +COMMON_MAPPINGS: Final = ( + mapping(rust_span="ocr", python_frame=r"ocr/main\.py:\d+ a?ocr$"), + mapping(rust_span="prepare_ocr_call", python_frame=r"ocr/main\.py:\d+ _prepare_ocr_request$"), + mapping(rust_span="ocr_provider_config", python_frame=r"ProviderConfigManager\.get_provider_ocr_config$"), + mapping(rust_span="supported_ocr_params", python_frame=r"get_supported_ocr_params$"), + mapping(rust_span="map_ocr_params", python_frame=r"(? RouteFixture: + response: Final = json.dumps( + { + "pages": [{"index": 0, "markdown": "hello"}], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1}, + } + ).encode() + return RouteFixture( + kwargs={ + "model": model, + "document": document or {"type": "document_url", "document_url": "https://example.com/document.pdf"}, + **({"optional_params": {"pages": [0]}} if engine == "rust" else {"pages": [0]}), + }, + provider_responses=( + RecordedHttpResponse.from_bytes( + 200, (HttpHeader(name="content-type", value="application/json"),), response + ), + ), + ) + + +def _mistral_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return _fixture(engine, "mistral/mistral-ocr-latest") + + +def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture: + return _fixture( + engine, + "azure_ai/pixtral-12b-2409", + {"type": "image_url", "image_url": "data:image/png;base64,aGVsbG8="}, + ) + + +def _vertex_fixture(engine: Engine, _base_url: str) -> RouteFixture: + fixture: Final = _fixture( + engine, + "vertex_ai/mistral-ocr-maas", + {"type": "image_url", "image_url": "data:image/png;base64,aGVsbG8="}, + ) + vertex: Final = {"vertex_project": "trace-project", "vertex_location": "us-central1"} + optional_params: Final = cast(dict[str, object], fixture.kwargs.get("optional_params", {})) + return RouteFixture( + kwargs={ + **fixture.kwargs, + **({"optional_params": {**optional_params, **vertex}} if engine == "rust" else vertex), + }, + provider_responses=fixture.provider_responses, + ) + + +def _vertex_deepseek_fixture(engine: Engine, _base_url: str) -> RouteFixture: + vertex: Final = {"vertex_project": "trace-project", "vertex_location": "us-central1"} + return RouteFixture( + kwargs={ + "model": "vertex_ai/deepseek-ocr-maas", + "document": {"type": "image_url", "image_url": "data:image/png;base64,aGVsbG8="}, + **({"optional_params": vertex} if engine == "rust" else vertex), + }, + provider_responses=( + RecordedHttpResponse.from_bytes( + 200, + (HttpHeader(name="content-type", value="application/json"),), + json.dumps( + { + "choices": [{"message": {"role": "assistant", "content": "hello"}}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1}, + } + ).encode(), + ), + ), + ) + + +def _azure_document_intelligence_fixture(engine: Engine, base_url: str) -> RouteFixture: + completed: Final = json.dumps( + { + "status": "succeeded", + "analyzeResult": { + "content": "hello", + "pages": [ + { + "pageNumber": 1, + "width": 8.5, + "height": 11, + "unit": "inch", + "lines": [{"content": "hello"}], + } + ], + }, + } + ).encode() + return RouteFixture( + kwargs={ + "model": "azure_ai/doc-intelligence/prebuilt-read", + "document": { + "type": "document_url", + "document_url": "data:application/pdf;base64,aGVsbG8=", + }, + **({"optional_params": {"pages": [0]}} if engine == "rust" else {"pages": [0]}), + }, + provider_responses=( + RecordedHttpResponse.from_bytes( + 202, + ( + HttpHeader(name="content-type", value="application/json"), + HttpHeader(name="operation-location", value=f"{base_url}/operations/trace"), + ), + b"{}", + ), + RecordedHttpResponse.from_bytes( + 200, + (HttpHeader(name="content-type", value="application/json"),), + completed, + ), + ), + ) + + +VERTEX_COMMON_MAPPINGS: Final = ( + *COMMON_MAPPINGS[:7], + mapping( + rust_span="transform_ocr_request", + python_frame=( + r"VertexAIOCRConfig\.(?:async_)?transform_ocr_request$" + r"|MistralOCRConfig\.transform_ocr_request$" + ), + ), + COMMON_MAPPINGS[-1], +) +VERTEX_SYNC_MAPPINGS: Final = ( + *VERTEX_COMMON_MAPPINGS, + mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.ocr$"), + mapping(span="python_transform_ocr_response_wrapper", python_frame=r"BaseLLMHTTPHandler\._transform_ocr_response$"), + mapping(rust_span="transform_ocr_response", python_frame=r"MistralOCRConfig\.transform_ocr_response$"), +) +VERTEX_ASYNC_MAPPINGS: Final = ( + *VERTEX_COMMON_MAPPINGS, + mapping(span="python_ocr_wrapper", python_frame=r"BaseLLMHTTPHandler\.ocr$"), + mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.async_ocr$"), + mapping(rust_span="transform_ocr_response", python_frame=r"MistralOCRConfig\.transform_ocr_response$"), +) + +DEEPSEEK_COMMON_MAPPINGS: Final = ( + mapping(rust_span="ocr", python_frame=r"ocr/main\.py:\d+ a?ocr$"), + mapping(rust_span="prepare_ocr_call", python_frame=r"ocr/main\.py:\d+ _prepare_ocr_request$"), + mapping(rust_span="ocr_provider_config", python_frame=r"ProviderConfigManager\.get_provider_ocr_config$"), + mapping(rust_span="supported_ocr_params", python_frame=r"get_supported_ocr_params$"), + mapping(rust_span="map_ocr_params", python_frame=r"(? bytes: + with io.BytesIO() as buffer: + with wave.open(buffer, "wb") as audio: + audio.setnchannels(1) + audio.setsampwidth(2) + audio.setframerate(16000) + audio.writeframes(b"\x00\x00" * 1600) + return buffer.getvalue() + + +def _fixture(engine: Engine, _base_url: str) -> RouteFixture: + credentials: Final = { + "aws_access_key_id": "test-access", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-east-1", + } + audio: Final = _audio_bytes() + payload: Final = ( + {"audio": {"data": base64.b64encode(audio).decode(), "format": "wav"}, "optional_params": credentials} + if engine == "rust" + else {"file": ("sample.wav", audio, "audio/wav"), **credentials} + ) + response: Final = json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": "hello"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 2, "outputTokens": 3, "totalTokens": 5}, + } + ).encode() + return RouteFixture( + kwargs={"model": "bedrock/mistral.voxtral-mini-3b-2507", **payload}, + provider_responses=( + RecordedHttpResponse.from_bytes( + 200, (HttpHeader(name="content-type", value="application/json"),), response + ), + ), + ) + + +SPEC: Final = RouteSpec( + "transcription", + ("transcription", "atranscription"), + ("transcription", "atranscription"), + _fixture, +) +TRACE_SUITE: Final = TraceSuite( + route=SPEC, + scenarios=( + TraceScenario( + name="bedrock", + fixture=_fixture, + mappings=MAPPINGS, + sync_mappings=SYNC_MAPPINGS, + async_mappings=ASYNC_MAPPINGS, + ), + ), +) diff --git a/tests/rust-python-harness/strategies/trace_parity/strategy.json b/tests/rust-python-harness/strategies/trace_parity/strategy.json deleted file mode 100644 index 9b67d8570cc..00000000000 --- a/tests/rust-python-harness/strategies/trace_parity/strategy.json +++ /dev/null @@ -1,33 +0,0 @@ -{ - "order": 20, - "id": "trace_parity", - "label": "Trace parity", - "description": "Compare mapped operations, call counts, and required execution ordering.", - "functions": { - "ocr": { - "coverage": "planned", - "selectors": [] - }, - "messages": { - "coverage": "planned", - "selectors": [] - }, - "chat_completions": { - "coverage": "planned", - "selectors": [] - }, - "responses": { - "coverage": "planned", - "selectors": [] - }, - "count_tokens": { - "coverage": "planned", - "selectors": [] - }, - "transcription": { - "coverage": "planned", - "selectors": [] - } - }, - "gateway": {} -} diff --git a/tests/rust-python-harness/strategies/trace_parity/test_reporting.py b/tests/rust-python-harness/strategies/trace_parity/test_reporting.py new file mode 100644 index 00000000000..68264064c92 --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/test_reporting.py @@ -0,0 +1,251 @@ +from __future__ import annotations + +from collections.abc import Sequence +from typing import Final, Literal + +import pytest + +from ...shared.reporting.models import CaseResult, Coverage, HarnessCase, ResultArtifact, RunStatus +from ...shared.reporting.strategy import ModuleCaseSpec, NotImplementedCaseSpec +from ...shared.tracing.steps import PipelineStep, TraceContract, TraceMapping, mapping +from . import reporting +from .reporting import TRACE_COMPARISON_ARTIFACT, TraceComparisonArtifact, render_trace_results + +MAPPINGS: Final = ( + mapping(rust_span="ocr", python_frame=r"ocr/main\.py:\d+ a?ocr$"), + mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$"), +) + + +def _result(comparison: TraceComparisonArtifact) -> CaseResult: + case: Final = HarnessCase( + strategy_id="trace_parity", + strategy_label="Trace parity", + sdk_function=comparison.sdk_function, + spec=ModuleCaseSpec(coverage=Coverage.PARTIAL, module="example"), + surface=comparison.surface, + ) + result: Final = CaseResult(case=case) + nodeid: Final = f"trace:sdk:{comparison.sdk_function}:{comparison.scenario}:{comparison.mode}" + result.collected.add(nodeid) + result.record( + nodeid, + RunStatus.PASSED, + artifacts=(ResultArtifact(TRACE_COMPARISON_ARTIFACT, comparison.model_dump_json()),), + ) + return result + + +def _comparison( + python: tuple[PipelineStep, ...], + rust: tuple[PipelineStep, ...], + *, + mappings: Sequence[TraceMapping] = MAPPINGS, + rust_error: str | None = None, +) -> TraceComparisonArtifact: + return TraceComparisonArtifact.from_traces( + surface="sdk", + sdk_function="ocr", + scenario="default", + mode="sync", + mappings=mappings, + contract=TraceContract(), + python=python, + rust=rust, + python_unmatched=796, + rust_error=rust_error, + ) + + +def _events(*items: tuple[str, int, str | None]) -> tuple[PipelineStep, ...]: + parents: dict[int, int] = {} + steps: list[PipelineStep] = [] + for event_id, (span, depth, raw) in enumerate(items): + parent_id = parents.get(depth - 1) if depth else None + steps.append(PipelineStep(event_id, parent_id, span, raw if raw is not None else span)) + parents[depth] = event_id + return tuple(steps) + + +def test_renderer_shows_matching_python_and_rust_paths() -> None: + rust: Final = _events(("ocr", 0, None), ("http_request", 1, None)) + python: Final = _events( + ("ocr", 0, "ocr/main.py:88 aocr"), + ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"), + ) + + section: Final = render_trace_results((_result(_comparison(python, rust)),))[0] + report: Final = "\n\n".join(section.blocks) + + assert section.title == "SDK trace comparisons" + assert "Case: ocr" in report + assert "PYTHON (2 steps)\n1 aocr (ocr/main.py:88)\n2 AsyncHTTPHandler.post (http_handler.py:673)" in report + assert "RUST (2 steps)\nocr -> 1 aocr\n http_request -> 2 AsyncHTTPHandler.post" in report + assert "Mapping (identifier -> span)" not in report + assert "Trace: MATCH" in report + assert "Same steps, order, and nesting" in report + assert "Unseen mappings:" not in report + + +def test_renderer_reports_mappings_that_matched_nothing() -> None: + events: Final = _events(("ocr", 0, None)) + + section: Final = render_trace_results((_result(_comparison(events, events)),))[0] + report: Final = "\n\n".join(section.blocks) + + assert "Unseen mappings: http_request" in report + assert "Contract: FAIL" in report + + +def test_renderer_numbers_repeated_span_occurrences() -> None: + mappings: Final = (MAPPINGS[0], MAPPINGS[1]) + rust: Final = _events(("ocr", 0, None), ("http_request", 1, None), ("http_request", 1, None)) + python: Final = _events( + ("ocr", 0, "ocr/main.py:88 aocr"), + ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"), + ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"), + ) + + report: Final = "\n\n".join(render_trace_results((_result(_comparison(python, rust, mappings=mappings)),))[0].blocks) + + assert "http_request#2" in report + + +def test_renderer_accepts_declared_engine_specific_steps() -> None: + mappings: Final = ( + *MAPPINGS[:1], + mapping(span="python_prepare", python_frame=r"python_prepare$"), + mapping(rust_span="rust_prepare"), + ) + python: Final = _events(("ocr", 0, None), ("python_prepare", 1, "prep.py:1 python_prepare")) + rust: Final = _events(("ocr", 0, None), ("rust_prepare", 1, None)) + + section: Final = render_trace_results((_result(_comparison(python, rust, mappings=mappings)),))[0] + report: Final = "\n\n".join(section.blocks) + + assert "2 python_prepare (prep.py:1) [python only]" in report + assert "rust_prepare -> [rust only]" in report + assert "Trace: MATCH" in report + assert "Contract: PASS" in report + + +def test_unavailable_check_reports_mode_from_nodeid() -> None: + case: Final = HarnessCase( + strategy_id="trace_parity", + strategy_label="Trace parity", + sdk_function="ocr", + spec=ModuleCaseSpec(coverage=Coverage.PARTIAL, module="example"), + surface="sdk", + ) + result: Final = CaseResult(case=case) + result.collected.add("trace:sdk:ocr:default:sync") + result.record("trace:sdk:ocr:default:sync", RunStatus.ERROR) + + section: Final = render_trace_results((result,))[0] + report: Final = "\n\n".join(section.blocks) + + assert "Case: ocr" in report + assert "Scenario: default / Mode: sync" in report + assert "Trace: NOT AVAILABLE\nTest outcome: error" in report + assert "unknown mode" not in report + + +def test_renderer_keeps_collected_trace_when_one_engine_errors() -> None: + python: Final = _events( + ("ocr", 0, "ocr/main.py:88 aocr"), + ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"), + ) + + section: Final = render_trace_results( + (_result(_comparison(python, (), rust_error="rust: native Rust bridge must include the trace-parity feature")),) + )[0] + report: Final = "\n\n".join(section.blocks) + + assert "PYTHON (2 steps)\n1 aocr (ocr/main.py:88) [python only]" in report + assert "Rust error: rust: native Rust bridge must include the trace-parity feature" in report + assert "hint: rebuild the native bridge with the trace-parity feature" in report + assert "Contract: FAIL" in report + + +def test_renderer_groups_all_modes_under_one_case_header() -> None: + case: Final = HarnessCase( + strategy_id="trace_parity", + strategy_label="Trace parity", + sdk_function="ocr", + spec=ModuleCaseSpec(coverage=Coverage.PARTIAL, module="example"), + surface="sdk", + ) + result: Final = CaseResult(case=case) + events: Final = _events(("ocr", 0, None)) + modes: Final[tuple[Literal["sync", "async"], ...]] = ("sync", "async") + for mode in modes: + nodeid = f"trace:sdk:ocr:default:{mode}" + result.collected.add(nodeid) + comparison = TraceComparisonArtifact.from_traces( + surface="sdk", + sdk_function="ocr", + scenario="default", + mode=mode, + mappings=MAPPINGS, + contract=TraceContract(), + python=events, + rust=events, + python_unmatched=0, + ) + result.record( + nodeid, + RunStatus.PASSED, + artifacts=(ResultArtifact(TRACE_COMPARISON_ARTIFACT, comparison.model_dump_json()),), + ) + + section: Final = render_trace_results((result,))[0] + + assert len(section.blocks) == 1 + report: Final = section.blocks[0] + assert report.count("Case: ocr") == 1 + assert "Scenario: default / Mode: sync" in report + assert "Scenario: default / Mode: async" in report + + +def test_renderer_colors_every_trace_line_in_a_terminal(monkeypatch: pytest.MonkeyPatch) -> None: + rust: Final = _events(("ocr", 0, None), ("http_request", 1, None)) + python: Final = _events( + ("ocr", 0, "ocr/main.py:88 aocr"), + ("http_request", 1, "http_handler.py:673 AsyncHTTPHandler.post"), + ) + monkeypatch.setattr(reporting.sys.stdout, "isatty", lambda: True) + monkeypatch.delenv("NO_COLOR", raising=False) + + section: Final = render_trace_results((_result(_comparison(python, rust)),))[0] + report: Final = "\n\n".join(section.blocks) + + assert "\033[36mPYTHON\033[0m (2 steps)" in report + assert "\033[36m1 aocr (ocr/main.py:88)\033[0m" in report + assert "\033[33mRUST\033[0m (2 steps)" in report + assert "\033[33mocr\033[0m -> \033[36m1 aocr\033[0m" in report + assert "\033[33mhttp_request\033[0m -> \033[36m2 AsyncHTTPHandler.post\033[0m" in report + + +def test_renderer_groups_cases_and_unavailable_entries_by_surface() -> None: + events: Final = _events(("ocr", 0, None)) + gateway_results: Final = tuple( + CaseResult( + case=HarnessCase( + strategy_id="trace_parity", + strategy_label="Trace parity", + sdk_function=sdk_function, + spec=NotImplementedCaseSpec(reason=f"No {sdk_function} case is registered."), + surface="gateway", + ), + status=RunStatus.NOT_IMPLEMENTED, + ) + for sdk_function in ("ocr", "messages") + ) + + sections: Final = render_trace_results((_result(_comparison(events, events)), *gateway_results)) + + assert tuple(section.title for section in sections) == ("SDK trace comparisons", "GATEWAY trace comparisons") + gateway_report: Final = "\n\n".join(sections[1].blocks) + assert gateway_report.count("Not implemented") == 1 + assert "- ocr: No ocr case is registered." in gateway_report + assert "- messages: No messages case is registered." in gateway_report diff --git a/tests/rust-python-harness/strategies/trace_parity/test_runner.py b/tests/rust-python-harness/strategies/trace_parity/test_runner.py new file mode 100644 index 00000000000..0fcc5860ff2 --- /dev/null +++ b/tests/rust-python-harness/strategies/trace_parity/test_runner.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +from typing import Final + +from ...shared.reporting.models import Coverage, HarnessCase, HarnessRun, RunStatus, SdkFunction, Surface +from ...shared.reporting.strategy import ModuleCaseSpec +from ...shared.tracing.steps import Engine +from .models import GatewayRouteSpec, RouteFixture, RouteSpec, TraceScenario, TraceSuite +from .runner import run_trace_mode, scenario_nodeids, validate_trace_suite + + +def _fixture(_engine: Engine, _base_url: str) -> RouteFixture: + return RouteFixture(kwargs={}, provider_responses=()) + + +def _case(*, surface: Surface = "sdk", function: SdkFunction = "ocr") -> HarnessCase: + return HarnessCase( + strategy_id="trace_parity", + strategy_label="Trace parity", + sdk_function=function, + spec=ModuleCaseSpec(coverage=Coverage.PARTIAL, module="example"), + surface=surface, + ) + + +def test_scenario_filtering_and_occurrence_node_ids() -> None: + suite: Final = TraceSuite( + route=RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture), + scenarios=( + TraceScenario("one", _fixture, (), modes=("sync", "async")), + TraceScenario("two", _fixture, (), modes=("async",)), + ), + ) + case: Final = _case() + + nodes: Final = scenario_nodeids(suite, case, frozenset({"two"})) + + assert tuple(nodeid for _, _, nodeid in nodes) == ("trace:sdk:ocr:two:async",) + + +def test_scenario_validation_rejects_duplicate_and_unsafe_names() -> None: + route: Final = RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture) + duplicate: Final = TraceSuite( + route=route, + scenarios=(TraceScenario("same", _fixture, ()), TraceScenario("same", _fixture, ())), + ) + unsafe: Final = TraceSuite(route=route, scenarios=(TraceScenario("bad:name", _fixture, ()),)) + case: Final = _case() + + assert validate_trace_suite(duplicate, case) is not None + assert validate_trace_suite(unsafe, case) is not None + + +def test_scenario_validation_rejects_invalid_modes_and_route_registration() -> None: + invalid_modes: Final = TraceSuite( + route=RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture), + scenarios=(TraceScenario("invalid", _fixture, (), modes=("sync", "sync")),), + ) + wrong_function: Final = TraceSuite( + route=RouteSpec("messages", ("create", "acreate"), ("messages", "amessages"), _fixture), + scenarios=(TraceScenario("one", _fixture, ()),), + ) + wrong_surface: Final = TraceSuite( + route=GatewayRouteSpec("ocr"), + scenarios=(TraceScenario("one", _fixture, ()),), + ) + case: Final = _case() + + assert "unique sync/async modes" in (validate_trace_suite(invalid_modes, case) or "") + assert "does not match case function" in (validate_trace_suite(wrong_function, case) or "") + assert "must use RouteSpec" in (validate_trace_suite(wrong_surface, case) or "") + + +def test_invalid_route_dispatch_records_harness_error() -> None: + case: Final = _case() + run: Final = HarnessRun.from_cases((case,)) + result: Final = run.results[case.key] + suite: Final = TraceSuite( + route=GatewayRouteSpec("ocr"), + scenarios=(TraceScenario("one", _fixture, (), modes=("sync",)),), + ) + nodeid: Final = "trace:sdk:ocr:one:sync" + + run_trace_mode(run, result, suite, suite.scenarios[0], "sync", "sdk", nodeid, lambda _: None) + + assert result.outcomes[nodeid] is RunStatus.ERROR + assert run.failures == [(nodeid, "gateway route cannot run on the sdk surface")] diff --git a/tests/rust-python-harness/strategies/unit_tests/README.md b/tests/rust-python-harness/strategies/unit_tests/README.md deleted file mode 100644 index bd37072ae4b..00000000000 --- a/tests/rust-python-harness/strategies/unit_tests/README.md +++ /dev/null @@ -1,7 +0,0 @@ -# Unit tests - -Run independently with `uv run python -m tests.rust-python-harness.strategies.unit_tests.runner --plain`. Configure a `unit_suite` for each mapped API in `strategy.json` - -The runner combines mapping validation, Python tests in separate verified backend processes, and Cargo tests. It reports missing and ambiguous counterparts. Native Rust tests and existing Python tests stay in their original locations - -See [the suite format](../../README.md#configure-cases) for configuration. No complete API mapping is configured yet diff --git a/tests/rust-python-harness/strategies/unit_tests/__init__.py b/tests/rust-python-harness/strategies/unit_tests/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/rust-python-harness/strategies/unit_tests/ledgers/ocr/ocr_test_ledger.json b/tests/rust-python-harness/strategies/unit_tests/ledgers/ocr/ocr_test_ledger.json deleted file mode 100644 index 1ceb79b52bc..00000000000 --- a/tests/rust-python-harness/strategies/unit_tests/ledgers/ocr/ocr_test_ledger.json +++ /dev/null @@ -1,211 +0,0 @@ -{ - "sdk_function": "ocr", - "python_scope": [ - "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", - "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", - "tests/test_litellm/ocr/test_rust_bridge.py", - "tests/test_litellm/ocr/test_ocr_file_input.py", - "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", - "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", - "tests/test_litellm/ocr/test_ocr_native_format.py", - "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", - "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py" - ], - "rust_scope": [ - "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", - "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", - "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", - "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", - "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", - "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs" - ], - "entries": [ - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_should_encode_azure_document_intelligence_model_id", "status": "unmapped", "reason": "model-id URL percent-encoding has no Rust test; Rust only tests pages/features query building"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_should_reject_dot_segment_azure_document_intelligence_model_id", "status": "unmapped", "reason": "model-id dot-segment validation has no Rust test"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_transform_ocr_response_preserves_azure_native_fields", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "document_intelligence_response_normalizes_pages", "justification": "both assert page markdown, dimension (inch-to-pixel) normalization, and usage_info.pages_processed from the same Azure succeeded response shape"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_async_transform_ocr_response_preserves_azure_native_fields", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "document_intelligence_response_normalizes_pages", "justification": "async twin of the sync case above, same underlying transform is exercised on the Rust side"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_transform_ocr_response_tolerates_missing_native_fields", "status": "unmapped", "reason": "tables/keyValuePairs absence tolerance is not asserted by the Rust response test"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_transform_ocr_response_non_succeeded_status_raises", "status": "unmapped", "reason": "no Rust test asserts on a non-succeeded Azure DI status"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_get_supported_ocr_params_includes_features", "status": "unmapped", "reason": "supported-params list content has no Rust equivalent for Azure"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_transform_ocr_response_native_format_carries_raw_operation", "status": "unmapped", "reason": "native req_format raw-operation passthrough is not tested in Rust"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_async_transform_ocr_response_native_format_carries_raw_operation", "status": "unmapped", "reason": "native req_format raw-operation passthrough is not tested in Rust"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_transform_ocr_response_default_format_omits_raw_operation", "status": "unmapped", "reason": "req_format gating of raw-operation output has no Rust test"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_map_ocr_params_passes_through_req_format", "status": "unmapped", "reason": "req_format passthrough in map_ocr_params has no Rust test"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_map_ocr_params_rejects_unknown_req_format_as_bad_request", "status": "unmapped", "reason": "req_format validation error path has no Rust test"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_get_complete_url_omits_req_format_query_param", "status": "unmapped", "reason": "no Rust test asserts req_format is excluded from the built URL"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_map_ocr_params_features", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "document_intelligence_url_normalizes_features", "justification": "both normalize comma-separated feature names and whitespace; Python does this during parameter mapping and Rust during URL construction"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_map_ocr_params_empty_features_list_omitted", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "document_intelligence_url_omits_empty_feature_list", "justification": "both omit empty feature lists from the outgoing request; Python removes the parameter and Rust omits the query field"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_map_ocr_params_invalid_features_raises", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "document_intelligence_url_rejects_invalid_features", "justification": "both reject malformed feature values, including query injection, empty strings, and objects before sending the request"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_get_complete_url_appends_features_query", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "document_intelligence_url_normalizes_features", "justification": "both assert the selected feature names appear in the outgoing features query parameter"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_get_complete_url_combines_pages_and_features", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "document_intelligence_url_combines_pages_and_feature_list", "justification": "both combine zero-based pages [0, 1, 2] with keyValuePairs and languages into pages=1,2,3 and features=keyValuePairs,languages"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_validate_environment_uses_subscription_key", "status": "unmapped", "reason": "Python-side header derivation from litellm_params; Rust's poll test only checks the header is present, not how it was resolved"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_validate_environment_falls_back_to_entra_token", "status": "unmapped", "reason": "Entra bearer-token fallback logic has no Rust test"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestIsAzureDocumentIntelligenceModel::test_matches_doc_intelligence_route", "status": "unmapped", "reason": "model-route string matching is Python-only dispatch logic"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestIsAzureDocumentIntelligenceModel::test_matches_documentintelligence_and_is_case_insensitive", "status": "unmapped", "reason": "model-route string matching is Python-only dispatch logic"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestIsAzureDocumentIntelligenceModel::test_does_not_match_mistral_route", "status": "unmapped", "reason": "model-route string matching is Python-only dispatch logic"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestDocIntelligenceApiBaseResolution::test_generic_azure_ai_base_does_not_hijack_doc_intelligence", "status": "unmapped", "reason": "api_base resolution from the secret manager runs before the Rust bridge is called, no Rust test exists for it"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestDocIntelligenceApiBaseResolution::test_explicit_api_base_is_honoured_for_doc_intelligence", "status": "unmapped", "reason": "api_base precedence resolution is Python-only"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestDocIntelligenceApiBaseResolution::test_generic_azure_ai_base_still_applies_to_mistral_ocr", "status": "unmapped", "reason": "api_base precedence resolution is Python-only"}, - - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_use_litellm_rust_toggles_flag", "status": "unmapped", "reason": "bridge-plumbing: Python-side feature-flag toggle, no Rust equivalent"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_env_var_enables_rust_ocr", "status": "unmapped", "reason": "bridge-plumbing: Python-side env-var flag gating"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_explicit_false_overrides_process_enable", "status": "unmapped", "reason": "Python request-level Rust opt-out overrides the process flag before any Rust implementation runs"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_load_rust_ocr_returns_injected_impl", "status": "unmapped", "reason": "bridge-plumbing: dependency-injection test hook, not provider behavior"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_native_bridge_loader_returns_none_when_extension_absent", "status": "unmapped", "reason": "bridge-plumbing: native-extension import/loader fallback"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_native_bridge_loader_caches_absent_extension", "status": "unmapped", "reason": "bridge-plumbing: loader caching behavior"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_native_bridge_available_reflects_loader", "status": "unmapped", "reason": "bridge-plumbing: loader availability check"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_load_rust_aocr_returns_injected_impl", "status": "unmapped", "reason": "bridge-plumbing: dependency-injection test hook"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_toggle_without_ocr_arg_preserves_injected_impl", "status": "unmapped", "reason": "bridge-plumbing: injected-impl state retention regression"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_explicit_ocr_none_clears_injected_impl", "status": "unmapped", "reason": "bridge-plumbing: injected-impl clearing behavior"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_load_rust_ocr_none_when_extension_absent", "status": "unmapped", "reason": "bridge-plumbing: degrade path when the native extension is missing"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_load_rust_ocr_uses_compiled_extension", "status": "unmapped", "reason": "bridge-plumbing: native module resolution"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_timeout_to_seconds_handles_float_timeout_and_none", "status": "unmapped", "reason": "bridge-plumbing: Python-side timeout normalization helper"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_bridge_wrapper_forwards_prepared_args_and_wraps_response", "status": "unmapped", "reason": "bridge-plumbing: wrapper argument forwarding, asserted against a fake bridge not the real Rust code"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response", "status": "unmapped", "reason": "bridge-plumbing: async wrapper argument forwarding"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_run_rust_ocr_prepares_request_and_wraps_response", "status": "unmapped", "reason": "bridge-plumbing: request preparation and response wrapping in Python"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_run_rust_ocr_resolves_key_via_secret_manager_when_missing", "status": "unmapped", "reason": "secret-manager: API key resolution happens in Python before the bridge is invoked"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_run_rust_ocr_prefers_explicit_key_over_resolver", "status": "unmapped", "reason": "secret-manager: key precedence resolution"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_run_rust_ocr_uses_provider_api_key_env_var", "status": "unmapped", "reason": "secret-manager: provider-specific env var name resolution"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_prepare_rust_ocr_call_forwards_vertex_routing_metadata", "status": "unmapped", "reason": "secret-manager: vertex routing metadata merge happens in Python"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager", "status": "unmapped", "reason": "secret-manager: vertex project/location resolution"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager", "status": "unmapped", "reason": "secret-manager: azure_ai api_base resolution"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint", "status": "unmapped", "reason": "secret-manager: doc-intelligence endpoint resolution"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_run_rust_ocr_runs_pre_call_logging", "status": "unmapped", "reason": "bridge-plumbing: Python logging-object pre_call invocation"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_routes_to_rust_when_enabled", "status": "unmapped", "reason": "bridge-plumbing: routing to a fake bridge, not the real Rust transform"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_routes_azure_ai_to_rust_when_enabled", "status": "unmapped", "reason": "bridge-plumbing: provider-prefix stripping before routing"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_rust_path_converts_file_document_before_bridge", "status": "unmapped", "reason": "file-normalization: raw-bytes-to-data-URI conversion happens in Python before the bridge call"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_exception_type_uses_resolved_provider_context", "status": "unmapped", "reason": "bridge-plumbing: Python exception-type mapping on bridge failure"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_aocr_routes_to_async_rust_when_enabled", "status": "unmapped", "reason": "bridge-plumbing: async routing to a fake bridge"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_aocr_exception_type_uses_resolved_provider_context", "status": "unmapped", "reason": "bridge-plumbing: async exception-type mapping on bridge failure"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_does_not_route_to_rust_when_disabled", "status": "unmapped", "reason": "bridge-plumbing: Python control flow for the toggle-disabled branch, no Rust-owned behavior runs"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_falls_back_to_python_when_bridge_unavailable", "status": "unmapped", "reason": "bridge-plumbing: Python-only fallback when the compiled Rust extension is absent, Rust cannot test its own absence"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_forwards_timeout_to_rust", "status": "unmapped", "reason": "bridge-plumbing: asserts the Python call site forwards a timeout kwarg, Rust receives an already-constructed request"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_passes_default_request_timeout_to_rust", "status": "unmapped", "reason": "bridge-plumbing: asserts the Python call site supplies a default timeout kwarg, no Rust equivalent"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_provider_configs_expose_api_key_env_vars", "status": "unmapped", "reason": "asserts per-provider get_api_key_env_var() strings; the closest Rust test (ocr_dispatch_supports_migrated_providers) asserts provider dispatch/param resolution instead, not API key env var names"}, - - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_pdf_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection is Python-only preprocessing before the bridge call"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_png_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_jpg_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_jpeg_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_gif_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_webp_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_tiff_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_tif_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_bmp_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_be_case_insensitive", "status": "unmapped", "reason": "file-normalization: MIME detection case handling"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_fallback_for_unknown_extension", "status": "unmapped", "reason": "file-normalization: MIME detection fallback"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_pdf_pathlib_path_to_document_url", "status": "unmapped", "reason": "file-normalization: local-path-to-data-URI conversion happens in Python"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_image_pathlib_path_to_image_url", "status": "unmapped", "reason": "file-normalization: local-path-to-data-URI conversion"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_reject_bare_str_path", "status": "unmapped", "reason": "file-normalization: arbitrary-file-read guard on bare str paths"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_pathlib_path", "status": "unmapped", "reason": "file-normalization: local-path-to-data-URI conversion"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_raw_bytes", "status": "unmapped", "reason": "file-normalization: raw-bytes-to-data-URI conversion"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_raw_bytes_with_explicit_mime_type", "status": "unmapped", "reason": "file-normalization: explicit MIME override on raw bytes"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_raw_bytes_with_image_mime_type", "status": "unmapped", "reason": "file-normalization: explicit MIME override on raw bytes"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_file_like_object", "status": "unmapped", "reason": "file-normalization: file-like-object conversion"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_file_like_object_with_name", "status": "unmapped", "reason": "file-normalization: file-like-object name-based MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_missing_file_field", "status": "unmapped", "reason": "file-normalization: missing-field validation"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_nonexistent_pathlib_path", "status": "unmapped", "reason": "file-normalization: missing-file validation"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_empty_file", "status": "unmapped", "reason": "file-normalization: empty-file validation"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_unsupported_type", "status": "unmapped", "reason": "file-normalization: unsupported input type validation"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_invalid_mime_type", "status": "unmapped", "reason": "file-normalization: MIME-type injection validation"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_override_mime_type_for_pathlib_path", "status": "unmapped", "reason": "file-normalization: explicit MIME override precedence"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_build_document_url_for_pdf", "status": "unmapped", "reason": "file-normalization: multipart upload conversion"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_build_image_url_for_png", "status": "unmapped", "reason": "file-normalization: multipart upload conversion"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_build_image_url_for_jpeg", "status": "unmapped", "reason": "file-normalization: multipart upload conversion"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_detect_mime_from_filename_when_content_type_is_octet_stream", "status": "unmapped", "reason": "file-normalization: filename-based MIME fallback"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_detect_mime_from_filename_when_content_type_is_none", "status": "unmapped", "reason": "file-normalization: filename-based MIME fallback"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_fallback_to_octet_stream_for_unknown", "status": "unmapped", "reason": "file-normalization: default MIME fallback"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_preserve_base64_content_correctly", "status": "unmapped", "reason": "file-normalization: binary round-trip through base64"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_strip_mime_parameters_from_content_type", "status": "unmapped", "reason": "file-normalization: content-type parameter stripping"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_strip_mime_parameters_with_multiple_params", "status": "unmapped", "reason": "file-normalization: content-type parameter stripping"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestProxySecurityGuard::test_should_reject_file_type_document_in_json_body", "status": "unmapped", "reason": "proxy-layer JSON-body file-type guard, a different mechanism than Rust's URL-fetch SSRF guard"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestProxySecurityGuard::test_should_accept_document_url_type_in_json_body", "status": "unmapped", "reason": "proxy-layer JSON-body parsing"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestProxySecurityGuard::test_should_raise_on_invalid_json_body", "status": "unmapped", "reason": "proxy-layer JSON-body parsing error path"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestProxySecurityGuard::test_should_ignore_document_form_field_injection", "status": "unmapped", "reason": "proxy-layer multipart form-field injection guard, a different mechanism than Rust's URL-fetch SSRF guard"}, - - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestGetSupportedOcrParams::test_extract_header_in_supported_params", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "extract_header_is_a_supported_ocr_param", "justification": "both assert extract_header appears in the Mistral supported OCR params list"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestGetSupportedOcrParams::test_extract_footer_in_supported_params", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "extract_footer_is_a_supported_ocr_param", "justification": "both assert extract_footer appears in the Mistral supported OCR params list"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestGetSupportedOcrParams::test_existing_params_still_present", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "existing_ocr_params_remain_supported", "justification": "both assert the previously supported params are still present in the supported list"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestMapOcrParams::test_extract_header_passed_through", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "map_ocr_params_forwards_extract_header", "justification": "both assert extract_header alone survives map_ocr_params unchanged"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestMapOcrParams::test_extract_footer_passed_through", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "map_ocr_params_forwards_extract_footer", "justification": "both assert extract_footer alone survives map_ocr_params unchanged"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestMapOcrParams::test_extract_header_and_footer_together", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "map_ocr_params_forwards_extract_header_and_footer", "justification": "both assert header and footer passed together are both forwarded with their given values"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestMapOcrParams::test_unknown_param_is_dropped", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "map_ocr_params_drops_unknown_params", "justification": "both assert an unrecognized param key is dropped while a known one is kept"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestNewSupportedParams::test_new_param_in_supported_list", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "new_ocr_params_are_supported", "justification": "both assert each OCR4 param is in the supported list"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestNewParamsMapOcr::test_new_param_passed_through", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "map_ocr_params_forwards_new_ocr_params", "justification": "both assert each OCR4 param/value pair survives map_ocr_params unchanged"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestTransformOcrRequest::test_param_included_in_request_body", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "transform_ocr_request_includes_each_optional_param", "justification": "both assert each optional param value lands in the built request body alongside model/document with no files"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestTransformOcrRequest::test_multiple_new_params_together", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "transform_ocr_request_includes_multiple_new_params", "justification": "both assert multiple OCR4 params passed together all land in the same request body"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestTransformOcrResponseOcr4Fields::test_blocks_and_confidence_scores_preserved", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "transform_ocr_response_preserves_blocks_and_confidence_scores", "justification": "both assert blocks and confidence_scores survive the OCR response transform on the returned page"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestTransformOcrResponseOcr4Fields::test_ocr4_fields_survive_model_dump", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "transform_ocr_response_preserves_ocr4_page_fields", "justification": "both assert tables, hyperlinks, header and footer survive the OCR response transform on the returned page"}, - - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_model_info_ocr4_price", "status": "unmapped", "reason": "cost-calc: pricing/model-info lookup is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr4_cost_scales_with_pages", "status": "unmapped", "reason": "cost-calc: per-page pricing math is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_pricing_entry", "status": "unmapped", "reason": "cost-calc: cost-map JSON entry validation is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_model_info_price", "status": "unmapped", "reason": "cost-calc: pricing/model-info lookup is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_cost_scales_with_pages", "status": "unmapped", "reason": "cost-calc: per-page pricing math is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_bills_ocr_and_annotation_pages_at_their_own_rates", "status": "unmapped", "reason": "cost-calc: mixed-rate billing math is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_bills_annotation_only_response", "status": "unmapped", "reason": "cost-calc: annotation-only billing math is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_bills_annotation_pages_when_pages_processed_missing", "status": "unmapped", "reason": "cost-calc: fallback billing math is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_azure_doc_ai_annotation_pages_fall_back_to_ocr_rate", "status": "unmapped", "reason": "cost-calc: fallback billing math is Python-only"}, - - {"python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", "python_test": "test_rust_ocr_serves_default_format", "status": "unmapped", "reason": "request-format gating decision is made in Python before the Rust bridge is ever invoked"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", "python_test": "test_rust_ocr_skipped_for_native_format", "status": "unmapped", "reason": "request-format gating decision is made in Python before the Rust bridge is ever invoked"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", "python_test": "test_native_format_rejected_for_provider_without_support_as_bad_request", "status": "unmapped", "reason": "provider-support validation for req_format happens in Python"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", "python_test": "test_unknown_format_rejected_for_provider_without_support_as_bad_request", "status": "unmapped", "reason": "req_format validation error path is Python-only"}, - - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestHandlerDiscovery::test_handler_discovered_for_ocr", "status": "unmapped", "reason": "guardrail-translation handler discovery is a Python proxy-layer concern"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestHandlerDiscovery::test_handler_discovered_for_aocr", "status": "unmapped", "reason": "guardrail-translation handler discovery is a Python proxy-layer concern"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestInputProcessing::test_process_document_url", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestInputProcessing::test_process_image_url", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestInputProcessing::test_process_no_document", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestInputProcessing::test_process_invalid_document", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestInputProcessing::test_input_blocking_guardrail", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_process_single_page", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_process_multiple_pages", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_process_empty_pages", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_process_page_with_empty_markdown", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_process_preserves_page_metadata", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_output_blocking_guardrail", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestPIIMaskingScenario::test_pii_masking_in_ocr_pages", "status": "unmapped", "reason": "PII redaction in the translation handler has no Rust equivalent"}, - - {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_read_req_format_from_header", "status": "unmapped", "reason": "proxy-layer header parsing has no Rust equivalent"}, - {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_prefer_body_req_format_over_header", "status": "unmapped", "reason": "proxy-layer body-vs-header precedence has no Rust equivalent"}, - {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_omit_req_format_when_header_absent", "status": "unmapped", "reason": "proxy-layer parsing has no Rust equivalent"}, - {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_reject_unknown_req_format", "status": "unmapped", "reason": "proxy-layer validation has no Rust equivalent"}, - {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_return_native_payload_with_litellm_response_headers", "status": "unmapped", "reason": "proxy-layer response construction has no Rust equivalent"}, - {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_return_normalized_response_when_no_native_payload", "status": "unmapped", "reason": "proxy-layer response construction has no Rust equivalent"} - ], - "rust_only_tests": [ - {"rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "document_intelligence_maps_features", "reason": "Rust retains the feature list while filtering unsupported parameters; Python normalizes the list to a string during mapping"}, - {"rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "document_intelligence_url_normalizes_zero_based_pages", "reason": "Python covers ascending page indices with features, but has no dedicated test for deduplicating and sorting page indices"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", "rust_test": "rust_custom_logger_reads_failure_payload_for_non_ocr_call_type", "reason": "exercises the non-OCR (acompletion) call-type branch of the logger; the OCR branch is covered separately by rust_custom_logger_reads_success_payload_for_ocr"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", "rust_test": "no_callback_fast_path_dispatches_nothing", "reason": "Rust-only fast-path optimization test for when zero callbacks are registered; Python has no equivalent no-op dispatch path"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", "rust_test": "with_standard_logging_payload_keeps_top_level_fields_in_sync", "reason": "Rust-internal builder-method invariant, Python has no equivalent internal builder"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", "rust_test": "blocks_private_and_metadata_ips", "reason": "SSRF IP-blocking helper has no Python unit test; Python relies on the proxy-layer JSON/form guards instead"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", "rust_test": "convert_document_url_rejects_loopback_fetch", "reason": "URL-fetch SSRF protection is Rust-gateway-only"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", "rust_test": "convert_document_url_leaves_data_uri_untouched", "reason": "URL-fetch SSRF protection is Rust-gateway-only"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "truncate_error_body_passes_short_strings_through", "reason": "Rust-gateway error-body truncation helper has no Python counterpart"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "truncate_error_body_caps_long_payloads", "reason": "Rust-gateway error-body truncation helper has no Python counterpart"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "truncate_error_body_does_not_split_multibyte_chars", "reason": "Rust-gateway error-body truncation helper has no Python counterpart"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "ocr_dispatch_supports_migrated_providers", "reason": "Rust-internal provider-config dispatch table has no equivalent Python unit test"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "string_headers_accepts_string_values", "reason": "Rust-gateway header-coercion helper has no Python counterpart"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "auth_header_detection_is_case_insensitive", "reason": "Rust-gateway header-detection helper has no Python counterpart"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "ocr_lifecycle_runs_pre_during_and_success_hooks", "reason": "full gateway-level guardrail-hook-plus-HTTP-lifecycle test with no Python equivalent at this integration scope"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "ocr_lifecycle_runs_failure_hook_on_provider_error", "reason": "full gateway-level failure-hook-plus-HTTP-lifecycle test with no Python equivalent at this integration scope"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "ocr_lifecycle_pre_call_block_skips_provider_socket", "reason": "full gateway-level pre-call-block-plus-socket-skip test with no Python equivalent at this integration scope"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "ocr_does_not_duplicate_authorization_header_when_header_is_supplied", "reason": "outgoing HTTP header dedup at the Rust gateway has no Python counterpart"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "document_intelligence_poll_uses_resolved_subscription_key", "reason": "full Azure DI poll-loop integration test with no Python equivalent at this scope"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "string_headers_rejects_non_string_values", "reason": "Rust-gateway header-coercion error path has no Python counterpart"}, - {"rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "azure_ai_reuses_mistral_body_transform", "reason": "Rust-internal delegation-to-Mistral-transform implementation detail, no Python test asserts this delegation"}, - {"rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "document_intelligence_request_uses_base64_source_for_data_uri", "reason": "no Python test asserts on the base64Source request body shape"}, - {"rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", "rust_test": "vertex_mistral_url_uses_project_location_and_model", "reason": "vertex OCR support has no Python unit test coverage yet"}, - {"rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", "rust_test": "vertex_mistral_reuses_mistral_body_transform", "reason": "vertex OCR support has no Python unit test coverage yet"}, - {"rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", "rust_test": "vertex_deepseek_request_uses_ocr_endpoint_shape", "reason": "vertex OCR support has no Python unit test coverage yet"}, - {"rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", "rust_test": "vertex_deepseek_response_wraps_markdown_content", "reason": "vertex OCR support has no Python unit test coverage yet"}, - {"rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "transform_ocr_request_rejects_non_object_document", "reason": "non-object document rejection has no dedicated Python unit test"}, - {"rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "transform_ocr_response_normalizes_mistral_json", "reason": "Python's response tests target OCR4-specific fields only, none asserts the same base normalization this Rust test checks"}, - {"rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "complete_url_defaults_and_dedupes_v1", "reason": "URL-building/defaulting for Mistral has no Python unit test"}, - {"rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "resolve_api_key_prefers_param_then_env", "reason": "API key resolution precedence at the Rust provider-config layer has no Python unit test"}, - {"rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "resolve_api_key_errors_when_absent", "reason": "API key resolution error path at the Rust provider-config layer has no Python unit test"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", "rust_test": "rust_custom_logger_reads_success_payload_for_ocr", "reason": "Rust-internal custom-logger dispatch for OCR payloads has no Python unit test at this layer"} - ] -} diff --git a/tests/rust-python-harness/strategies/unit_tests/mapping_validator.py b/tests/rust-python-harness/strategies/unit_tests/mapping_validator.py deleted file mode 100644 index d805311e488..00000000000 --- a/tests/rust-python-harness/strategies/unit_tests/mapping_validator.py +++ /dev/null @@ -1,173 +0,0 @@ -from __future__ import annotations - -from collections import Counter -from collections.abc import Sequence -from dataclasses import dataclass -from pathlib import Path -from typing import Final - -from pydantic import BaseModel, ConfigDict - -from ...shared.parity.ledger import TestLedger, load_ledger -from .python_runner import enumerate_python_tests -from .rust_runner import enumerate_rust_tests - - -class TestMapping(BaseModel): - model_config = ConfigDict(frozen=True, extra="forbid") - - python: str - rust: str - - -@dataclass(frozen=True, slots=True) -class MappingReport: - pairs: tuple[TestMapping, ...] - problems: tuple[str, ...] - - -def _name(node: str) -> str: - return node.rsplit("::", 1)[-1].split("[", 1)[0] - - -def validate_mapping( - python_tests: Sequence[str], - rust_tests: Sequence[str], - annotations: Sequence[TestMapping] = (), -) -> MappingReport: - explicit_problems: Final = ( - *(f"missing Python counterpart: {pair.python}" for pair in annotations if pair.python not in python_tests), - *(f"missing Rust counterpart: {pair.rust}" for pair in annotations if pair.rust not in rust_tests), - *( - f"ambiguous Python annotation: {name}" - for name, count in Counter(p.python for p in annotations).items() - if count > 1 - ), - *( - f"ambiguous Rust annotation: {name}" - for name, count in Counter(p.rust for p in annotations).items() - if count > 1 - ), - ) - explicit_python: Final = {pair.python for pair in annotations} - candidates: Final = { - python: tuple(rust for rust in rust_tests if _name(python) == _name(rust)) - for python in python_tests - if python not in explicit_python - } - pairs: Final = ( - *annotations, - *(TestMapping(python=python, rust=matches[0]) for python, matches in candidates.items() if len(matches) == 1), - ) - problems: Final = ( - *explicit_problems, - *(f"missing Rust counterpart: {python}" for python, matches in candidates.items() if not matches), - *( - f"ambiguous Rust counterparts: {python}: {matches}" - for python, matches in candidates.items() - if len(matches) > 1 - ), - *( - f"ambiguous Python counterparts: {rust}" - for rust, count in Counter(pair.rust for pair in pairs).items() - if count > 1 - ), - *(f"missing Python counterpart: {rust}" for rust in rust_tests if rust not in {pair.rust for pair in pairs}), - *(("no Python tests collected",) if not python_tests else ()), - *(("no Rust tests collected",) if not rust_tests else ()), - ) - return MappingReport(pairs, problems) - - -REPO_ROOT = Path(__file__).resolve().parents[4] -LEDGER_ROOT = Path(__file__).parent / "ledgers" - - -def ledger_path_for(sdk_function: str) -> Path: - return LEDGER_ROOT / sdk_function / f"{sdk_function}_test_ledger.json" - - -@dataclass(frozen=True, slots=True) -class AuditReport: - missing_python_tests: tuple[str, ...] - stale_python_tests: tuple[str, ...] - missing_rust_tests: tuple[str, ...] - stale_rust_tests: tuple[str, ...] - - @property - def is_clean(self) -> bool: - return not ( - self.missing_python_tests - or self.stale_python_tests - or self.missing_rust_tests - or self.stale_rust_tests - ) - - -def _ledger_python_tests_by_file(ledger: TestLedger) -> dict[str, set[str]]: - grouping: dict[str, set[str]] = {path: set() for path in ledger.python_scope} - for entry in ledger.entries: - grouping.setdefault(entry.python_file, set()).add(entry.python_test) - return grouping - - -def _ledger_rust_tests_by_file(ledger: TestLedger) -> dict[str, set[str]]: - grouping: dict[str, set[str]] = {path: set() for path in ledger.rust_scope} - for entry in ledger.entries: - if entry.status == "mapped": - grouping.setdefault(entry.rust_file, set()).add(entry.rust_test) - for rust_only in ledger.rust_only_tests: - grouping.setdefault(rust_only.rust_file, set()).add(rust_only.rust_test) - return grouping - - -def audit_ledger(ledger: TestLedger, repo_root: Path = REPO_ROOT) -> AuditReport: - missing_python: list[str] = [] - stale_python: list[str] = [] - for python_file, ledger_tests in _ledger_python_tests_by_file(ledger).items(): - actual_tests = enumerate_python_tests(repo_root, python_file) - for missing in sorted(ledger_tests - actual_tests): - missing_python.append(f"{python_file}:{missing}") - for stale in sorted(actual_tests - ledger_tests): - stale_python.append(f"{python_file}:{stale}") - - missing_rust: list[str] = [] - stale_rust: list[str] = [] - for rust_file, ledger_tests in _ledger_rust_tests_by_file(ledger).items(): - actual_tests = enumerate_rust_tests(repo_root, rust_file) - for missing in sorted(ledger_tests - actual_tests): - missing_rust.append(f"{rust_file}:{missing}") - for stale in sorted(actual_tests - ledger_tests): - stale_rust.append(f"{rust_file}:{stale}") - - return AuditReport( - missing_python_tests=tuple(missing_python), - stale_python_tests=tuple(stale_python), - missing_rust_tests=tuple(missing_rust), - stale_rust_tests=tuple(stale_rust), - ) - - -@dataclass(frozen=True, slots=True) -class FunctionReport: - sdk_function: str - ledger: TestLedger | None - audit: AuditReport | None - - @property - def has_ledger(self) -> bool: - return self.ledger is not None - - @property - def is_clean(self) -> bool: - return self.audit is None or self.audit.is_clean - - -def build_function_report(sdk_function: str, repo_root: Path = REPO_ROOT) -> FunctionReport: - path = ledger_path_for(sdk_function) - if not path.exists(): - return FunctionReport(sdk_function=sdk_function, ledger=None, audit=None) - ledger = load_ledger(path) - return FunctionReport( - sdk_function=sdk_function, ledger=ledger, audit=audit_ledger(ledger, repo_root) - ) diff --git a/tests/rust-python-harness/strategies/unit_tests/runner.py b/tests/rust-python-harness/strategies/unit_tests/runner.py deleted file mode 100644 index d10fa364d1c..00000000000 --- a/tests/rust-python-harness/strategies/unit_tests/runner.py +++ /dev/null @@ -1,91 +0,0 @@ -from __future__ import annotations - -import sys -from collections.abc import Sequence -from pathlib import Path -from time import monotonic -from typing import Final - -from pydantic import BaseModel, ConfigDict - -from ...shared.reporting.models import HarnessCase, HarnessRun, RunStatus -from ...shared.reporting.pytest_runner import UpdateCallback -from .mapping_validator import TestMapping, validate_mapping -from .python_runner import BackendSpec, compare_python_runs, run_python_tests -from .rust_runner import run_rust_tests - - -class UnitSuite(BaseModel): - model_config = ConfigDict(frozen=True, extra="forbid") - - python_selectors: tuple[str, ...] - cargo_manifest: str - cargo_package: str - cargo_filter: str - backend: BackendSpec - mappings: tuple[TestMapping, ...] = () - - -def run_suite(suite: UnitSuite, repo_root: Path, pytest_args: Sequence[str] = ()) -> tuple[str, ...]: - if not suite.python_selectors or not suite.cargo_filter: - return ("unit suites must select Python tests and a focused Cargo filter",) - python: Final = run_python_tests(suite.python_selectors, repo_root, "python", suite.backend, pytest_args) - rust_python: Final = run_python_tests(suite.python_selectors, repo_root, "rust", suite.backend, pytest_args) - inventory: Final = run_rust_tests( - repo_root / suite.cargo_manifest, suite.cargo_package, suite.cargo_filter, collect_only=True - ) - mapping: Final = validate_mapping(python.tests, inventory.tests, suite.mappings) - rust: Final = run_rust_tests(repo_root / suite.cargo_manifest, suite.cargo_package, suite.cargo_filter) - return ( - *compare_python_runs(python, rust_python), - *mapping.problems, - *(("native Rust tests did not all pass",) if set(inventory.tests) != set(rust.tests) else ()), - *((inventory.output,) if inventory.exit_code else ()), - *((rust.output,) if rust.exit_code else ()), - ) - - -def run( - cases: Sequence[HarnessCase], - repo_root: Path, - on_update: UpdateCallback, - pytest_args: Sequence[str] = (), -) -> tuple[int, HarnessRun]: - report: Final = HarnessRun.from_cases(cases) - for case in cases: - result: Final = report.results[case.key] - if case.unit_suite is None: - result.finalize() - continue - nodeid: Final = f"unit-suite:{case.unit_suite}" - result.collected.add(nodeid) - result.status = RunStatus.RUNNING - on_update(report) - try: - suite: Final = UnitSuite.model_validate_json((repo_root / case.unit_suite).read_text()) - problems: Final = run_suite(suite, repo_root, pytest_args) - except (OSError, ValueError) as error: - result.record(nodeid, RunStatus.ERROR) - report.failures.append((nodeid, str(error))) - continue - result.record(nodeid, RunStatus.FAILED if problems else RunStatus.PASSED) - report.failures.extend((nodeid, problem) for problem in problems) - on_update(report) - report.finished_at = monotonic() - on_update(report) - return int( - any( - result.status in {RunStatus.ERROR, RunStatus.FAILED, RunStatus.MISSING} - for result in report.results.values() - ) - ), report - - -def main(argv: Sequence[str] | None = None) -> int: - from ...cli import main as harness_main - - return harness_main(argv, strategy_id="unit_tests") - - -if __name__ == "__main__": - raise SystemExit(main(sys.argv[1:])) diff --git a/tests/rust-python-harness/strategies/unit_tests/rust_runner.py b/tests/rust-python-harness/strategies/unit_tests/rust_runner.py deleted file mode 100644 index b24034199b5..00000000000 --- a/tests/rust-python-harness/strategies/unit_tests/rust_runner.py +++ /dev/null @@ -1,53 +0,0 @@ -from __future__ import annotations - -import re -import subprocess -from dataclasses import dataclass -from pathlib import Path -from typing import Final - - -@dataclass(frozen=True, slots=True) -class RustReport: - tests: tuple[str, ...] - exit_code: int - output: str - - -def run_rust_tests(manifest: Path, package: str, test_filter: str, *, collect_only: bool = False) -> RustReport: - command: Final = ( - "cargo", - "test", - "--manifest-path", - str(manifest), - "--package", - package, - "--lib", - test_filter, - "--", - *(("--list",) if collect_only else ("--format=pretty",)), - ) - try: - result: Final = subprocess.run(command, capture_output=True, text=True, check=False, timeout=600) - except (OSError, subprocess.TimeoutExpired) as error: - return RustReport((), 1, str(error)) - tests: Final = ( - tuple(line.removesuffix(": test") for line in result.stdout.splitlines() if line.endswith(": test")) - if collect_only - else tuple( - line.removeprefix("test ").removesuffix(" ... ok") - for line in result.stdout.splitlines() - if line.startswith("test ") and line.endswith(" ... ok") - ) - ) - return RustReport(tests, result.returncode, result.stdout + result.stderr) - - -_RUST_TEST_PATTERN = re.compile( - r"#\[(?:test|tokio::test)\][^\n]*\n(?:[^\n]*\n)*?\s*(?:async\s+)?fn\s+(\w+)\s*\(" -) - - -def enumerate_rust_tests(repo_root: Path, relative_path: str) -> frozenset[str]: - source = (repo_root / relative_path).read_text(encoding="utf-8") - return frozenset(match.group(1) for match in _RUST_TEST_PATTERN.finditer(source)) diff --git a/tests/rust-python-harness/strategies/unit_tests/strategy.json b/tests/rust-python-harness/strategies/unit_tests/strategy.json deleted file mode 100644 index 7ae5d9c22c6..00000000000 --- a/tests/rust-python-harness/strategies/unit_tests/strategy.json +++ /dev/null @@ -1,32 +0,0 @@ -{ - "order": 30, - "id": "unit_tests", - "label": "Unit tests", - "description": "Validate Python/Rust test mappings and compare isolated Python runs alongside native Cargo tests.", - "functions": { - "ocr": { - "coverage": "planned", - "selectors": [] - }, - "messages": { - "coverage": "planned", - "selectors": [] - }, - "chat_completions": { - "coverage": "planned", - "selectors": [] - }, - "responses": { - "coverage": "planned", - "selectors": [] - }, - "count_tokens": { - "coverage": "planned", - "selectors": [] - }, - "transcription": { - "coverage": "planned", - "selectors": [] - } - } -} diff --git a/tests/rust-python-harness/strategies/unit_tests/test_mapping_validator.py b/tests/rust-python-harness/strategies/unit_tests/test_mapping_validator.py deleted file mode 100644 index 25c63faf3a7..00000000000 --- a/tests/rust-python-harness/strategies/unit_tests/test_mapping_validator.py +++ /dev/null @@ -1,38 +0,0 @@ -from __future__ import annotations - -import pytest - -from .mapping_validator import TestMapping as Mapping, validate_mapping - - -def test_matches_names_and_explicit_annotations() -> None: - report = validate_mapping( - ("tests/test_api.py::test_decode", "tests/test_api.py::test_error"), - ("api::test_decode", "api::preserves_error"), - (Mapping(python="tests/test_api.py::test_error", rust="api::preserves_error"),), - ) - assert report.problems == () - assert {(pair.python, pair.rust) for pair in report.pairs} == { - ("tests/test_api.py::test_decode", "api::test_decode"), - ("tests/test_api.py::test_error", "api::preserves_error"), - } - - -@pytest.mark.parametrize( - ("python", "rust", "message"), - ( - (("test_decode",), (), "missing Rust counterpart"), - ((), ("test_decode",), "missing Python counterpart"), - (("test_decode",), ("one::test_decode", "two::test_decode"), "ambiguous Rust counterparts"), - (("one::test_decode", "two::test_decode"), ("test_decode",), "ambiguous Python counterparts"), - ), -) -def test_reports_missing_and_ambiguous_counterparts( - python: tuple[str, ...], rust: tuple[str, ...], message: str -) -> None: - assert any(message in problem for problem in validate_mapping(python, rust).problems) - - -def test_rejects_stale_annotations_even_when_names_match() -> None: - report = validate_mapping(("test_decode",), ("test_decode",), (Mapping(python="test_decode", rust="removed"),)) - assert "missing Rust counterpart: removed" in report.problems diff --git a/tests/rust-python-harness/strategies/unit_tests/test_runner.py b/tests/rust-python-harness/strategies/unit_tests/test_runner.py deleted file mode 100644 index c652d6e12b1..00000000000 --- a/tests/rust-python-harness/strategies/unit_tests/test_runner.py +++ /dev/null @@ -1,59 +0,0 @@ -from __future__ import annotations - -import json -import shutil -from pathlib import Path -from typing import Final - -import pytest - -from ...shared.reporting.models import Coverage, HarnessCase, RunStatus -from .runner import run - - -@pytest.mark.skipif(shutil.which("cargo") is None, reason="Cargo is required for the combined unit strategy") -def test_combines_mapping_backend_comparison_and_cargo_results(tmp_path: Path, monkeypatch) -> None: - monkeypatch.setenv("PYTHONPATH", str(Path(__file__).resolve().parents[4])) - monkeypatch.setenv("PYTEST_DISABLE_PLUGIN_AUTOLOAD", "1") - (tmp_path / "pytest.ini").write_text("[pytest]\n") - (tmp_path / "backend_probe.py").write_text( - "import os\ndef selected():\n return 'rust' if os.environ['TEST_USE_RUST'] == '1' else 'python'\n" - ) - (tmp_path / "test_api.py").write_text("def test_decode():\n assert int('42') == 42\n") - (tmp_path / "Cargo.toml").write_text( - '[package]\nname = "combined-check"\nversion = "0.1.0"\nedition = "2021"\n[workspace]\n' - ) - (tmp_path / "src").mkdir() - (tmp_path / "src/lib.rs").write_text('#[test] fn test_decode() { assert_eq!("42".parse::().unwrap(), 42); }\n') - suite: Final = { - "python_selectors": ("test_api.py",), - "cargo_manifest": "Cargo.toml", - "cargo_package": "combined-check", - "cargo_filter": "test_decode", - "backend": {"environment_variable": "TEST_USE_RUST", "probe": "backend_probe:selected"}, - } - (tmp_path / "suite.json").write_text(json.dumps(suite)) - case: Final = HarnessCase( - strategy_id="unit_tests", - strategy_label="Unit tests", - sdk_function="ocr", - coverage=Coverage.COMPLETE, - selectors=(), - unit_suite="suite.json", - ) - code, report = run((case,), tmp_path, lambda _: None) - assert code == 0, report.failures - assert report.results[case.key].status is RunStatus.PASSED - (tmp_path / "suite.json").write_text( - json.dumps({**suite, "mappings": [{"python": "test_api.py::test_decode", "rust": "removed"}]}) - ) - failed_code, failed_report = run((case,), tmp_path, lambda _: None) - assert failed_code == 1 - assert failed_report.results[case.key].status is RunStatus.FAILED - assert any("missing Rust counterpart: removed" in detail for _, detail in failed_report.failures) - - (tmp_path / "suite.json").write_text(json.dumps(suite)) - (tmp_path / "src/lib.rs").write_text("#[test] #[ignore] fn test_decode() {}\n") - skipped_code, skipped_report = run((case,), tmp_path, lambda _: None) - assert skipped_code == 1 - assert any("native Rust tests did not all pass" in detail for _, detail in skipped_report.failures) diff --git a/tests/rust-python-harness/strategies/unit_tests/test_rust_runner.py b/tests/rust-python-harness/strategies/unit_tests/test_rust_runner.py deleted file mode 100644 index aeeb7f602f5..00000000000 --- a/tests/rust-python-harness/strategies/unit_tests/test_rust_runner.py +++ /dev/null @@ -1,27 +0,0 @@ -from __future__ import annotations - -import shutil -from pathlib import Path -from typing import Final - -import pytest - -from .rust_runner import run_rust_tests - - -@pytest.mark.skipif(shutil.which("cargo") is None, reason="Cargo is required for native runner integration") -def test_collects_and_runs_native_tests_and_propagates_failure(tmp_path: Path) -> None: - manifest: Final = tmp_path / "Cargo.toml" - manifest.write_text('[package]\nname = "harness-runner-check"\nversion = "0.1.0"\nedition = "2021"\n[workspace]\n') - (tmp_path / "src").mkdir() - source: Final = tmp_path / "src/lib.rs" - source.write_text("#[test] fn test_parity() { assert_eq!(2 + 2, 4); }\n") - inventory: Final = run_rust_tests(manifest, "harness-runner-check", "test_parity", collect_only=True) - assert inventory.exit_code == 0, inventory.output - assert inventory.tests == ("test_parity",) - passing: Final = run_rust_tests(manifest, "harness-runner-check", "test_parity") - assert passing.exit_code == 0, passing.output - source.write_text("#[test] fn test_parity() { assert_eq!(2 + 2, 5); }\n") - failed: Final = run_rust_tests(manifest, "harness-runner-check", "test_parity") - assert failed.exit_code != 0 - assert "test_parity" in failed.output diff --git a/tests/rust-python-harness/strategies/unit_tests_mapping/AGENTS.md b/tests/rust-python-harness/strategies/unit_tests_mapping/AGENTS.md new file mode 100644 index 00000000000..379d1443f33 --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests_mapping/AGENTS.md @@ -0,0 +1,13 @@ +# What this is + +Validates that unit tests covering traced Python behavior have semantic counterparts among colocated Rust unit tests + +# How it works + +Trace parity runs representative public API scenarios and records the Python and Rust functions reached, including their source files and lines. The OCR contract selects the behavior-level trace spans that require parity and excludes shared infrastructure such as generic HTTP transport + +For Python, those traced functions define the denominator. Static references and explicit includes create a safe pytest discovery universe, then a pytest profiler keeps only tests that actually execute at least one selected function. Static matches do not count by themselves. Parametrized pytest cases are collapsed to one logical test function in the mapping report. Explicit includes and exclusions cover dynamic callers or intentional harness behavior that static discovery cannot express reliably + +For Rust, each traced function identifies its source file and module. If that source file has a colocated `#[cfg(test)] mod tests`, the harness inventories that module for the configured Rust target. Rust test names are therefore derived from traced implementation files, not from a hand-maintained list of OCR test modules + +The Python-to-Rust mappings remain explicit because equivalent behavior often has different test boundaries and names in each SDK. Host-only exclusions require a reason. The report validates both against the live inventories, then shows mapped, excluded, and unmapped Python tests plus Rust-only tests diff --git a/tests/rust-python-harness/strategies/unit_tests_mapping/__init__.py b/tests/rust-python-harness/strategies/unit_tests_mapping/__init__.py new file mode 100644 index 00000000000..4d857c01ed0 --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests_mapping/__init__.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +from functools import partial +from pathlib import Path +from typing import Final + +from ...shared.reporting.models import SDK_FUNCTIONS, Coverage +from ...shared.reporting.strategy import ( + CaseDefinition, + NotImplementedCaseSpec, + RunnerArgumentDefinition, + StrategyDefinition, + SuiteCaseSpec, +) +from ...shared.unit_runners.suite_runner import run_suites +from .mappings import UNIT_TEST_CONTRACTS +from .reporting import render_mapping_results +from .runner import run_suite + + +CASES: Final[tuple[CaseDefinition, ...]] = ( + *( + CaseDefinition( + sdk_function, + SuiteCaseSpec(coverage=Coverage.COMPLETE, suite=sdk_function) + if sdk_function in UNIT_TEST_CONTRACTS + else NotImplementedCaseSpec(reason=f"No {sdk_function} unit-test mapping is registered."), + ) + for sdk_function in SDK_FUNCTIONS + ), +) + +STRATEGY: Final = StrategyDefinition( + id="unit_tests_mapping", + order=30, + label="Unit test mapping", + description="Validate Python/Rust unit-test mappings against collected test inventories.", + directory=Path(__file__).parent, + runnable_spec=SuiteCaseSpec, + cases=CASES, + run=partial(run_suites, suites=UNIT_TEST_CONTRACTS, execute=run_suite), + render=render_mapping_results, + runner_argument=RunnerArgumentDefinition( + option="--detail", + metavar="MODE", + help="show individual test names; any value enables full detail", + ), +) diff --git a/tests/rust-python-harness/strategies/unit_tests_mapping/cases/__init__.py b/tests/rust-python-harness/strategies/unit_tests_mapping/cases/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests_mapping/cases/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py b/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py new file mode 100644 index 00000000000..dc3167017d9 --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py @@ -0,0 +1,418 @@ +from __future__ import annotations + +from typing import Final + +from ....shared.unit_runners.rust_runner import RustTarget, RustTestIdentity +from ..contracts import ( + MappingExclusionSpec, + MappingSpec, + PythonFunctionDiscoverySpec, + RustTestFamily, + RustUnitSpec, + TestMapping, + UnitParityExclusionSpec, + UnitParitySpec, + UnitTestContract, +) + +_CORE_TARGET: Final = RustTarget(package="litellm-core", name="litellm_core", kind="lib") +_GATEWAY_TARGET: Final = RustTarget( + package="litellm-ai-gateway", + name="litellm_ai_gateway", + kind="lib", +) +_AZURE_OCR_TESTS: Final = "providers::azure_ai::ocr::transformation::tests" +_MISTRAL_OCR_TESTS: Final = "providers::mistral::ocr::transformation::tests" +_VERTEX_OCR_TESTS: Final = "providers::vertex_ai::ocr::transformation::tests" +_REDUCTO_OCR_TESTS: Final = "providers::reducto::ocr::tests" +_GATEWAY_OCR_TESTS: Final = "ocr::tests" +_GATEWAY_PREPARE_OCR_TESTS: Final = "ocr::prepare::tests" + + +def _rust_test(target: RustTarget, module: str, test: str) -> RustTestIdentity: + return RustTestIdentity(target=target, name=f"{module}::{test}") + + +def _rust_family(target: RustTarget, module: str, test: str) -> RustTestFamily: + return RustTestFamily(target=target, name=f"{module}::{test}") + + +def _test_mappings(target: RustTarget, module: str, pairs: tuple[tuple[str, str], ...]) -> tuple[TestMapping, ...]: + return tuple(TestMapping(python=python, rust=_rust_test(target, module, test)) for python, test in pairs) + + +_AZURE_TRANSFORM_FILE: Final = "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py" +_AZURE_PAGES_FILE: Final = "tests/ocr_tests/test_ocr_azure_document_intelligence.py" +_AZURE_BASE_FILE: Final = "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py" +_RUST_BRIDGE_FILE: Final = "tests/test_litellm/ocr/test_rust_bridge.py" + +_AZURE_PORT_MAPPINGS: Final = _test_mappings( + _CORE_TARGET, + _AZURE_OCR_TESTS, + ( + ( + f"{_AZURE_TRANSFORM_FILE}::test_should_encode_azure_document_intelligence_model_id", + "azure_document_intelligence_model_id_is_encoded", + ), + ( + f"{_AZURE_TRANSFORM_FILE}::test_should_reject_dot_segment_azure_document_intelligence_model_id", + "azure_document_intelligence_dot_segment_model_id_is_rejected", + ), + ( + f"{_AZURE_TRANSFORM_FILE}::test_async_transform_ocr_response_preserves_azure_native_fields", + "document_intelligence_async_response_preserves_normalized_fields", + ), + ( + f"{_AZURE_TRANSFORM_FILE}::test_transform_ocr_response_tolerates_missing_native_fields", + "document_intelligence_response_tolerates_missing_native_fields", + ), + ( + f"{_AZURE_TRANSFORM_FILE}::test_transform_ocr_response_non_succeeded_status_raises", + "document_intelligence_non_succeeded_status_is_rejected", + ), + ( + f"{_AZURE_TRANSFORM_FILE}::test_get_supported_ocr_params_includes_features", + "document_intelligence_supported_params_include_features", + ), + ( + f"{_AZURE_TRANSFORM_FILE}::test_transform_ocr_response_native_format_carries_raw_operation", + "document_intelligence_native_format_carries_raw_operation", + ), + ( + f"{_AZURE_TRANSFORM_FILE}::test_async_transform_ocr_response_native_format_carries_raw_operation", + "document_intelligence_async_native_format_carries_raw_operation", + ), + ( + f"{_AZURE_TRANSFORM_FILE}::test_map_ocr_params_rejects_unknown_req_format_as_bad_request", + "document_intelligence_rejects_unknown_req_format", + ), + ( + f"{_AZURE_TRANSFORM_FILE}::test_get_complete_url_omits_req_format_query_param", + "document_intelligence_url_omits_req_format", + ), + ( + f"{_AZURE_TRANSFORM_FILE}::test_validate_environment_uses_subscription_key", + "document_intelligence_validate_environment_uses_subscription_key", + ), + ( + f"{_AZURE_TRANSFORM_FILE}::test_validate_environment_falls_back_to_entra_token", + "document_intelligence_validate_environment_falls_back_to_entra_token", + ), + ( + f"{_AZURE_PAGES_FILE}::TestAzureDocumentIntelligencePagesParam::test_get_supported_ocr_params_includes_pages_and_features", + "document_intelligence_supported_params_include_pages_features_and_req_format", + ), + ( + f"{_AZURE_PAGES_FILE}::TestAzureDocumentIntelligencePagesParam::test_map_ocr_params_mistral_zero_based_int_list", + "document_intelligence_maps_zero_based_page_list", + ), + ( + f"{_AZURE_PAGES_FILE}::TestAzureDocumentIntelligencePagesParam::test_map_ocr_params_dedupes_and_sorts", + "document_intelligence_page_mapping_dedupes_and_sorts", + ), + ( + f"{_AZURE_PAGES_FILE}::TestAzureDocumentIntelligencePagesParam::test_map_ocr_params_empty_list_omits_pages", + "document_intelligence_page_mapping_omits_empty_list", + ), + ( + f"{_AZURE_PAGES_FILE}::TestAzureDocumentIntelligencePagesParam::test_map_ocr_params_azure_native_string_range", + "document_intelligence_page_mapping_accepts_native_range", + ), + ( + f"{_AZURE_PAGES_FILE}::TestAzureDocumentIntelligencePagesParam::test_map_ocr_params_azure_native_string_with_spaces_stripped", + "document_intelligence_page_mapping_strips_spaces", + ), + ( + f"{_AZURE_PAGES_FILE}::TestAzureDocumentIntelligencePagesParam::test_map_ocr_params_list_of_string_tokens", + "document_intelligence_page_mapping_accepts_string_tokens", + ), + ( + f"{_AZURE_PAGES_FILE}::TestAzureDocumentIntelligencePagesParam::test_map_ocr_params_invalid_string_raises", + "document_intelligence_page_mapping_rejects_invalid_string", + ), + ( + f"{_AZURE_PAGES_FILE}::TestAzureDocumentIntelligencePagesParam::test_map_ocr_params_negative_index_raises", + "document_intelligence_page_mapping_rejects_negative_index", + ), + ( + f"{_AZURE_PAGES_FILE}::TestAzureDocumentIntelligencePagesParam::test_map_ocr_params_bool_list_raises", + "document_intelligence_page_mapping_rejects_bool_list", + ), + ( + f"{_AZURE_PAGES_FILE}::TestAzureDocumentIntelligencePagesParam::test_map_ocr_params_unsupported_type_raises", + "document_intelligence_page_mapping_rejects_unsupported_type", + ), + ( + f"{_AZURE_PAGES_FILE}::TestAzureDocumentIntelligencePagesParam::test_get_complete_url_appends_pages_query", + "document_intelligence_url_appends_pages_query", + ), + ( + f"{_AZURE_PAGES_FILE}::TestAzureDocumentIntelligencePagesParam::test_get_complete_url_no_pages_when_optional_params_empty", + "document_intelligence_url_has_no_pages_when_params_are_empty", + ), + ( + f"{_AZURE_PAGES_FILE}::TestAzureDocumentIntelligencePagesParam::test_transform_ocr_request_does_not_put_pages_in_body", + "document_intelligence_request_keeps_pages_out_of_body", + ), + ( + f"{_AZURE_PAGES_FILE}::TestAzureDocumentIntelligencePagesParam::test_end_to_end_mistral_shape_to_azure_query", + "document_intelligence_mistral_pages_flow_to_query_only", + ), + ( + "tests/test_litellm/llms/azure_ai/test_azure_ai_entra_auth.py::test_ocr_authenticates_with_entra_token", + "azure_ai_ocr_authenticates_with_entra_token", + ), + ( + f"{_AZURE_BASE_FILE}::TestDocIntelligenceApiBaseResolution::test_generic_azure_ai_base_does_not_hijack_doc_intelligence", + "document_intelligence_endpoint_ignores_generic_azure_ai_base", + ), + ( + f"{_AZURE_BASE_FILE}::TestDocIntelligenceApiBaseResolution::test_explicit_api_base_is_honoured_for_doc_intelligence", + "document_intelligence_endpoint_honors_explicit_api_base", + ), + ( + f"{_AZURE_BASE_FILE}::TestDocIntelligenceApiBaseResolution::test_generic_azure_ai_base_still_applies_to_mistral_ocr", + "azure_ai_mistral_ocr_uses_generic_api_base", + ), + ), +) + +_REDUCTO_PORT_MAPPINGS: Final = _test_mappings( + _CORE_TARGET, + _REDUCTO_OCR_TESTS, + ( + ( + "tests/test_litellm/llms/reducto/test_parse_v3.py::test_parse_v3_reducto_id_passthrough_skips_upload", + "test_parse_v3_reducto_id_passthrough_skips_upload", + ), + ( + "tests/test_litellm/llms/reducto/test_parse_legacy.py::test_parse_legacy_wraps_enhance_under_options", + "test_parse_legacy_wraps_enhance_under_options", + ), + ( + "tests/test_litellm/llms/reducto/test_upload.py::test_parse_v3_image_data_uri_upload_uses_image_mime", + "test_parse_v3_image_data_uri_upload_uses_image_mime", + ), + ( + "tests/test_litellm/llms/reducto/test_upload.py::test_parse_v3_uses_programmatic_api_key_over_env", + "test_parse_v3_uses_programmatic_api_key_over_env", + ), + ), +) + +_REDUCTO_GATEWAY_MAPPING: Final = TestMapping( + python="tests/test_litellm/llms/reducto/test_parse_v3.py::test_parse_v3_file_upload_and_response_mapping", + rust=_rust_test(_GATEWAY_TARGET, _GATEWAY_OCR_TESTS, "reducto_file_upload_then_parse_maps_response"), +) + +_GATEWAY_PORT_MAPPINGS: Final = _test_mappings( + _GATEWAY_TARGET, + _GATEWAY_PREPARE_OCR_TESTS, + ( + ( + "tests/test_litellm/ocr/test_ocr_native_format.py::test_native_format_rejected_for_provider_without_support_as_bad_request", + "native_format_rejected_for_provider_without_support_as_bad_request", + ), + ( + "tests/test_litellm/ocr/test_ocr_native_format.py::test_unknown_format_rejected_for_provider_without_support_as_bad_request", + "unknown_format_rejected_for_provider_without_support_as_bad_request", + ), + ), +) + +_HOST_ONLY_BRIDGE_EXCLUSIONS: Final = tuple( + MappingExclusionSpec(nodeid=f"{_RUST_BRIDGE_FILE}::{test}", reason=reason) + for test, reason in ( + ("test_ocr_routes_to_rust_when_enabled", "Python selects and invokes the native bridge."), + ("test_ocr_routes_azure_ai_to_rust_when_enabled", "Python resolves provider arguments before the bridge."), + ("test_ocr_rust_path_converts_file_document_before_bridge", "Python converts file inputs before the bridge."), + ( + "test_ocr_exception_type_uses_resolved_provider_context", + "Python wraps bridge exceptions into public errors.", + ), + ("test_aocr_routes_to_async_rust_when_enabled", "Python selects and invokes the async native bridge."), + ("test_aocr_exception_type_uses_resolved_provider_context", "Python wraps async bridge exceptions."), + ("test_ocr_forwards_timeout_to_rust", "Python converts and forwards explicit timeouts."), + ("test_ocr_passes_default_request_timeout_to_rust", "Python supplies its process-level default timeout."), + ("test_ocr_falls_back_to_python_when_bridge_unavailable", "Python owns fallback when the extension is absent."), + ) +) + +_FAMILY_PORT_MAPPINGS: Final = ( + TestMapping( + python=f"{_AZURE_TRANSFORM_FILE}::test_transform_ocr_response_default_format_omits_raw_operation", + rust=_rust_family( + _CORE_TARGET, + _AZURE_OCR_TESTS, + "document_intelligence_default_format_omits_raw_operation", + ), + ), + TestMapping( + python=f"{_AZURE_TRANSFORM_FILE}::test_map_ocr_params_passes_through_req_format", + rust=_rust_family(_CORE_TARGET, _AZURE_OCR_TESTS, "document_intelligence_maps_req_format"), + ), + TestMapping( + python="tests/ocr_tests/test_ocr_vertex_ai.py::test_deepseek_request_uses_single_provider_namespace", + rust=_rust_family( + _CORE_TARGET, + _VERTEX_OCR_TESTS, + "vertex_deepseek_request_uses_single_provider_namespace", + ), + ), + TestMapping( + python="tests/test_litellm/llms/reducto/test_upload.py::test_parse_v3_rejects_plain_http_urls", + rust=_rust_family(_CORE_TARGET, _REDUCTO_OCR_TESTS, "test_parse_v3_rejects_plain_http_urls"), + ), +) + + +OCR_CONTRACT: Final = UnitTestContract( + mapping=MappingSpec( + python_functions=PythonFunctionDiscoverySpec( + trace_module="tests.rust-python-harness.strategies.trace_parity.sdk.ocr.case", + trace_spans=( + "ocr", + "prepare_ocr_call", + "ocr_provider_config", + "supported_ocr_params", + "map_ocr_params", + "validate_environment", + "complete_url", + "transform_ocr_request", + "execute_ocr_provider_call", + "transform_ocr_response", + "poll_document_intelligence", + ), + search_roots=("tests",), + exclude_roots=( + "tests/e2e", + "tests/ocr_tests/test_ocr_mistral.py", + "tests/rust-python-harness", + ), + includes=( + "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "tests/test_litellm/llms/mistral/ocr", + "tests/test_litellm/llms/ocr", + "tests/test_litellm/ocr", + "tests/test_litellm/proxy/ocr_endpoints", + ), + exclusions=( + "tests/ocr_tests/test_ocr_azure_document_intelligence.py::TestAzureDocumentIntelligenceOCR", + "tests/ocr_tests/test_ocr_vertex_ai.py::TestVertexAIMistralOCR", + "tests/ocr_tests/test_ocr_vertex_ai.py::TestVertexAIDeepSeekOCR", + ), + ), + rust_targets=(_CORE_TARGET, _GATEWAY_TARGET), + mappings=( + TestMapping( + python="tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py::test_transform_ocr_response_preserves_azure_native_fields", + rust=_rust_test(_CORE_TARGET, _AZURE_OCR_TESTS, "document_intelligence_response_normalizes_pages"), + ), + TestMapping( + python="tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py::test_map_ocr_params_features", + rust=_rust_family(_CORE_TARGET, _AZURE_OCR_TESTS, "document_intelligence_maps_features"), + ), + TestMapping( + python="tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py::test_map_ocr_params_empty_features_list_omitted", + rust=_rust_test(_CORE_TARGET, _AZURE_OCR_TESTS, "document_intelligence_url_omits_empty_feature_list"), + ), + TestMapping( + python="tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py::test_map_ocr_params_invalid_features_raises", + rust=_rust_family( + _CORE_TARGET, + _AZURE_OCR_TESTS, + "document_intelligence_mapping_rejects_invalid_features", + ), + ), + TestMapping( + python="tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py::test_get_complete_url_appends_features_query", + rust=_rust_test(_CORE_TARGET, _AZURE_OCR_TESTS, "document_intelligence_url_normalizes_features"), + ), + TestMapping( + python="tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py::test_get_complete_url_combines_pages_and_features", + rust=_rust_test( + _CORE_TARGET, _AZURE_OCR_TESTS, "document_intelligence_url_combines_pages_and_feature_list" + ), + ), + TestMapping( + python="tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py::TestGetSupportedOcrParams::test_extract_header_in_supported_params", + rust=_rust_test(_CORE_TARGET, _MISTRAL_OCR_TESTS, "extract_header_is_a_supported_ocr_param"), + ), + TestMapping( + python="tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py::TestGetSupportedOcrParams::test_extract_footer_in_supported_params", + rust=_rust_test(_CORE_TARGET, _MISTRAL_OCR_TESTS, "extract_footer_is_a_supported_ocr_param"), + ), + TestMapping( + python="tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py::TestGetSupportedOcrParams::test_existing_params_still_present", + rust=_rust_test(_CORE_TARGET, _MISTRAL_OCR_TESTS, "existing_ocr_params_remain_supported"), + ), + TestMapping( + python="tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py::TestMapOcrParams::test_extract_header_passed_through", + rust=_rust_test(_CORE_TARGET, _MISTRAL_OCR_TESTS, "map_ocr_params_forwards_extract_header"), + ), + TestMapping( + python="tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py::TestMapOcrParams::test_extract_footer_passed_through", + rust=_rust_test(_CORE_TARGET, _MISTRAL_OCR_TESTS, "map_ocr_params_forwards_extract_footer"), + ), + TestMapping( + python="tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py::TestMapOcrParams::test_extract_header_and_footer_together", + rust=_rust_test(_CORE_TARGET, _MISTRAL_OCR_TESTS, "map_ocr_params_forwards_extract_header_and_footer"), + ), + TestMapping( + python="tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py::TestMapOcrParams::test_unknown_param_is_dropped", + rust=_rust_test(_CORE_TARGET, _MISTRAL_OCR_TESTS, "map_ocr_params_drops_unknown_params"), + ), + TestMapping( + python="tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py::TestNewSupportedParams::test_new_param_in_supported_list", + rust=_rust_test(_CORE_TARGET, _MISTRAL_OCR_TESTS, "new_ocr_params_are_supported"), + ), + TestMapping( + python="tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py::TestNewParamsMapOcr::test_new_param_passed_through", + rust=_rust_test(_CORE_TARGET, _MISTRAL_OCR_TESTS, "map_ocr_params_forwards_new_ocr_params"), + ), + TestMapping( + python="tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py::TestTransformOcrRequest::test_param_included_in_request_body", + rust=_rust_test(_CORE_TARGET, _MISTRAL_OCR_TESTS, "transform_ocr_request_includes_each_optional_param"), + ), + TestMapping( + python="tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py::TestTransformOcrRequest::test_multiple_new_params_together", + rust=_rust_test(_CORE_TARGET, _MISTRAL_OCR_TESTS, "transform_ocr_request_includes_multiple_new_params"), + ), + TestMapping( + python="tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py::TestTransformOcrResponseOcr4Fields::test_blocks_and_confidence_scores_preserved", + rust=_rust_test( + _CORE_TARGET, _MISTRAL_OCR_TESTS, "transform_ocr_response_preserves_blocks_and_confidence_scores" + ), + ), + TestMapping( + python="tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py::TestTransformOcrResponseOcr4Fields::test_ocr4_fields_survive_model_dump", + rust=_rust_test(_CORE_TARGET, _MISTRAL_OCR_TESTS, "transform_ocr_response_preserves_ocr4_page_fields"), + ), + *_AZURE_PORT_MAPPINGS, + *_REDUCTO_PORT_MAPPINGS, + _REDUCTO_GATEWAY_MAPPING, + *_GATEWAY_PORT_MAPPINGS, + *_FAMILY_PORT_MAPPINGS, + ), + exclusions=_HOST_ONLY_BRIDGE_EXCLUSIONS, + require_complete=True, + ), + unit_parity=UnitParitySpec( + python_selectors=( + "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "tests/test_litellm/llms/mistral/ocr", + "tests/test_litellm/llms/ocr", + "tests/test_litellm/ocr", + ), + exclusions=( + UnitParityExclusionSpec( + nodeid="tests/test_litellm/ocr/test_rust_bridge.py::test_use_litellm_rust_toggles_flag", + reason="This test asserts the process-level backend flag selected by the parity runner.", + ), + ), + ), + rust=RustUnitSpec( + cargo_manifest="litellm-rust/Cargo.toml", + cargo_filter="ocr", + ), +) diff --git a/tests/rust-python-harness/strategies/unit_tests_mapping/contracts.py b/tests/rust-python-harness/strategies/unit_tests_mapping/contracts.py new file mode 100644 index 00000000000..a8f309cc8f3 --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests_mapping/contracts.py @@ -0,0 +1,220 @@ +from __future__ import annotations + +from collections import Counter +from typing import Final, Literal + +from pydantic import BaseModel, ConfigDict, field_validator, model_validator +from typing_extensions import Self + +from ...shared.tracing.pytest_usage import PythonFunctionReference +from ...shared.unit_runners.rust_runner import RustTarget, RustTestIdentity, RustTestScope + + +class _ContractModel(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + +def _clean_unique(values: tuple[str, ...], field: str) -> tuple[str, ...]: + cleaned: Final = tuple(value.strip().rstrip("/") for value in values) + if not cleaned or any(not value for value in cleaned): + raise ValueError(f"{field} must contain non-empty paths") + duplicates: Final = tuple(value for value, count in Counter(cleaned).items() if count > 1) + if duplicates: + raise ValueError(f"{field} contains duplicates: {sorted(duplicates)}") + return cleaned + + +def _selector_contains(parent: str, child: str) -> bool: + return child == parent or child.startswith(f"{parent}/") + + +class RustTestFamily(_ContractModel): + kind: Literal["family"] = "family" + target: RustTarget + name: str + + @field_validator("name") + @classmethod + def validate_name(cls, value: str) -> str: + stripped: Final = value.strip() + if not stripped or stripped.endswith("::"): + raise ValueError("must be a non-empty Rust test base name") + return stripped + + @property + def key(self) -> str: + return f"{self.target.key}::{self.name}::case_*" + + def contains(self, identity: RustTestIdentity) -> bool: + return identity.target == self.target and identity.name.startswith(f"{self.name}::case_") + + +class TestMapping(_ContractModel): + python: str + rust: RustTestIdentity | RustTestFamily + + @field_validator("python") + @classmethod + def validate_python_nodeid(cls, value: str) -> str: + stripped: Final = value.strip() + if "::" not in stripped: + raise ValueError("must be a source path and test name separated by '::'") + return stripped + + +class PythonFunctionDiscoverySpec(_ContractModel): + functions: tuple[PythonFunctionReference, ...] = () + trace_module: str | None = None + trace_spans: tuple[str, ...] = () + search_roots: tuple[str, ...] + exclude_roots: tuple[str, ...] = () + includes: tuple[str, ...] = () + exclusions: tuple[str, ...] = () + + @field_validator("search_roots") + @classmethod + def validate_search_roots(cls, value: tuple[str, ...]) -> tuple[str, ...]: + return _clean_unique(value, "python function search_roots") + + @field_validator("exclude_roots") + @classmethod + def validate_exclude_roots(cls, value: tuple[str, ...]) -> tuple[str, ...]: + if not value: + return () + return _clean_unique(value, "python function exclude_roots") + + @model_validator(mode="after") + def validate_functions(self) -> Self: + if bool(self.functions) == bool(self.trace_module): + raise ValueError("python function discovery needs exactly one function list or trace module") + if self.trace_module is not None and not self.trace_spans: + raise ValueError("trace-derived Python function discovery needs trace_spans") + if not self.functions: + return self + keys: Final = tuple(f"{function.module}:{function.qualname}" for function in self.functions) + duplicates: Final = tuple(key for key, count in Counter(keys).items() if count > 1) + if duplicates: + raise ValueError(f"python function discovery contains duplicates: {sorted(duplicates)}") + return self + + +class UnitParityExclusionSpec(_ContractModel): + nodeid: str + reason: str + + @field_validator("nodeid", "reason") + @classmethod + def validate_fields(cls, value: str) -> str: + stripped: Final = value.strip() + if not stripped: + raise ValueError("must be a non-empty string") + return stripped + + +class MappingExclusionSpec(_ContractModel): + nodeid: str + reason: str + + @field_validator("nodeid", "reason") + @classmethod + def validate_fields(cls, value: str) -> str: + stripped: Final = value.strip() + if not stripped: + raise ValueError("must be a non-empty string") + return stripped + + +class MappingSpec(_ContractModel): + python_selectors: tuple[str, ...] = () + python_functions: PythonFunctionDiscoverySpec | None = None + rust_scope: tuple[RustTestScope, ...] = () + rust_targets: tuple[RustTarget, ...] = () + mappings: tuple[TestMapping, ...] + exclusions: tuple[MappingExclusionSpec, ...] = () + require_complete: bool = False + + @field_validator("python_selectors") + @classmethod + def validate_python_selectors(cls, value: tuple[str, ...]) -> tuple[str, ...]: + if not value: + return () + return _clean_unique(value, "python_selectors") + + @model_validator(mode="after") + def validate_rust_scope(self) -> Self: + if bool(self.python_selectors) == bool(self.python_functions): + raise ValueError("mapping needs exactly one Python selector or function-discovery scope") + targets: Final = tuple(scope.target.key for scope in self.rust_scope) + duplicates: Final = tuple(target for target, count in Counter(targets).items() if count > 1) + if duplicates: + raise ValueError(f"rust_scope contains duplicate targets: {sorted(duplicates)}") + target_names: Final = tuple(target.name for target in self.rust_targets) + duplicate_names: Final = tuple(name for name, count in Counter(target_names).items() if count > 1) + if duplicate_names: + raise ValueError(f"rust_targets contains duplicate names: {sorted(duplicate_names)}") + exclusion_nodeids: Final = tuple(exclusion.nodeid for exclusion in self.exclusions) + duplicate_exclusions: Final = tuple(nodeid for nodeid, count in Counter(exclusion_nodeids).items() if count > 1) + if duplicate_exclusions: + raise ValueError(f"mapping exclusions contain duplicate nodeids: {sorted(duplicate_exclusions)}") + return self + + +class UnitParitySpec(_ContractModel): + python_selectors: tuple[str, ...] + exclusions: tuple[UnitParityExclusionSpec, ...] = () + + @field_validator("python_selectors") + @classmethod + def validate_python_selectors(cls, value: tuple[str, ...]) -> tuple[str, ...]: + return _clean_unique(value, "unit parity python_selectors") + + @model_validator(mode="after") + def validate_exclusions(self) -> Self: + nodeids: Final = tuple(exclusion.nodeid for exclusion in self.exclusions) + duplicates: Final = tuple(nodeid for nodeid, count in Counter(nodeids).items() if count > 1) + if duplicates: + raise ValueError(f"unit parity exclusions contain duplicate nodeids: {sorted(duplicates)}") + return self + + +class RustUnitSpec(_ContractModel): + cargo_manifest: str + cargo_filter: str + cargo_package: str | None = None + + @field_validator("cargo_manifest", "cargo_filter") + @classmethod + def validate_required_fields(cls, value: str) -> str: + stripped: Final = value.strip() + if not stripped: + raise ValueError("must be a non-empty string") + return stripped + + @field_validator("cargo_package") + @classmethod + def validate_package(cls, value: str | None) -> str | None: + if value is None: + return None + stripped: Final = value.strip() + if not stripped: + raise ValueError("must be a non-empty string when provided") + return stripped + + +class UnitTestContract(_ContractModel): + mapping: MappingSpec + unit_parity: UnitParitySpec + rust: RustUnitSpec + + @model_validator(mode="after") + def validate_unit_parity_scope(self) -> Self: + if not self.mapping.python_selectors: + return self + unknown: Final = tuple( + selector + for selector in self.unit_parity.python_selectors + if not any(_selector_contains(parent, selector) for parent in self.mapping.python_selectors) + ) + if unknown: + raise ValueError(f"unit parity selectors must be contained in mapping selectors: {sorted(unknown)}") + return self diff --git a/tests/rust-python-harness/strategies/unit_tests_mapping/mapping_report.py b/tests/rust-python-harness/strategies/unit_tests_mapping/mapping_report.py new file mode 100644 index 00000000000..a5fd92e449d --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests_mapping/mapping_report.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +from collections import Counter +from collections.abc import Callable, Sequence +from typing import Final + +from pydantic import BaseModel, ConfigDict + +from .mapping_validator import MappingReport + + +class MappingReportArtifact(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + report: MappingReport + detailed: bool = False + + +def _group_counts(nodeids: Sequence[str], owner: Callable[[str], str]) -> tuple[str, ...]: + counts: Final = Counter(owner(nodeid) for nodeid in nodeids) + width: Final = max((len(str(count)) for count in counts.values()), default=1) + return tuple( + f" {count:>{width}} {name}" for name, count in sorted(counts.items(), key=lambda item: (-item[1], item[0])) + ) + + +def _python_file(nodeid: str) -> str: + return nodeid.partition("::")[0] + + +def _rust_module(nodeid: str) -> str: + return nodeid.rpartition("::")[0] + + +def _details(nodeids: Sequence[str], owner: Callable[[str], str]) -> tuple[str, ...]: + owners: Final = tuple(sorted(frozenset(owner(nodeid) for nodeid in nodeids))) + return tuple( + line + for name in owners + for line in ( + f" {name}", + *(f" {nodeid.removeprefix(f'{name}::')}" for nodeid in nodeids if owner(nodeid) == name), + ) + ) + + +def _contract_errors(report: MappingReport) -> tuple[str, ...]: + return ( + *(f" Missing Python test: {nodeid}" for nodeid in report.missing_python_tests), + *(f" Missing Rust test: {nodeid}" for nodeid in report.missing_rust_tests), + *(f" Python test mapped more than once: {nodeid}" for nodeid in report.duplicate_python_mappings), + *(f" Rust test mapped more than once: {nodeid}" for nodeid in report.duplicate_rust_mappings), + *(f" Missing mapping exclusion: {nodeid}" for nodeid in report.invalid_mapping_exclusions), + *(f" Python test is both mapped and excluded: {nodeid}" for nodeid in report.mapped_and_excluded_python_tests), + *(f" Missing unit-parity exclusion: {nodeid}" for nodeid in report.invalid_unit_parity_exclusions), + ) + + +def mapping_report_lines(report: MappingReport, *, detailed: bool = False) -> tuple[str, ...]: + unmapped_count: Final = len(report.unmapped_python_tests) + excluded_count: Final = len(report.excluded_python_tests) + excluded_percentage: Final = ( + 0.0 if not report.total_count else round(100.0 * excluded_count / report.total_count, 1) + ) + unmapped_percentage: Final = ( + 0.0 if not report.total_count else round(100.0 * unmapped_count / report.total_count, 1) + ) + rust_total: Final = len(report.rust_tests) + rust_only_count: Final = len(report.rust_only_tests) + rust_mapped_count: Final = rust_total - rust_only_count + contract_errors: Final = _contract_errors(report) + detail_lines: Final = ( + ( + "", + "Unmapped Python test details", + *_details(report.unmapped_python_tests, _python_file), + "", + "Excluded Python test details", + *_details(report.excluded_python_tests, _python_file), + "", + "Rust-only test details", + *_details(report.rust_only_tests, _rust_module), + ) + if detailed + else () + ) + return ( + f"Contract: {'PASS' if report.is_valid else 'FAIL'}", + "", + "Python coverage", + f" Mapped {report.mapped_count:>3} / {report.total_count} ({report.percentage}%)", + f" Excluded {excluded_count:>3} / {report.total_count} ({excluded_percentage}%)", + f" Unmapped {unmapped_count:>3} / {report.total_count} ({unmapped_percentage}%)", + "", + "Rust inventory", + f" Mapped {rust_mapped_count:>3} / {rust_total}", + f" Rust-only {rust_only_count:>3} / {rust_total}", + "", + f"Unmapped Python tests by file ({unmapped_count})", + *_group_counts(report.unmapped_python_tests, _python_file), + "", + f"Excluded Python tests by file ({excluded_count})", + *_group_counts(report.excluded_python_tests, _python_file), + "", + f"Rust-only tests by module ({rust_only_count})", + *_group_counts(report.rust_only_tests, _rust_module), + *(("", "Contract errors", *contract_errors) if contract_errors else ()), + *detail_lines, + ) diff --git a/tests/rust-python-harness/strategies/unit_tests_mapping/mapping_validator.py b/tests/rust-python-harness/strategies/unit_tests_mapping/mapping_validator.py new file mode 100644 index 00000000000..98ea0b02e68 --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests_mapping/mapping_validator.py @@ -0,0 +1,298 @@ +from __future__ import annotations + +import importlib +from collections import Counter, defaultdict +from collections.abc import Callable, Sequence +from pathlib import Path +from typing import Final, TypeAlias + +from pydantic import BaseModel, ConfigDict + +from ...shared.tracing.pytest_usage import ( + PythonFunctionIdentity, + RustFunctionIdentity, + candidate_test_files, + collect_python_function_tests, +) +from ...shared.tracing.steps import pipeline_projection +from ...shared.unit_runners.python_runner import collect_python_tests, contract_nodeid +from ...shared.unit_runners.rust_runner import RustTarget, RustTestIdentity, RustTestScope, enumerate_rust_tests +from .contracts import PythonFunctionDiscoverySpec, RustTestFamily, TestMapping, UnitTestContract + +PythonInventory: TypeAlias = Callable[[Sequence[str], Path], frozenset[str]] +RustInventory: TypeAlias = Callable[[Path, tuple[RustTestScope, ...]], frozenset[RustTestIdentity]] + + +def _trace_functions( + spec: PythonFunctionDiscoverySpec, +) -> tuple[tuple[PythonFunctionIdentity, ...], tuple[RustFunctionIdentity, ...]]: + from ..trace_parity.models import RouteSpec, TraceExecutionFailure, TraceSuite + from ..trace_parity.sdk.execution import collect_trace + + if spec.trace_module is None: + return () + module: Final = importlib.import_module(spec.trace_module) + suite: Final = getattr(module, "TRACE_SUITE", None) + if not isinstance(suite, TraceSuite) or not isinstance(suite.route, RouteSpec): + raise ValueError(f"{spec.trace_module} must export an SDK TRACE_SUITE") + python_functions: Final[dict[str, PythonFunctionIdentity]] = {} + rust_functions: Final[dict[str, RustFunctionIdentity]] = {} + for scenario in suite.scenarios: + for mode in scenario.modes: + route: Final = RouteSpec( + route=suite.route.route, + python_entrypoints=suite.route.python_entrypoints, + rust_entrypoints=suite.route.rust_entrypoints, + fixture=scenario.fixture, + ) + python_trace: Final = collect_trace(route, "python", asynchronous=mode == "async") + rust_trace: Final = collect_trace(route, "rust", asynchronous=mode == "async") + if isinstance(python_trace, TraceExecutionFailure): + raise ValueError(f"Python trace discovery failed for {scenario.name}/{mode}: {python_trace.message}") + if isinstance(rust_trace, TraceExecutionFailure): + raise ValueError(f"Rust trace discovery failed for {scenario.name}/{mode}: {rust_trace.message}") + mappings: Final = scenario.mappings_for(mode) + python_projection: Final = pipeline_projection("python", python_trace, mappings) + rust_projection: Final = pipeline_projection("rust", rust_trace, mappings) + for step in python_projection.steps: + if step.span in spec.trace_spans: + function: Final = PythonFunctionIdentity.from_trace(step.raw) + python_functions[function.raw] = function + for step in rust_projection.steps: + if step.span in spec.trace_spans: + function: Final = RustFunctionIdentity.from_trace(step.raw) + rust_functions[step.raw] = function + if not python_functions or not rust_functions: + raise ValueError(f"Python trace discovery found no functions for spans: {', '.join(spec.trace_spans)}") + return ( + tuple(python_functions[key] for key in sorted(python_functions)), + tuple(rust_functions[key] for key in sorted(rust_functions)), + ) + + +def collect_python_function_inventory( + spec: PythonFunctionDiscoverySpec, + repo_root: Path, + traced_functions: Sequence[PythonFunctionIdentity] = (), +) -> frozenset[str]: + source_root: Final = repo_root / "litellm" + functions: Final = ( + tuple(reference.resolve(source_root) for reference in spec.functions) + if spec.functions + else tuple(traced_functions) + ) + discovered: Final = candidate_test_files( + functions, + spec.search_roots, + repo_root, + exclude_roots=spec.exclude_roots, + ) + selectors: Final = tuple(dict.fromkeys((*discovered, *spec.includes))) + if not selectors: + raise ValueError("Python function discovery found no candidate test files") + report: Final = collect_python_function_tests( + functions, + selectors, + repo_root, + source_root=source_root, + exclusions=spec.exclusions, + ) + if report.exit_code or report.problems: + details: Final = "\n".join(report.problems) or f"pytest exited with code {report.exit_code}" + raise ValueError(f"Python function test discovery failed:\n{details}") + return frozenset(contract_nodeid(nodeid) for usage in report.usages for nodeid in usage.tests) + + +def _colocated_rust_scope(mappings: Sequence[TestMapping]) -> tuple[RustTestScope, ...]: + modules_by_target: Final[dict[str, set[str]]] = defaultdict(set) + targets: Final[dict[str, RustTarget]] = {} + for item in mappings: + module, separator, _ = item.rust.name.partition("::tests::") + if not separator: + raise ValueError(f"Rust unit test is not colocated in a tests module: {item.rust.key}") + target_key: Final = item.rust.target.key + targets[target_key] = item.rust.target + modules_by_target[target_key].add(f"{module}::tests") + return tuple( + RustTestScope( + target=targets[target_key], + modules=tuple(sorted(modules_by_target[target_key])), + ) + for target_key in sorted(targets) + ) + + +def _traced_rust_scope( + functions: Sequence[RustFunctionIdentity], + targets: Sequence[RustTarget], + repo_root: Path, +) -> tuple[RustTestScope, ...]: + targets_by_name: Final = {target.name: target for target in targets} + modules_by_target: Final[dict[str, set[str]]] = defaultdict(set) + for function in functions: + crate: Final = function.module_path.partition("::")[0] + target: Final = targets_by_name.get(crate) + if target is None: + continue + source_candidates: Final = ( + repo_root / "litellm-rust" / function.file, + repo_root / function.file, + ) + source: Final = next((path for path in source_candidates if path.is_file()), None) + if source is None: + raise ValueError(f"Traced Rust source does not exist: {function.file}") + contents: Final = source.read_text() + if "mod tests" in contents and "#[cfg(test)]" in contents: + modules_by_target[target.key].add(function.test_module) + selected_targets: Final = {target.key: target for target in targets} + scopes: Final = tuple( + RustTestScope(target=selected_targets[key], modules=tuple(sorted(modules))) + for key, modules in sorted(modules_by_target.items()) + if modules + ) + if not scopes: + raise ValueError("Traced Rust functions have no colocated test modules") + return scopes + + +def _merge_rust_scopes(scopes: Sequence[RustTestScope]) -> tuple[RustTestScope, ...]: + targets: Final = {scope.target.key: scope.target for scope in scopes} + modules: Final[dict[str, set[str]]] = defaultdict(set) + features: Final[dict[str, set[str]]] = defaultdict(set) + default_features: Final[dict[str, bool]] = {} + for scope in scopes: + modules[scope.target.key].update(scope.modules) + features[scope.target.key].update(scope.features) + default_features[scope.target.key] = default_features.get(scope.target.key, True) and scope.default_features + return tuple( + RustTestScope( + target=targets[key], + modules=tuple( + sorted( + module + for module in modules[key] + if not any(module.startswith(f"{parent}::") for parent in modules[key]) + ) + ), + features=tuple(sorted(features[key])), + default_features=default_features[key], + ) + for key in sorted(targets) + ) + + +def _owned_rust_tests( + rust: RustTestIdentity | RustTestFamily, + inventory: frozenset[RustTestIdentity], +) -> frozenset[RustTestIdentity]: + if isinstance(rust, RustTestFamily): + return frozenset(identity for identity in inventory if rust.contains(identity)) + return frozenset((rust,)) if rust in inventory else frozenset() + + +class MappingReport(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + python_tests: tuple[str, ...] + rust_tests: tuple[str, ...] + mapped_python_tests: tuple[str, ...] + excluded_python_tests: tuple[str, ...] + unmapped_python_tests: tuple[str, ...] + rust_only_tests: tuple[str, ...] + missing_python_tests: tuple[str, ...] + missing_rust_tests: tuple[str, ...] + duplicate_python_mappings: tuple[str, ...] + duplicate_rust_mappings: tuple[str, ...] + invalid_mapping_exclusions: tuple[str, ...] + mapped_and_excluded_python_tests: tuple[str, ...] + invalid_unit_parity_exclusions: tuple[str, ...] + + @property + def mapped_count(self) -> int: + return len(self.mapped_python_tests) + + @property + def total_count(self) -> int: + return len(self.python_tests) + + @property + def percentage(self) -> float: + return 0.0 if not self.total_count else round(100.0 * self.mapped_count / self.total_count, 1) + + @property + def is_valid(self) -> bool: + return not ( + self.missing_python_tests + or self.missing_rust_tests + or self.duplicate_python_mappings + or self.duplicate_rust_mappings + or self.invalid_mapping_exclusions + or self.mapped_and_excluded_python_tests + or self.invalid_unit_parity_exclusions + ) + + +def audit_mapping( + contract: UnitTestContract, + repo_root: Path, + *, + python_inventory: PythonInventory = collect_python_tests, + rust_inventory: RustInventory = enumerate_rust_tests, +) -> MappingReport: + mapping: Final = contract.mapping + traced_python: tuple[PythonFunctionIdentity, ...] = () + traced_rust: tuple[RustFunctionIdentity, ...] = () + if mapping.python_functions is not None and mapping.python_functions.trace_module is not None: + traced_python, traced_rust = _trace_functions(mapping.python_functions) + python_tests: Final = ( + collect_python_function_inventory(mapping.python_functions, repo_root, traced_python) + if mapping.python_functions is not None + else python_inventory(mapping.python_selectors, repo_root) + ) + unit_parity_tests: Final = python_inventory(contract.unit_parity.python_selectors, repo_root) + traced_scope: Final = _traced_rust_scope(traced_rust, mapping.rust_targets, repo_root) if traced_rust else () + rust_scope: Final = _merge_rust_scopes( + (*mapping.rust_scope, *traced_scope, *_colocated_rust_scope(mapping.mappings)) + ) + rust_tests: Final = rust_inventory(repo_root, rust_scope) + mapped_python: Final = frozenset(item.python for item in mapping.mappings) + excluded_python: Final = frozenset(exclusion.nodeid for exclusion in mapping.exclusions) + rust_ownership: Final = tuple((item.rust, _owned_rust_tests(item.rust, rust_tests)) for item in mapping.mappings) + mapped_rust: Final = frozenset(identity for _, identities in rust_ownership for identity in identities) + duplicate_python: Final = tuple( + sorted(nodeid for nodeid, count in Counter(item.python for item in mapping.mappings).items() if count > 1) + ) + duplicate_exact_rust: Final = frozenset( + identity.key + for identity, count in Counter( + item.rust for item in mapping.mappings if isinstance(item.rust, RustTestIdentity) + ).items() + if count > 1 + ) + duplicate_owned_rust: Final = frozenset( + identity.key + for identity, count in Counter(identity for _, identities in rust_ownership for identity in identities).items() + if count > 1 + ) + duplicate_rust: Final = tuple(sorted(duplicate_exact_rust | duplicate_owned_rust)) + return MappingReport( + python_tests=tuple(sorted(python_tests)), + rust_tests=tuple(sorted(identity.key for identity in rust_tests)), + mapped_python_tests=tuple(sorted(python_tests & mapped_python)), + excluded_python_tests=tuple(sorted((python_tests & excluded_python) - mapped_python)), + unmapped_python_tests=tuple(sorted(python_tests - mapped_python - excluded_python)), + rust_only_tests=tuple(sorted(identity.key for identity in rust_tests - mapped_rust)), + missing_python_tests=tuple(sorted(mapped_python - python_tests)), + missing_rust_tests=tuple(sorted(rust.key for rust, identities in rust_ownership if not identities)), + duplicate_python_mappings=duplicate_python, + duplicate_rust_mappings=duplicate_rust, + invalid_mapping_exclusions=tuple(sorted(excluded_python - python_tests)), + mapped_and_excluded_python_tests=tuple(sorted(mapped_python & excluded_python)), + invalid_unit_parity_exclusions=tuple( + sorted( + exclusion.nodeid + for exclusion in contract.unit_parity.exclusions + if exclusion.nodeid not in unit_parity_tests + ) + ), + ) diff --git a/tests/rust-python-harness/strategies/unit_tests_mapping/mappings.py b/tests/rust-python-harness/strategies/unit_tests_mapping/mappings.py new file mode 100644 index 00000000000..efb5b2a644a --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests_mapping/mappings.py @@ -0,0 +1,11 @@ +from __future__ import annotations + +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +from ...shared.reporting.models import SdkFunction +from .cases.ocr import OCR_CONTRACT +from .contracts import UnitTestContract + +UNIT_TEST_CONTRACTS: Final[Mapping[SdkFunction, UnitTestContract]] = MappingProxyType({"ocr": OCR_CONTRACT}) diff --git a/tests/rust-python-harness/strategies/unit_tests_mapping/reporting.py b/tests/rust-python-harness/strategies/unit_tests_mapping/reporting.py new file mode 100644 index 00000000000..d4bce7bc768 --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests_mapping/reporting.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +from collections.abc import Sequence +from typing import Final + +from pydantic import ValidationError + +from ...shared.reporting.models import CaseResult +from ...shared.reporting.rendering import ReportSection, render_case_outcome +from .mapping_report import MappingReportArtifact, mapping_report_lines +from .runner import MAPPING_REPORT_ARTIFACT + + +def _render_artifact(body: str) -> str: + try: + artifact: Final = MappingReportArtifact.model_validate_json(body) + except ValidationError as error: + return f"Mapping report artifact is invalid: {error}" + return "\n".join(mapping_report_lines(artifact.report, detailed=artifact.detailed)) + + +def _render_result(result: CaseResult) -> str: + reports: Final = tuple( + _render_artifact(artifact.body) + for artifacts in result.artifacts.values() + for artifact in artifacts + if artifact.kind == MAPPING_REPORT_ARTIFACT + ) + if reports: + return "\n".join((f"Case: {result.case.display_name}", *reports)) + return render_case_outcome(result) + + +def render_mapping_results(results: Sequence[CaseResult]) -> tuple[ReportSection, ...]: + blocks: Final = tuple(_render_result(result) for result in results) + return (ReportSection("Python/Rust unit-test mappings", blocks or ("No mapping cases selected",)),) diff --git a/tests/rust-python-harness/strategies/unit_tests_mapping/runner.py b/tests/rust-python-harness/strategies/unit_tests_mapping/runner.py new file mode 100644 index 00000000000..540edca9385 --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests_mapping/runner.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +from collections.abc import Sequence +from pathlib import Path +from typing import Final + +from ...shared.native_build import ensure_trace_bridge +from ...shared.reporting.models import ResultArtifact +from ...shared.unit_runners.python_runner import collect_python_tests +from ...shared.unit_runners.rust_runner import enumerate_rust_tests +from ...shared.unit_runners.suite_runner import SuiteExecution +from .contracts import UnitTestContract +from .mapping_report import MappingReportArtifact +from .mapping_validator import PythonInventory, RustInventory, audit_mapping + +MAPPING_REPORT_ARTIFACT: Final = "mapping_report" + + +def _audit_problems(artifact: MappingReportArtifact) -> tuple[str, ...]: + report: Final = artifact.report + return ( + *(f"mapped Python test does not exist: {nodeid}" for nodeid in report.missing_python_tests), + *(f"mapped Rust test does not exist: {nodeid}" for nodeid in report.missing_rust_tests), + *(f"Python test has multiple mappings: {nodeid}" for nodeid in report.duplicate_python_mappings), + *(f"Rust test has multiple mappings: {nodeid}" for nodeid in report.duplicate_rust_mappings), + *(f"mapping exclusion does not exist: {nodeid}" for nodeid in report.invalid_mapping_exclusions), + *(f"Python test is both mapped and excluded: {nodeid}" for nodeid in report.mapped_and_excluded_python_tests), + *(f"unit parity exclusion does not exist: {nodeid}" for nodeid in report.invalid_unit_parity_exclusions), + ) + + +def run_suite( + contract: UnitTestContract, + repo_root: Path, + runner_args: Sequence[str] = (), + *, + python_inventory: PythonInventory = collect_python_tests, + rust_inventory: RustInventory = enumerate_rust_tests, +) -> SuiteExecution: + if contract.mapping.python_functions is not None and contract.mapping.python_functions.trace_module is not None: + bridge_error: Final = ensure_trace_bridge(repo_root) + if bridge_error is not None: + return SuiteExecution(problems=(bridge_error,)) + artifact: Final = MappingReportArtifact( + report=audit_mapping( + contract, + repo_root, + python_inventory=python_inventory, + rust_inventory=rust_inventory, + ), + detailed=bool(runner_args), + ) + completeness_problems: Final = ( + tuple(f"Python test has no Rust mapping: {nodeid}" for nodeid in artifact.report.unmapped_python_tests) + if contract.mapping.require_complete + else () + ) + return SuiteExecution( + problems=(*_audit_problems(artifact), *completeness_problems), + artifacts=(ResultArtifact(MAPPING_REPORT_ARTIFACT, artifact.model_dump_json()),), + ) diff --git a/tests/rust-python-harness/strategies/unit_tests_mapping/test_mapping_validator.py b/tests/rust-python-harness/strategies/unit_tests_mapping/test_mapping_validator.py new file mode 100644 index 00000000000..6635a0eb522 --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests_mapping/test_mapping_validator.py @@ -0,0 +1,314 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Final + +import pytest +from pydantic import ValidationError + +from ...shared.unit_runners.rust_runner import RustTarget, RustTestIdentity, RustTestScope +from .contracts import ( + MappingExclusionSpec, + MappingSpec, + RustTestFamily, + RustUnitSpec, + UnitParityExclusionSpec, + UnitParitySpec, + UnitTestContract, +) +from .contracts import TestMapping as MappingPair +from .mapping_validator import audit_mapping + +_TARGET: Final = RustTarget(package="example", name="example", kind="lib") +_SCOPE: Final = RustTestScope(target=_TARGET, modules=("api::tests",)) +_PYTHON_TESTS: Final = frozenset(("test_api.py::test_decode", "test_api.py::test_unmapped")) +_RUST_TEST: Final = RustTestIdentity(target=_TARGET, name="api::tests::decodes") +_RUST_ONLY: Final = RustTestIdentity(target=_TARGET, name="api::tests::rust_only") +_RUST_TESTS: Final = frozenset((_RUST_TEST, _RUST_ONLY)) + + +def _python_inventory(*_: object) -> frozenset[str]: + return _PYTHON_TESTS + + +def _rust_inventory(*_: object) -> frozenset[RustTestIdentity]: + return _RUST_TESTS + + +def _contract(*mappings: MappingPair, exclusions: tuple[UnitParityExclusionSpec, ...] = ()) -> UnitTestContract: + return UnitTestContract( + mapping=MappingSpec( + python_selectors=("test_api.py",), + rust_scope=(_SCOPE,), + mappings=mappings, + ), + unit_parity=UnitParitySpec(python_selectors=("test_api.py",), exclusions=exclusions), + rust=RustUnitSpec(cargo_manifest="Cargo.toml", cargo_filter="api"), + ) + + +def _mapping_exclusion(nodeid: str) -> MappingExclusionSpec: + return MappingExclusionSpec(nodeid=nodeid, reason="Python bridge availability is host-only") + + +def test_derives_mapping_status_from_live_inventories(tmp_path: Path) -> None: + contract: Final = _contract(MappingPair(python="test_api.py::test_decode", rust=_RUST_TEST)) + + report: Final = audit_mapping( + contract, tmp_path, python_inventory=_python_inventory, rust_inventory=_rust_inventory + ) + + assert report.is_valid + assert report.mapped_python_tests == ("test_api.py::test_decode",) + assert report.unmapped_python_tests == ("test_api.py::test_unmapped",) + assert report.rust_only_tests == (_RUST_ONLY.key,) + assert report.percentage == 50.0 + + +def test_reports_stale_and_duplicate_mappings(tmp_path: Path) -> None: + removed: Final = RustTestIdentity(target=_TARGET, name="api::tests::removed") + contract: Final = _contract( + MappingPair(python="test_api.py::removed", rust=removed), + MappingPair(python="test_api.py::removed", rust=_RUST_TEST), + ) + + report: Final = audit_mapping( + contract, tmp_path, python_inventory=_python_inventory, rust_inventory=_rust_inventory + ) + + assert not report.is_valid + assert report.missing_python_tests == ("test_api.py::removed",) + assert report.missing_rust_tests == (removed.key,) + assert report.duplicate_python_mappings == ("test_api.py::removed",) + + +def test_reports_duplicate_rust_mapping_and_invalid_exclusion(tmp_path: Path) -> None: + contract: Final = _contract( + MappingPair(python="test_api.py::test_decode", rust=_RUST_TEST), + MappingPair(python="test_api.py::test_unmapped", rust=_RUST_TEST), + exclusions=(UnitParityExclusionSpec(nodeid="test_api.py::removed", reason="Removed test"),), + ) + + report: Final = audit_mapping( + contract, tmp_path, python_inventory=_python_inventory, rust_inventory=_rust_inventory + ) + + assert not report.is_valid + assert report.duplicate_rust_mappings == (_RUST_TEST.key,) + assert report.invalid_unit_parity_exclusions == ("test_api.py::removed",) + + +def test_excludes_host_only_python_test_from_unmapped_inventory(tmp_path: Path) -> None: + partial: Final = _contract(MappingPair(python="test_api.py::test_decode", rust=_RUST_TEST)) + contract: Final = partial.model_copy( + update={ + "mapping": partial.mapping.model_copy( + update={"exclusions": (_mapping_exclusion("test_api.py::test_unmapped"),)} + ) + } + ) + + report: Final = audit_mapping( + contract, tmp_path, python_inventory=_python_inventory, rust_inventory=_rust_inventory + ) + + assert report.is_valid + assert report.excluded_python_tests == ("test_api.py::test_unmapped",) + assert report.unmapped_python_tests == () + + +def test_reports_missing_and_mapped_mapping_exclusions(tmp_path: Path) -> None: + partial: Final = _contract(MappingPair(python="test_api.py::test_decode", rust=_RUST_TEST)) + contract: Final = partial.model_copy( + update={ + "mapping": partial.mapping.model_copy( + update={ + "exclusions": ( + _mapping_exclusion("test_api.py::test_decode"), + _mapping_exclusion("test_api.py::removed"), + ) + } + ) + } + ) + + report: Final = audit_mapping( + contract, tmp_path, python_inventory=_python_inventory, rust_inventory=_rust_inventory + ) + + assert not report.is_valid + assert report.invalid_mapping_exclusions == ("test_api.py::removed",) + assert report.mapped_and_excluded_python_tests == ("test_api.py::test_decode",) + + +def test_resolves_rstest_family_to_generated_cases(tmp_path: Path) -> None: + first_case: Final = RustTestIdentity(target=_TARGET, name="api::tests::decodes::case_1_png") + second_case: Final = RustTestIdentity(target=_TARGET, name="api::tests::decodes::case_2_pdf") + family: Final = RustTestFamily(target=_TARGET, name="api::tests::decodes") + contract: Final = _contract(MappingPair(python="test_api.py::test_decode", rust=family)) + + report: Final = audit_mapping( + contract, + tmp_path, + python_inventory=_python_inventory, + rust_inventory=lambda *_: frozenset((first_case, second_case)), + ) + + assert report.is_valid + assert report.mapped_python_tests == ("test_api.py::test_decode",) + assert report.missing_rust_tests == () + + +def test_reports_missing_rstest_family(tmp_path: Path) -> None: + family: Final = RustTestFamily(target=_TARGET, name="api::tests::decodes") + contract: Final = _contract(MappingPair(python="test_api.py::test_decode", rust=family)) + + report: Final = audit_mapping( + contract, tmp_path, python_inventory=_python_inventory, rust_inventory=_rust_inventory + ) + + assert not report.is_valid + assert report.missing_rust_tests == (family.key,) + + +def test_reports_concrete_test_owned_by_exact_and_family_mappings(tmp_path: Path) -> None: + generated: Final = RustTestIdentity(target=_TARGET, name="api::tests::decodes::case_1_png") + family: Final = RustTestFamily(target=_TARGET, name="api::tests::decodes") + contract: Final = _contract( + MappingPair(python="test_api.py::test_decode", rust=family), + MappingPair(python="test_api.py::test_unmapped", rust=generated), + ) + + report: Final = audit_mapping( + contract, + tmp_path, + python_inventory=_python_inventory, + rust_inventory=lambda *_: frozenset((generated,)), + ) + + assert not report.is_valid + assert report.duplicate_rust_mappings == (generated.key,) + + +def test_rstest_family_cases_are_not_rust_only(tmp_path: Path) -> None: + generated: Final = RustTestIdentity(target=_TARGET, name="api::tests::decodes::case_1_png") + unrelated: Final = RustTestIdentity(target=_TARGET, name="api::tests::rust_only") + family: Final = RustTestFamily(target=_TARGET, name="api::tests::decodes") + contract: Final = _contract(MappingPair(python="test_api.py::test_decode", rust=family)) + + report: Final = audit_mapping( + contract, + tmp_path, + python_inventory=_python_inventory, + rust_inventory=lambda *_: frozenset((generated, unrelated)), + ) + + assert report.rust_only_tests == (unrelated.key,) + + +def test_merges_configured_and_colocated_rust_scopes(tmp_path: Path) -> None: + support_test: Final = RustTestIdentity(target=_TARGET, name="support::tests::rust_only") + configured_scope: Final = RustTestScope( + target=_TARGET, + modules=("support::tests",), + features=("mock",), + default_features=False, + ) + expected_scope: Final = RustTestScope( + target=_TARGET, + modules=("api::tests", "support::tests"), + features=("mock",), + default_features=False, + ) + contract: Final = UnitTestContract( + mapping=MappingSpec( + python_selectors=("test_api.py",), + rust_scope=(configured_scope,), + mappings=(MappingPair(python="test_api.py::test_decode", rust=_RUST_TEST),), + ), + unit_parity=UnitParitySpec(python_selectors=("test_api.py",)), + rust=RustUnitSpec(cargo_manifest="Cargo.toml", cargo_filter="api"), + ) + + def assert_merged_scope(_: Path, scopes: tuple[RustTestScope, ...]) -> frozenset[RustTestIdentity]: + assert scopes == (expected_scope,) + return frozenset((_RUST_TEST, support_test)) + + report: Final = audit_mapping( + contract, + tmp_path, + python_inventory=_python_inventory, + rust_inventory=assert_merged_scope, + ) + + assert report.is_valid + assert report.rust_only_tests == (support_test.key,) + + +def test_merged_rust_scope_removes_modules_contained_by_parent(tmp_path: Path) -> None: + expected_scope: Final = RustTestScope(target=_TARGET, modules=("api",)) + contract: Final = UnitTestContract( + mapping=MappingSpec( + python_selectors=("test_api.py",), + rust_scope=(expected_scope,), + mappings=(MappingPair(python="test_api.py::test_decode", rust=_RUST_TEST),), + ), + unit_parity=UnitParitySpec(python_selectors=("test_api.py",)), + rust=RustUnitSpec(cargo_manifest="Cargo.toml", cargo_filter="api"), + ) + + def assert_parent_scope(_: Path, scopes: tuple[RustTestScope, ...]) -> frozenset[RustTestIdentity]: + assert scopes == (expected_scope,) + return frozenset((_RUST_TEST,)) + + report: Final = audit_mapping( + contract, + tmp_path, + python_inventory=_python_inventory, + rust_inventory=assert_parent_scope, + ) + + assert report.is_valid + + +def test_accepts_descendant_unit_parity_selector() -> None: + contract: Final = UnitTestContract( + mapping=MappingSpec(python_selectors=("tests/api",), rust_scope=(_SCOPE,), mappings=()), + unit_parity=UnitParitySpec(python_selectors=("tests/api/test_ocr.py",)), + rust=RustUnitSpec(cargo_manifest="Cargo.toml", cargo_filter="api"), + ) + + assert contract.unit_parity.python_selectors == ("tests/api/test_ocr.py",) + + +@pytest.mark.parametrize( + "mapping_selectors,parity_selectors", + (((), ("tests/api",)), (("tests/api", "tests/api"), ("tests/api",)), (("tests/api",), ("tests/chat",))), +) +def test_rejects_invalid_selector_contracts( + mapping_selectors: tuple[str, ...], parity_selectors: tuple[str, ...] +) -> None: + with pytest.raises(ValidationError): + UnitTestContract( + mapping=MappingSpec(python_selectors=mapping_selectors, rust_scope=(_SCOPE,), mappings=()), + unit_parity=UnitParitySpec(python_selectors=parity_selectors), + rust=RustUnitSpec(cargo_manifest="Cargo.toml", cargo_filter="api"), + ) + + +def test_rejects_duplicate_scopes_and_exclusions() -> None: + exclusion: Final = UnitParityExclusionSpec(nodeid="test_api.py::test_skip", reason="Backend assertion") + with pytest.raises(ValidationError, match="duplicate targets"): + MappingSpec(python_selectors=("test_api.py",), rust_scope=(_SCOPE, _SCOPE), mappings=()) + with pytest.raises(ValidationError, match="duplicate nodeids"): + UnitParitySpec(python_selectors=("test_api.py",), exclusions=(exclusion, exclusion)) + mapping_exclusion: Final = _mapping_exclusion("test_api.py::test_skip") + with pytest.raises(ValidationError, match="mapping exclusions contain duplicate nodeids"): + MappingSpec( + python_selectors=("test_api.py",), + rust_scope=(_SCOPE,), + mappings=(), + exclusions=(mapping_exclusion, mapping_exclusion), + ) + with pytest.raises(ValidationError, match="must be a non-empty string"): + MappingExclusionSpec(nodeid="test_api.py::test_skip", reason=" ") diff --git a/tests/rust-python-harness/strategies/unit_tests_mapping/test_reporting.py b/tests/rust-python-harness/strategies/unit_tests_mapping/test_reporting.py new file mode 100644 index 00000000000..36e18a9d109 --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests_mapping/test_reporting.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +from typing import Final + +from ...shared.reporting.models import CaseResult, Coverage, HarnessCase, ResultArtifact, RunStatus +from ...shared.reporting.strategy import SuiteCaseSpec +from .mapping_report import MappingReportArtifact +from .mapping_validator import MappingReport +from .reporting import render_mapping_results +from .runner import MAPPING_REPORT_ARTIFACT + + +def _report(*, invalid: bool = False, excluded: bool = False) -> MappingReport: + return MappingReport( + python_tests=("test_api.py::test_decode", "test_api.py::test_unmapped"), + rust_tests=("example/lib/example::api::tests::decodes", "example/lib/example::api::tests::rust_only"), + mapped_python_tests=("test_api.py::test_decode",), + excluded_python_tests=(("test_api.py::test_unmapped",) if excluded else ()), + unmapped_python_tests=(() if excluded else ("test_api.py::test_unmapped",)), + rust_only_tests=("example/lib/example::api::tests::rust_only",), + missing_python_tests=("test_api.py::removed",) if invalid else (), + missing_rust_tests=(), + duplicate_python_mappings=(), + duplicate_rust_mappings=(), + invalid_mapping_exclusions=(), + mapped_and_excluded_python_tests=(), + invalid_unit_parity_exclusions=(), + ) + + +def _result(body: str) -> CaseResult: + case: Final = HarnessCase( + strategy_id="unit_tests_mapping", + strategy_label="Unit test mapping", + sdk_function="ocr", + spec=SuiteCaseSpec(coverage=Coverage.COMPLETE, suite="ocr"), + ) + result: Final = CaseResult(case=case) + result.record( + "suite:unit_tests_mapping:ocr:ocr", + RunStatus.PASSED, + artifacts=(ResultArtifact(MAPPING_REPORT_ARTIFACT, body),), + ) + return result + + +def test_renderer_preserves_summary_and_detailed_output() -> None: + summary: Final = MappingReportArtifact(report=_report()).model_dump_json() + detailed: Final = MappingReportArtifact(report=_report(), detailed=True).model_dump_json() + + summary_text: Final = "\n".join(render_mapping_results((_result(summary),))[0].blocks) + detailed_text: Final = "\n".join(render_mapping_results((_result(detailed),))[0].blocks) + + assert "Mapped 1 / 2 (50.0%)" in summary_text + assert "Unmapped Python test details" not in summary_text + assert "Unmapped Python test details\n test_api.py\n test_unmapped" in detailed_text + assert "Rust-only test details" in detailed_text + + +def test_renderer_shows_contract_errors() -> None: + body: Final = MappingReportArtifact(report=_report(invalid=True)).model_dump_json() + rendered: Final = "\n".join(render_mapping_results((_result(body),))[0].blocks) + + assert "Contract: FAIL" in rendered + assert "Missing Python test: test_api.py::removed" in rendered + + +def test_renderer_distinguishes_excluded_python_tests() -> None: + body: Final = MappingReportArtifact(report=_report(excluded=True), detailed=True).model_dump_json() + rendered: Final = "\n".join(render_mapping_results((_result(body),))[0].blocks) + + assert "Excluded 1 / 2 (50.0%)" in rendered + assert "Unmapped 0 / 2 (0.0%)" in rendered + assert "Excluded Python test details\n test_api.py\n test_unmapped" in rendered + + +def test_renderer_handles_empty_inventory_and_malformed_artifact() -> None: + empty: Final = MappingReport( + python_tests=(), + rust_tests=(), + mapped_python_tests=(), + excluded_python_tests=(), + unmapped_python_tests=(), + rust_only_tests=(), + missing_python_tests=(), + missing_rust_tests=(), + duplicate_python_mappings=(), + duplicate_rust_mappings=(), + invalid_mapping_exclusions=(), + mapped_and_excluded_python_tests=(), + invalid_unit_parity_exclusions=(), + ) + empty_text: Final = "\n".join( + render_mapping_results((_result(MappingReportArtifact(report=empty).model_dump_json()),))[0].blocks + ) + invalid_text: Final = "\n".join(render_mapping_results((_result("not-json"),))[0].blocks) + + assert "Mapped 0 / 0 (0.0%)" in empty_text + assert "Mapping report artifact is invalid:" in invalid_text diff --git a/tests/rust-python-harness/strategies/unit_tests_mapping/test_runner.py b/tests/rust-python-harness/strategies/unit_tests_mapping/test_runner.py new file mode 100644 index 00000000000..2b14c716e1d --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests_mapping/test_runner.py @@ -0,0 +1,166 @@ +from __future__ import annotations + +from functools import partial +from pathlib import Path +from typing import Final + +from ...shared.reporting.models import Coverage, HarnessCase, RunStatus +from ...shared.reporting.strategy import SuiteCaseSpec +from ...shared.unit_runners.rust_runner import RustTarget, RustTestIdentity, RustTestScope +from ...shared.unit_runners.suite_runner import run_suites +from .contracts import ( + MappingExclusionSpec, + MappingSpec, + RustUnitSpec, + TestMapping as MappingPair, + UnitParitySpec, + UnitTestContract, +) +from .mapping_report import MappingReportArtifact +from .runner import MAPPING_REPORT_ARTIFACT, run_suite + +_TARGET: Final = RustTarget(package="example", name="example", kind="lib") +_RUST_TEST: Final = RustTestIdentity(target=_TARGET, name="api::tests::decodes") +_RUST_ONLY: Final = RustTestIdentity(target=_TARGET, name="api::tests::rust_only") + + +def _python_inventory(*_: object) -> frozenset[str]: + return frozenset(("test_api.py::test_decode", "test_api.py::test_unmapped")) + + +def _rust_inventory(*_: object) -> frozenset[RustTestIdentity]: + return frozenset((_RUST_TEST, _RUST_ONLY)) + + +def _contract(mapping: MappingPair) -> UnitTestContract: + return UnitTestContract( + mapping=MappingSpec( + python_selectors=("test_api.py",), + rust_scope=(RustTestScope(target=_TARGET, modules=("api::tests",)),), + mappings=(mapping,), + ), + unit_parity=UnitParitySpec(python_selectors=("test_api.py",)), + rust=RustUnitSpec(cargo_manifest="Cargo.toml", cargo_filter="api"), + ) + + +def _case() -> HarnessCase: + return HarnessCase( + strategy_id="unit_tests_mapping", + strategy_label="Unit test mapping", + sdk_function="ocr", + spec=SuiteCaseSpec(coverage=Coverage.COMPLETE, suite="ocr"), + ) + + +def test_reports_structured_mapping_status_without_running_tests(tmp_path: Path) -> None: + contract: Final = _contract(MappingPair(python="test_api.py::test_decode", rust=_RUST_TEST)) + case: Final = _case() + + code, report = run_suites( + (case,), + tmp_path, + lambda _: None, + suites={"ocr": contract}, + execute=partial( + run_suite, + python_inventory=_python_inventory, + rust_inventory=_rust_inventory, + ), + ) + + result: Final = report.results[case.key] + artifacts: Final = tuple( + artifact + for values in result.artifacts.values() + for artifact in values + if artifact.kind == MAPPING_REPORT_ARTIFACT + ) + parsed: Final = MappingReportArtifact.model_validate_json(artifacts[0].body) + assert code == 0, report.failures + assert result.status is RunStatus.PASSED + assert parsed.report.mapped_count == 1 + assert parsed.report.total_count == 2 + assert not parsed.detailed + + +def test_fails_when_a_mapping_target_is_missing(tmp_path: Path) -> None: + missing: Final = RustTestIdentity(target=_TARGET, name="api::tests::missing") + contract: Final = _contract(MappingPair(python="test_api.py::test_decode", rust=missing)) + case: Final = _case() + + code, report = run_suites( + (case,), + tmp_path, + lambda _: None, + suites={"ocr": contract}, + execute=partial( + run_suite, + python_inventory=_python_inventory, + rust_inventory=_rust_inventory, + ), + ) + + assert code == 1 + assert report.results[case.key].status is RunStatus.FAILED + assert any("mapped Rust test does not exist" in detail for _, detail in report.failures) + + +def test_required_complete_mapping_fails_for_unmapped_python_test(tmp_path: Path) -> None: + partial: Final = _contract(MappingPair(python="test_api.py::test_decode", rust=_RUST_TEST)) + contract: Final = partial.model_copy( + update={"mapping": partial.mapping.model_copy(update={"require_complete": True})} + ) + + execution: Final = run_suite( + contract, + tmp_path, + python_inventory=_python_inventory, + rust_inventory=_rust_inventory, + ) + + assert execution.problems == ("Python test has no Rust mapping: test_api.py::test_unmapped",) + + +def test_required_complete_mapping_accepts_host_only_exclusion(tmp_path: Path) -> None: + partial: Final = _contract(MappingPair(python="test_api.py::test_decode", rust=_RUST_TEST)) + contract: Final = partial.model_copy( + update={ + "mapping": partial.mapping.model_copy( + update={ + "require_complete": True, + "exclusions": ( + MappingExclusionSpec( + nodeid="test_api.py::test_unmapped", + reason="Python bridge availability is host-only", + ), + ), + } + ) + } + ) + + execution: Final = run_suite( + contract, + tmp_path, + python_inventory=_python_inventory, + rust_inventory=_rust_inventory, + ) + artifact: Final = MappingReportArtifact.model_validate_json(execution.artifacts[0].body) + + assert execution.problems == () + assert artifact.report.excluded_python_tests == ("test_api.py::test_unmapped",) + + +def test_detail_argument_is_stored_in_artifact(tmp_path: Path) -> None: + contract: Final = _contract(MappingPair(python="test_api.py::test_decode", rust=_RUST_TEST)) + execution: Final = run_suite( + contract, + tmp_path, + ("full",), + python_inventory=_python_inventory, + rust_inventory=_rust_inventory, + ) + artifact: Final = MappingReportArtifact.model_validate_json(execution.artifacts[0].body) + + assert artifact.detailed diff --git a/tests/rust-python-harness/strategies/unit_tests_parity/AGENTS.md b/tests/rust-python-harness/strategies/unit_tests_parity/AGENTS.md new file mode 100644 index 00000000000..ccab6b1ff12 --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests_parity/AGENTS.md @@ -0,0 +1 @@ +Runs the existing litellm Python unit tests with LITELLM_RUST=0 and LITELLM_RUST=1 in separate processes and requires the two runs to match, including on failures. diff --git a/tests/rust-python-harness/strategies/unit_tests_parity/__init__.py b/tests/rust-python-harness/strategies/unit_tests_parity/__init__.py new file mode 100644 index 00000000000..0067bf6dfe5 --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests_parity/__init__.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +from collections.abc import Mapping +from functools import partial +from pathlib import Path +from types import MappingProxyType +from typing import Final + +from ...shared.reporting.models import SDK_FUNCTIONS, Coverage, SdkFunction +from ...shared.reporting.strategy import ( + CaseDefinition, + NotImplementedCaseSpec, + RunnerArgumentDefinition, + StrategyDefinition, + SuiteCaseSpec, +) +from ...shared.unit_runners.suite_runner import run_suites +from ..unit_tests_mapping.mappings import UNIT_TEST_CONTRACTS +from .reporting import render_unit_parity_results +from .runner import UnitParityExclusion, UnitParitySuite, run_suite + + +UNIT_PARITY_SUITES: Final[Mapping[SdkFunction, UnitParitySuite]] = MappingProxyType( + { + sdk_function: UnitParitySuite( + python_selectors=contract.unit_parity.python_selectors, + exclusions=tuple( + UnitParityExclusion( + nodeid=exclusion.nodeid, + reason=exclusion.reason, + ) + for exclusion in contract.unit_parity.exclusions + ), + ) + for sdk_function, contract in UNIT_TEST_CONTRACTS.items() + } +) + + +CASES: Final[tuple[CaseDefinition, ...]] = ( + *( + CaseDefinition( + sdk_function, + SuiteCaseSpec(coverage=Coverage.COMPLETE, suite=sdk_function) + if sdk_function in UNIT_PARITY_SUITES + else NotImplementedCaseSpec(reason=f"No {sdk_function} unit-test parity suite is registered."), + ) + for sdk_function in SDK_FUNCTIONS + ), +) + +STRATEGY: Final = StrategyDefinition( + id="unit_tests_parity", + order=31, + label="Unit test parity", + description=( + "Run existing Python unit tests with LITELLM_RUST disabled and enabled and require matching outcomes." + ), + directory=Path(__file__).parent, + runnable_spec=SuiteCaseSpec, + cases=CASES, + run=partial(run_suites, suites=UNIT_PARITY_SUITES, execute=run_suite), + render=render_unit_parity_results, + runner_argument=RunnerArgumentDefinition( + option="--pytest-arg", + help="append an argument to both Python and Rust-backed pytest runs", + ), +) diff --git a/tests/rust-python-harness/strategies/unit_tests_parity/reporting.py b/tests/rust-python-harness/strategies/unit_tests_parity/reporting.py new file mode 100644 index 00000000000..339e893c60d --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests_parity/reporting.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +from collections.abc import Sequence +from typing import Final + +from ...shared.reporting.models import CaseResult +from ...shared.reporting.rendering import ReportSection, render_case_outcome + + +def render_unit_parity_results(results: Sequence[CaseResult]) -> tuple[ReportSection, ...]: + blocks: Final = tuple(render_case_outcome(result) for result in results) + return (ReportSection("Python backend parity outcomes", blocks or ("No unit-parity cases selected",)),) diff --git a/tests/rust-python-harness/strategies/unit_tests_parity/runner.py b/tests/rust-python-harness/strategies/unit_tests_parity/runner.py new file mode 100644 index 00000000000..5a3a70ea03c --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests_parity/runner.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +from collections.abc import Sequence +from pathlib import Path +from typing import Final + +from pydantic import BaseModel, ConfigDict + +from ...shared.unit_runners.python_runner import BackendSpec, compare_python_runs, run_python_tests +from ...shared.unit_runners.suite_runner import SuiteExecution + +BACKEND: Final = BackendSpec(environment_variable="LITELLM_RUST") + + +class UnitParityExclusion(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + nodeid: str + reason: str + + +class UnitParitySuite(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + python_selectors: tuple[str, ...] + exclusions: tuple[UnitParityExclusion, ...] = () + + +def run_suite(suite: UnitParitySuite, repo_root: Path, pytest_args: Sequence[str] = ()) -> SuiteExecution: + if not suite.python_selectors: + return SuiteExecution(problems=("unit parity suites must select Python tests",)) + deselections: Final = tuple(f"--deselect={exclusion.nodeid}" for exclusion in suite.exclusions) + args: Final = (*pytest_args, *deselections) + python: Final = run_python_tests(suite.python_selectors, repo_root, "python", BACKEND, args) + rust: Final = run_python_tests(suite.python_selectors, repo_root, "rust", BACKEND, args) + return SuiteExecution(problems=compare_python_runs(python, rust)) diff --git a/tests/rust-python-harness/strategies/unit_tests_parity/test_runner.py b/tests/rust-python-harness/strategies/unit_tests_parity/test_runner.py new file mode 100644 index 00000000000..a4a9524c85f --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests_parity/test_runner.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Final + +from ...shared.reporting.models import Coverage, HarnessCase, HarnessRun, RunStatus +from ...shared.reporting.strategy import SuiteCaseSpec +from ...shared.unit_runners.suite_runner import run_suites +from .runner import UnitParityExclusion, UnitParitySuite, run_suite + + +def _write_tests(tmp_path: Path, *, mismatch: bool = False, failing: bool = False) -> None: + (tmp_path / "pytest.ini").write_text("[pytest]\n") + (tmp_path / "test_api.py").write_text( + "import os\n" + "def test_decode():\n assert int('42') == 42\n" + + ("def test_backend():\n assert os.environ['LITELLM_RUST'] == '0'\n" if mismatch else "") + + ("def test_fails():\n assert False\n" if failing else "") + ) + + +def _case() -> HarnessCase: + return HarnessCase( + strategy_id="unit_tests_parity", + strategy_label="Unit test parity", + sdk_function="ocr", + spec=SuiteCaseSpec(coverage=Coverage.COMPLETE, suite="ocr"), + ) + + +def _run(case: HarnessCase, tmp_path: Path, suite: UnitParitySuite) -> tuple[int, HarnessRun]: + return run_suites((case,), tmp_path, lambda _: None, suites={"ocr": suite}, execute=run_suite) + + +def test_passes_when_both_backends_agree(tmp_path: Path) -> None: + _write_tests(tmp_path) + case: Final = _case() + + code, report = _run(case, tmp_path, UnitParitySuite(python_selectors=("test_api.py",))) + + assert code == 0, report.failures + assert report.results[case.key].status is RunStatus.PASSED + + +def test_passes_when_both_backends_fail_identically(tmp_path: Path) -> None: + _write_tests(tmp_path, failing=True) + case: Final = _case() + + code, report = _run(case, tmp_path, UnitParitySuite(python_selectors=("test_api.py",))) + + assert code == 0, report.failures + assert report.results[case.key].status is RunStatus.PASSED + + +def test_fails_when_backend_outcomes_differ(tmp_path: Path) -> None: + _write_tests(tmp_path, mismatch=True) + case: Final = _case() + + code, report = _run(case, tmp_path, UnitParitySuite(python_selectors=("test_api.py",))) + + assert code == 1 + assert report.results[case.key].status is RunStatus.FAILED + assert any("Python/Rust test outcomes differ" in detail for _, detail in report.failures) + assert any("Python only: test_api.py::test_backend [call] passed" in detail for _, detail in report.failures) + assert any("Rust only: test_api.py::test_backend [call] failed" in detail for _, detail in report.failures) + + +def test_excludes_tests_whose_contract_is_the_backend_flag(tmp_path: Path) -> None: + _write_tests(tmp_path, mismatch=True) + suite: Final = UnitParitySuite( + python_selectors=("test_api.py",), + exclusions=( + UnitParityExclusion( + nodeid="test_api.py::test_backend", + reason="The test intentionally asserts which backend is selected.", + ), + ), + ) + case: Final = _case() + + code, report = _run(case, tmp_path, suite) + + assert code == 0, report.failures + assert report.results[case.key].status is RunStatus.PASSED diff --git a/tests/rust-python-harness/strategies/unit_tests_rust/AGENTS.md b/tests/rust-python-harness/strategies/unit_tests_rust/AGENTS.md new file mode 100644 index 00000000000..250da763530 --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests_rust/AGENTS.md @@ -0,0 +1 @@ +Runs the focused native Cargo test suite for each mapped API. diff --git a/tests/rust-python-harness/strategies/unit_tests_rust/__init__.py b/tests/rust-python-harness/strategies/unit_tests_rust/__init__.py new file mode 100644 index 00000000000..8114e12ab96 --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests_rust/__init__.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +from collections.abc import Mapping +from functools import partial +from pathlib import Path +from types import MappingProxyType +from typing import Final + +from ...shared.reporting.models import SDK_FUNCTIONS, Coverage, SdkFunction +from ...shared.reporting.strategy import ( + CaseDefinition, + NotImplementedCaseSpec, + StrategyDefinition, + SuiteCaseSpec, +) +from ...shared.unit_runners.suite_runner import run_suites +from ..unit_tests_mapping.mappings import UNIT_TEST_CONTRACTS +from .reporting import render_rust_unit_results +from .runner import RustSuite, run_suite + + +RUST_SUITES: Final[Mapping[SdkFunction, RustSuite]] = MappingProxyType( + { + sdk_function: RustSuite( + cargo_manifest=contract.rust.cargo_manifest, + cargo_filter=contract.rust.cargo_filter, + cargo_package=contract.rust.cargo_package, + ) + for sdk_function, contract in UNIT_TEST_CONTRACTS.items() + } +) + + +CASES: Final[tuple[CaseDefinition, ...]] = ( + *( + CaseDefinition( + sdk_function, + SuiteCaseSpec(coverage=Coverage.COMPLETE, suite=sdk_function) + if sdk_function in RUST_SUITES + else NotImplementedCaseSpec(reason=f"No focused {sdk_function} Rust unit suite is registered."), + ) + for sdk_function in SDK_FUNCTIONS + ), +) + +STRATEGY: Final = StrategyDefinition( + id="unit_tests_rust", + order=32, + label="Unit test Rust", + description="Run the focused native Cargo test suite for each mapped API.", + directory=Path(__file__).parent, + runnable_spec=SuiteCaseSpec, + cases=CASES, + run=partial(run_suites, suites=RUST_SUITES, execute=run_suite), + render=render_rust_unit_results, +) diff --git a/tests/rust-python-harness/strategies/unit_tests_rust/reporting.py b/tests/rust-python-harness/strategies/unit_tests_rust/reporting.py new file mode 100644 index 00000000000..575fa5e8cd1 --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests_rust/reporting.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +from collections.abc import Sequence +from typing import Final + +from ...shared.reporting.models import CaseResult +from ...shared.reporting.rendering import ReportSection, render_case_outcome + + +def render_rust_unit_results(results: Sequence[CaseResult]) -> tuple[ReportSection, ...]: + blocks: Final = tuple(render_case_outcome(result) for result in results) + return (ReportSection("Native Rust unit-test outcomes", blocks or ("No Rust unit-test cases selected",)),) diff --git a/tests/rust-python-harness/strategies/unit_tests_rust/runner.py b/tests/rust-python-harness/strategies/unit_tests_rust/runner.py new file mode 100644 index 00000000000..601b5a3c96b --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests_rust/runner.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from collections.abc import Sequence +from pathlib import Path + +from pydantic import BaseModel, ConfigDict + +from ...shared.unit_runners.rust_runner import run_rust_tests +from ...shared.unit_runners.suite_runner import SuiteExecution + + +class RustSuite(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + cargo_manifest: str + cargo_package: str | None = None + cargo_filter: str + + +def run_suite(suite: RustSuite, repo_root: Path, pytest_args: Sequence[str] = ()) -> SuiteExecution: + del pytest_args + if not suite.cargo_filter: + return SuiteExecution(problems=("rust suites must configure a focused Cargo filter",)) + inventory = run_rust_tests( + repo_root / suite.cargo_manifest, suite.cargo_package, suite.cargo_filter, collect_only=True + ) + rust = run_rust_tests(repo_root / suite.cargo_manifest, suite.cargo_package, suite.cargo_filter) + return SuiteExecution( + problems=( + *(("native Rust tests did not all pass",) if set(inventory.tests) != set(rust.tests) else ()), + *((inventory.output,) if inventory.exit_code else ()), + *((rust.output,) if rust.exit_code else ()), + ) + ) diff --git a/tests/rust-python-harness/strategies/unit_tests_rust/test_runner.py b/tests/rust-python-harness/strategies/unit_tests_rust/test_runner.py new file mode 100644 index 00000000000..e151308699a --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests_rust/test_runner.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +import shutil +from collections.abc import Callable +from pathlib import Path +from typing import Final + +import pytest + +from ...shared.reporting.models import Coverage, HarnessCase, RunStatus +from ...shared.reporting.strategy import SuiteCaseSpec +from ...shared.unit_runners.suite_runner import run_suites +from .runner import RustSuite, run_suite + + +@pytest.mark.skipif(shutil.which("cargo") is None, reason="Cargo is required for the native unit strategy") +def test_runs_cargo_tests_and_propagates_ignored_or_failing_tests( + tmp_path: Path, + cargo_project: Callable[[str, str], Path], +) -> None: + cargo_project("rust-unit-check", '#[test] fn test_decode() { assert_eq!("42".parse::().unwrap(), 42); }\n') + rust_root: Final = tmp_path / "litellm-rust" + rust_root.mkdir() + (tmp_path / "Cargo.toml").rename(rust_root / "Cargo.toml") + (tmp_path / "src").rename(rust_root / "src") + suite: Final = RustSuite( + cargo_manifest="litellm-rust/Cargo.toml", + cargo_filter="test_decode", + ) + case: Final = HarnessCase( + strategy_id="unit_tests_rust", + strategy_label="Unit test Rust", + sdk_function="ocr", + spec=SuiteCaseSpec(coverage=Coverage.COMPLETE, suite="ocr"), + ) + + code, report = run_suites((case,), tmp_path, lambda _: None, suites={"ocr": suite}, execute=run_suite) + + assert code == 0, report.failures + assert report.results[case.key].status is RunStatus.PASSED + + (rust_root / "src/lib.rs").write_text("#[test] #[ignore] fn test_decode() {}\n") + ignored_code, ignored_report = run_suites( + (case,), tmp_path, lambda _: None, suites={"ocr": suite}, execute=run_suite + ) + + assert ignored_code == 1 + assert any("native Rust tests did not all pass" in detail for _, detail in ignored_report.failures) + + (rust_root / "src/lib.rs").write_text("#[test] fn test_decode() { assert_eq!(2 + 2, 5); }\n") + failed_code, failed_report = run_suites((case,), tmp_path, lambda _: None, suites={"ocr": suite}, execute=run_suite) + + assert failed_code == 1 + assert failed_report.results[case.key].status is RunStatus.FAILED + assert any("test_decode" in detail for _, detail in failed_report.failures) diff --git a/tests/sdk_function_trace/README.md b/tests/sdk_function_trace/README.md deleted file mode 100644 index d3a3b654aea..00000000000 --- a/tests/sdk_function_trace/README.md +++ /dev/null @@ -1,30 +0,0 @@ -# SDK function tracing - -The compare runner executes the same SDK calls through the Python engine and the Rust native bridge against a local HTTP provider fixture, then prints their pipeline trees side by side. Matching calls align on the same row in green; Python-only calls are blue, Rust-only calls yellow, and reordered calls red. Gaps preserve execution order and each column retains its own nesting. A comparison column labels every row even without color. Colors are enabled in terminals unless `NO_COLOR` is set. A difference summary follows (shared step order, python-only steps, rust-only steps). Each invocation must issue exactly one HTTP request. It requires the LiteLLM Python dependencies and the native extension built with tracing support - -From the repository root, using the project's Python environment: - -```bash -uv run python -m tests.sdk_function_trace.compare -uv run python -m tests.sdk_function_trace.compare --route ocr -uv run python -m tests.sdk_function_trace.compare --route ocr --sync -uv run python -m tests.sdk_function_trace.compare --route all --both --check -``` - -Calls default to async; use `--sync` for synchronous calls or `--both` for the complete matrix. Python sync Messages raises `not implemented for sync calls`; only that exact failure is marked `SKIP`, and the runner still executes Rust sync Messages and subsequent routes. Bedrock transcription has no independent Python provider implementation: its Python trace covers SDK dispatch into Rust - -Both engines are projected onto a shared per-route step table (`steps.py`): canonical names such as `transform_ocr_request` map Python functions (`MistralOCRConfig.transform_ocr_request`) and Rust spans (`transform_ocr_request`) to the same label. Only the first occurrence of each step is kept. Python indentation uses each event's actual frame ancestors and the nearest already displayed ancestor, so returned helpers and coroutine resumptions do not create false parents. Rust indentation uses instrumented span ancestry. Unmatched Rust span names pass through unchanged. `--full` prints every captured runtime event; validation still uses projected steps - -Every report checks required stage presence and dependency order. Provider lookup must precede request transformation, which must precede HTTP, followed by response transformation. The handler must precede HTTP; parameter mapping and supported-parameter checks must precede request transformation. Environment validation and URL construction, where mapped, must precede HTTP. Python transcription is checked only through native dispatch. `--check` also requires identical canonical step sequences for comparable routes and exits nonzero for missing, extra, or reordered steps, or an unexpected call failure, after finishing all selected cases - -Individual stage checks are separate from cross-language `step parity`. Passing stage checks cannot override a failing step comparison. Bedrock transcription and Python sync Messages report `UNAVAILABLE` for cross-language parity because they lack an independent Python execution to compare. Absolute nesting depth is not a cross-language gate: async Python Messages dispatches its handler onto another thread. See `route-comparison.md` for the audited matrix and remaining contract limitations - -The Python runner uses the existing `profile_python` / `sys.setprofile` collector, selecting executed code under the installed `litellm` source directory instead of maintaining a function-name allowlist. It prints source locations and qualified function names, including repeated calls. Coroutine resumptions are counted once per invocation. It profiles the current thread and threads created during the call, including the fresh async executor. Existing worker threads are not retroactively profiled; background Python calls may appear, and indentation follows selected Python stack ancestors within each thread - -The Rust runner calls the compiled PyO3 SDK entrypoints with `trace=True`. The existing `FunctionTrace` subscriber collects `#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]` spans for the route entrypoint, preparation, provider lookup, HTTP handler, and selected provider transformations. The shared `http_request` helper instruments the existing Rust send operation without changing clients, timeouts, signing, or error mapping. Function names come from the actual functions. `WithSubscriber` attaches the collector to each future across async polls. Arguments and provider payloads are not recorded in trace events. Uninstrumented functions do not appear; this is scoped instrumentation, not an exhaustive native call graph - -Tracing is opt-in: native calls without `trace=True` keep their original response shape. Traced calls return `{"response": ..., "trace": [{"function": ..., "depth": ...}]}`. The runners print only trace events. Missing native support or empty traces fail instead of falling back to source searching. The old `--repo`, `--signatures`, and `--calls` options are removed - -`profile_python(functions)` still supports direct function references for focused parity checks. `assert_function_trace_parity` compares selected Python events with Rust events supplied by an executable scenario. Successful stage checks prove the declared pipeline ran in a valid dependency order for this fixture; they do not assert identical function contracts, request bodies, responses, streaming behavior, or live-provider correctness - -Build the extension with `maturin develop` in the project's virtual environment. Then run either command above to get the executed function order diff --git a/tests/sdk_function_trace/__init__.py b/tests/sdk_function_trace/__init__.py deleted file mode 100644 index da62b8041f6..00000000000 --- a/tests/sdk_function_trace/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -from tests.sdk_function_trace.harness import ( - TraceScenario, - TraceStep, - assert_function_trace_parity, -) -from tests.sdk_function_trace.profiler import FunctionTraceEvent - -__all__ = [ - "FunctionTraceEvent", - "TraceScenario", - "TraceStep", - "assert_function_trace_parity", -] diff --git a/tests/sdk_function_trace/compare.py b/tests/sdk_function_trace/compare.py deleted file mode 100644 index 941c1b6e067..00000000000 --- a/tests/sdk_function_trace/compare.py +++ /dev/null @@ -1,45 +0,0 @@ -from __future__ import annotations - -import argparse -import os -import sys -from typing import Final - -from tests.sdk_function_trace.fixtures import ROUTES -from tests.sdk_function_trace.report import compare, render - - -def _run(route: str, asynchronous: bool, *, full: bool, colorize: bool) -> bool: - comparison: Final = compare(route, asynchronous=asynchronous) - sys.stdout.write(render(comparison, full=full, colorize=colorize)) - return comparison.passed - - -def main() -> None: - parser: Final = argparse.ArgumentParser(description="Compare Python and Rust SDK pipeline steps per route") - parser.add_argument("--route", choices=("all", *ROUTES), default="all") - mode: Final = parser.add_mutually_exclusive_group() - mode.add_argument("--async", dest="asynchronous", action="store_true", default=True) - mode.add_argument("--sync", dest="asynchronous", action="store_false") - mode.add_argument("--both", action="store_true", help="run async and sync for every selected route") - parser.add_argument( - "--check", action="store_true", help="exit nonzero for missing, extra, or reordered comparable steps" - ) - parser.add_argument( - "--full", action="store_true", help="print every captured runtime event instead of pipeline steps" - ) - args: Final = parser.parse_args() - os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "True") - colorize: Final = sys.stdout.isatty() and "NO_COLOR" not in os.environ - results: Final = tuple( - _run(selected, selected_mode, full=args.full, colorize=colorize) - for selected in ROUTES - if args.route in ("all", selected) - for selected_mode in ((True, False) if args.both else (args.asynchronous,)) - ) - if args.check and not all(results): - raise SystemExit(1) - - -if __name__ == "__main__": - main() diff --git a/tests/sdk_function_trace/fixtures.py b/tests/sdk_function_trace/fixtures.py deleted file mode 100644 index 47bbe839627..00000000000 --- a/tests/sdk_function_trace/fixtures.py +++ /dev/null @@ -1,200 +0,0 @@ -from __future__ import annotations - -import base64 -import io -import json -import wave -from collections.abc import Callable -from dataclasses import dataclass -from typing import Final, Protocol, cast - -from tests.sdk_function_trace.mock_provider import MockProviderResponse -from tests.sdk_function_trace.steps import Engine - -ANTHROPIC_MODEL: Final = "claude-sonnet-5" -OCR_MODEL: Final = "mistral-ocr-latest" -AUDIO_MODEL: Final = "mistral.voxtral-mini-3b-2507" - - -class SdkCall(Protocol): - def __call__(self, **kwargs: object) -> object: ... - - -@dataclass(frozen=True, slots=True) -class Fixture: - kwargs: dict[str, object] - provider_response: MockProviderResponse - - -@dataclass(frozen=True, slots=True) -class RouteSpec: - label: str - python_entrypoints: tuple[str, str] - rust_entrypoints: tuple[str, str] - fixture: Callable[[Engine], Fixture] - - -@dataclass(frozen=True, slots=True) -class Invocation: - function: SdkCall - kwargs: dict[str, object] - provider_response: MockProviderResponse - label: str - - -def audio_bytes() -> bytes: - with io.BytesIO() as buffer: - with wave.open(buffer, "wb") as audio: - audio.setnchannels(1) - audio.setsampwidth(2) - audio.setframerate(16000) - audio.writeframes(b"\x00\x00" * 1600) - return buffer.getvalue() - - -def _anthropic_message_response() -> MockProviderResponse: - body: Final = { - "id": "msg_trace", - "type": "message", - "role": "assistant", - "model": ANTHROPIC_MODEL, - "content": [{"type": "text", "text": "hello"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 2, "output_tokens": 3}, - } - return MockProviderResponse(200, (("content-type", "application/json"),), json.dumps(body).encode()) - - -def _conversation() -> dict[str, object]: - return {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 16} - - -def _ocr_fixture(engine: Engine) -> Fixture: - return Fixture( - kwargs={ - "model": f"mistral/{OCR_MODEL}", - "document": {"type": "document_url", "document_url": "https://example.com/document.pdf"}, - **({"optional_params": {"pages": [0]}} if engine == "rust" else {"pages": [0]}), - }, - provider_response=MockProviderResponse( - 200, - (("content-type", "application/json"),), - json.dumps( - { - "pages": [{"index": 0, "markdown": "hello"}], - "model": OCR_MODEL, - "usage_info": {"pages_processed": 1}, - } - ).encode(), - ), - ) - - -def _chat_completions_fixture(engine: Engine) -> Fixture: - conversation: Final = _conversation() - payload: Final = ( - {"messages": conversation["messages"], "optional_params": {"max_tokens": 16}} - if engine == "rust" - else conversation - ) - return Fixture( - kwargs={"model": f"anthropic/{ANTHROPIC_MODEL}", **payload}, - provider_response=_anthropic_message_response(), - ) - - -def _messages_fixture(engine: Engine) -> Fixture: - conversation: Final = _conversation() - payload: Final = {"body": {**conversation, "model": ANTHROPIC_MODEL}} if engine == "rust" else conversation - return Fixture( - kwargs={"model": f"anthropic/{ANTHROPIC_MODEL}", **payload}, - provider_response=_anthropic_message_response(), - ) - - -def _transcription_fixture(engine: Engine) -> Fixture: - credentials: Final = { - "aws_access_key_id": "test-access", - "aws_secret_access_key": "test-secret", - "aws_region_name": "us-east-1", - } - payload: Final = ( - { - "audio": {"data": base64.b64encode(audio_bytes()).decode(), "format": "wav"}, - "optional_params": credentials, - } - if engine == "rust" - else {"file": ("sample.wav", audio_bytes(), "audio/wav"), **credentials} - ) - return Fixture( - kwargs={"model": f"bedrock/{AUDIO_MODEL}", **payload}, - provider_response=MockProviderResponse( - 200, - (("content-type", "application/json"),), - json.dumps( - { - "output": {"message": {"role": "assistant", "content": [{"text": "hello"}]}}, - "stopReason": "end_turn", - "usage": {"inputTokens": 2, "outputTokens": 3, "totalTokens": 5}, - } - ).encode(), - ), - ) - - -ROUTE_SPECS: Final[dict[str, RouteSpec]] = { - "chat_completions": RouteSpec( - label="anthropic", - python_entrypoints=("completion", "acompletion"), - rust_entrypoints=("chat_completions", "achat_completions"), - fixture=_chat_completions_fixture, - ), - "audio_transcription": RouteSpec( - label="bedrock (Rust-only provider; Python trace covers SDK dispatch)", - python_entrypoints=("transcription", "atranscription"), - rust_entrypoints=("transcription", "atranscription"), - fixture=_transcription_fixture, - ), - "messages": RouteSpec( - label="anthropic", - python_entrypoints=("create", "acreate"), - rust_entrypoints=("messages", "amessages"), - fixture=_messages_fixture, - ), - "ocr": RouteSpec( - label="mistral", - python_entrypoints=("ocr", "aocr"), - rust_entrypoints=("ocr", "aocr"), - fixture=_ocr_fixture, - ), -} - -ROUTES: Final = tuple(ROUTE_SPECS) - - -def sdk_invocation(route: str, *, engine: Engine, asynchronous: bool) -> Invocation: - import litellm - from litellm.anthropic_interface import messages as sdk_messages - from litellm.rust_bridge import get_native_bridge - - rust: Final = engine == "rust" - bridge: Final = get_native_bridge() if rust else None - if rust and bridge is None: - raise RuntimeError("Build the native extension first: maturin develop") - spec: Final = ROUTE_SPECS.get(route) - if spec is None: - raise ValueError(f"Unknown route: {route}") - fixture: Final = spec.fixture(engine) - owner: Final = bridge if rust else (sdk_messages if route == "messages" else litellm) - entrypoint: Final = (spec.rust_entrypoints if rust else spec.python_entrypoints)[int(asynchronous)] - return Invocation( - function=cast(SdkCall, getattr(owner, entrypoint)), - kwargs={ - **fixture.kwargs, - "api_key": "test-key", - **({"trace": True, "timeout_seconds": 5} if rust else {"timeout": 5}), - }, - provider_response=fixture.provider_response, - label=spec.label, - ) diff --git a/tests/sdk_function_trace/harness.py b/tests/sdk_function_trace/harness.py deleted file mode 100644 index 8f707402449..00000000000 --- a/tests/sdk_function_trace/harness.py +++ /dev/null @@ -1,39 +0,0 @@ -from __future__ import annotations - -from collections.abc import Callable, Sequence -from dataclasses import dataclass -from types import FunctionType -from typing import Final, cast - -from tests.sdk_function_trace.profiler import FunctionTraceEvent, profile_python - - -@dataclass(frozen=True, slots=True) -class TraceStep: - function: FunctionType - depth: int - - -@dataclass(frozen=True, slots=True) -class TraceScenario: - steps: tuple[TraceStep, ...] - invoke_python: Callable[[], object] - invoke_rust: Callable[[], Sequence[FunctionTraceEvent]] - - -def assert_function_trace_parity(scenario: TraceScenario) -> None: - expected: Final = tuple( - FunctionTraceEvent(function=step.function.__name__, depth=step.depth) for step in scenario.steps - ) - functions: Final = cast(tuple[FunctionType, ...], tuple(step.function for step in scenario.steps)) - with profile_python(functions) as profiler: - scenario.invoke_python() - python_trace: Final = tuple(profiler.events) - rust_trace: Final = tuple(scenario.invoke_rust()) - - if python_trace != expected: - raise AssertionError(f"Python function trace differs: {python_trace!r} != {expected!r}") - if rust_trace != expected: - raise AssertionError(f"Rust function trace differs: {rust_trace!r} != {expected!r}") - if python_trace != rust_trace: - raise AssertionError(f"Python and Rust function traces differ: {python_trace!r} != {rust_trace!r}") diff --git a/tests/sdk_function_trace/mock_provider.py b/tests/sdk_function_trace/mock_provider.py deleted file mode 100644 index 37eca665586..00000000000 --- a/tests/sdk_function_trace/mock_provider.py +++ /dev/null @@ -1,67 +0,0 @@ -from __future__ import annotations - -from collections.abc import Generator -from contextlib import contextmanager -from dataclasses import dataclass -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from threading import Lock, Thread -from typing import Final, cast - - -@dataclass(frozen=True, slots=True) -class MockProviderResponse: - status_code: int - headers: tuple[tuple[str, str], ...] - body: bytes - - -class _MockProviderServer(ThreadingHTTPServer): - def __init__(self, response: MockProviderResponse) -> None: - super().__init__(("127.0.0.1", 0), _MockProviderHandler) - self.response: Final = response - self._request_count = 0 - self._request_count_lock: Final = Lock() - - def record_request(self) -> None: - with self._request_count_lock: - self._request_count += 1 - - @property - def request_count(self) -> int: - with self._request_count_lock: - return self._request_count - - -class _MockProviderHandler(BaseHTTPRequestHandler): - protocol_version = "HTTP/1.1" - - def do_POST(self) -> None: - content_length: Final = int(self.headers.get("content-length", "0")) - self.rfile.read(content_length) - server: Final = cast(_MockProviderServer, self.server) - server.record_request() - self.send_response(server.response.status_code) - for name, value in server.response.headers: - self.send_header(name, value) - self.send_header("content-length", str(len(server.response.body))) - self.end_headers() - self.wfile.write(server.response.body) - - def log_message(self, format: str, *args: object) -> None: # noqa: A002 # matches BaseHTTPRequestHandler - pass - - -@contextmanager -def mock_provider(response: MockProviderResponse) -> Generator[str]: - server: Final = _MockProviderServer(response) - thread: Final = Thread(target=server.serve_forever, daemon=True) - thread.start() - host, port = cast(tuple[str, int], server.server_address) - try: - yield f"http://{host}:{port}" - finally: - server.shutdown() - server.server_close() - thread.join() - if server.request_count != 1: - raise AssertionError(f"expected one provider request, received {server.request_count}") diff --git a/tests/sdk_function_trace/ocr-comparison.md b/tests/sdk_function_trace/ocr-comparison.md deleted file mode 100644 index d252480e218..00000000000 --- a/tests/sdk_function_trace/ocr-comparison.md +++ /dev/null @@ -1,59 +0,0 @@ -# OCR Python and Rust comparison - -Audited implementation revision: `edcba483b2`. The implementations do not match in function contracts, call structure, or all tested response behavior. This audit changes the source listing coverage, not OCR runtime behavior - -Run both source listings from the repository root: - -```bash -python3 tests/sdk_function_trace/list_python_steps.py --route ocr --signatures --calls -uv run tests/sdk_function_trace/list_rust_steps.py --route ocr --signatures --calls -``` - -Both cover Mistral, Azure AI Mistral, Azure Document Intelligence, Vertex Mistral, and Vertex DeepSeek. Listings show declarations and source call sites, not executed traces - -## Function contracts - -Comparing Python `BaseOCRConfig` with Rust `OcrProviderConfig`, omitting `self` and language-specific ownership details: - -| Python | Rust | Difference | -| --- | --- | --- | -| `get_supported_ocr_params(model)` | `supported_ocr_params()` | Name and model argument | -| `get_api_key_env_var()` | No corresponding method | Missing contract | -| `map_ocr_params(non_default_params, optional_params, model)` | `map_ocr_params(non_default_params)` | Missing accumulator and model | -| `validate_environment(headers, model, api_key, api_base, litellm_params, **kwargs)` | Separate auth/key/header helpers | Different contract | -| `get_complete_url(api_base, model, optional_params, litellm_params, **kwargs)` | `complete_url(api_base, model, optional_params, env_lookup)` | Name and context | -| `transform_ocr_request(model, document, optional_params, headers, **kwargs)` | `transform_ocr_request(model, document, optional_params)` | Missing headers and extra context | -| `async_transform_ocr_request(...)` | No corresponding method | Missing async override | -| `transform_ocr_response(model, raw_response, logging_obj, **kwargs)` | `transform_ocr_response(model, response_json)` | Missing HTTP metadata, logging and extra context | -| `async_transform_ocr_response(...)` | No corresponding method | Missing async override | -| `get_error_class(error_message, status_code, headers)` | Central Rust error mapping | Different contract | - -Python's default mapper returns the supplied `optional_params`; Rust's filters `non_default_params`. Provider overrides must also be compared - -Python maps parameters during SDK preparation, before HTTP-handler environment validation and URL construction. Rust resolves auth and URL before mapping parameters in `prepare_provider_request`. Python has async provider transforms; both native entrypoints execute the same Rust async route using synchronous transform hooks, with polling and document downloading in gateway helpers - -The native bindings also accept `optional_params` and `timeout_seconds`, while the Python SDK accepts `**kwargs` and `timeout`. Public SDK calls with Rust enabled still execute Python preparation before entering Rust, so matching SDK responses would not prove matching standalone Rust steps - -## Runtime results - -Built the native extension from the audited source using `cargo build -p litellm-python-bridge --features extension-module --offline`. Supplied that build's functions through `use_litellm_rust` dependency injection. Ran public `litellm.ocr` and `litellm.aocr` with Rust disabled and enabled against identical local HTTP response fixtures, requiring one request per invocation - -Successful `model_dump()` results and failure exception classes were compared. These checks cover Mistral response outcomes only, not request equality, error messages, live providers, or every execution branch - -| Mistral response fixture | Sync | Async | Observation | -| --- | --- | --- | --- | -| Valid page/model/usage | Match | Match | Same normalized response | -| Model omitted | Match | Match | Both use the requested model | -| `model: null` | Different | Different | Python rejects; Rust uses the requested model | -| `pages: null` | Different | Different | Python rejects; Rust returns an empty array | -| Invalid page element | Match | Match | Both reject during response validation | - -Six of ten fixture/mode comparisons match, four differ. Rust's Mistral response transform conflates missing values with explicit nulls through `as_array`/`as_str` fallbacks. Python preserves explicit nulls into response validation, which rejects them - -## Other provider gaps found in source - -Azure Document Intelligence's Python configuration supports `pages`, `features`, and `req_format`; Rust lists only `pages`. Python normalizes parameters before URL construction; Rust normalizes pages during URL construction - -Python preserves Azure `content`, `tables`, and `keyValuePairs`, and supports retaining the native operation payload. Rust's `OcrResponseData` has no corresponding fields, and its Azure transform does not preserve those values - -Azure and Vertex async document transforms and Azure polling also use different helper contracts. Their runtime equivalence was not tested in this audit diff --git a/tests/sdk_function_trace/profiler.py b/tests/sdk_function_trace/profiler.py deleted file mode 100644 index c71c74ab0d3..00000000000 --- a/tests/sdk_function_trace/profiler.py +++ /dev/null @@ -1,76 +0,0 @@ -from __future__ import annotations - -import sys -import threading -from collections.abc import Generator, Sequence -from contextlib import contextmanager -from dataclasses import dataclass -from pathlib import Path -from types import CodeType, FrameType, FunctionType -from typing import Final - - -@dataclass(frozen=True, slots=True) -class FunctionTraceEvent: - function: str - depth: int - ancestors: tuple[str, ...] | None = None - - -class PythonProfiler: - def __init__(self, functions: Sequence[FunctionType], source_root: Path | None = None) -> None: - self._source_root: Final = str(source_root.resolve()) + "/" if source_root is not None else None - self._names_by_code: Final = {function.__code__: function.__name__ for function in functions} - self._seen_frames: Final[set[FrameType]] = set() - self.events: Final[list[FunctionTraceEvent]] = [] - - def __call__(self, frame: FrameType, event: str, _arg: object) -> None: - if event != "call" or frame in self._seen_frames: - return - function_name: Final = self.function_name(frame.f_code) - if function_name is None: - return - ancestors: Final = tuple( - name for ancestor in _frame_ancestors(frame) if (name := self.function_name(ancestor.f_code)) is not None - ) - self._seen_frames.add(frame) - self.events.append( - FunctionTraceEvent( - function=function_name, - depth=len(ancestors), - ancestors=ancestors if self._source_root is not None else None, - ) - ) - - def function_name(self, code: CodeType) -> str | None: - if self._source_root is None: - return self._names_by_code.get(code) - if not code.co_filename.startswith(self._source_root): - return None - relative: Final = code.co_filename.removeprefix(self._source_root) - return f"{relative}:{code.co_firstlineno} {getattr(code, 'co_qualname', code.co_name)}" - - -def _frame_ancestors(frame: FrameType) -> Generator[FrameType]: - ancestor: Final = frame.f_back - if ancestor is not None: - yield ancestor - yield from _frame_ancestors(ancestor) - - -@contextmanager -def profile_python( - functions: Sequence[FunctionType] = (), *, source_root: Path | None = None, threads: bool = False -) -> Generator[PythonProfiler]: - profiler: Final = PythonProfiler(functions, source_root) - previous_thread: Final = threading.getprofile() - if threads: - threading.setprofile(profiler) - previous: Final = sys.getprofile() - sys.setprofile(profiler) - try: - yield profiler - finally: - sys.setprofile(previous) - if threads: - threading.setprofile(previous_thread) diff --git a/tests/sdk_function_trace/report.py b/tests/sdk_function_trace/report.py deleted file mode 100644 index 9b654e571f8..00000000000 --- a/tests/sdk_function_trace/report.py +++ /dev/null @@ -1,175 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import Final - -from tests.sdk_function_trace.fixtures import ROUTE_SPECS -from tests.sdk_function_trace.profiler import FunctionTraceEvent -from tests.sdk_function_trace.runtime import ( - TraceDiff, - TraceFailed, - TraceOk, - TraceRun, - TraceSkipped, - attempt_trace, - trace_diff, -) -from tests.sdk_function_trace.steps import Engine, pipeline_issues, pipeline_steps -from tests.sdk_function_trace.table import format_trace_table - -_PYTHON_ONLY_COLOR: Final = "\033[34m" -_RUST_ONLY_COLOR: Final = "\033[33m" -_RESET: Final = "\033[0m" - -_ENGINE_COLOR: Final[dict[Engine, str]] = {"python": _PYTHON_ONLY_COLOR, "rust": _RUST_ONLY_COLOR} - - -@dataclass(frozen=True, slots=True) -class EngineReport: - engine: Engine - run: TraceRun - events: tuple[FunctionTraceEvent, ...] - steps: tuple[FunctionTraceEvent, ...] - issues: tuple[str, ...] - - -@dataclass(frozen=True, slots=True) -class Comparison: - route: str - label: str - asynchronous: bool - engines: tuple[EngineReport, ...] - diff: TraceDiff - - @property - def comparable(self) -> bool: - return self.route != "audio_transcription" and all(isinstance(report.run, TraceOk) for report in self.engines) - - @property - def passed(self) -> bool: - return ( - (not self.comparable or self.diff.matches) - and not any(report.issues for report in self.engines) - and all(not isinstance(report.run, TraceFailed) for report in self.engines) - ) - - -def _events(run: TraceRun) -> tuple[FunctionTraceEvent, ...]: - match run: - case TraceOk(events=events): - return events - case TraceSkipped() | TraceFailed(): - return () - - -def _engine_report(route: str, engine: Engine, run: TraceRun) -> EngineReport: - events: Final = _events(run) - steps: Final = pipeline_steps(route, engine, events) - issues: Final = pipeline_issues(route, engine, steps) if isinstance(run, TraceOk) else () - return EngineReport(engine=engine, run=run, events=events, steps=steps, issues=issues) - - -def compare(route: str, *, asynchronous: bool) -> Comparison: - runs: Final = { - engine: attempt_trace(route, engine=engine, asynchronous=asynchronous) for engine in ("python", "rust") - } - engines: Final = tuple(_engine_report(route, engine, run) for engine, run in runs.items()) - return Comparison( - route=route, - label=ROUTE_SPECS[route].label, - asynchronous=asynchronous, - engines=engines, - diff=trace_diff(engines[0].steps, engines[1].steps), - ) - - -def _tree_line(event: FunctionTraceEvent, only: frozenset[str], marker: str, color: str, *, colorize: bool) -> str: - line: Final = f"{' ' * event.depth}{event.function}" + (f" {marker}" if event.function in only else "") - return f"{color}{line}{_RESET}\n" if colorize and event.function in only else f"{line}\n" - - -def _tree_lines( - events: tuple[FunctionTraceEvent, ...], - only: frozenset[str], - marker: str, - color: str, - *, - colorize: bool, -) -> tuple[str, ...]: - return tuple(_tree_line(event, only, marker, color, colorize=colorize) for event in events) - - -def _engine_lines( - report: EngineReport, diff: TraceDiff, *, comparable: bool, full: bool, colorize: bool -) -> tuple[str, ...]: - match report.run: - case TraceSkipped(reason=reason): - return (f"{report.engine}: SKIP ({reason})\n\n",) - case TraceFailed(reason=reason): - return (f"{report.engine}: FAIL ({reason})\n\n",) - case TraceOk(): - shown: Final = report.events if full else report.steps - only: Final = ( - () if full or not comparable else (diff.python_only if report.engine == "python" else diff.rust_only) - ) - return ( - f"{report.engine} ({len(shown)} steps)\n\n", - *_tree_lines( - shown, - frozenset(only), - f"<- {report.engine} only", - _ENGINE_COLOR[report.engine], - colorize=colorize, - ), - "\n", - ) - - -def _parity_lines(comparison: Comparison) -> tuple[str, ...]: - if not comparison.comparable: - if comparison.route == "audio_transcription": - return ("step parity: UNAVAILABLE (Bedrock transcription has no independent Python implementation)\n",) - return ("step parity: UNAVAILABLE (both engines must complete)\n",) - diff: Final = comparison.diff - order: Final = "the same" if diff.shared_order_matches else "a different" - return ( - "diff\n\n", - f"shared steps appear in {order} order\n", - f"python-only: {', '.join(diff.python_only) or 'none'}\n", - f"rust-only: {', '.join(diff.rust_only) or 'none'}\n\n", - f"step parity: {'PASS' if diff.matches else 'FAIL'}\n", - ) - - -def _stage_lines(comparison: Comparison) -> tuple[str, ...]: - return tuple( - f"{report.engine} " - f"{'SDK dispatch only' if comparison.route == 'audio_transcription' and report.engine == 'python' else 'pipeline'}: " - f"{'FAIL: ' + '; '.join(report.issues) if report.issues else 'PASS'}\n" - for report in comparison.engines - if isinstance(report.run, TraceOk) - ) - - -def render(comparison: Comparison, *, full: bool, colorize: bool) -> str: - mode: Final = "async" if comparison.asynchronous else "sync" - traces: Final = ( - (format_trace_table(comparison.engines[0].steps, comparison.engines[1].steps, colorize=colorize) + "\n\n",) - if not full and all(isinstance(report.run, TraceOk) for report in comparison.engines) - else tuple( - line - for report in comparison.engines - for line in _engine_lines( - report, comparison.diff, comparable=comparison.comparable, full=full, colorize=colorize - ) - ) - ) - return "".join( - ( - f"route: {comparison.route} provider: {comparison.label} mode: {mode}\n\n", - *traces, - *_parity_lines(comparison), - *_stage_lines(comparison), - "Each successful invocation issued exactly one local provider request\n\n", - ) - ) diff --git a/tests/sdk_function_trace/route-comparison.md b/tests/sdk_function_trace/route-comparison.md deleted file mode 100644 index 009d3544d05..00000000000 --- a/tests/sdk_function_trace/route-comparison.md +++ /dev/null @@ -1,26 +0,0 @@ -# SDK route trace audit - -Run the four native HTTP route families in both modes from the repository root: - -```bash -uv run python -m tests.sdk_function_trace.compare --route all --both --check -``` - -The local fixture matrix on 2026-09-02 completed 15 successful engine invocations and one expected skip. Every successful invocation issued exactly one local HTTP request. All five comparable route/mode pairs have identical canonical steps in the same order, with no Python-only or Rust-only steps - -| Route | Python async | Python sync | Rust async | Rust sync | -| --- | --- | --- | --- | --- | -| Chat completions, Anthropic | Pass | Pass | Pass | Pass | -| Messages, Anthropic | Pass | Unsupported, skipped | Pass | Pass | -| OCR, Mistral | Pass | Pass | Pass | Pass | -| Audio transcription, Bedrock | Dispatch only | Dispatch only | Pass | Pass | - -The same canonical step sequence ran in sync and async for each engine with both modes available. Bedrock transcription's Python SDK delegates to Rust, so its two successful calls do not establish independent provider parity. Realtime and Responses WebSockets are outside this HTTP fixture runner - -Chat and OCR also have identical projected nesting in both modes. Async Messages has the same helper nesting beneath its handler, but Python starts that handler on a worker thread, so it appears as a second root. The comparison preserves this physical thread boundary and checks step order independently of absolute depth - -Rust now resolves chat providers and supported parameters before entering its handler. Chat and Messages validate the environment and transform requests inside their handlers. Messages builds the final URL after transformation. OCR resolves its config and maps supported parameters during preparation, then validates credentials, builds the URL, and transforms the request inside its handler. Its during-call guardrails still run before HTTP, within the provider-call lifecycle phase - -The environment hooks execute credential and header validation. Chat's supported-parameter hooks return OpenAI names paired with provider names and feed the existing request acceptance checks. The direct Rust API still accepts provider-mapped parameters, and its supported subset is smaller than Python's. Matching the pipeline does not establish identical parameter contracts - -`--check` now fails if either comparable engine has missing, extra, or reordered canonical steps, even if its individual stage checks pass. Bedrock transcription and sync Messages report `UNAVAILABLE` for cross-language parity; native execution is still checked. Passing establishes step coverage and order for one non-streaming fixture per route, not complete request, response, error, or provider parity. The previously recorded OCR response gaps remain in `ocr-comparison.md` diff --git a/tests/sdk_function_trace/runtime.py b/tests/sdk_function_trace/runtime.py deleted file mode 100644 index d5bf15694bc..00000000000 --- a/tests/sdk_function_trace/runtime.py +++ /dev/null @@ -1,128 +0,0 @@ -from __future__ import annotations - -import asyncio -import os -from collections.abc import Awaitable, Generator -from contextlib import contextmanager -from dataclasses import dataclass -from pathlib import Path -from typing import Final, cast -from unittest.mock import patch - -from pydantic import BaseModel, ConfigDict - -from tests.sdk_function_trace.fixtures import Invocation, sdk_invocation -from tests.sdk_function_trace.mock_provider import mock_provider -from tests.sdk_function_trace.profiler import FunctionTraceEvent, profile_python -from tests.sdk_function_trace.steps import Engine - - -class TraceEventPayload(BaseModel): - model_config = ConfigDict(strict=True, extra="forbid") - function: str - depth: int - - -class TraceResponsePayload(BaseModel): - model_config = ConfigDict(strict=True, extra="forbid") - response: object - trace: tuple[TraceEventPayload, ...] | list[TraceEventPayload] - - -@contextmanager -def _python_engine() -> Generator[None]: - from litellm.rust_bridge import ocr as ocr_bridge - - previous_ocr: Final = ocr_bridge.rust_ocr_enabled() - with patch.dict(os.environ, {"LITELLM_RUST": "false"}): - ocr_bridge.use_litellm_rust(False) - try: - yield - finally: - ocr_bridge.use_litellm_rust(previous_ocr) - - -def _invoke(case: Invocation, api_base: str, *, asynchronous: bool) -> object: - async def invoke_async() -> object: - return await cast("Awaitable[object]", case.function(**case.kwargs, api_base=api_base)) - - if asynchronous: - return asyncio.run(invoke_async()) - return case.function(**case.kwargs, api_base=api_base) - - -def collect(case: Invocation, api_base: str, *, engine: Engine, asynchronous: bool) -> tuple[FunctionTraceEvent, ...]: - import litellm - - if engine == "rust": - payload: Final = TraceResponsePayload.model_validate(_invoke(case, api_base, asynchronous=asynchronous)) - return tuple(FunctionTraceEvent(event.function, event.depth) for event in payload.trace) - with profile_python(source_root=Path(litellm.__file__).parent, threads=True) as profiler: - _invoke(case, api_base, asynchronous=asynchronous) - return tuple(profiler.events) - - -def run_trace(route: str, *, engine: Engine, asynchronous: bool = False) -> tuple[FunctionTraceEvent, ...]: - case: Final = sdk_invocation(route, engine=engine, asynchronous=asynchronous) - with _python_engine(), mock_provider(case.provider_response) as api_base: - events: Final = collect(case, api_base, engine=engine, asynchronous=asynchronous) - if not events: - raise RuntimeError(f"No runtime events for {route}; rebuild the native extension with tracing support") - return events - - -@dataclass(frozen=True, slots=True) -class TraceOk: - events: tuple[FunctionTraceEvent, ...] - - -@dataclass(frozen=True, slots=True) -class TraceSkipped: - reason: str - - -@dataclass(frozen=True, slots=True) -class TraceFailed: - reason: str - - -TraceRun = TraceOk | TraceSkipped | TraceFailed - - -def attempt_trace(route: str, *, engine: Engine, asynchronous: bool) -> TraceRun: - try: - return TraceOk(run_trace(route, engine=engine, asynchronous=asynchronous)) - except Exception as error: - reason: Final = f"{type(error).__name__}: {error}" - if ( - route == "messages" - and engine == "python" - and not asynchronous - and isinstance(error, ValueError) - and str(error) == "anthropic_messages_handler is not implemented for sync calls" - ): - return TraceSkipped(reason) - return TraceFailed(reason) - - -@dataclass(frozen=True, slots=True) -class TraceDiff: - python_only: tuple[str, ...] - rust_only: tuple[str, ...] - shared_order_matches: bool - - @property - def matches(self) -> bool: - return not self.python_only and not self.rust_only and self.shared_order_matches - - -def trace_diff(python: tuple[FunctionTraceEvent, ...], rust: tuple[FunctionTraceEvent, ...]) -> TraceDiff: - python_names: Final = {event.function for event in python} - rust_names: Final = {event.function for event in rust} - shared_python: Final = tuple(event.function for event in python if event.function in rust_names) - shared_rust: Final = tuple(event.function for event in rust if event.function in python_names) - return TraceDiff( - python_only=tuple(event.function for event in python if event.function not in rust_names), - rust_only=tuple(event.function for event in rust if event.function not in python_names), - shared_order_matches=bool(shared_python) and shared_python == shared_rust, - ) diff --git a/tests/sdk_function_trace/steps.py b/tests/sdk_function_trace/steps.py deleted file mode 100644 index bb50d4ebe57..00000000000 --- a/tests/sdk_function_trace/steps.py +++ /dev/null @@ -1,181 +0,0 @@ -from __future__ import annotations - -import re -from collections.abc import Sequence -from dataclasses import dataclass -from functools import reduce -from typing import Final, Literal - -from tests.sdk_function_trace.profiler import FunctionTraceEvent - -Engine = Literal["python", "rust"] - - -@dataclass(frozen=True, slots=True) -class Step: - name: str - python: re.Pattern[str] | None - rust: str | None - - -def _step(name: str, python: str | None = None, rust: str | None = None) -> Step: - return Step(name, re.compile(python) if python is not None else None, rust) - - -_POST: Final = r"AsyncHTTPHandler\.post$|HTTPHandler\.post$" - -STEPS: Final[dict[str, tuple[Step, ...]]] = { - "ocr": ( - _step("ocr", r"ocr/main\.py:\d+ a?ocr$", "ocr"), - _step("prepare_ocr_call", r"ocr/main\.py:\d+ _prepare_ocr_request$", "prepare_ocr_call"), - _step("get_provider_ocr_config", r"ProviderConfigManager\.get_provider_ocr_config$", "ocr_provider_config"), - _step("supported_ocr_params", r"get_supported_ocr_params$", "supported_ocr_params"), - _step("map_ocr_params", r"(? tuple[str, ...]: - names: Final = tuple(event.function for event in events) - required: Final = tuple(step.name for step in STEPS[route] if getattr(step, engine) is not None) - missing: Final = tuple(f"missing {name}" for name in required if name not in names) - provider: Final = next(name for name in required if name.startswith("get_provider_")) - handler: Final = next(name for name in required if name.startswith("execute_")) - dispatch_only: Final = route == "audio_transcription" and engine == "python" - request: Final = next( - (name for name in required if name.startswith("transform_") and name.endswith("request")), handler - ) - response: Final = next( - (name for name in required if name.startswith("transform_") and name.endswith("response")), handler - ) - phases: Final = ( - (route, "map_transcription_params", provider, handler) - if dispatch_only - else (route, provider, request, "http_request", response) - ) - extra_edges: Final = ( - () - if dispatch_only - else ( - (handler, "http_request"), - *((name, request) for name in required if name.startswith(("map_", "supported_"))), - *((name, "http_request") for name in ("validate_environment", "complete_url") if name in required), - ) - ) - edges: Final = (*zip(phases, phases[1:]), *extra_edges) - return missing + tuple( - f"{before} must precede {after}" - for before, after in edges - if before in names and after in names and names.index(before) >= names.index(after) - ) - - -def _canonical_name(route: str, engine: Engine, function: str) -> str | None: - for step in STEPS[route]: - if engine == "python": - if step.python is not None and step.python.search(function): - return step.name - elif step.rust is not None and function == step.rust: - return step.name - return function if engine == "rust" else None - - -@dataclass(frozen=True, slots=True) -class _Projection: - shown: tuple[FunctionTraceEvent, ...] = () - stack: tuple[tuple[int, int], ...] = () - seen: frozenset[str] = frozenset() - - -def _project(route: str, engine: Engine, state: _Projection, event: FunctionTraceEvent) -> _Projection: - stack: Final = tuple(pair for pair in state.stack if event.depth > pair[0]) - name: Final = _canonical_name(route, engine, event.function) - if name is None or name in state.seen: - return _Projection(state.shown, stack, state.seen) - depth: Final = ( - next( - ( - kept.depth + 1 - for ancestor in event.ancestors - for kept in state.shown - if kept.function == _canonical_name(route, engine, ancestor) - ), - 0, - ) - if event.ancestors is not None - else stack[-1][1] + 1 - if stack - else 0 - ) - return _Projection( - state.shown + (FunctionTraceEvent(function=name, depth=depth),), - stack + ((event.depth, depth),), - state.seen | {name}, - ) - - -def pipeline_steps(route: str, engine: Engine, events: Sequence[FunctionTraceEvent]) -> tuple[FunctionTraceEvent, ...]: - projection: Final = reduce(lambda state, event: _project(route, engine, state, event), events, _Projection()) - return projection.shown diff --git a/tests/sdk_function_trace/table.py b/tests/sdk_function_trace/table.py deleted file mode 100644 index 2124d7e3faf..00000000000 --- a/tests/sdk_function_trace/table.py +++ /dev/null @@ -1,72 +0,0 @@ -from __future__ import annotations - -from collections.abc import Iterator -from difflib import SequenceMatcher -from typing import Final - -from tests.sdk_function_trace.profiler import FunctionTraceEvent - - -def _aligned_rows( - python: tuple[FunctionTraceEvent, ...], rust: tuple[FunctionTraceEvent, ...] -) -> Iterator[tuple[FunctionTraceEvent | None, FunctionTraceEvent | None]]: - matcher: Final = SequenceMatcher( - a=tuple(event.function for event in python), - b=tuple(event.function for event in rust), - autojunk=False, - ) - for tag, python_start, python_end, rust_start, rust_end in matcher.get_opcodes(): - if tag == "equal": - yield from zip(python[python_start:python_end], rust[rust_start:rust_end]) - else: - yield from ((event, None) for event in python[python_start:python_end]) - yield from ((None, event) for event in rust[rust_start:rust_end]) - - -def _label(event: FunctionTraceEvent | None) -> str: - return f"{' ' * event.depth}{event.function}" if event is not None else "" - - -def _status( - python: FunctionTraceEvent | None, - rust: FunctionTraceEvent | None, - python_names: frozenset[str], - rust_names: frozenset[str], -) -> tuple[str, str]: - if python is not None and rust is not None: - return "match", "\033[32m" - if python is not None: - return ("reordered", "\033[31m") if python.function in rust_names else ("python only", "\033[34m") - if rust is not None: - return ("reordered", "\033[31m") if rust.function in python_names else ("rust only", "\033[33m") - return "", "" - - -def format_trace_table( - python: tuple[FunctionTraceEvent, ...], - rust: tuple[FunctionTraceEvent, ...], - *, - colorize: bool, -) -> str: - python_header: Final = f"python ({len(python)} steps)" - rust_header: Final = f"rust ({len(rust)} steps)" - python_width: Final = max(len(python_header), *(len(_label(event)) for event in python), 0) - rust_width: Final = max(len(rust_header), *(len(_label(event)) for event in rust), 0) - python_names: Final = frozenset(event.function for event in python) - rust_names: Final = frozenset(event.function for event in rust) - border: Final = f"+-{'-' * python_width}-+-{'-' * rust_width}-+-------------+" - rows: Final = tuple( - f"{color}{line}\033[0m" if colorize else line - for left, right in _aligned_rows(python, rust) - for status, color in (_status(left, right, python_names, rust_names),) - for line in (f"| {_label(left):<{python_width}} | {_label(right):<{rust_width}} | {status:<11} |",) - ) - return "\n".join( - ( - border, - f"| {python_header:<{python_width}} | {rust_header:<{rust_width}} | {'comparison':<11} |", - border, - *rows, - border, - ) - ) diff --git a/tests/sdk_function_trace/test_mock_provider.py b/tests/sdk_function_trace/test_mock_provider.py deleted file mode 100644 index 88d7d5392d0..00000000000 --- a/tests/sdk_function_trace/test_mock_provider.py +++ /dev/null @@ -1,33 +0,0 @@ -from __future__ import annotations - -from contextlib import ExitStack -from typing import Final -from urllib.error import HTTPError -from urllib.request import Request, urlopen - -import pytest - -from tests.sdk_function_trace.mock_provider import MockProviderResponse, mock_provider - - -def test_mock_provider_preserves_error_response() -> None: - response: Final = MockProviderResponse(429, (("retry-after", "2"),), b'{"error":"rate limited"}') - with mock_provider(response) as api_base: - with pytest.raises(HTTPError) as error: - urlopen(Request(api_base, data=b"{}"), timeout=5) - with error.value as received: - assert received.code == 429 - assert received.headers["retry-after"] == "2" - assert received.read() == response.body - - -@pytest.mark.parametrize("request_count", [0, 2]) -def test_mock_provider_rejects_missing_or_duplicate_requests(request_count: int) -> None: - response: Final = MockProviderResponse(200, (), b"{}") - with ExitStack() as stack: - api_base: Final = stack.enter_context(mock_provider(response)) - for _ in range(request_count): - with urlopen(Request(api_base, data=b"{}"), timeout=5) as received: - assert received.read() == response.body - with pytest.raises(AssertionError, match=f"expected one provider request, received {request_count}"): - stack.close() diff --git a/tests/sdk_function_trace/test_profiler.py b/tests/sdk_function_trace/test_profiler.py deleted file mode 100644 index 10a266fb1e8..00000000000 --- a/tests/sdk_function_trace/test_profiler.py +++ /dev/null @@ -1,167 +0,0 @@ -from __future__ import annotations - -import asyncio -import sys -from pathlib import Path -from types import FunctionType -from typing import Final, cast - -import pytest - -from tests.sdk_function_trace import ( - FunctionTraceEvent, - TraceScenario, - TraceStep, - assert_function_trace_parity, -) -from tests.sdk_function_trace.profiler import profile_python - - -class First: - @staticmethod - def run() -> None: - return None - - -class Second: - @staticmethod - def run() -> None: - return None - - -def test_profiler_matches_code_objects_and_keeps_repeated_calls() -> None: - with profile_python((First.run,)) as profiler: - Second.run() - First.run() - First.run() - - assert profiler.events == [ - FunctionTraceEvent(function="run", depth=0), - FunctionTraceEvent(function="run", depth=0), - ] - - -def test_profiler_records_selected_function_nesting_depth() -> None: - class Nested: - @staticmethod - def run() -> None: - First.run() - - with profile_python((Nested.run, First.run)) as profiler: - Nested.run() - - assert profiler.events == [ - FunctionTraceEvent(function="run", depth=0), - FunctionTraceEvent(function="run", depth=1), - ] - - -def test_profiler_restores_previous_profiler_after_failure() -> None: - previous: Final = sys.getprofile() - - with profile_python((First.run,)) as outer: - with pytest.raises(RuntimeError, match="stop"): - with profile_python((Second.run,)): - raise RuntimeError("stop") - assert sys.getprofile() is outer - First.run() - - assert sys.getprofile() is previous - assert outer.events == [FunctionTraceEvent(function="run", depth=0)] - - -def test_profiler_does_not_count_coroutine_resumption_as_another_call() -> None: - async def suspended() -> None: - await asyncio.sleep(0) - First.run() - await asyncio.sleep(0) - - with profile_python((suspended, First.run)) as profiler: - asyncio.run(suspended()) - - assert profiler.events == [ - FunctionTraceEvent(function="suspended", depth=0), - FunctionTraceEvent(function="run", depth=1), - ] - - -def test_source_profiler_records_real_frame_ancestry() -> None: - def outer() -> None: - First.run() - - with profile_python(source_root=Path(__file__).parent) as profiler: - outer() - Second.run() - - outer_event, first_event, second_event = ( - event for event in profiler.events if event.function.startswith("test_profiler.py:") - ) - assert first_event.ancestors is not None - assert outer_event.function in first_event.ancestors - assert second_event.ancestors is not None - assert outer_event.function not in second_event.ancestors - - -@pytest.mark.parametrize( - "rust_trace", - [ - (), - (FunctionTraceEvent(function="renamed", depth=0),), - (FunctionTraceEvent(function="run", depth=1),), - (FunctionTraceEvent(function="run", depth=0),) * 2, - ], - ids=["missing", "renamed", "wrong-depth", "extra-call"], -) -def test_harness_rejects_rust_function_trace_drift(rust_trace: tuple[FunctionTraceEvent, ...]) -> None: - with pytest.raises(AssertionError, match="Rust function trace differs"): - assert_function_trace_parity( - TraceScenario( - steps=(TraceStep(cast(FunctionType, First.run), depth=0),), - invoke_python=First.run, - invoke_rust=lambda: rust_trace, - ) - ) - - -def test_harness_rejects_python_function_trace_drift() -> None: - with pytest.raises(AssertionError, match="Python function trace differs"): - assert_function_trace_parity( - TraceScenario( - steps=(TraceStep(cast(FunctionType, First.run), depth=0),), - invoke_python=Second.run, - invoke_rust=lambda: (FunctionTraceEvent(function="run", depth=0),), - ) - ) - - -def test_harness_accepts_matching_traces() -> None: - assert_function_trace_parity( - TraceScenario( - steps=(TraceStep(cast(FunctionType, First.run), depth=0),), - invoke_python=First.run, - invoke_rust=lambda: (FunctionTraceEvent(function="run", depth=0),), - ) - ) - - -def test_harness_rejects_reordered_calls() -> None: - def begin() -> None: - return None - - def finish() -> None: - return None - - with pytest.raises(AssertionError, match="Rust function trace differs"): - assert_function_trace_parity( - TraceScenario( - steps=( - TraceStep(cast(FunctionType, begin), depth=0), - TraceStep(cast(FunctionType, finish), depth=0), - ), - invoke_python=lambda: (begin(), finish()), - invoke_rust=lambda: ( - FunctionTraceEvent(function="finish", depth=0), - FunctionTraceEvent(function="begin", depth=0), - ), - ) - ) diff --git a/tests/sdk_function_trace/test_runtime.py b/tests/sdk_function_trace/test_runtime.py deleted file mode 100644 index 015cba55083..00000000000 --- a/tests/sdk_function_trace/test_runtime.py +++ /dev/null @@ -1,47 +0,0 @@ -from __future__ import annotations - -from typing import Final - -import pytest - -from tests.sdk_function_trace.runtime import ( - TraceFailed, - TraceSkipped, - attempt_trace, - run_trace, - trace_diff, -) -from tests.sdk_function_trace.steps import pipeline_issues, pipeline_steps - - -def test_sync_messages_records_the_known_python_limitation() -> None: - result: Final = attempt_trace("messages", engine="python", asynchronous=False) - - assert isinstance(result, TraceSkipped) - assert result.reason == "ValueError: anthropic_messages_handler is not implemented for sync calls" - - -def test_unexpected_call_failure_is_not_skipped() -> None: - result: Final = attempt_trace("unknown", engine="python", asynchronous=False) - - assert isinstance(result, TraceFailed) - assert result.reason == "ValueError: Unknown route: unknown" - - -@pytest.mark.parametrize( - ("route", "asynchronous"), - (("chat_completions", False), ("chat_completions", True), ("messages", True), ("ocr", False), ("ocr", True)), -) -def test_compiled_routes_match_python_steps(route: str, asynchronous: bool) -> None: - from litellm.rust_bridge import get_native_bridge - - if get_native_bridge() is None: - pytest.skip("build the native bridge to run executed route parity") - python: Final = pipeline_steps(route, "python", run_trace(route, engine="python", asynchronous=asynchronous)) - rust: Final = pipeline_steps(route, "rust", run_trace(route, engine="rust", asynchronous=asynchronous)) - - assert pipeline_issues(route, "python", python) == () - assert pipeline_issues(route, "rust", rust) == () - assert trace_diff(python, rust).matches - if route != "messages": - assert python == rust diff --git a/tests/sdk_function_trace/test_steps.py b/tests/sdk_function_trace/test_steps.py deleted file mode 100644 index b5432951187..00000000000 --- a/tests/sdk_function_trace/test_steps.py +++ /dev/null @@ -1,244 +0,0 @@ -from __future__ import annotations - -from typing import Final - -import pytest - -from tests.sdk_function_trace.profiler import FunctionTraceEvent -from tests.sdk_function_trace.runtime import trace_diff -from tests.sdk_function_trace.steps import pipeline_issues, pipeline_steps - - -def test_python_ocr_projection_keeps_pipeline_and_drops_noise() -> None: - events: Final = ( - FunctionTraceEvent("utils.py:1747 client..wrapper_async", 0), - FunctionTraceEvent("ocr/main.py:331 aocr", 1), - FunctionTraceEvent("ocr/main.py:70 _prepare_ocr_request", 2), - FunctionTraceEvent("litellm_core_utils/get_llm_provider_logic.py:142 get_llm_provider", 3), - FunctionTraceEvent("utils.py:9303 ProviderConfigManager.get_provider_ocr_config", 3), - FunctionTraceEvent("llms/mistral/ocr/transformation.py:34 MistralOCRConfig.get_supported_ocr_params", 4), - FunctionTraceEvent("llms/mistral/ocr/transformation.py:72 MistralOCRConfig.map_ocr_params", 4), - FunctionTraceEvent("llms/mistral/ocr/transformation.py:34 MistralOCRConfig.get_supported_ocr_params", 5), - FunctionTraceEvent("llms/custom_httpx/llm_http_handler.py:1705 BaseLLMHTTPHandler.async_ocr", 2), - FunctionTraceEvent("llms/mistral/ocr/transformation.py:94 MistralOCRConfig.validate_environment", 4), - FunctionTraceEvent("llms/mistral/ocr/transformation.py:124 MistralOCRConfig.get_complete_url", 4), - FunctionTraceEvent("llms/base_llm/ocr/transformation.py:209 BaseOCRConfig.async_transform_ocr_request", 5), - FunctionTraceEvent("llms/mistral/ocr/transformation.py:149 MistralOCRConfig.transform_ocr_request", 6), - FunctionTraceEvent("llms/custom_httpx/http_handler.py:654 AsyncHTTPHandler.post", 6), - FunctionTraceEvent("llms/base_llm/ocr/transformation.py:255 BaseOCRConfig.async_transform_ocr_response", 4), - FunctionTraceEvent("llms/mistral/ocr/transformation.py:200 MistralOCRConfig.transform_ocr_response", 5), - FunctionTraceEvent("cost_calculator.py:1874 ocr_cost", 6), - ) - - assert pipeline_steps("ocr", "python", events) == ( - FunctionTraceEvent("ocr", 0), - FunctionTraceEvent("prepare_ocr_call", 1), - FunctionTraceEvent("get_provider_ocr_config", 2), - FunctionTraceEvent("supported_ocr_params", 3), - FunctionTraceEvent("map_ocr_params", 3), - FunctionTraceEvent("execute_ocr_provider_call", 1), - FunctionTraceEvent("validate_environment", 2), - FunctionTraceEvent("complete_url", 2), - FunctionTraceEvent("transform_ocr_request", 3), - FunctionTraceEvent("http_request", 3), - FunctionTraceEvent("transform_ocr_response", 2), - ) - - -def test_rust_ocr_projection_reuses_step_names_and_keeps_unknown_spans() -> None: - events: Final = ( - FunctionTraceEvent("ocr", 0), - FunctionTraceEvent("prepare_ocr_call", 1), - FunctionTraceEvent("map_ocr_params", 2), - FunctionTraceEvent("supported_ocr_params", 3), - FunctionTraceEvent("map_ocr_params", 2), - FunctionTraceEvent("transform_ocr_request", 2), - FunctionTraceEvent("execute_ocr_provider_call", 2), - FunctionTraceEvent("transform_ocr_response", 3), - FunctionTraceEvent("new_uninstrumented_span", 3), - ) - - assert pipeline_steps("ocr", "rust", events) == ( - FunctionTraceEvent("ocr", 0), - FunctionTraceEvent("prepare_ocr_call", 1), - FunctionTraceEvent("map_ocr_params", 2), - FunctionTraceEvent("supported_ocr_params", 3), - FunctionTraceEvent("transform_ocr_request", 2), - FunctionTraceEvent("execute_ocr_provider_call", 2), - FunctionTraceEvent("transform_ocr_response", 3), - FunctionTraceEvent("new_uninstrumented_span", 3), - ) - - -def test_projection_resets_depth_on_thread_root() -> None: - events: Final = ( - FunctionTraceEvent("main.py:387 acompletion", 1), - FunctionTraceEvent("llms/anthropic/chat/handler.py:255 AnthropicChatCompletion.acompletion_function", 2), - FunctionTraceEvent( - "llms/anthropic/experimental_pass_through/messages/handler.py:416 anthropic_messages_handler", 0 - ), - FunctionTraceEvent( - "llms/anthropic/experimental_pass_through/messages/transformation.py:575" - " AnthropicMessagesConfig.transform_anthropic_messages_request", - 4, - ), - ) - - assert pipeline_steps("chat_completions", "python", events) == ( - FunctionTraceEvent("chat_completions", 0), - FunctionTraceEvent("execute_chat_completions_provider_call", 1), - ) - assert pipeline_steps("messages", "python", events) == ( - FunctionTraceEvent("execute_messages_provider_call", 0), - FunctionTraceEvent("transform_request", 1), - ) - - -@pytest.mark.parametrize("function", ("completion", "completion_function", "acompletion_function")) -def test_chat_projection_includes_sync_and_async_handlers(function: str) -> None: - events: Final = (FunctionTraceEvent(f"llms/anthropic/chat/handler.py:100 AnthropicChatCompletion.{function}", 0),) - - assert pipeline_steps("chat_completions", "python", events) == ( - FunctionTraceEvent("execute_chat_completions_provider_call", 0), - ) - - -def test_trace_diff_reports_no_difference_for_identical_steps() -> None: - steps: Final = ( - FunctionTraceEvent("ocr", 0), - FunctionTraceEvent("transform_ocr_request", 1), - ) - - diff: Final = trace_diff(steps, steps) - - assert diff.python_only == () - assert diff.rust_only == () - assert diff.shared_order_matches - assert diff.matches - - -def test_trace_diff_reports_exclusive_steps_and_reordered_shared_steps() -> None: - python: Final = ( - FunctionTraceEvent("ocr", 0), - FunctionTraceEvent("supported_ocr_params", 1), - FunctionTraceEvent("map_ocr_params", 1), - FunctionTraceEvent("http_request", 2), - ) - rust: Final = ( - FunctionTraceEvent("ocr", 0), - FunctionTraceEvent("map_ocr_params", 1), - FunctionTraceEvent("supported_ocr_params", 2), - FunctionTraceEvent("transform_ocr_response", 2), - ) - - diff: Final = trace_diff(python, rust) - - assert diff.python_only == ("http_request",) - assert diff.rust_only == ("transform_ocr_response",) - assert not diff.shared_order_matches - assert not diff.matches - - -def test_trace_diff_does_not_claim_empty_or_disjoint_traces_match() -> None: - assert not trace_diff((), ()).shared_order_matches - assert not trace_diff((FunctionTraceEvent("ocr", 0),), (FunctionTraceEvent("messages", 0),)).shared_order_matches - - -def test_projection_uses_actual_ancestors_after_coroutine_resumption() -> None: - entrypoint: Final = "main.py:387 acompletion" - handler: Final = "llms/anthropic/chat/handler.py:255 AnthropicChatCompletion.acompletion_function" - events: Final = ( - FunctionTraceEvent(entrypoint, 0, ()), - FunctionTraceEvent(handler, 1, (entrypoint,)), - FunctionTraceEvent("utils.py:100 unrelated_worker", 0, ()), - FunctionTraceEvent("llms/anthropic/chat/transformation.py:100 transform_response", 1, (handler,)), - ) - - assert pipeline_steps("chat_completions", "python", events) == ( - FunctionTraceEvent("chat_completions", 0), - FunctionTraceEvent("execute_chat_completions_provider_call", 1), - FunctionTraceEvent("transform_response", 2), - ) - - -def test_projection_does_not_nest_siblings_under_a_returned_config_lookup() -> None: - events: Final = ( - FunctionTraceEvent("main.py:387 completion", 0), - FunctionTraceEvent("utils.py:100 ProviderConfigManager.get_provider_chat_config", 1), - FunctionTraceEvent("utils.py:200 unrelated_helper", 1), - FunctionTraceEvent("llms/anthropic/chat/transformation.py:100 transform_request", 2), - ) - - assert pipeline_steps("chat_completions", "python", events) == ( - FunctionTraceEvent("chat_completions", 0), - FunctionTraceEvent("get_provider_chat_config", 1), - FunctionTraceEvent("transform_request", 1), - ) - - -CHAT_RUST_STEPS: Final = ( - "chat_completions", - "get_provider_chat_config", - "supported_openai_params", - "execute_chat_completions_provider_call", - "validate_environment", - "transform_request", - "http_request", - "transform_response", -) - - -@pytest.mark.parametrize("missing", CHAT_RUST_STEPS) -def test_pipeline_check_rejects_missing_stages(missing: str) -> None: - steps: Final = tuple(FunctionTraceEvent(name, 0) for name in CHAT_RUST_STEPS if name != missing) - - assert f"missing {missing}" in pipeline_issues("chat_completions", "rust", steps) - - -def test_pipeline_check_rejects_http_before_request_transformation() -> None: - steps: Final = tuple( - FunctionTraceEvent(name, 0) - for name in ( - "chat_completions", - "get_provider_chat_config", - "supported_openai_params", - "execute_chat_completions_provider_call", - "validate_environment", - "http_request", - "transform_request", - "transform_response", - ) - ) - - assert "transform_request must precede http_request" in pipeline_issues("chat_completions", "rust", steps) - - -def test_step_parity_rejects_different_handler_boundaries_even_with_valid_stages() -> None: - rust: Final = tuple(FunctionTraceEvent(name, 0) for name in CHAT_RUST_STEPS) - python: Final = tuple( - FunctionTraceEvent(name, 0) - for name in ( - "chat_completions", - "get_provider_chat_config", - "supported_openai_params", - "validate_environment", - "transform_request", - "execute_chat_completions_provider_call", - "http_request", - "transform_response", - ) - ) - - assert not trace_diff(python, rust).shared_order_matches - assert not trace_diff(python, rust).matches - assert pipeline_issues("chat_completions", "python", python) == () - assert pipeline_issues("chat_completions", "rust", rust) == () - - -def test_step_parity_rejects_an_exclusive_helper_with_matching_shared_order() -> None: - rust: Final = tuple(FunctionTraceEvent(name, 0) for name in CHAT_RUST_STEPS) - python: Final = (*rust, FunctionTraceEvent("unmatched_helper", 0)) - diff: Final = trace_diff(python, rust) - - assert diff.shared_order_matches - assert not diff.matches diff --git a/tests/sdk_function_trace/test_table.py b/tests/sdk_function_trace/test_table.py deleted file mode 100644 index c2341a391a9..00000000000 --- a/tests/sdk_function_trace/test_table.py +++ /dev/null @@ -1,67 +0,0 @@ -from __future__ import annotations - -import re -from typing import Final - -from tests.sdk_function_trace.profiler import FunctionTraceEvent -from tests.sdk_function_trace.table import format_trace_table - - -def test_table_aligns_matches_after_missing_steps_and_preserves_indentation() -> None: - python: Final = ( - FunctionTraceEvent("ocr", 0), - FunctionTraceEvent("python_helper", 1), - FunctionTraceEvent("http_request", 2), - ) - rust: Final = ( - FunctionTraceEvent("ocr", 0), - FunctionTraceEvent("rust_helper", 1), - FunctionTraceEvent("http_request", 1), - ) - output: Final = format_trace_table(python, rust, colorize=False) - rows: Final = tuple(line.split("|")[1:-1] for line in output.splitlines() if line.startswith("|")) - - assert tuple(tuple(cell.strip() for cell in row) for row in rows) == ( - ("python (3 steps)", "rust (3 steps)", "comparison"), - ("ocr", "ocr", "match"), - ("python_helper", "", "python only"), - ("", "rust_helper", "rust only"), - ("http_request", "http_request", "match"), - ) - assert rows[-1][0].startswith(" http_request") - assert rows[-1][1].startswith(" http_request") - assert len({len(line) for line in output.splitlines()}) == 1 - assert "\033[" not in output - - -def test_table_marks_reordered_calls_and_keeps_both_execution_orders() -> None: - python: Final = tuple(FunctionTraceEvent(name, 0) for name in ("ocr", "map", "validate", "http")) - rust: Final = tuple(FunctionTraceEvent(name, 0) for name in ("ocr", "validate", "map", "http")) - output: Final = format_trace_table(python, rust, colorize=True) - plain: Final = re.sub(r"\033\[[0-9;]*m", "", output) - rows: Final = tuple(line.split("|")[1:-1] for line in plain.splitlines() if line.startswith("|"))[1:] - - assert tuple(row[0].strip() for row in rows if row[0].strip()) == tuple(event.function for event in python) - assert tuple(row[1].strip() for row in rows if row[1].strip()) == tuple(event.function for event in rust) - assert plain.count("reordered") == 2 - assert output.count("\033[31m") == 2 - assert "only" not in output - - -def test_table_colors_match_and_exclusive_rows_without_changing_alignment() -> None: - python: Final = (FunctionTraceEvent("ocr", 0), FunctionTraceEvent("python_helper", 1)) - rust: Final = (FunctionTraceEvent("ocr", 0), FunctionTraceEvent("rust_helper", 1)) - colored: Final = format_trace_table(python, rust, colorize=True) - - assert re.sub(r"\033\[[0-9;]*m", "", colored) == format_trace_table(python, rust, colorize=False) - assert next(line for line in colored.splitlines() if "match" in line).startswith("\033[32m") - assert next(line for line in colored.splitlines() if "python only" in line).startswith("\033[34m") - assert next(line for line in colored.splitlines() if "rust only" in line).startswith("\033[33m") - - -def test_table_handles_empty_traces() -> None: - output: Final = format_trace_table((), (), colorize=False) - - assert "python (0 steps)" in output - assert "rust (0 steps)" in output - assert "match" not in output diff --git a/tests/test_litellm/caching/test_caching.py b/tests/test_litellm/caching/test_caching.py index 955b0e531bc..4d0ec0fb677 100644 --- a/tests/test_litellm/caching/test_caching.py +++ b/tests/test_litellm/caching/test_caching.py @@ -4,7 +4,7 @@ import re import pytest from litellm.caching.caching import Cache -from litellm.types.caching import LiteLLMCacheType +from litellm.types.caching import LiteLLMCacheType, SemanticCacheScope from litellm.types.utils import Embedding, EmbeddingResponse, Usage @@ -80,12 +80,13 @@ def test_get_per_item_prompt_tokens_distributes_with_remainder(): assert per_item == [4, 3, 3] -def _semantic_cache(): +def _semantic_cache(**cache_kwargs): return Cache( type=LiteLLMCacheType.VALKEY_SEMANTIC, host="localhost", port="6379", similarity_threshold=0.8, + **cache_kwargs, ) @@ -139,6 +140,76 @@ def test_semantic_cache_key_isolates_tenants(): assert key_a != key_team +_SEMANTICALLY_IDENTICAL_PROMPTS = ( + [{"role": "user", "content": "What color is the sky?"}], + [{"role": "user", "content": "Tell me the colour of the daytime sky."}], +) + + +def _end_user_keys(cache, metadata_field, *end_user_ids): + return [ + cache.get_cache_key( + model="gpt-4o-mini", + messages=messages, + **{metadata_field: {"user_api_key": "hash-A", "user_api_key_end_user_id": end_user_id}}, + ) + for messages, end_user_id in zip(_SEMANTICALLY_IDENTICAL_PROMPTS, end_user_ids) + ] + + +@pytest.mark.parametrize("metadata_field", ["metadata", "litellm_metadata"]) +def test_semantic_cache_key_shares_bucket_across_end_users_by_default(metadata_field): + key_alice, key_bob = _end_user_keys(_semantic_cache(), metadata_field, "alice", "bob") + assert key_alice == key_bob + + +@pytest.mark.parametrize("metadata_field", ["metadata", "litellm_metadata"]) +def test_semantic_cache_key_isolates_end_users_under_end_user_scope(metadata_field): + cache = _semantic_cache(semantic_cache_scope="end_user") + key_alice, key_bob = _end_user_keys(cache, metadata_field, "alice", "bob") + key_alice_again, _ = _end_user_keys(cache, metadata_field, "alice", "alice") + assert key_alice != key_bob + assert key_alice == key_alice_again + + +def test_semantic_cache_key_end_user_scope_without_end_user_falls_back_to_key_scope(): + cache = _semantic_cache(semantic_cache_scope=SemanticCacheScope.END_USER) + messages = [{"role": "user", "content": "What color is the sky?"}] + key_scope_only = cache.get_cache_key(model="gpt-4o-mini", messages=messages, metadata={"user_api_key": "hash-A"}) + end_user_absent = cache.get_cache_key( + model="gpt-4o-mini", + messages=messages, + metadata={"user_api_key": "hash-A", "user_api_key_end_user_id": None}, + ) + other_key = cache.get_cache_key(model="gpt-4o-mini", messages=messages, metadata={"user_api_key": "hash-B"}) + key_alice, _ = _end_user_keys(cache, "metadata", "alice", "alice") + default_scope_key = _semantic_cache().get_cache_key( + model="gpt-4o-mini", messages=messages, metadata={"user_api_key": "hash-A"} + ) + assert key_scope_only == end_user_absent == default_scope_key + assert key_scope_only != other_key + assert key_scope_only != key_alice + + +def test_semantic_cache_key_reads_tenant_identity_from_litellm_metadata(): + cache = _semantic_cache() + messages = [{"role": "user", "content": "What color is the sky?"}] + key_a = cache.get_cache_key(model="gpt-4o-mini", messages=messages, litellm_metadata={"user_api_key": "hash-A"}) + key_b = cache.get_cache_key(model="gpt-4o-mini", messages=messages, litellm_metadata={"user_api_key": "hash-B"}) + key_a_in_litellm_params = cache.get_cache_key( + model="gpt-4o-mini", + messages=messages, + litellm_params={"litellm_metadata": {"user_api_key": "hash-A"}}, + ) + assert key_a != key_b + assert key_a == key_a_in_litellm_params + + +def test_semantic_cache_scope_rejects_unknown_value(): + with pytest.raises(ValueError, match="'team' is not a valid SemanticCacheScope"): + _semantic_cache(semantic_cache_scope="team") + + def test_semantic_cache_key_still_separates_models_and_params(): cache = _semantic_cache() messages = [{"role": "user", "content": "hi"}] diff --git a/tests/test_litellm/caching/test_redis_cluster_node_isolation.py b/tests/test_litellm/caching/test_redis_cluster_node_isolation.py index c16ceec8c31..9ebc3673ab6 100644 --- a/tests/test_litellm/caching/test_redis_cluster_node_isolation.py +++ b/tests/test_litellm/caching/test_redis_cluster_node_isolation.py @@ -5,20 +5,26 @@ CLIENT PAUSE) showed 100% of concurrent commands to the other two, untouched nod stalling for the full pause duration before this fix, and zero after -- these tests pin the same behavior at the unit level so it can run without a live Redis Cluster.""" +import asyncio from typing import TYPE_CHECKING -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, Mock, patch import pytest from redis.exceptions import ( + AskError, BusyLoadingError, ClusterDownError, + ClusterError, MaxConnectionsError, MovedError, + TryAgainError, ) from redis.exceptions import ( ConnectionError as RedisConnectionError, ) -from redis.exceptions import TimeoutError as RedisTimeoutError +from redis.exceptions import ( + TimeoutError as RedisTimeoutError, +) from litellm.caching.redis_cluster_node_isolation import ( get_litellm_async_redis_cluster_class, @@ -39,10 +45,35 @@ class _NodeClassWithoutPerConnectionRecovery: class _FakeClusterNode: def __init__(self, name: str, raises: Exception | None = None, response: object = None) -> None: self.name = name - self.execute_command = AsyncMock(side_effect=raises, return_value=response) + + async def execute_command(*args: object, **kwargs: object) -> object: + await asyncio.sleep(0) + if raises is not None: + raise raises + return response + + self.execute_command = AsyncMock(side_effect=execute_command) self.disconnect = AsyncMock() +class _Fake8xRedisCluster: + def __init__(self) -> None: + self._initialize = False + + async def _execute_command( + self, target_node: _FakeClusterNode, *args: object, **kwargs: object + ) -> object: + try: + return await target_node.execute_command(*args, **kwargs) + except (RedisConnectionError, RedisTimeoutError): + self._initialize = True + await asyncio.sleep(0) + raise + + async def aclose(self) -> None: + self._initialize = True + + class _FakeNodesManager: def __init__(self, node_to_return: _FakeClusterNode) -> None: self._moved_exception: object = None @@ -68,31 +99,198 @@ def _build_cluster_instance() -> "_AsyncRedisClusterType": return instance -def test_per_connection_recovery_redis_py_gets_the_unmodified_upstream_class() -> None: - """Regression (redis-py 8.x): when upstream ClusterNode already recovers a node-level - connection error per-connection, the factory must NOT install the copied override, - whose node.disconnect() also kills connections other coroutines are mid-operation on.""" - from redis.asyncio.cluster import RedisCluster - +def _build_8x_cluster_instance() -> _Fake8xRedisCluster: cluster_cls = get_litellm_async_redis_cluster_class( - cluster_node_class=_NodeClassWithPerConnectionRecovery + cluster_node_class=_NodeClassWithPerConnectionRecovery, + base_cluster_class=_Fake8xRedisCluster, + ) + return cluster_cls() + + +def test_unverified_redis_version_logs_warning(caplog: pytest.LogCaptureFixture) -> None: + import redis + + with patch.object(redis, "__version__", "8.0.1"): + get_litellm_async_redis_cluster_class(cluster_node_class=_NodeClassWithoutPerConnectionRecovery) + + assert "not in the set this cluster-teardown-storm fix was verified against" in caplog.text + + +@pytest.mark.asyncio +async def test_single_timeout_does_not_request_topology_reinit() -> None: + error = RedisTimeoutError("timeout") + target_node = _FakeClusterNode("node-a") + target_node.execute_command.side_effect = error + instance = _build_8x_cluster_instance() + + with pytest.raises(RedisTimeoutError) as exc_info: + await instance._execute_command(target_node, "GET", "k") + + assert exc_info.value is error + assert instance._initialize is False + + +@pytest.mark.asyncio +async def test_connection_error_preserves_upstream_topology_reinit() -> None: + error = RedisConnectionError("connection error") + target_node = _FakeClusterNode("node-a") + target_node.execute_command.side_effect = error + instance = _build_8x_cluster_instance() + + with pytest.raises(RedisConnectionError) as exc_info: + await instance._execute_command(target_node, "GET", "k") + + assert exc_info.value is error + assert instance._initialize is True + + +@pytest.mark.asyncio +async def test_three_consecutive_timeouts_request_topology_reinit_and_reset_counter() -> None: + errors = [ + RedisTimeoutError("timeout-1"), + RedisTimeoutError("timeout-2"), + RedisTimeoutError("timeout-3"), + ] + fourth_error = RedisTimeoutError("timeout-4") + target_node = _FakeClusterNode("node-a") + target_node.execute_command.side_effect = [*errors, fourth_error] + instance = _build_8x_cluster_instance() + + for error in errors: + with pytest.raises(RedisTimeoutError) as exc_info: + await instance._execute_command(target_node, "GET", "k") + assert exc_info.value is error + + assert instance._initialize is True + instance._initialize = False + + with pytest.raises(RedisTimeoutError) as exc_info: + await instance._execute_command(target_node, "GET", "k") + + assert exc_info.value is fourth_error + assert instance._initialize is False + + +@pytest.mark.asyncio +async def test_success_resets_consecutive_timeout_counter() -> None: + errors = [RedisTimeoutError("timeout-1"), RedisTimeoutError("timeout-2")] + final_error = RedisTimeoutError("timeout-3") + target_node = _FakeClusterNode("node-a") + target_node.execute_command.side_effect = [*errors, b"value", final_error] + instance = _build_8x_cluster_instance() + + for error in errors: + with pytest.raises(RedisTimeoutError) as exc_info: + await instance._execute_command(target_node, "GET", "k") + assert exc_info.value is error + assert instance._initialize is False + + result = await instance._execute_command(target_node, "GET", "k") + assert result == b"value" + assert instance._initialize is False + + with pytest.raises(RedisTimeoutError) as exc_info: + await instance._execute_command(target_node, "GET", "k") + + assert exc_info.value is final_error + assert instance._initialize is False + + +@pytest.mark.asyncio +async def test_timeout_counters_are_per_node() -> None: + node_a_errors = [RedisTimeoutError("node-a-1"), RedisTimeoutError("node-a-2")] + node_b_error = RedisTimeoutError("node-b-1") + node_a = _FakeClusterNode("node-a") + node_b = _FakeClusterNode("node-b") + node_a.execute_command.side_effect = node_a_errors + node_b.execute_command.side_effect = node_b_error + instance = _build_8x_cluster_instance() + + for target_node, error in ( + (node_a, node_a_errors[0]), + (node_b, node_b_error), + (node_a, node_a_errors[1]), + ): + with pytest.raises(RedisTimeoutError) as exc_info: + await instance._execute_command(target_node, "GET", "k") + assert exc_info.value is error + + assert instance._initialize is False + + +@pytest.mark.asyncio +async def test_timeout_does_not_clear_concurrent_topology_reinit_request() -> None: + error = RedisTimeoutError("timeout") + instance = _build_8x_cluster_instance() + + async def request_reinit(*args: object, **kwargs: object) -> object: + await instance.aclose() + raise error + + target_node = _FakeClusterNode("node-a") + target_node.execute_command.side_effect = request_reinit + + with pytest.raises(RedisTimeoutError) as exc_info: + await instance._execute_command(target_node, "GET", "k") + + assert exc_info.value is error + assert instance._initialize is True + + +@pytest.mark.asyncio +async def test_tolerated_timeout_does_not_erase_concurrent_connection_error_reinit() -> None: + instance = _build_8x_cluster_instance() + failing_node = _FakeClusterNode("node-a", raises=RedisConnectionError("gone")) + slow_node = _FakeClusterNode("node-b", raises=RedisTimeoutError("slow")) + + results = await asyncio.gather( + instance._execute_command(failing_node, "GET", "a"), + instance._execute_command(slow_node, "GET", "b"), + return_exceptions=True, ) - assert cluster_cls is RedisCluster + assert isinstance(results[0], RedisConnectionError) + assert isinstance(results[1], RedisTimeoutError) + assert instance._initialize is True -def test_pre_recovery_redis_py_still_gets_the_node_isolation_override() -> None: - """Old redis-py (5.x) responds to a node-level error with a full-cluster aclose(), - so those versions must keep litellm's per-node isolation override.""" - from redis.asyncio.cluster import RedisCluster +@pytest.mark.asyncio +async def test_overlapping_tolerated_timeouts_do_not_request_topology_reinit() -> None: + instance = _build_8x_cluster_instance() + node_a = _FakeClusterNode("node-a", raises=RedisTimeoutError("slow-a")) + node_b = _FakeClusterNode("node-b", raises=RedisTimeoutError("slow-b")) - cluster_cls = get_litellm_async_redis_cluster_class( - cluster_node_class=_NodeClassWithoutPerConnectionRecovery + results = await asyncio.gather( + instance._execute_command(node_a, "GET", "a"), + instance._execute_command(node_b, "GET", "b"), + return_exceptions=True, ) - assert cluster_cls is not RedisCluster - assert issubclass(cluster_cls, RedisCluster) - assert "_execute_command" in cluster_cls.__dict__ + assert all(isinstance(result, RedisTimeoutError) for result in results) + assert instance._initialize is False + + +@pytest.mark.asyncio +async def test_tolerated_timeout_does_not_clear_pending_reinit() -> None: + instance = _build_8x_cluster_instance() + instance._initialize = True + target_node = _FakeClusterNode("node-a", raises=RedisTimeoutError("slow")) + + with pytest.raises(RedisTimeoutError): + await instance._execute_command(target_node, "GET", "k") + + assert instance._initialize is True + + +@pytest.mark.asyncio +async def test_success_returns_value_without_topology_reinit() -> None: + target_node = _FakeClusterNode("node-a", response=b"value") + instance = _build_8x_cluster_instance() + + result = await instance._execute_command(target_node, "GET", "k") + + assert result == b"value" + assert instance._initialize is False @pytest.mark.asyncio @@ -110,6 +308,48 @@ async def test_node_level_error_resets_only_that_node_not_the_whole_client(error instance.aclose.assert_not_awaited() +@pytest.mark.asyncio +async def test_moved_error_retries_without_full_reinit_before_threshold() -> None: + moved_error = MovedError("1 127.0.0.1:7001") + target_node = _FakeClusterNode("node-a") + target_node.execute_command = AsyncMock(side_effect=[moved_error, b"value"]) + instance = _build_cluster_instance() + instance.RedisClusterRequestTTL = 2 + instance.nodes_manager = _FakeNodesManager(node_to_return=target_node) + instance._determine_slot = AsyncMock(return_value=0) + + result = await instance._execute_command(target_node, "GET", "k") + + assert result == b"value" + assert instance.nodes_manager._moved_exception is moved_error + instance.aclose.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_ask_error_sends_asking_and_retries_on_redirected_node() -> None: + ask_error = AskError("0 127.0.0.1:7001") + target_node = _FakeClusterNode("node-a") + target_node.execute_command = AsyncMock(side_effect=[ask_error, None, b"value"]) + instance = _build_cluster_instance() + instance.RedisClusterRequestTTL = 2 + instance.get_node = Mock(return_value=target_node) + + result = await instance._execute_command(target_node, "GET", "k") + + assert result == b"value" + instance.get_node.assert_called_once_with(node_name="127.0.0.1:7001") + + +@pytest.mark.asyncio +async def test_try_again_error_exhausts_ttl() -> None: + target_node = _FakeClusterNode("node-a", raises=TryAgainError("try again")) + instance = _build_cluster_instance() + instance.RedisClusterRequestTTL = 2 + + with pytest.raises(ClusterError): + await instance._execute_command(target_node, "GET", "k") + + @pytest.mark.asyncio async def test_successful_command_touches_neither_disconnect_nor_aclose() -> None: target_node = _FakeClusterNode("node-a", response=b"v") diff --git a/tests/test_litellm/integrations/azure_storage/test_azure_storage.py b/tests/test_litellm/integrations/azure_storage/test_azure_storage.py index a96eae0f9c3..6e1dab4a71a 100644 --- a/tests/test_litellm/integrations/azure_storage/test_azure_storage.py +++ b/tests/test_litellm/integrations/azure_storage/test_azure_storage.py @@ -42,6 +42,7 @@ def workload_identity_env_vars(monkeypatch): "AZURE_STORAGE_ENDPOINT_SUFFIX", "AZURE_CLIENT_SECRET", "AZURE_CREDENTIAL", + "AZURE_TOKEN_CREDENTIALS", "AZURE_SCOPE", ): monkeypatch.delenv(unset, raising=False) @@ -206,10 +207,28 @@ def test_default_chain_provider_is_storage_scoped_and_built_once_per_process(): assert first() == "chain-token" mock_builder.assert_called_once_with( azure_scope="https://storage.azure.com/.default", - azure_credential=AzureCredentialType.DefaultAzureCredential, + azure_credential=AzureCredentialType.DeploymentIdentityCredential, ) +def test_storage_chain_reaches_only_the_identities_a_deployment_carries(workload_identity_env_vars): + """ + The chain runs on a server, where a developer sign-in is a person and not the deployment, so + the storage token must come from workload identity or managed identity or from nothing + """ + _cached_credential_chain_token_provider.cache_clear() + with patch("azure.identity.get_bearer_token_provider", return_value=lambda: "chain-token") as bearer: + _cached_credential_chain_token_provider() + _cached_credential_chain_token_provider.cache_clear() + + bearer.assert_called_once() + with bearer.call_args.args[0] as chain: + assert {type(link).__name__ for link in chain.credentials} == { + "WorkloadIdentityCredential", + "ManagedIdentityCredential", + } + + @pytest.mark.asyncio async def test_chain_tokens_are_read_from_the_provider_on_every_refresh( workload_identity_env_vars, diff --git a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py index a182656e4a8..25a961c3413 100644 --- a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py +++ b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py @@ -17,7 +17,7 @@ import pytest import litellm from litellm import completion, acompletion from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler -from litellm.llms.snowflake.chat.transformation import SnowflakeConfig +from litellm.llms.snowflake.chat.transformation import SnowflakeConfig, SnowflakeStreamingHandler from litellm.types.utils import ModelResponse @@ -114,8 +114,7 @@ class TestSnowflakeToolTransformation: ) assert transformed_request["tool_choice"] == value, ( - f"tool_choice='{value}' should pass through unchanged, " - f"got {transformed_request['tool_choice']}" + f"tool_choice='{value}' should pass through unchanged, got {transformed_request['tool_choice']}" ) def test_transform_response_with_tool_calls(self): @@ -159,9 +158,7 @@ class TestSnowflakeToolTransformation: headers={"Content-Type": "application/json"}, ) - model_response = ModelResponse( - choices=[litellm.Choices(index=0, message=litellm.Message())] - ) + model_response = ModelResponse(choices=[litellm.Choices(index=0, message=litellm.Message())]) logging_obj = MagicMock() @@ -232,9 +229,7 @@ class TestSnowflakeToolTransformation: headers={"Content-Type": "application/json"}, ) - model_response = ModelResponse( - choices=[litellm.Choices(index=0, message=litellm.Message())] - ) + model_response = ModelResponse(choices=[litellm.Choices(index=0, message=litellm.Message())]) logging_obj = MagicMock() @@ -280,9 +275,7 @@ class TestSnowflakeToolTransformation: headers={"Content-Type": "application/json"}, ) - model_response = ModelResponse( - choices=[litellm.Choices(index=0, message=litellm.Message())] - ) + model_response = ModelResponse(choices=[litellm.Choices(index=0, message=litellm.Message())]) logging_obj = MagicMock() @@ -300,10 +293,7 @@ class TestSnowflakeToolTransformation: # Verify standard response works assert isinstance(result, ModelResponse) - assert ( - result.choices[0].message.content - == "Hello! I'm doing well, thank you for asking." - ) + assert result.choices[0].message.content == "Hello! I'm doing well, thank you for asking." def test_get_supported_openai_params_includes_tools(self): """ @@ -318,6 +308,385 @@ class TestSnowflakeToolTransformation: assert "max_tokens" in supported_params +class TestSnowflakeCortexClaudeFixes: + def setup_method(self): + self.config = SnowflakeConfig() + + @staticmethod + def _transform(messages, optional_params=None): + return SnowflakeConfig().transform_request( + model="snowflake/claude-sonnet-4-6", + messages=messages, + optional_params=optional_params or {}, + litellm_params={}, + headers={}, + ) + + def test_thinking_is_offered_on_every_claude_model(self): + """Cortex documents extended thinking (budget_tokens) for Claude generally, so a + 4.6-only gate would silently drop it on the models that do support it.""" + for model in ( + "snowflake/claude-sonnet-4-6", + "snowflake/claude-sonnet-4-5", + "snowflake/claude-3-7-sonnet", + "snowflake/claude-4-opus", + ): + assert "thinking" in self.config.get_supported_openai_params(model), model + assert "thinking" not in self.config.get_supported_openai_params("snowflake/llama3.1-70b") + + def test_system_blocks_preserve_cache_control_and_strip_ttl(self): + body = self._transform( + [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are helpful", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + }, + {"role": "user", "content": "hi"}, + ] + ) + assert body["system"] == [{"type": "text", "text": "You are helpful", "cache_control": {"type": "ephemeral"}}] + + def test_direct_system_param_is_normalized(self): + body = self._transform( + [{"role": "user", "content": "hi"}], + {"system": [{"type": "text", "text": "direct", "cache_control": {"type": "ephemeral", "ttl": "1h"}}]}, + ) + assert body["system"] == [{"type": "text", "text": "direct", "cache_control": {"type": "ephemeral"}}] + + def test_message_and_tool_cache_control_are_normalized(self): + body = self._transform( + [ + { + "role": "user", + "content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral", "ttl": "1h"}}], + } + ], + { + "tools": [ + { + "name": "f", + "input_schema": {"type": "object", "properties": {}}, + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ] + }, + ) + assert body["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral"} + assert body["tools"][0]["cache_control"] == {"type": "ephemeral"} + + def test_extra_body_message_override_is_normalized(self): + body = self._transform( + [{"role": "user", "content": "original"}], + { + "extra_body": { + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "override", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + } + ] + } + }, + ) + assert body["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral"} + + def test_image_blocks_are_converted_to_anthropic_source(self): + body = self._transform( + [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,ZmFrZQ==", "format": "image/jpeg"}, + } + ], + } + ] + ) + assert body["messages"][0]["content"] == [ + {"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": "ZmFrZQ=="}} + ] + + def test_tool_result_image_list_is_converted(self): + body = self._transform( + [ + {"role": "user", "content": "look"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "read", "arguments": "{}"}} + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,ZmFrZQ=="}}], + }, + ] + ) + assert body["messages"][2]["content"][0]["content"] == [ + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "ZmFrZQ=="}} + ] + + def test_tool_result_preserves_cache_control(self): + """A cache breakpoint the bridge puts on a tool message must survive onto the tool_result.""" + for tool_content in ("done", [{"type": "text", "text": "done"}]): + body = self._transform( + [ + {"role": "user", "content": "look"}, + { + "role": "tool", + "tool_call_id": "call_1", + "content": tool_content, + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + }, + ] + ) + tool_result = body["messages"][1]["content"][0] + assert tool_result["cache_control"] == {"type": "ephemeral"}, tool_content + + def test_pdf_data_uri_becomes_a_document_block(self): + """A bridged pdf data URI is a document block; forwarding it as an image is malformed.""" + body = self._transform( + [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": "data:application/pdf;base64,ZmFrZQ=="}}, + ], + } + ] + ) + assert body["messages"][0]["content"] == [ + { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": "ZmFrZQ=="}, + } + ] + + def test_multipart_tool_result_preserves_text_and_converts_image(self): + body = self._transform( + [ + {"role": "user", "content": "look"}, + { + "role": "tool", + "tool_call_id": "call_1", + "content": [ + {"type": "text", "text": "first"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,ZmFrZQ=="}}, + {"type": "text", "text": "last"}, + ], + }, + ] + ) + assert body["messages"][1]["content"][0]["content"] == [ + {"type": "text", "text": "first"}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "ZmFrZQ=="}}, + {"type": "text", "text": "last"}, + ] + + def test_plain_text_tool_result_remains_string(self): + body = self._transform( + [{"role": "user", "content": "look"}, {"role": "tool", "tool_call_id": "call_1", "content": "done"}] + ) + assert body["messages"][1]["content"][0]["content"] == "done" + + def test_anthropic_tool_schema_strips_only_top_level_schema_key(self): + tools = [ + { + "name": "f", + "input_schema": {"$schema": "schema", "type": "object", "properties": {"$schema": {"type": "string"}}}, + } + ] + body = self._transform([{"role": "user", "content": "hi"}], {"tools": tools}) + schema = body["tools"][0]["input_schema"] + assert "$schema" not in schema + assert "$schema" in schema["properties"] + + def test_tool_schema_strips_only_top_level_schema_key(self): + tools = [ + { + "type": "function", + "function": { + "name": "f", + "parameters": { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": {"$schema": {"type": "string"}}, + }, + }, + } + ] + body = self._transform([{"role": "user", "content": "hi"}], {"tools": tools}) + schema = body["tools"][0]["input_schema"] + assert "$schema" not in schema + assert "$schema" in schema["properties"] + + def test_streaming_tool_identity_is_emitted_only_on_start(self): + handler = SnowflakeStreamingHandler(streaming_response=[], sync_stream=True) + start = handler.chunk_parser( + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "tool_use", "id": "tool_1", "name": "read"}, + } + ) + first_delta = handler.chunk_parser( + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '{"path":'}, + } + ) + second_delta = handler.chunk_parser( + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '"/tmp"}'}, + } + ) + + def _tool_call(chunk): + return chunk.choices[0].delta.tool_calls[0] + + assert _tool_call(start).id == "tool_1" + assert _tool_call(start).function.name == "read" + assert _tool_call(first_delta).id is None + assert _tool_call(first_delta).function.name is None + assert _tool_call(second_delta).id is None + assert _tool_call(second_delta).function.name is None + assert _tool_call(first_delta).function.arguments == '{"path":' + assert _tool_call(second_delta).function.arguments == '"/tmp"}' + + def test_signed_thinking_blocks_lead_the_assistant_turn(self): + """Multi-turn tool use with thinking only works if the signed block is echoed back first.""" + body = self._transform( + [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": None, + "thinking_blocks": [ + {"type": "thinking", "thinking": "391", "signature": "Eto"}, + {"type": "thinking", "thinking": "unsigned"}, + ], + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "read", "arguments": "{}"}} + ], + }, + ] + ) + blocks = body["messages"][1]["content"] + assert blocks[0] == {"type": "thinking", "thinking": "391", "signature": "Eto"} + assert [b["type"] for b in blocks] == ["thinking", "tool_use"] + + def test_signed_thinking_blocks_lead_a_plain_text_assistant_turn(self): + """A thinking response without a tool call must also round-trip on the next request.""" + body = self._transform( + [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": "391", + "thinking_blocks": [{"type": "thinking", "thinking": "391", "signature": "Eto"}], + }, + {"role": "user", "content": "continue"}, + ] + ) + assert body["messages"][1] == { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "391", "signature": "Eto"}, + {"type": "text", "text": "391"}, + ], + } + + def test_signed_thinking_blocks_preserve_list_content(self): + """Cached assistant text reaches this transform as a content list, not a string.""" + body = self._transform( + [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [{"type": "text", "text": "391", "cache_control": {"type": "ephemeral"}}], + "thinking_blocks": [{"type": "thinking", "thinking": "391", "signature": "Eto"}], + }, + {"role": "user", "content": "continue"}, + ] + ) + assert body["messages"][1] == { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "391", "signature": "Eto"}, + {"type": "text", "text": "391", "cache_control": {"type": "ephemeral"}}, + ], + } + + def test_thinking_only_assistant_turn_sends_no_empty_text_block(self): + """Anthropic-shaped APIs reject empty text blocks, so a content-less thinking turn is thinking only.""" + body = self._transform( + [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": None, + "thinking_blocks": [{"type": "thinking", "thinking": "391", "signature": "Eto"}], + }, + {"role": "user", "content": "continue"}, + ] + ) + assert body["messages"][1]["content"] == [{"type": "thinking", "thinking": "391", "signature": "Eto"}] + + def test_streaming_surfaces_thinking_and_prompt_cache_usage(self): + """Cortex streams thinking deltas, signatures and cache counts; all must reach the caller.""" + handler = SnowflakeStreamingHandler(streaming_response=[], sync_stream=True) + handler.chunk_parser( + { + "type": "message_start", + "message": {"usage": {"input_tokens": 18, "cache_creation_input_tokens": 1323}}, + } + ) + thinking = handler.chunk_parser( + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "thinking_delta", "thinking": "391"}, + } + ) + signature = handler.chunk_parser( + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "signature_delta", "signature": "Eto"}, + } + ) + final = handler.chunk_parser( + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"output_tokens": 8, "cache_read_input_tokens": 1323}, + } + ) + + assert thinking.choices[0].delta.reasoning_content == "391" + assert signature.choices[0].delta.thinking_blocks[0]["signature"] == "Eto" + assert final.usage.prompt_tokens_details.cached_tokens == 1323 + + class TestSnowFlakeCompletion: model_name = "mistral" @@ -380,10 +749,7 @@ class TestSnowFlakeCompletion: # PAT key was used post_kwargs = mock_post.call_args_list[-1][1] assert "xxxxx" in post_kwargs["headers"]["Authorization"] - assert ( - post_kwargs["headers"]["X-Snowflake-Authorization-Token-Type"] - == "PROGRAMMATIC_ACCESS_TOKEN" - ) + assert post_kwargs["headers"]["X-Snowflake-Authorization-Token-Type"] == "PROGRAMMATIC_ACCESS_TOKEN" # account id was used assert "AAAA-BBBB" in post_kwargs["url"] @@ -495,9 +861,7 @@ class TestSnowflakeChatCompletion: ) mock_post.assert_called_once() else: - with patch.object( - AsyncHTTPHandler, "post", new_callable=AsyncMock, return_value=mock_resp - ) as mock_post: + with patch.object(AsyncHTTPHandler, "post", new_callable=AsyncMock, return_value=mock_resp) as mock_post: response = asyncio.run( acompletion( model="snowflake/mistral-7b", @@ -580,8 +944,4 @@ class TestSnowflakeChatCompletion: chunks_received = asyncio.run(_run()) assert len(chunks_received) > 0 - content = "".join( - c.choices[0].delta.content - for c in chunks_received - if c.choices[0].delta.content - ) + content = "".join(c.choices[0].delta.content for c in chunks_received if c.choices[0].delta.content) diff --git a/tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py b/tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py index fb21e2e6f6b..7970f7771fc 100644 --- a/tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py +++ b/tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py @@ -338,7 +338,7 @@ class TestAnthropicConfigRequest: litellm_params={}, headers={}, ) - assert body["system"] == "You are helpful." + assert body["system"] == [{"type": "text", "text": "You are helpful."}] assert all(m["role"] != "system" for m in body["messages"]) assert body["messages"][0] == {"role": "user", "content": "Hello"} @@ -422,6 +422,64 @@ class TestAnthropicConfigResponse: assert result.usage.completion_tokens == 5 assert result.usage.total_tokens == 15 + def test_prompt_cache_usage_is_surfaced(self): + """Cortex reports cache creation/read counts; dropping them hides caching and bills cached input at full price.""" + raw = httpx.Response( + 200, + json={ + "id": "msg_1", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "hi"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 18, "cache_creation_input_tokens": 1323, "cache_read_input_tokens": 0}, + }, + ) + result = self.cfg.transform_response( + model="snowflake/claude-sonnet-4-6", + raw_response=raw, + model_response=ModelResponse(), + logging_obj=_mock_logging(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert result.usage.prompt_tokens == 1341 + assert result.usage.prompt_tokens_details.cache_creation_tokens == 1323 + assert result.usage.prompt_tokens_details.cached_tokens == 0 + + def test_thinking_block_and_signature_are_preserved(self): + """The signature must survive so a client can echo the thinking block on the next turn.""" + raw = httpx.Response( + 200, + json={ + "id": "msg_1", + "model": "claude-sonnet-4-6", + "content": [ + {"type": "thinking", "thinking": "391", "signature": "Eto"}, + {"type": "text", "text": "391"}, + ], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 5}, + }, + ) + result = self.cfg.transform_response( + model="snowflake/claude-sonnet-4-6", + raw_response=raw, + model_response=ModelResponse(), + logging_obj=_mock_logging(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + message = result.choices[0].message + assert message.content == "391" + assert message.reasoning_content == "391" + assert message.thinking_blocks[0]["signature"] == "Eto" + def test_stop_reason_end_turn_maps_to_stop(self): raw = _make_anthropic_response() result = self.cfg.transform_response( diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 0764aec7185..4afb8303d03 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -297,6 +297,25 @@ def test_native_bridge_loader_caches_absent_extension(monkeypatch): assert attempts == 1 +def test_native_bridge_loader_reset_forces_relookup(monkeypatch): + real_import = builtins.__import__ + attempts = 0 + + def fake_import(name, globals=None, locals=None, fromlist=(), level=0): + nonlocal attempts + if name == "litellm.rust_bridge" and "_native" in fromlist: + attempts += 1 + raise ImportError + return real_import(name, globals, locals, fromlist, level) + + monkeypatch.setattr(builtins, "__import__", fake_import) + + assert rust_bridge_loader.get_native_bridge() is None + rust_bridge_loader.reset_native_bridge_cache() + assert rust_bridge_loader.get_native_bridge() is None + assert attempts == 2 + + def test_native_bridge_available_reflects_loader(monkeypatch): fake_module = types.ModuleType("litellm.rust_bridge._native") monkeypatch.setattr(rust_bridge_loader, "get_native_bridge", lambda: fake_module) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py index 7b82e004f37..5d6d47b8c38 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py @@ -7,6 +7,7 @@ round-trips exactly, a store failure never escapes into the login path, and a sa rotation re-encrypts stored rows like the sibling per-user credential tables. """ +import asyncio import json import os import time @@ -16,8 +17,10 @@ import jwt as pyjwt import pytest from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + _ASSERTION_CACHE, AssertionStoreUnavailable, DbSSOAssertionStore, + SSOAssertionCache, assertion_from_sso_login, ema_assertion_retention_enabled, fetch_sso_identity_assertion, @@ -25,7 +28,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_s retain_sso_identity_assertion_for_ema, rotate_sso_identity_assertions_master_key, ) -from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper, encrypt_value_helper from litellm.types.mcp import MCPAuth SALT_KEY = "test-salt-key-for-sso-assertion-tests-1234" @@ -38,6 +41,11 @@ def _set_salt_key(monkeypatch): monkeypatch.setenv("LITELLM_SALT_KEY", SALT_KEY) +@pytest.fixture(autouse=True) +def _flush_assertion_cache(): + _ASSERTION_CACHE.flush() + + def _make_id_token(exp_offset: int = 3600, iss: str = ISSUER) -> str: return pyjwt.encode( {"iss": iss, "sub": "u1", "exp": int(time.time()) + exp_offset}, @@ -52,9 +60,7 @@ def _make_prisma(stored: dict, db_has_id_jag_server: bool = False): ``db_has_id_jag_server`` drives the retention gate's authoritative DB fallback; it is wired explicitly so the gate never reads a truthy bare MagicMock.""" prisma = MagicMock() - prisma.db.litellm_mcpservertable.find_first = AsyncMock( - return_value=MagicMock() if db_has_id_jag_server else None - ) + prisma.db.litellm_mcpservertable.find_first = AsyncMock(return_value=MagicMock() if db_has_id_jag_server else None) async def _upsert(where, data): stored[where["user_id"]] = data["update"]["assertion_b64"] @@ -235,6 +241,78 @@ async def test_persist_overwrites_previous_login(): assert fetched.refresh_token is not None +@pytest.mark.asyncio +async def test_fetch_serves_second_read_from_cache_without_db_read(): + stored = {} + prisma = _make_prisma(stored) + cache = SSOAssertionCache() + token = _make_id_token() + assertion = assertion_from_sso_login(token, "rt_1") + with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary + await persist_sso_identity_assertion("user-a", assertion, cache=cache) + first = await fetch_sso_identity_assertion("user-a", cache=cache) + second = await fetch_sso_identity_assertion("user-a", cache=cache) + assert first is not None + assert second is not None + assert first.id_token.get_secret_value() == token + assert second.id_token.get_secret_value() == token + prisma.db.litellm_ssoidentityassertion.find_unique.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_persist_busts_cache_so_relogin_is_visible_immediately(): + stored = {} + prisma = _make_prisma(stored) + cache = SSOAssertionCache() + first_token = _make_id_token(exp_offset=100) + second_token = _make_id_token(exp_offset=7200) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary + await persist_sso_identity_assertion("user-a", assertion_from_sso_login(first_token, None), cache=cache) + first = await fetch_sso_identity_assertion("user-a", cache=cache) + await persist_sso_identity_assertion("user-a", assertion_from_sso_login(second_token, "rt_new"), cache=cache) + second = await fetch_sso_identity_assertion("user-a", cache=cache) + assert first is not None + assert second is not None + assert first.id_token.get_secret_value() == first_token + assert second.id_token.get_secret_value() == second_token + + +@pytest.mark.asyncio +async def test_fetch_racing_a_relogin_does_not_cache_the_previous_assertion(): + stored = {} + prisma = _make_prisma(stored) + cache = SSOAssertionCache() + first_token = _make_id_token(exp_offset=100) + second_token = _make_id_token(exp_offset=7200) + db_read_started = asyncio.Event() + relogin_done = asyncio.Event() + unpaused_find_unique = prisma.db.litellm_ssoidentityassertion.find_unique + + async def _paused_find_unique(where): + row = await unpaused_find_unique(where=where) + db_read_started.set() + await relogin_done.wait() + return row + + prisma.db.litellm_ssoidentityassertion.find_unique = _paused_find_unique + with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary + await persist_sso_identity_assertion("user-a", assertion_from_sso_login(first_token, None), cache=cache) + racing_fetch = asyncio.create_task(fetch_sso_identity_assertion("user-a", cache=cache)) + await db_read_started.wait() + await persist_sso_identity_assertion( + "user-a", + assertion_from_sso_login(second_token, "rt_new"), + cache=cache, + ) + relogin_done.set() + raced = await racing_fetch + after = await fetch_sso_identity_assertion("user-a", cache=cache) + assert raced is not None + assert after is not None + assert raced.id_token.get_secret_value() == first_token + assert after.id_token.get_secret_value() == second_token + + @pytest.mark.asyncio async def test_fetch_missing_row_returns_none(): prisma = _make_prisma({}) @@ -242,6 +320,18 @@ async def test_fetch_missing_row_returns_none(): assert await fetch_sso_identity_assertion("nobody") is None +@pytest.mark.asyncio +async def test_fetch_does_not_cache_a_missing_row(): + prisma = _make_prisma({}) + cache = SSOAssertionCache() + with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary + first = await fetch_sso_identity_assertion("nobody", cache=cache) + second = await fetch_sso_identity_assertion("nobody", cache=cache) + assert first is None + assert second is None + assert prisma.db.litellm_ssoidentityassertion.find_unique.await_count == 2 + + @pytest.mark.asyncio async def test_fetch_undecryptable_row_returns_none(): prisma = _make_prisma({"user-a": "not-an-encrypted-blob"}) @@ -251,13 +341,37 @@ async def test_fetch_undecryptable_row_returns_none(): @pytest.mark.asyncio async def test_fetch_unparseable_payload_returns_none(): - from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper - prisma = _make_prisma({"user-a": encrypt_value_helper("]]not json")}) with patch("litellm.proxy.proxy_server.prisma_client", prisma): assert await fetch_sso_identity_assertion("user-a") is None +@pytest.mark.asyncio +async def test_cached_assertion_expires_after_ttl(): + stored = {} + prisma = _make_prisma(stored) + cache = SSOAssertionCache(ttl_seconds=1) + first_token = _make_id_token(exp_offset=100) + second_token = _make_id_token(exp_offset=7200) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary + await persist_sso_identity_assertion("user-a", assertion_from_sso_login(first_token, None), cache=cache) + first = await fetch_sso_identity_assertion("user-a", cache=cache) + await persist_sso_identity_assertion( + "user-a", + assertion_from_sso_login(second_token, "rt_new"), + cache=SSOAssertionCache(), + ) + cached = await fetch_sso_identity_assertion("user-a", cache=cache) + time.sleep(1.1) + expired = await fetch_sso_identity_assertion("user-a", cache=cache) + assert first is not None + assert cached is not None + assert expired is not None + assert first.id_token.get_secret_value() == first_token + assert cached.id_token.get_secret_value() == first_token + assert expired.id_token.get_secret_value() == second_token + + @pytest.mark.asyncio async def test_retain_noop_when_no_id_jag_server(): stored = {} @@ -357,6 +471,24 @@ async def test_db_store_converts_a_driver_failure_into_assertion_store_unavailab await DbSSOAssertionStore().fetch("alice") +@pytest.mark.asyncio +async def test_db_store_uses_injected_cache(): + stored = {} + prisma = _make_prisma(stored) + cache = SSOAssertionCache() + token = _make_id_token() + with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary + await persist_sso_identity_assertion("alice", assertion_from_sso_login(token, None), cache=cache) + store = DbSSOAssertionStore(cache=cache) + first = await store.fetch("alice") + second = await store.fetch("alice") + assert first is not None + assert second is not None + assert first.id_token.get_secret_value() == token + assert second.id_token.get_secret_value() == token + prisma.db.litellm_ssoidentityassertion.find_unique.assert_awaited_once() + + @pytest.mark.asyncio async def test_db_store_returns_none_for_a_user_with_no_stored_assertion(): """An absent row stays an absence, not an outage, so a user who never signed in still gets the 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 0fd35e674b7..9a6815a61e5 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 @@ -8201,6 +8201,129 @@ class TestPreemptive401ModeAware: await self._run(delegate, self.LITELLM_KEY_HEADERS, has_stored_token=False) +class TestSingleServerPreflightReachesIdJag: + """The connect-time preflight is what turns a credential failure into an HTTP status the client + can read. An oauth2_id_jag server has to reach it: its subject comes from the assertion stored at + SSO login, so the failure is decided before any IdP call and there is nothing later in the session + that can report it (tools/list degrades to an empty list, tools/call to 'tool not found').""" + + def _id_jag_server(self) -> MCPServer: + return MCPServer( + server_id="id-idjag", + name="idjag", + alias="idjag", + server_name="idjag", + url="https://idjag.test/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_id_jag, + client_id="gateway-client", + client_secret="gateway-secret", + token_exchange_endpoint="https://org-idp.test/oauth2/token", + id_jag_resource_token_endpoint="https://resource-as.test/oauth2/token", + mcp_info={"server_name": "idjag"}, + ) + + async def _run(self, server: MCPServer, mcp_servers: list[str], preflight: AsyncMock) -> None: + from litellm.proxy._experimental.mcp_server import server as server_module + + with ( + patch.object( # test-quality-ok: route wiring must use the manager's configured server + server_module.global_mcp_server_manager, + "get_mcp_server_by_name", + return_value=server, + ), + patch.object( # test-quality-ok: route wiring must invoke the manager preflight + server_module.global_mcp_server_manager, + "preflight_token_exchange", + preflight, + ), + patch.object( # test-quality-ok: allowed-set resolution needs the DB; the test controls its answer + server_module, "_get_allowed_mcp_servers", AsyncMock(return_value=[server]) + ), + ): + await server_module._raise_preemptive_401_for_unauthenticated_servers( + scope={"type": "http", "method": "POST", "path": "/mcp/idjag", "headers": []}, + mcp_servers=mcp_servers, + oauth2_headers={"Authorization": "Bearer sk-litellm-virtual-key"}, + mcp_server_auth_headers=None, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-virtual-key", user_id="u-1"), + client_ip=None, + ) + + @pytest.mark.asyncio + async def test_id_jag_single_server_route_surfaces_the_preflight_status(self): + """The 412 the preflight raises must propagate out of connect, not be swallowed.""" + server = self._id_jag_server() + preflight = AsyncMock(side_effect=HTTPException(status_code=412, detail="no stored assertion")) + + with pytest.raises(HTTPException) as exc: + await self._run(server, ["idjag"], preflight) + + assert exc.value.status_code == 412 + assert preflight.await_args.kwargs["server"] is server + + @pytest.mark.asyncio + async def test_token_exchange_without_a_bearer_still_challenges_and_never_pre_flights(self): + """The already-shipped OBO path must be untouched by the call site dropping its mode test. + A token_exchange server with no inbound bearer has nothing to exchange, so it still gets the + RFC 9728 discovery challenge from the block above and the preflight is never reached; pushing + a subject-less exchange through the resolver would turn that challenge into some other status + and strand a client that only had to SSO and retry.""" + from litellm.proxy._experimental.mcp_server import server as server_module + + token_exchange = MCPServer( + server_id="id-obo", + name="obo", + alias="obo", + server_name="obo", + url="https://obo.test/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_token_exchange, + token_exchange_endpoint="https://idp.test/oauth2/token", + client_id="cid", + client_secret="csec", + mcp_info={"server_name": "obo"}, + ) + preflight = AsyncMock() + + with ( + patch.object( # test-quality-ok: route wiring must use the manager's configured server + server_module.global_mcp_server_manager, + "get_mcp_server_by_name", + return_value=token_exchange, + ), + patch.object( # test-quality-ok: route wiring must invoke the manager preflight + server_module.global_mcp_server_manager, + "preflight_token_exchange", + preflight, + ), + pytest.raises(HTTPException) as exc, + ): + await server_module._raise_preemptive_401_for_unauthenticated_servers( + scope={"type": "http", "method": "POST", "path": "/mcp/obo", "headers": []}, + mcp_servers=["obo"], + oauth2_headers=None, + mcp_server_auth_headers=None, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-virtual-key", user_id="u-1"), + client_ip=None, + ) + + assert exc.value.status_code == 401 + headers = exc.value.headers or {} + assert "resource_metadata" in (headers.get("WWW-Authenticate") or headers.get("www-authenticate") or "") + preflight.assert_not_awaited() + + @pytest.mark.asyncio + async def test_id_jag_multi_server_route_still_absorbs_the_failure(self): + """The aggregate contract is unchanged: with more than one target the preflight does not run, + so one server with no stored assertion cannot fail the whole connect.""" + preflight = AsyncMock(side_effect=HTTPException(status_code=412, detail="no stored assertion")) + + await self._run(self._id_jag_server(), ["idjag", "other"], preflight) + + preflight.assert_not_awaited() + + def _make_obo_server(alias: str) -> MCPServer: return MCPServer( server_id=f"id-{alias}", 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 482f779bcb8..02b1a19081a 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 @@ -2616,6 +2616,165 @@ class TestMCPServerManager: await manager.preflight_token_exchange(server=server, oauth2_headers=None, user_api_key_auth=None) assert resolved == ["good-subject"] + def _id_jag_server(self, server_id: str) -> "MCPServer": + return MCPServer( + server_id=server_id, + name=f"{server_id}-server", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_id_jag, + client_id="gateway-client", + client_secret="gateway-secret", + token_exchange_endpoint="https://org-idp.example/oauth2/token", + id_jag_resource_token_endpoint="https://resource-as.example/oauth2/token", + ) + + @pytest.mark.asyncio + async def test_preflight_id_jag_surfaces_missing_assertion_as_a_plain_412(self): + """ID-JAG's missing/expired-assertion precondition must reach the client as a 412 whose body + names the fix, at the transport edge. Without the preflight the session opens and the caller + gets a 200 with an empty tool list and then 'tool not found', which is not what happened. + 412 is a precondition, not an RFC 9728 discovery challenge, so it carries no + WWW-Authenticate: there is nothing for the client to discover and retry against.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError + + summary = ( + "ID-JAG requires an IdP identity assertion for this user and none is stored. " + "Sign in through LiteLLM SSO so the gateway captures one." + ) + + class _FakeProvider: + async def resolve_credentials(self, subject, server): + return Error(CredError.of_precondition_required(summary)) + + manager = MCPServerManager(cred_provider=_FakeProvider()) + + with pytest.raises(HTTPException) as exc_info: + await manager.preflight_token_exchange( + server=self._id_jag_server("id-jag-preflight-412"), + oauth2_headers=None, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-virtual-key", user_id="u-1"), + ) + assert exc_info.value.status_code == 412 + assert summary in exc_info.value.detail + assert not (exc_info.value.headers or {}) + + @pytest.mark.asyncio + async def test_preflight_id_jag_surfaces_assertion_store_outage_as_503(self): + """A store outage is the other failure the session would swallow, and it is a different + answer than 412: the user has nothing to fix by signing in again.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError + + class _FakeProvider: + async def resolve_credentials(self, subject, server): + return Error(CredError.of_upstream_unavailable("assertion store unreachable")) + + manager = MCPServerManager(cred_provider=_FakeProvider()) + + with pytest.raises(HTTPException) as exc_info: + await manager.preflight_token_exchange( + server=self._id_jag_server("id-jag-preflight-503"), + oauth2_headers=None, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-virtual-key", user_id="u-1"), + ) + assert exc_info.value.status_code == 503 + + @pytest.mark.asyncio + async def test_preflight_id_jag_preflights_litellm_key_and_skips_identity_bearer(self): + """ID-JAG preflights when Authorization carries a LiteLLM key, but skips a caller identity + bearer that the session passes through unchanged.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + StaticHeaderAuth, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok + + subjects = [] + + class _FakeProvider: + async def resolve_credentials(self, subject, server): + subjects.append( + ( + subject.subject_id, + subject.inbound_token.get_secret_value() if subject.inbound_token else None, + ) + ) + return Ok(StaticHeaderAuth("Bearer minted-id-jag", header_name="Authorization")) + + manager = MCPServerManager(cred_provider=_FakeProvider()) + caller = UserAPIKeyAuth(api_key="sk-litellm-virtual-key", user_id="u-1") + + await manager.preflight_token_exchange( + server=self._id_jag_server("id-jag-preflight-key"), + oauth2_headers={"Authorization": "Bearer sk-litellm-virtual-key"}, + raw_headers={"authorization": "Bearer sk-litellm-virtual-key"}, + user_api_key_auth=caller, + ) + assert subjects == [("u-1", None)] + + await manager.preflight_token_exchange( + server=self._id_jag_server("id-jag-preflight-identity"), + oauth2_headers={"Authorization": "Bearer caller-idp-id-token"}, + raw_headers={ + "x-litellm-api-key": "Bearer sk-admission-key", + "authorization": "Bearer caller-idp-id-token", + }, + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="u-1"), + ) + assert subjects == [("u-1", None)] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "server_fields", + [ + {"auth_type": MCPAuth.none}, + {"auth_type": MCPAuth.api_key, "authentication_token": "static-upstream-key"}, + {"auth_type": MCPAuth.bearer_token, "authentication_token": "static-upstream-key"}, + { + "auth_type": MCPAuth.oauth2, + "oauth2_flow": "client_credentials", + "client_id": "cid", + "client_secret": "csec", + "token_url": "https://idp.example.com/token", + }, + {"auth_type": MCPAuth.true_passthrough}, + ], + ) + async def test_preflight_resolves_nothing_for_a_mode_that_does_not_pre_flight(self, server_fields): + """The manager is the only thing deciding which modes pre-flight, so it has to reject every + other mode itself. The single-server call site no longer tests the mode before calling, so a + mode that falls through here would start resolving its credential a second time, at connect, + for flows that never had a connect-time resolution at all.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError + + calls = [] + + class _FakeProvider: + async def resolve_credentials(self, subject, server): + calls.append(server.server_id) + return Error(CredError.of_misconfigured("the preflight must never get here")) + + manager = MCPServerManager(cred_provider=_FakeProvider()) + server = MCPServer( + server_id="not-pre-flighted", + name="not-pre-flighted-server", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + **server_fields, + ) + + assert ( + await manager.preflight_token_exchange( + server=server, + oauth2_headers={"Authorization": "Bearer sk-litellm-virtual-key"}, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-virtual-key", user_id="u-1"), + ) + is None + ) + assert calls == [] + @pytest.mark.asyncio @pytest.mark.parametrize( "authorization", @@ -2638,7 +2797,9 @@ class TestMCPServerManager: resolved: Final[list[str | None]] = [] class _FakeProvider: - async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Ok[StaticHeaderAuth, CredError]: + async def resolve_credentials( + self, subject: Subject, server: ServerSpec + ) -> Ok[StaticHeaderAuth, CredError]: resolved.append(subject.inbound_token.get_secret_value() if subject.inbound_token else None) return Ok(StaticHeaderAuth("Bearer MINTED", header_name="Authorization")) @@ -2664,7 +2825,9 @@ class TestMCPServerManager: resolved: Final[list[str | None]] = [] class _FakeProvider: - async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Ok[StaticHeaderAuth, CredError]: + async def resolve_credentials( + self, subject: Subject, server: ServerSpec + ) -> Ok[StaticHeaderAuth, CredError]: resolved.append(subject.inbound_token.get_secret_value() if subject.inbound_token else None) return Ok(StaticHeaderAuth("Bearer MINTED", header_name="Authorization")) diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py index c83ba142011..f809fadc879 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py @@ -125,11 +125,12 @@ class TestBlockedResponseUsage: mock_logging.post_call_failure_hook.assert_awaited_once() -class TestProxyExceptionPassthrough: +class TestProxyExceptionAnthropicEnvelope: @pytest.mark.asyncio - async def test_anthropic_response_reraises_proxy_exception_unwrapped(self): - """A 400 ProxyException from request validation must surface as-is, - not be re-wrapped into a code-500 ProxyException.""" + async def test_anthropic_response_maps_proxy_exception_to_anthropic_envelope(self): + """LIT-6468: a 400 ProxyException from request validation must surface as + Anthropic's documented {"type": "error", "error": {...}} envelope with the + original status and message, not the OpenAI {"error": {...}} envelope.""" import litellm.proxy.anthropic_endpoints.endpoints as ep import litellm.proxy.proxy_server as proxy_server from litellm.proxy._types import ProxyErrorTypes, ProxyException @@ -140,6 +141,8 @@ class TestProxyExceptionPassthrough: param="metadata", code=400, ) + request = MagicMock() + request.headers = {"x-request-id": "req_test_6468"} with ( patch.object(ep, "_read_request_body", new=AsyncMock(return_value={})), @@ -151,30 +154,61 @@ class TestProxyExceptionPassthrough: patch.object(proxy_server, "proxy_logging_obj") as mock_logging, ): mock_logging.post_call_failure_hook = AsyncMock() - with pytest.raises(ProxyException) as exc_info: - await ep.anthropic_response( - fastapi_response=MagicMock(), - request=MagicMock(), - user_api_key_dict=MagicMock(), - ) + response = await ep.anthropic_response( + fastapi_response=MagicMock(), + request=request, + user_api_key_dict=MagicMock(), + ) - assert exc_info.value is exc - assert exc_info.value.code == "400" - assert exc_info.value.param == "metadata" + assert response.status_code == 400 + body = json.loads(response.body) + assert body == { + "type": "error", + "error": { + "type": "invalid_request_error", + "message": "Invalid type for 'metadata': expected an object, but got a string instead.", + }, + "request_id": "req_test_6468", + } mock_logging.post_call_failure_hook.assert_awaited_once() + @pytest.mark.asyncio + async def test_anthropic_response_maps_429_to_rate_limit_error(self): + """The Anthropic error type follows the status code (429 -> rate_limit_error), + and a code-less exception falls back to 500 api_error.""" + import litellm.proxy.anthropic_endpoints.endpoints as ep + from litellm.proxy._types import ProxyException + + request = MagicMock() + request.headers = {} + + response = ep._anthropic_error_json_response( + ProxyException(message="Rate limit exceeded", type="rate_limit_error", param=None, code=429), + request, + ) + assert response.status_code == 429 + assert json.loads(response.body)["error"]["type"] == "rate_limit_error" + + fallback = ep._anthropic_error_json_response( + ProxyException(message="boom", type="None", param=None, code=None), + request, + ) + assert fallback.status_code == 500 + assert json.loads(fallback.body)["error"]["type"] == "api_error" + class TestHttpExceptionDictDetail: @pytest.mark.asyncio async def test_anthropic_response_serializes_dict_detail_http_exception(self): - """LIT-6466: a post_call guardrail's HTTPException(detail=) must - surface with a clean message plus provider_specific_fields, matching - /v1/chat/completions and /v1/responses, not the str() of the exception.""" + """LIT-6466 + LIT-6468: a post_call guardrail's HTTPException(detail=) + must surface as Anthropic's {"type": "error", "error": {...}} envelope with + the guardrail's clean message plus provider_specific_fields, not the str() + of the exception and not the OpenAI envelope.""" from fastapi import HTTPException import litellm.proxy.anthropic_endpoints.endpoints as ep import litellm.proxy.proxy_server as proxy_server - from litellm.proxy._types import ProxyException, UserAPIKeyAuth + from litellm.proxy._types import UserAPIKeyAuth detail = { "error": "Content blocked: keyword 'kumquat' detected", @@ -182,6 +216,8 @@ class TestHttpExceptionDictDetail: "guardrail": "keyword-block", } exc = HTTPException(status_code=400, detail=detail) + request = MagicMock() + request.headers = {} with ( patch.object(ep, "_read_request_body", new=AsyncMock(return_value={})), # test-quality-ok: endpoint reads the body via a module function; no injection seam @@ -193,17 +229,19 @@ class TestHttpExceptionDictDetail: patch.object(proxy_server, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global imported at call time; no injection seam ): mock_logging.post_call_failure_hook = AsyncMock() - with pytest.raises(ProxyException) as exc_info: - await ep.anthropic_response( - fastapi_response=MagicMock(), - request=MagicMock(), - user_api_key_dict=UserAPIKeyAuth(), - ) + response = await ep.anthropic_response( + fastapi_response=MagicMock(), + request=request, + user_api_key_dict=UserAPIKeyAuth(), + ) - assert exc_info.value.message == "Content blocked: keyword 'kumquat' detected" - assert "{'error'" not in exc_info.value.message - assert exc_info.value.provider_specific_fields == detail - assert exc_info.value.code == "400" + assert response.status_code == 400 + body = json.loads(response.body) + assert body["type"] == "error" + assert body["error"]["type"] == "invalid_request_error" + assert body["error"]["message"] == "Content blocked: keyword 'kumquat' detected" + assert "{'error'" not in body["error"]["message"] + assert body["error"]["provider_specific_fields"] == detail mock_logging.post_call_failure_hook.assert_awaited_once() @@ -215,7 +253,7 @@ class TestFailureHookRequestData: handler must pass that replaced dict, not the raw request body dict.""" import litellm.proxy.anthropic_endpoints.endpoints as ep import litellm.proxy.proxy_server as proxy_server - from litellm.proxy._types import ProxyException, UserAPIKeyAuth + from litellm.proxy._types import UserAPIKeyAuth captured = {} @@ -224,18 +262,23 @@ class TestFailureHookRequestData: captured["processor_data"] = self.data raise RuntimeError("provider timeout") + request = MagicMock() + request.headers = {} + with ( patch.object(ep, "_read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet"})), patch.object(ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fake_process), patch.object(proxy_server, "proxy_logging_obj") as mock_logging, ): mock_logging.post_call_failure_hook = AsyncMock() - with pytest.raises(ProxyException): - await ep.anthropic_response( - fastapi_response=MagicMock(), - request=MagicMock(), - user_api_key_dict=UserAPIKeyAuth(), - ) + response = await ep.anthropic_response( + fastapi_response=MagicMock(), + request=request, + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert response.status_code == 500 + assert json.loads(response.body)["error"]["message"] == "provider timeout" hook_request_data = mock_logging.post_call_failure_hook.await_args.kwargs["request_data"] assert hook_request_data is captured["processor_data"] diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index 62073f4bf51..d58683fd1e5 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -754,6 +754,84 @@ def test_expand_wildcard_invalid_litellm_params_passthrough(): assert result == [deployment] +def test_get_complete_model_list_excludes_wildcard_routes_by_default(): + """Regression (LIT-4108): a wildcard with a matching router deployment leaked into /v1/models.""" + from litellm import Router + from litellm.proxy.auth.model_checks import get_complete_model_list + + router = Router( + model_list=[ + { + "model_name": "bedrock/*", + "litellm_params": {"model": "bedrock/*"}, + }, + { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4"}, + }, + ] + ) + + result = get_complete_model_list( + key_models=[], + team_models=[], + proxy_model_list=["bedrock/*", "gpt-4"], + user_model=None, + infer_model_from_keys=False, + return_wildcard_routes=False, + llm_router=router, + ) + + assert "bedrock/*" not in result + assert "gpt-4" in result + assert any(m.startswith("bedrock/") for m in result) + + +def test_get_complete_model_list_excludes_wildcard_routes_without_router(): + from litellm.proxy.auth.model_checks import get_complete_model_list + + result = get_complete_model_list( + key_models=[], + team_models=[], + proxy_model_list=["bedrock/*", "gpt-4"], + user_model=None, + infer_model_from_keys=False, + return_wildcard_routes=False, + llm_router=None, + ) + + assert "bedrock/*" not in result + assert "gpt-4" in result + assert any(m.startswith("bedrock/") for m in result) + + +def test_get_complete_model_list_includes_wildcard_routes_when_requested(): + from litellm import Router + from litellm.proxy.auth.model_checks import get_complete_model_list + + router = Router( + model_list=[ + { + "model_name": "bedrock/*", + "litellm_params": {"model": "bedrock/*"}, + }, + ] + ) + + result = get_complete_model_list( + key_models=[], + team_models=[], + proxy_model_list=["bedrock/*"], + user_model=None, + infer_model_from_keys=False, + return_wildcard_routes=True, + llm_router=router, + ) + + assert result.count("bedrock/*") == 1 + assert any(m.startswith("bedrock/") and m != "bedrock/*" for m in result) + + def test_add_known_models_refreshes_models_by_provider_for_wildcard_expansion(): """models_by_provider was a frozen import-time snapshot of set unions, so cost map reloads (which call add_known_models) never reached wildcard expansion until a diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index e3f71692c78..0e90c107865 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -1238,6 +1238,18 @@ def test_health_liveness_endpoint(proxy_client): print(f"\n/health/liveness response time: {duration_ms:.2f}ms") +def test_health_backlog_includes_admission_control_stats(proxy_client): + response = proxy_client.get("/health/backlog") + + assert response.status_code == 200, response.text + assert set(response.json()) == { + "in_flight_requests", + "admitted_requests", + "queued_requests", + "rejected_requests", + } + + def test_health_readiness(proxy_client): """ Test /health/readiness endpoint. diff --git a/tests/test_litellm/proxy/list_api/test_list_framework.py b/tests/test_litellm/proxy/list_api/test_list_framework.py index 6ed3ab369c2..ab7ef229e1a 100644 --- a/tests/test_litellm/proxy/list_api/test_list_framework.py +++ b/tests/test_litellm/proxy/list_api/test_list_framework.py @@ -25,6 +25,7 @@ from litellm.proxy.list_api.list_framework import ( SortKey, Within, build_query_plan, + handle_facet, handle_list, order_by_sql, where_sql, @@ -892,3 +893,107 @@ def test_the_facet_page_shapes_are_untouched_by_page_mode(): assert set(links) == {"self", "prev", "next"} assert links["next"] == "/management/v1/budgets?q=ac&page=3" + + +# ------------------------------------------------------- facet request handling + + +class RecordingFacetExecutor: + """Records the one call `handle_facet` is allowed to make, so a rejected request + can be shown never to have reached it.""" + + def __init__(self, values: tuple[str, ...] = ()) -> None: + self.values = values + self.field: str | None = None + self.where: tuple[object, ...] | None = None + + async def distinct(self, field: str, where: tuple[object, ...]) -> Sequence[str]: + self.field = field + self.where = where + return self.values + + +async def _facet_problem(query: str, spec: ListSpec[BudgetRow, BudgetOut] | None = None) -> ProblemDetail: + executor = RecordingFacetExecutor(values=("a", "b")) + with pytest.raises(ManagementProblem) as raised: + await handle_facet( + spec=spec or _spec(), + executor=executor, + request=_request(query), + caller=CALLER, + field="created_by", + ) + assert executor.field is None, "a rejected facet request still queried the executor" + return raised.value.problem + + +@pytest.mark.asyncio +async def test_a_facet_conjoins_the_scope_with_the_callers_filters(): + """The scope is the one predicate a caller cannot drop, so a facet has to add to it + rather than replace it: otherwise a dropdown lists values from rows the caller + cannot see in the table.""" + spec = _spec(scope=lambda caller: ScopeWhere(where=(Compare(field="created_by", op="eq", value="caller-1"),))) + executor = RecordingFacetExecutor(values=("caller-1",)) + + response = await handle_facet( + spec=spec, + executor=executor, + request=_request("filter[max_budget][gte]=5&q=ac"), + caller=CALLER, + field="created_by", + ) + + assert tuple(response.data) == ("caller-1",) + assert executor.where == ( + Compare(field="created_by", op="eq", value="caller-1"), + Compare(field="max_budget", op="gte", value=5.0), + AnyOf( + clauses=( + Compare(field="budget_id", op="contains", value="ac"), + Compare(field="created_by", op="contains", value="ac"), + ) + ), + ) + + +@pytest.mark.asyncio +async def test_a_denied_scope_on_a_facet_never_reaches_the_executor(): + """A 200 with an empty list would read as "no such values" rather than "not yours".""" + problem = await _facet_problem("", spec=_spec(scope=lambda caller: ScopeDenied(reason="nope"))) + + assert problem.status == 403 + assert problem.type == f"{PROBLEM_TYPE_BASE}forbidden" + + +@pytest.mark.asyncio +async def test_a_facet_rejects_a_filter_operator_its_spec_does_not_offer(): + problem = await _facet_problem("filter[created_by][gte]=x") + + assert problem.status == 400 + assert "gte" in problem.detail + + +@pytest.mark.asyncio +async def test_a_facet_rejects_a_repeated_query_parameter(): + problem = await _facet_problem("page=1&page=2") + + assert problem.type == f"{PROBLEM_TYPE_BASE}duplicate-query-parameter" + assert "page" in problem.detail + + +@pytest.mark.asyncio +@pytest.mark.parametrize("query", ("page=0", "page=one")) +async def test_a_facet_rejects_a_page_that_is_not_a_positive_integer(query: str): + problem = await _facet_problem(query) + + assert problem.status == 400 + assert problem.type == f"{PROBLEM_TYPE_BASE}invalid-query-parameter" + assert "'page'" in problem.detail + + +@pytest.mark.asyncio +async def test_a_facet_rejects_a_page_size_that_is_not_a_positive_integer(): + problem = await _facet_problem("page_size=0") + + assert problem.status == 400 + assert "'page_size'" in problem.detail diff --git a/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_endpoints.py b/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_endpoints.py index 4e063dd0c5b..86f9aeafc08 100644 --- a/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_endpoints.py @@ -225,3 +225,71 @@ def test_compute_overall_action_all_passed(): def test_compute_overall_action_empty(): assert _compute_overall_action([]) == "passed" + + +class TestEnrichPolicyTemplateStreamKeepalive: + async def _collect_endpoint_body(self, monkeypatch, interval, delay=0.3) -> tuple[list[bytes], dict]: + import asyncio + from unittest.mock import MagicMock + + import litellm + import litellm.proxy.management_endpoints.policy_endpoints.endpoints as policy_endpoints + import litellm.proxy.proxy_server as proxy_server + from fastapi.responses import StreamingResponse + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.policy_endpoints.endpoints import ( + EnrichTemplateRequest, + enrich_policy_template_stream, + ) + + monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", interval) + + async def _name_chunks(): + await asyncio.sleep(delay) + chunk = MagicMock() + chunk.choices = [MagicMock()] + chunk.choices[0].delta.content = "Rival Air\n" + yield chunk + + class SlowRouter: + async def acompletion(self, **kwargs): + return _name_chunks() + + async def _no_variations(competitors, model): + return {} + + monkeypatch.setattr(proxy_server, "llm_router", SlowRouter()) + monkeypatch.setattr(policy_endpoints, "_generate_competitor_variations", _no_variations) + + response = await enrich_policy_template_stream( + data=EnrichTemplateRequest( + template_id="competitor-mention-detection", + parameters={"brand_name": "Acme"}, + model="gpt-5.4-mini", + ), + request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert isinstance(response, StreamingResponse) + chunks = [chunk if isinstance(chunk, bytes) else chunk.encode() async for chunk in response.body_iterator] + return chunks, dict(response.headers) + + @pytest.mark.asyncio + async def test_endpoint_pings_while_competitor_discovery_is_still_running(self, monkeypatch): + chunks, headers = await self._collect_endpoint_body(monkeypatch, interval=0.05) + + assert headers["content-type"].startswith("text/event-stream") + assert headers["cache-control"] == "no-cache" + assert headers["x-accel-buffering"] == "no" + assert chunks[0] == b": ping\n\n" + assert chunks.count(b": ping\n\n") >= 3 + assert b'data: {"type": "competitor", "name": "Rival Air"}\n\n' in chunks + assert chunks[-1].startswith(b'data: {"type": "done"') + + @pytest.mark.asyncio + async def test_endpoint_stream_is_untouched_while_keepalives_are_unconfigured(self, monkeypatch): + chunks, _ = await self._collect_endpoint_body(monkeypatch, interval=None, delay=0.15) + + assert b": ping\n\n" not in chunks + assert chunks[0] == b'data: {"type": "competitor", "name": "Rival Air"}\n\n' + assert chunks[-1].startswith(b'data: {"type": "done"') diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index 1697b77b99a..60f9a1a55e2 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -19,11 +19,13 @@ from litellm.proxy._types import ( NewUserResponse, ProxyErrorTypes, ProxyException, + UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.scim.scim_v2 import ( SCIMRosterSyncError, UserProvisionerHelpers, _apply_group_patch_updates, + _create_user_if_not_exists, _extract_group_member_ids, _extract_ids_from_path_filter, _handle_group_membership_changes, @@ -37,8 +39,8 @@ from litellm.proxy.management_endpoints.scim.scim_v2 import ( delete_group, delete_user, get_groups, - get_users, get_service_provider_config, + get_users, merge_placeholder, patch_group, patch_team_membership, @@ -304,6 +306,85 @@ async def test_create_user_uses_default_internal_user_params_role(mocker, monkey assert called_args.user_role == LitellmUserRoles.PROXY_ADMIN +def _mock_scim_create_user_deps(mocker: MockerFixture, scim_user: SCIMUser) -> AsyncMock: + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable into create_user + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable into create_user + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=scim_user), + ) + return mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable into create_user + "litellm.proxy.management_endpoints.scim.scim_v2.new_user", + AsyncMock(return_value=NewUserRequest(user_id=scim_user.userName)), + ) + + +@pytest.mark.asyncio +async def test_create_user_without_groups_defers_to_default_team(mocker: MockerFixture, monkeypatch): + """IdPs omit groups on POST /Users; teams must stay unset so new_user applies default_internal_user_params.teams""" + scim_user = SCIMUser( + schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + userName="new-user", + emails=[SCIMUserEmail(value="new@example.com")], + ) + monkeypatch.setattr( + "litellm.default_internal_user_params", + {"teams": [{"team_id": "default-team", "max_budget_in_team": 25}]}, + raising=False, + ) + new_user_mock = _mock_scim_create_user_deps(mocker, scim_user) + + await create_user(user=scim_user) + + assert new_user_mock.call_args.kwargs["data"].teams is None + assert new_user_mock.call_args.kwargs["user_api_key_dict"] == UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + + +@pytest.mark.asyncio +async def test_create_user_with_groups_keeps_idp_teams(mocker: MockerFixture, monkeypatch): + scim_user = SCIMUser( + schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + userName="new-user", + emails=[SCIMUserEmail(value="new@example.com")], + groups=[SCIMUserGroup(value="idp-team")], + ) + monkeypatch.setattr( + "litellm.default_internal_user_params", + {"teams": [{"team_id": "default-team", "max_budget_in_team": 25}]}, + raising=False, + ) + new_user_mock = _mock_scim_create_user_deps(mocker, scim_user) + + await create_user(user=scim_user) + + assert new_user_mock.call_args.kwargs["data"].teams == ["idp-team"] + + +@pytest.mark.asyncio +async def test_create_user_if_not_exists_defers_to_default_team(mocker: MockerFixture, monkeypatch): + monkeypatch.setattr( + "litellm.default_internal_user_params", + {"teams": [{"team_id": "default-team", "max_budget_in_team": 25}]}, + raising=False, + ) + new_user_mock = mocker.patch( # test-quality-ok: new_user is imported inside the helper, not injectable + "litellm.proxy.management_endpoints.internal_user_endpoints.new_user", + AsyncMock(return_value=NewUserResponse(user_id="group-user", key="k")), + ) + + created = await _create_user_if_not_exists(user_id="group-user") + + assert created is not None + assert new_user_mock.call_args.kwargs["data"].teams is None + assert new_user_mock.call_args.kwargs["user_api_key_dict"] == UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + + @pytest.mark.asyncio async def test_scim_create_user_respects_default_role_set_via_ui(mocker, monkeypatch): """ @@ -1176,6 +1257,67 @@ async def test_update_user_success(mocker): assert call_args[1]["data"]["teams"] == ["new-team"] +@pytest.mark.asyncio +@pytest.mark.parametrize("groups", [None, []], ids=["groups-omitted", "groups-empty"]) +async def test_update_user_without_groups_preserves_memberships_and_role(mocker, monkeypatch, groups): + """Okta profile PUTs carry no `groups` or `groups: []`; neither may drop teams (and their keys) or recompute role""" + from litellm.proxy.proxy_server import proxy_config + + async def mock_get_config(): + return {"litellm_settings": {"scim_admin_group": "litellm-admins"}} + + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False) + + existing_user = mocker.MagicMock() + existing_user.teams = ["litellm-admins", "engineering"] + existing_user.metadata = {} + + scim_user = SCIMUser( + schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + userName="okta-user", + name=SCIMUserName(familyName="Renamed", givenName="Okta"), + emails=[SCIMUserEmail(value="okta@example.com")], + **({} if groups is None else {"groups": groups}), + ) + response_scim_user = SCIMUser( + schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + id="okta-user", + userName="okta-user", + emails=[SCIMUserEmail(value="okta@example.com")], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={"user_id": "okta-user"}) + + mocker.patch( # test-quality-ok: update_user's collaborators are module-level, not injectable + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + mocker.patch( # test-quality-ok: update_user's collaborators are module-level, not injectable + "litellm.proxy.management_endpoints.scim.scim_v2._check_user_exists", + AsyncMock(return_value=existing_user), + ) + mocker.patch( # test-quality-ok: update_user's collaborators are module-level, not injectable + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=response_scim_user), + ) + patch_membership = mocker.patch( # test-quality-ok: roster writes are module-level, not injectable + "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", + AsyncMock(), + ) + + result = await update_user(user_id="okta-user", user=scim_user) + + assert result == response_scim_user + patch_membership.assert_not_awaited() + update_data = mock_prisma_client.db.litellm_usertable.update.call_args.kwargs["data"] + assert update_data["teams"] == ["litellm-admins", "engineering"] + assert "user_role" not in update_data + + @pytest.mark.asyncio async def test_update_user_not_found(mocker): """Should raise 404 when user doesn't exist""" diff --git a/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py index 3a32b3cc128..e5616a1975d 100644 --- a/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py +++ b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py @@ -466,3 +466,61 @@ class TestUsageAiChatServiceAccountGuard: is_admin=False, ) assert "Endpoint-level guard missing" in str(exc_info.value) + + +class TestUsageAiChatKeepalive: + async def _collect_endpoint_body(self, monkeypatch, interval, delay=0.3) -> tuple[list[bytes], dict]: + import asyncio + + import litellm + from fastapi.responses import StreamingResponse + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.usage_endpoints.endpoints import ( + ChatMessage, + UsageAIChatRequest, + usage_ai_chat, + ) + + monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", interval) + + async def slow_acompletion(**kwargs): + await asyncio.sleep(delay) + response = MagicMock() + response.choices = [MagicMock()] + response.choices[0].message.tool_calls = None + response.choices[0].message.content = "Total spend is $50.25" + return response + + with patch( # test-quality-ok: the stream calls the module-level litellm.acompletion directly; no injection seam + "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm.acompletion", + new=AsyncMock(side_effect=slow_acompletion), + ): + response = await usage_ai_chat( + data=UsageAIChatRequest(messages=[ChatMessage(role="user", content="hi")], model="gpt-4o-mini"), + request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert isinstance(response, StreamingResponse) + chunks = [chunk if isinstance(chunk, bytes) else chunk.encode() async for chunk in response.body_iterator] + return chunks, dict(response.headers) + + @pytest.mark.asyncio + async def test_endpoint_pings_while_the_planning_completion_is_still_running(self, monkeypatch): + chunks, headers = await self._collect_endpoint_body(monkeypatch, interval=0.05) + + assert headers["content-type"].startswith("text/event-stream") + assert headers["cache-control"] == "no-cache" + assert headers["x-accel-buffering"] == "no" + assert chunks[0].startswith(b'data: {"type": "status"') + assert chunks[1] == b": ping\n\n" + assert chunks.count(b": ping\n\n") >= 3 + assert b'"content": "Total spend is $50.25"' in b"".join(chunks) + assert chunks[-1] == b'data: {"type": "done"}\n\n' + + @pytest.mark.asyncio + async def test_endpoint_stream_is_untouched_while_keepalives_are_unconfigured(self, monkeypatch): + chunks, _ = await self._collect_endpoint_body(monkeypatch, interval=None, delay=0.15) + + assert b": ping\n\n" not in chunks + assert chunks[0].startswith(b'data: {"type": "status"') + assert chunks[-1] == b'data: {"type": "done"}\n\n' diff --git a/tests/test_litellm/proxy/middleware/test_admission_control_middleware.py b/tests/test_litellm/proxy/middleware/test_admission_control_middleware.py new file mode 100644 index 00000000000..f1ca13daa03 --- /dev/null +++ b/tests/test_litellm/proxy/middleware/test_admission_control_middleware.py @@ -0,0 +1,402 @@ +import asyncio +import json +from typing import Final + +import pytest +from starlette.middleware.base import BaseHTTPMiddleware +from starlette.types import ASGIApp, Message, Receive, Scope, Send + +from litellm.proxy.middleware.admission_control_middleware import ( + AdmissionControlMetrics, + AdmissionControlMiddleware, + AdmissionControlSettings, + AdmissionControlState, + AdmissionControlStats, + _parse_admission_control_settings, + create_prometheus_admission_metrics, + get_admission_control_settings, +) + + +@pytest.fixture +def state() -> AdmissionControlState: + return AdmissionControlState(lambda: None) + + +async def _call( + middleware: AdmissionControlMiddleware, + path: str = "/", + root_path: str = "", +) -> tuple[Message, ...]: + messages: Final[list[Message]] = [] + + async def receive() -> Message: + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(message: Message) -> None: + messages.append(message) + + scope: Final[Scope] = { + "type": "http", + "path": path, + "root_path": root_path, + "method": "GET", + "headers": [], + } + await middleware(scope, receive, send) + return tuple(messages) + + +def _handler_with_release( + started: asyncio.Event, + release: asyncio.Event, +) -> ASGIApp: + async def handler(scope: Scope, receive: Receive, send: Send) -> None: + started.set() + await release.wait() + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"ok", "more_body": False}) + + return handler + + +def test_is_not_base_http_middleware() -> None: + assert not issubclass(AdmissionControlMiddleware, BaseHTTPMiddleware) + + +@pytest.mark.asyncio +async def test_capacity_rejects_excess_and_releases_queued_request(state: AdmissionControlState) -> None: + started: Final = asyncio.Event() + release: Final = asyncio.Event() + middleware: Final = AdmissionControlMiddleware( + _handler_with_release(started, release), + lambda: AdmissionControlSettings(1, 1, 1.0), + state, + ) + + first: Final = asyncio.create_task(_call(middleware)) + await started.wait() + second: Final = asyncio.create_task(_call(middleware)) + await asyncio.sleep(0) + assert state.get_stats().queued == 1 + + third: Final = await _call(middleware) + assert third[0]["status"] == 503 + headers: Final = dict(third[0]["headers"]) + assert headers[b"retry-after"] == b"1" + assert headers[b"content-type"] == b"application/json" + assert json.loads(third[1]["body"])["error"] == { + "message": "Worker at capacity: 1 in-flight, 1 queued requests. Retry later.", + "type": "overloaded_error", + "code": "503", + } + assert state.get_stats().rejected_total == 1 + + release.set() + assert (await first)[0]["status"] == 200 + assert (await second)[0]["status"] == 200 + assert state.get_stats() == AdmissionControlStats(0, 0, 1) + + +@pytest.mark.asyncio +async def test_pending_waiter_is_not_skipped_after_admission_is_released(state: AdmissionControlState) -> None: + started: Final = asyncio.Event() + release: Final = asyncio.Event() + third_trigger: Final = asyncio.Event() + middleware: Final = AdmissionControlMiddleware( + _handler_with_release(started, release), + lambda: AdmissionControlSettings(1, 2, 1.0), + state, + ) + + first: Final = asyncio.create_task(_call(middleware)) + await started.wait() + second: Final = asyncio.create_task(_call(middleware)) + await asyncio.sleep(0) + + async def call_third() -> tuple[Message, ...]: + await third_trigger.wait() + return await _call(middleware) + + third: Final = asyncio.create_task(call_third()) + await asyncio.sleep(0) + release.set() + third_trigger.set() + await asyncio.sleep(0) + + assert state.get_stats().queued == 2 + await asyncio.gather(first, second, third) + + +@pytest.mark.asyncio +async def test_queue_timeout_rejects_and_decrements_queue(state: AdmissionControlState) -> None: + started: Final = asyncio.Event() + release: Final = asyncio.Event() + middleware: Final = AdmissionControlMiddleware( + _handler_with_release(started, release), + lambda: AdmissionControlSettings(1, 1, 0.05), + state, + ) + + first: Final = asyncio.create_task(_call(middleware)) + await started.wait() + start_time: Final = asyncio.get_running_loop().time() + second: Final = await _call(middleware) + elapsed: Final = asyncio.get_running_loop().time() - start_time + + assert second[0]["status"] == 503 + assert elapsed < 0.5 + assert state.get_stats().queued == 0 + assert state.get_stats().rejected_total == 1 + release.set() + await first + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("root_path", "probe_path"), + ( + ("", "/health/liveliness"), + ("/proxy", "/proxy/health/liveliness"), + ("/proxy", "/proxy/metrics"), + ), +) +async def test_exempt_path_passes_through_when_saturated( + state: AdmissionControlState, + root_path: str, + probe_path: str, +) -> None: + started: Final = asyncio.Event() + release: Final = asyncio.Event() + + async def handler(scope: Scope, receive: Receive, send: Send) -> None: + if scope["path"] == "/": + started.set() + await release.wait() + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"ok", "more_body": False}) + + middleware: Final = AdmissionControlMiddleware(handler, lambda: AdmissionControlSettings(1, 0, 1.0), state) + + first: Final = asyncio.create_task(_call(middleware)) + await started.wait() + health: Final = await _call(middleware, probe_path, root_path) + assert health[0]["status"] == 200 + blocked: Final = await _call(middleware, "/proxy/v1/chat/completions", root_path) + assert blocked[0]["status"] == 503 + lookalike: Final = await _call(middleware, "/proxyhealth/liveliness", "/proxy") + assert lookalike[0]["status"] == 503 + release.set() + await first + + +@pytest.mark.asyncio +async def test_non_http_scope_passes_through_when_saturated(state: AdmissionControlState) -> None: + seen: Final[list[str]] = [] + + async def handler(scope: Scope, receive: Receive, send: Send) -> None: + seen.append(scope["type"]) + + middleware: Final = AdmissionControlMiddleware(handler, lambda: AdmissionControlSettings(1, 0, 1.0), state) + state.record_admission() + + async def receive() -> Message: + return {"type": "lifespan.startup"} + + async def send(message: Message) -> None: + return None + + await middleware({"type": "lifespan"}, receive, send) + assert seen == ["lifespan"] + + +@pytest.mark.asyncio +async def test_none_settings_does_not_limit_concurrency() -> None: + active: Final = [0] + peak: Final = [0] + all_started: Final = asyncio.Event() + release: Final = asyncio.Event() + + async def handler(scope: Scope, receive: Receive, send: Send) -> None: + active[0] += 1 + peak[0] = max(peak[0], active[0]) + if active[0] == 3: + all_started.set() + await release.wait() + active[0] -= 1 + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"ok", "more_body": False}) + + middleware: Final = AdmissionControlMiddleware(handler, lambda: None, AdmissionControlState(lambda: None)) + requests: Final = tuple(asyncio.create_task(_call(middleware)) for _ in range(3)) + await all_started.wait() + assert peak[0] == 3 + release.set() + results: Final = await asyncio.gather(*requests) + assert tuple(result[0]["status"] for result in results) == (200, 200, 200) + + +@pytest.mark.asyncio +async def test_cancelling_queued_request_does_not_leak_counter(state: AdmissionControlState) -> None: + started: Final = asyncio.Event() + release: Final = asyncio.Event() + middleware: Final = AdmissionControlMiddleware( + _handler_with_release(started, release), + lambda: AdmissionControlSettings(1, 1, 1.0), + state, + ) + + first: Final = asyncio.create_task(_call(middleware)) + await started.wait() + queued: Final = asyncio.create_task(_call(middleware)) + await asyncio.sleep(0) + queued.cancel() + with pytest.raises(asyncio.CancelledError): + await queued + assert state.get_stats().queued == 0 + release.set() + await first + + +@pytest.mark.asyncio +async def test_streaming_response_holds_admission_until_final_body(state: AdmissionControlState) -> None: + first_chunk_sent: Final = asyncio.Event() + finish_stream: Final = asyncio.Event() + + async def handler(scope: Scope, receive: Receive, send: Send) -> None: + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"first", "more_body": True}) + first_chunk_sent.set() + await finish_stream.wait() + await send({"type": "http.response.body", "body": b"last", "more_body": False}) + + middleware: Final = AdmissionControlMiddleware( + handler, + lambda: AdmissionControlSettings(1, 1, 1.0), + state, + ) + first: Final = asyncio.create_task(_call(middleware)) + await first_chunk_sent.wait() + second: Final = asyncio.create_task(_call(middleware)) + await asyncio.sleep(0) + assert not second.done() + assert state.get_stats().queued == 1 + finish_stream.set() + assert (await first)[0]["status"] == 200 + assert (await second)[0]["status"] == 200 + assert state.get_stats().admitted == 0 + assert state.get_stats().queued == 0 + + +class _FakeGauge: + def __init__(self) -> None: + self.value = 0.0 + + def inc(self, amount: float = 1) -> None: + self.value += amount + + def dec(self, amount: float = 1) -> None: + self.value -= amount + + +class _FakeCounter: + def __init__(self) -> None: + self.by_reason: Final[dict[str, _FakeGauge]] = {} + + def labels(self, reason: str) -> _FakeGauge: + return self.by_reason.setdefault(reason, _FakeGauge()) + + +@pytest.mark.asyncio +async def test_metrics_track_admitted_queued_and_rejected() -> None: + admitted: Final = _FakeGauge() + queued: Final = _FakeGauge() + rejected: Final = _FakeCounter() + state: Final = AdmissionControlState( + lambda: AdmissionControlMetrics(admitted_gauge=admitted, queued_gauge=queued, rejected_counter=rejected) + ) + started: Final = asyncio.Event() + release: Final = asyncio.Event() + middleware: Final = AdmissionControlMiddleware( + _handler_with_release(started, release), + lambda: AdmissionControlSettings(1, 1, 0.05), + state, + ) + + first: Final = asyncio.create_task(_call(middleware)) + await started.wait() + second: Final = asyncio.create_task(_call(middleware)) + await asyncio.sleep(0) + assert (admitted.value, queued.value) == (1.0, 1.0) + await _call(middleware) + assert rejected.by_reason["queue_full"].value == 1.0 + await second + assert rejected.by_reason["queue_timeout"].value == 1.0 + release.set() + await first + assert (admitted.value, queued.value) == (0.0, 0.0) + + +def test_create_prometheus_admission_metrics_registers_named_metrics() -> None: + from prometheus_client import REGISTRY + + metrics: Final = create_prometheus_admission_metrics() + if metrics is not None: + metrics.admitted_gauge.inc() + metrics.queued_gauge.inc() + metrics.rejected_counter.labels(reason="queue_full").inc() + assert REGISTRY.get_sample_value("litellm_admission_admitted_requests") == 1.0 + assert REGISTRY.get_sample_value("litellm_admission_queued_requests") == 1.0 + assert REGISTRY.get_sample_value("litellm_admission_rejected_requests_total", {"reason": "queue_full"}) is not None + assert create_prometheus_admission_metrics() is None + + +@pytest.mark.parametrize( + ("settings", "expected"), + ( + ({}, None), + ({"max_in_flight_requests_per_worker": None}, None), + ({"max_in_flight_requests_per_worker": 0}, None), + ({"max_in_flight_requests_per_worker": "many"}, None), + ({"max_in_flight_requests_per_worker": 3, "max_queued_requests_per_worker": -1}, None), + ({"max_in_flight_requests_per_worker": 3, "admission_queue_timeout_seconds": 0}, None), + ({"max_in_flight_requests_per_worker": 3, "admission_queue_timeout_seconds": -0.5}, None), + ( + {"max_in_flight_requests_per_worker": 3, "max_queued_requests_per_worker": 0}, + AdmissionControlSettings(3, 0, 1.0), + ), + ( + {"max_in_flight_requests_per_worker": 3}, + AdmissionControlSettings(3, 3, 1.0), + ), + ( + { + "max_in_flight_requests_per_worker": 3, + "max_queued_requests_per_worker": 5, + "admission_queue_timeout_seconds": 0.25, + }, + AdmissionControlSettings(3, 5, 0.25), + ), + ), +) +def test_get_admission_control_settings( + settings: dict[str, object], + expected: AdmissionControlSettings | None, +) -> None: + assert get_admission_control_settings(settings) == expected + + +def test_invalid_admission_control_settings_logs_once(caplog: pytest.LogCaptureFixture) -> None: + _parse_admission_control_settings.cache_clear() + caplog.set_level("ERROR") + settings: Final = {"max_in_flight_requests_per_worker": [1]} + + assert get_admission_control_settings(settings) is None + assert get_admission_control_settings(settings) is None + + messages: Final = tuple( + record.message + for record in caplog.records + if record.message.startswith("Ignoring invalid admission control settings") + ) + assert len(messages) == 1 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 1d4b0264879..acb45038df0 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 @@ -319,9 +319,6 @@ class TestVertexAIPassThroughHandler: mock_handler.get_default_base_target_url.return_value = ( f"https://{test_location}-aiplatform.googleapis.com/" ) - mock_handler.update_base_target_url_with_credential_location = Mock( - return_value=f"https://{test_location}-aiplatform.googleapis.com/" - ) mock_get_handler.return_value = mock_handler # Mock create_pass_through_route to return a function that returns a mock response @@ -427,9 +424,6 @@ class TestVertexAIPassThroughHandler: mock_handler.get_default_base_target_url.return_value = ( "https://aiplatform.googleapis.com/" ) - mock_handler.update_base_target_url_with_credential_location = Mock( - return_value="https://aiplatform.googleapis.com/" - ) mock_get_handler.return_value = mock_handler # Mock create_pass_through_route to return a function that returns a mock response @@ -530,9 +524,6 @@ class TestVertexAIPassThroughHandler: mock_handler.get_default_base_target_url.return_value = ( f"https://{default_location}-aiplatform.googleapis.com/" ) - mock_handler.update_base_target_url_with_credential_location = Mock( - return_value=f"https://{default_location}-aiplatform.googleapis.com/" - ) mock_get_handler.return_value = mock_handler # Mock create_pass_through_route to return a function that returns a mock response @@ -1308,9 +1299,6 @@ class TestVertexAIDiscoveryPassThroughHandler: mock_handler.get_default_base_target_url.return_value = ( "https://discoveryengine.googleapis.com" ) - mock_handler.update_base_target_url_with_credential_location = Mock( - return_value="https://discoveryengine.googleapis.com" - ) mock_get_handler.return_value = mock_handler # Mock create_pass_through_route to return a function that returns a mock response @@ -3650,7 +3638,6 @@ class TestVertexRawPredictStreamingClassification: base_url = "https://us-east5-aiplatform.googleapis.com/" mock_handler = Mock() mock_handler.get_default_base_target_url.return_value = base_url - mock_handler.update_base_target_url_with_credential_location = Mock(return_value=base_url) module = "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints" with ( @@ -4234,6 +4221,126 @@ class TestVertexCredentiallessPassthroughVirtualKeyLeak: assert "sk-master-1234" not in " ".join(f"{name}:{value}" for name, value in forwarded.items()) +class TestVertexPassthroughDefaultLocationOnShortRoutes: + PROJECT = "test-project" + SHORT_ROUTE = "publishers/google/models/gemini-2.5-flash:generateContent" + + @staticmethod + def _forwarder() -> Mock: + return Mock(return_value=AsyncMock(return_value={"status": "success"})) + + async def _forward( + self, + monkeypatch, + endpoint: str, + default_config: dict | None, + headers: list[tuple[bytes, bytes]], + forwarder: Mock, + ) -> None: + from litellm.proxy.pass_through_endpoints.passthrough_endpoint_router import ( + PassthroughEndpointRouter, + ) + + async def receive(): + return {"type": "http.request", "body": b"{}", "more_body": False} + + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": f"/vertex_ai/{endpoint}", + "headers": headers, + "query_string": b"", + }, + receive=receive, + ) + router: Final = PassthroughEndpointRouter() + if default_config is not None: + router.set_default_vertex_config(dict(default_config)) + module: Final = "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints" + monkeypatch.setattr(f"{module}.passthrough_endpoint_router", router) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + mock_credentials: Final = Mock() + mock_credentials.token = "test-token" + caller: Final = UserAPIKeyAuth(api_key="test-key") + with ( + mock.patch( # test-quality-ok: the route mints its Google token through its own VertexBase, nothing injects the credential loader + "litellm.llms.vertex_ai.vertex_llm_base.VertexBase.load_auth", + return_value=(mock_credentials, self.PROJECT), + ), + mock.patch(f"{module}.create_pass_through_route", new=forwarder), + mock.patch(f"{module}.user_api_key_auth", new=AsyncMock(return_value=caller)), + ): + await vertex_proxy_route( + endpoint=endpoint, + request=request, + fastapi_response=Response(), + user_api_key_dict=caller, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("endpoint", "location", "expected_target"), + [ + ( + SHORT_ROUTE, + "global", + "https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/" + SHORT_ROUTE, + ), + ( + f"v1/{SHORT_ROUTE}", + "global", + "https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/" + SHORT_ROUTE, + ), + ( + f"v1beta1/{SHORT_ROUTE}", + "global", + "https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/" + SHORT_ROUTE, + ), + ( + SHORT_ROUTE, + "us-central1", + "https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/" + + SHORT_ROUTE, + ), + ], + ) + async def test_default_vertex_config_location_fills_routes_without_project_and_location( + self, monkeypatch, endpoint, location, expected_target + ): + forwarder: Final = self._forwarder() + await self._forward( + monkeypatch, + endpoint, + {"vertex_project": self.PROJECT, "vertex_location": location, "vertex_credentials": "test-creds"}, + [(b"content-type", b"application/json"), (b"authorization", b"Bearer test-key")], + forwarder, + ) + forwarded: Final = forwarder.call_args.kwargs + assert str(forwarded["target"]) == expected_target + assert forwarded["custom_headers"]["Authorization"] == "Bearer test-token" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("default_config", "headers"), + [ + (None, [(b"content-type", b"application/json"), (b"authorization", b"Bearer ya29.byo-google-oauth")]), + ( + {"vertex_project": PROJECT, "vertex_credentials": "test-creds"}, + [(b"content-type", b"application/json"), (b"authorization", b"Bearer test-key")], + ), + ], + ) + async def test_no_location_anywhere_is_a_400_not_a_500(self, monkeypatch, default_config, headers): + forwarder: Final = self._forwarder() + with pytest.raises(HTTPException) as raised: + await self._forward(monkeypatch, self.SHORT_ROUTE, default_config, headers, forwarder) + forwarder.assert_not_called() + assert raised.value.status_code == 400 + assert "/projects//locations//" in str(raised.value.detail) + assert "default_vertex_config" in str(raised.value.detail) + + class TestGetAzureAISearchIndexFromEndpoint: """The operable index is only the segment right after ``indexes``. @@ -5009,3 +5116,93 @@ class TestAzureRouterModelStreamingDispatch: assert result.status_code == 200 body = b"".join([chunk async for chunk in result.body_iterator]) assert body == upstream_body + + +class TestAzureRouterModelStreamingKeepalive: + async def _dispatch(self, monkeypatch, interval, headers_delay=0.0, body_delay=0.0) -> StreamingResponse: + import asyncio + + import litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm.passthrough.main import AsyncPassthroughStreamingResponse + + monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", interval) + + class _StallingBody(httpx.AsyncByteStream): + async def __aiter__(self): + await asyncio.sleep(body_delay) + yield b"data: hello\n\n" + + async def _upstream_response() -> httpx.Response: + await asyncio.sleep(headers_delay) + return httpx.Response( + 200, + headers={"content-type": "text/event-stream", "x-upstream": "kept"}, + stream=_StallingBody(), + request=httpx.Request("POST", "https://my-azure.openai.azure.com/openai/deployments/gpt-5/x"), + ) + + logging_obj = MagicMock() + logging_obj.async_flush_passthrough_collected_chunks = AsyncMock() + + class StreamingRouter: + async def allm_passthrough_route(self, **kwargs): + return await AsyncPassthroughStreamingResponse( + response=_upstream_response(), + litellm_logging_obj=logging_obj, + provider_config=MagicMock(), + ) + + async def fake_get_request_body(_request): + return {"model": "gpt-5", "stream": True} + + monkeypatch.setattr(proxy_server, "llm_router", StreamingRouter()) + monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) + monkeypatch.setattr(ep, "is_passthrough_request_using_router_model", lambda *a, **k: True) + + request = MagicMock(spec=Request) + request.method = "POST" + request.headers = {"content-type": "application/json"} + request.query_params = {} + + result = await azure_proxy_route( + endpoint="openai/deployments/gpt-5/chat/completions", + request=request, + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"), + ) + assert isinstance(result, StreamingResponse) + return result + + @pytest.mark.asyncio + async def test_pings_while_upstream_headers_are_still_pending(self, monkeypatch): + result = await self._dispatch(monkeypatch, interval=0.05, headers_delay=0.3) + + chunks = [chunk async for chunk in result.body_iterator] + + assert result.status_code == 200 + assert result.headers["x-accel-buffering"] == "no" + assert chunks[0] == b": ping\n\n" + assert chunks.count(b": ping\n\n") >= 3 + assert b"".join(chunks).endswith(b"data: hello\n\n") + + @pytest.mark.asyncio + async def test_pings_while_upstream_body_is_still_pending(self, monkeypatch): + result = await self._dispatch(monkeypatch, interval=0.05, body_delay=0.3) + + chunks = [chunk async for chunk in result.body_iterator] + + assert result.status_code == 200 + assert result.headers["x-upstream"] == "kept" + assert chunks[0] == b": ping\n\n" + assert chunks.count(b": ping\n\n") >= 3 + assert chunks[-1] == b"data: hello\n\n" + + @pytest.mark.asyncio + async def test_relays_upstream_bytes_untouched_while_keepalives_are_unconfigured(self, monkeypatch): + result = await self._dispatch(monkeypatch, interval=None, headers_delay=0.15, body_delay=0.15) + + chunks = [chunk async for chunk in result.body_iterator] + + assert result.headers["x-upstream"] == "kept" + assert chunks == [b"data: hello\n\n"] diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py index 6735f2a3780..e8fd5579631 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py @@ -4,6 +4,7 @@ import pytest from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + VertexAIPassThroughHandler, _base_vertex_proxy_route, _upstream_headers_for_vertex_route, ) @@ -20,6 +21,7 @@ async def test_vertex_passthrough_load_balancing(): mock_request = MagicMock() mock_response = MagicMock() mock_handler = MagicMock() + mock_handler.get_default_base_target_url.return_value = "https://test.url" # Mock the router mock_router = MagicMock() @@ -68,7 +70,6 @@ async def test_vertex_passthrough_load_balancing(): mock_pt_router.get_vertex_credentials.return_value = MagicMock() mock_prep_headers.return_value = ( {}, - "https://test.url", False, "test-project-lb", "us-central1-lb", @@ -290,12 +291,6 @@ async def test_vertex_passthrough_forwards_anthropic_beta_header(): mock_vertex_credentials.vertex_location = "us-central1" mock_vertex_credentials.vertex_credentials = "test-credentials" - # Create mock handler - mock_handler = MagicMock() - mock_handler.update_base_target_url_with_credential_location.return_value = ( - "https://us-central1-aiplatform.googleapis.com" - ) - with ( patch.object( VertexBase, @@ -313,7 +308,6 @@ async def test_vertex_passthrough_forwards_anthropic_beta_header(): # Call the function ( headers, - base_target_url, headers_passed_through, vertex_project, vertex_location, @@ -323,8 +317,6 @@ async def test_vertex_passthrough_forwards_anthropic_beta_header(): router_credentials=None, vertex_project="test-project", vertex_location="us-central1", - base_target_url="https://us-central1-aiplatform.googleapis.com", - get_vertex_pass_through_handler=mock_handler, user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-secret-key"), ) @@ -394,7 +386,6 @@ async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens( "content-type": "application/json", "Authorization": "Bearer vertex-access-token", }, - "https://aiplatform.googleapis.com", False, "test-project", "global", @@ -406,7 +397,7 @@ async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens( endpoint=f"{VERTEX_ANTHROPIC_MODELS_PREFIX}{model_segment}", request=MagicMock(), fastapi_response=MagicMock(), - get_vertex_pass_through_handler=MagicMock(), + get_vertex_pass_through_handler=VertexAIPassThroughHandler(), ) upstream_headers = mock_create_route.call_args.kwargs["custom_headers"] @@ -473,12 +464,6 @@ async def test_vertex_passthrough_does_not_forward_litellm_auth_token(): mock_vertex_credentials.vertex_location = "us-central1" mock_vertex_credentials.vertex_credentials = "test-credentials" - # Create mock handler - mock_handler = MagicMock() - mock_handler.update_base_target_url_with_credential_location.return_value = ( - "https://us-central1-aiplatform.googleapis.com" - ) - with ( patch.object( VertexBase, @@ -495,7 +480,6 @@ async def test_vertex_passthrough_does_not_forward_litellm_auth_token(): ( headers, - _base_target_url, _headers_passed_through, _vertex_project, _vertex_location, @@ -505,8 +489,6 @@ async def test_vertex_passthrough_does_not_forward_litellm_auth_token(): router_credentials=None, vertex_project="test-project", vertex_location="us-central1", - base_target_url="https://us-central1-aiplatform.googleapis.com", - get_vertex_pass_through_handler=mock_handler, user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-secret-key"), ) @@ -742,7 +724,6 @@ async def test_vertex_passthrough_custom_model_name_replaced_in_url(): mock_pt_router.get_vertex_credentials.return_value = MagicMock() mock_prep_headers.return_value = ( {}, - "https://global-aiplatform.googleapis.com", False, "nv-gcpllmgwit-20250411173346", "global", diff --git a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py index fdaad567f95..87e10ce7e8d 100644 --- a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py @@ -1819,3 +1819,92 @@ async def test_run_thread_stream_is_untouched_while_keepalives_are_unconfigured( assert not any(chunk.startswith(": ping") for chunk in chunks) assert chunks[-1] == "data: [DONE]\n\n" + + +# --------------------------------------------------------------------------- +# async_queue_request: SSE keepalives during the time-to-first-token +# --------------------------------------------------------------------------- + + +async def _queue_streaming(monkeypatch, interval, delay=0.3, fails_with=None): + _patch_logging_flags(monkeypatch) + monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", interval) + + router = MagicMock() + router.get_model_list.return_value = [] + + async def _schedule_after_the_scheduler_queue_drains(**kwargs): + await asyncio.sleep(delay) + if fails_with is not None: + raise fails_with + return _async_iter([_simple_chunk(content="queued reply")]) + + router.schedule_acompletion = _schedule_after_the_scheduler_queue_drains + monkeypatch.setattr(ps, "llm_router", router) + + request = MagicMock() + request.url = "http://testserver/queue/chat/completions" + request.method = "POST" + request.headers = {} + request.json = AsyncMock( + return_value={ + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + "priority": 0, + "stream": True, + } + ) + request.is_disconnected = AsyncMock(return_value=False) + + return await ps.async_queue_request( + request=request, + fastapi_response=Response(), + user_api_key_dict=_user_auth(), + ) + + +@pytest.mark.asyncio +async def test_queue_request_pings_while_the_scheduler_is_still_waiting(monkeypatch): + response = await _queue_streaming(monkeypatch, interval=0.05) + + assert isinstance(response, StreamingResponse) + assert response.headers["x-accel-buffering"] == "no" + chunks = [chunk async for chunk in response.body_iterator] + + assert chunks[0] == b": ping\n\n" + assert chunks.count(b": ping\n\n") >= 3 + assert b'"content":"queued reply"' in chunks[-2] + assert chunks[-1] == b"data: [DONE]\n\n" + + +@pytest.mark.asyncio +async def test_queue_request_audits_a_failure_that_arrives_after_the_first_ping(monkeypatch): + audited = [] + + async def _record_failure(*, user_api_key_dict, original_exception, request_data, **kwargs): + audited.append(original_exception) + return None + + monkeypatch.setattr(ps.proxy_logging_obj, "post_call_failure_hook", _record_failure) + + boom = RuntimeError("scheduler died after the wire was already open") + response = await _queue_streaming(monkeypatch, interval=0.05, fails_with=boom) + + assert isinstance(response, StreamingResponse) + chunks = [chunk async for chunk in response.body_iterator] + + assert chunks[0] == b": ping\n\n" + assert audited == [boom] + assert json.loads(chunks[-2].removeprefix(b"data: "))["error"]["code"] == "500" + assert chunks[-1] == b"data: [DONE]\n\n" + + +@pytest.mark.asyncio +async def test_queue_request_stream_is_untouched_while_keepalives_are_unconfigured(monkeypatch): + response = await _queue_streaming(monkeypatch, interval=None, delay=0.15) + + assert isinstance(response, StreamingResponse) + chunks = [chunk if isinstance(chunk, bytes) else chunk.encode() async for chunk in response.body_iterator] + + assert not any(chunk.startswith(b": ping") for chunk in chunks) + assert chunks[-1] == b"data: [DONE]\n\n" diff --git a/tests/test_litellm/proxy/public_endpoints/public_v1/test_model_hub.py b/tests/test_litellm/proxy/public_endpoints/public_v1/test_model_hub.py index 631e91dca11..de2d95e9f28 100644 --- a/tests/test_litellm/proxy/public_endpoints/public_v1/test_model_hub.py +++ b/tests/test_litellm/proxy/public_endpoints/public_v1/test_model_hub.py @@ -193,12 +193,12 @@ def test_sorting_by_a_numeric_field_puts_the_unset_ones_last_in_both_directions( def test_an_undeclared_sort_field_is_a_problem_naming_the_allowed_fields(monkeypatch): _publish(monkeypatch, _named(3)) - response = _get("sort=providers") + response = _get("sort=health_status") assert response.status_code == 400 assert response.headers["content-type"].startswith("application/problem+json") body = response.json() - assert "providers" in body["detail"] + assert "health_status" in body["detail"] assert body["allowed"] == [ "input_cost_per_token", "max_input_tokens", @@ -206,6 +206,9 @@ def test_an_undeclared_sort_field_is_a_problem_naming_the_allowed_fields(monkeyp "mode", "model_group", "output_cost_per_token", + "providers", + "rpm", + "tpm", ] @@ -347,3 +350,131 @@ def test_the_endpoint_it_supersedes_still_answers_with_its_bare_array(monkeypatc body = response.json() assert isinstance(body, list) assert [row["model_group"] for row in body] == ["model-000", "model-001", "model-002"] + + +FACET_PATHS = ("providers", "modes", "features") + + +def _facet(name: str, query: str = ""): + suffix = f"?{query}" if query else "" + return client.get(f"{MODEL_HUB_PATH}/{name}{suffix}") + + +def test_providers_filter_accepts_several_providers_at_once(monkeypatch): + """The hub's provider control is a multi-select, so the route has to OR the values.""" + _publish( + monkeypatch, + ( + _info("gpt-4", providers=("openai",)), + _info("claude", providers=("anthropic",)), + _info("mistral-large", providers=("mistral",)), + _info("router", providers=("openai", "anthropic")), + ), + ) + + response = _get("filter[providers][in]=openai,anthropic") + + assert response.status_code == 200, response.text + assert sorted(_groups(response)) == ["claude", "gpt-4", "router"] + + +def test_features_filter_matches_a_model_with_any_of_the_named_features(monkeypatch): + """Selecting two features widens the result set, the way the hub's multi-select always did.""" + _publish( + monkeypatch, + ( + _info("sees", supports_vision=True), + _info("calls", supports_function_calling=True), + _info("both", supports_vision=True, supports_function_calling=True), + _info("plain"), + ), + ) + + response = _get("filter[features][in]=vision,function_calling") + + assert response.status_code == 200, response.text + assert sorted(_groups(response)) == ["both", "calls", "sees"] + + +def test_a_single_feature_filter_selects_only_models_with_it(monkeypatch): + _publish(monkeypatch, (_info("sees", supports_vision=True), _info("plain"), _info("reasons", supports_reasoning=True))) + + assert _groups(_get("filter[features][in]=vision")) == ["sees"] + assert _groups(_get("filter[features][in]=reasoning")) == ["reasons"] + + +def test_providers_and_limits_are_sortable(monkeypatch): + """The hub sorted on these columns before it paged; they stay sortable now that the route orders.""" + _publish( + monkeypatch, + ( + _info("b-model", providers=("mistral",), rpm=10), + _info("a-model", providers=("anthropic",), rpm=30), + _info("c-model", providers=("openai",), rpm=20), + ), + ) + + assert _groups(_get("sort=providers")) == ["a-model", "b-model", "c-model"] + assert _groups(_get("sort=-rpm")) == ["a-model", "c-model", "b-model"] + + +@pytest.mark.parametrize("facet", FACET_PATHS) +def test_a_facet_serves_the_distinct_values_of_its_column(monkeypatch, facet): + _publish( + monkeypatch, + ( + _info("a", providers=("openai",), mode="chat", supports_vision=True), + _info("b", providers=("anthropic", "openai"), mode="embedding", supports_vision=True), + _info("c", providers=("mistral",), mode="chat"), + ), + ) + + response = _facet(facet) + + assert response.status_code == 200, response.text + assert response.json()["data"] == { + "providers": ["anthropic", "mistral", "openai"], + "modes": ["chat", "embedding"], + "features": ["vision"], + }[facet] + + +def test_a_facet_offers_only_values_the_table_can_show(monkeypatch): + """Section 12's reason for hanging facets off the resource: the dropdown matches the filtered table.""" + _publish( + monkeypatch, + ( + _info("chat-openai", providers=("openai",), mode="chat"), + _info("embed-cohere", providers=("cohere",), mode="embedding"), + ), + ) + + assert _facet("providers", "filter[mode][in]=chat").json()["data"] == ["openai"] + assert _facet("providers", "q=embed").json()["data"] == ["cohere"] + + +def test_a_facet_pages_and_reports_whether_more_remain(monkeypatch): + _publish(monkeypatch, tuple(_info(f"m-{index}", providers=(f"p-{index:02d}",)) for index in range(5))) + + first = _facet("providers", "page_size=2") + last = _facet("providers", "page=3&page_size=2") + + assert first.json()["data"] == ["p-00", "p-01"] + assert first.json()["meta"] == {"page": 1, "page_size": 2, "has_more": True} + assert last.json()["data"] == ["p-04"] + assert last.json()["meta"]["has_more"] is False + + +def test_a_facet_rejects_a_sort_it_does_not_offer(monkeypatch): + """Facet values are always ascending, so `sort` is not part of the facet contract.""" + _publish(monkeypatch, _named(3)) + + response = _facet("providers", "sort=-providers") + + assert response.status_code == 400 + assert response.json()["type"].endswith("unknown-query-parameter") + + +@pytest.mark.parametrize("facet", FACET_PATHS) +def test_a_facet_is_reachable_without_a_key(facet): + assert f"{MODEL_HUB_PATH}/{facet}" in LiteLLMRoutes.public_routes.value diff --git a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py index a176e91eaa4..832435711c6 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -11,7 +11,6 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi.testclient import TestClient - from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app @@ -324,6 +323,100 @@ def test_rag_query_stream_returns_event_stream(client_internal_user): assert "data: [DONE]" in response.text +def test_rag_query_stream_pings_while_retrieval_is_still_running(client_internal_user, monkeypatch): + import asyncio + + import litellm as litellm_module + + monkeypatch.setattr(litellm_module, "sse_keepalive_ping_interval_seconds", 0.05) + + async def slow_aquery(**kwargs): + await asyncio.sleep(0.3) + return await litellm_module.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "What is the codename?"}], + mock_response="The codename is AZURE-FALCON-42.", + stream=True, + api_key="test-key", + ) + + with ( + patch( # test-quality-ok: the handler calls the module-level litellm.aquery directly; no injection seam + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", + new=AsyncMock(side_effect=slow_aquery), + ), + patch("litellm.vector_store_registry", None), # test-quality-ok: proxy module global, no injection seam + patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: proxy module global, no injection seam + ): + response = client_internal_user.post( + "/v1/rag/query", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "What is the codename?"}], + "retrieval_config": { + "vector_store_id": "vs_test_123", + "custom_llm_provider": "openai", + }, + "stream": True, + }, + ) + + assert response.status_code == 200, response.text + assert response.headers.get("content-type", "").startswith("text/event-stream") + assert response.headers["x-accel-buffering"] == "no" + assert response.text.startswith(": ping\n\n") + assert response.text.count(": ping\n\n") >= 3 + assert '"object":"chat.completion.chunk"' in response.text + assert response.text.endswith("data: [DONE]\n\n") + + +def test_rag_query_stream_keeps_response_headers_when_retrieval_beats_the_keepalive( + client_internal_user, monkeypatch +): + import litellm as litellm_module + + monkeypatch.setattr(litellm_module, "sse_keepalive_ping_interval_seconds", 5) + + async def fast_aquery(**kwargs): + response = await litellm_module.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "What is the codename?"}], + mock_response="The codename is AZURE-FALCON-42.", + stream=True, + api_key="test-key", + ) + response._hidden_params["response_cost"] = 3.45e-06 + return response + + with ( + patch( # test-quality-ok: the handler calls the module-level litellm.aquery directly; no injection seam + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", + new=AsyncMock(side_effect=fast_aquery), + ), + patch("litellm.vector_store_registry", None), # test-quality-ok: proxy module global, no injection seam + patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: proxy module global, no injection seam + ): + response = client_internal_user.post( + "/v1/rag/query", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "What is the codename?"}], + "retrieval_config": { + "vector_store_id": "vs_test_123", + "custom_llm_provider": "openai", + }, + "stream": True, + }, + ) + + assert response.status_code == 200, response.text + assert response.headers.get("content-type", "").startswith("text/event-stream") + assert response.headers.get("x-litellm-response-cost") == "3.45e-06" + assert not response.text.startswith(": ping") + assert '"object":"chat.completion.chunk"' in response.text + assert response.text.endswith("data: [DONE]\n\n") + + def test_rag_query_merges_managed_store_params(client_internal_user): """ Regression: /v1/rag/query must consult the managed vector store registry diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 73a29afd9b9..329a33eb440 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -5,14 +5,12 @@ import hashlib import json import re from datetime import timezone +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException from fastapi.testclient import TestClient - -from unittest.mock import AsyncMock, MagicMock, patch - import litellm import litellm.proxy.proxy_server as ps @@ -3325,7 +3323,7 @@ def _compare_nested_dicts( return differences # Check for keys in actual but not in expected - for key in actual.keys(): + for key in actual: current_path = f"{path}.{key}" if path else key if current_path not in ignore_keys and key not in expected: differences.append(f"Extra key in actual: {current_path}") @@ -3495,24 +3493,22 @@ async def test_view_spend_logs_summarize_parameter(client, monkeypatch): # Return individual log entries when summarize=false return mock_spend_logs - async def group_by(self, *args, **kwargs): - # Return grouped data when summarize=true - # Simplified mock response for grouped data + async def query_raw(self, sql_query, *params): yesterday = datetime.datetime.now(timezone.utc) - timedelta(days=1) return [ { "api_key": "sk-test-key", "user": "test_user_1", "model": "gpt-3.5-turbo", - "startTime": yesterday.strftime("%Y-%m-%dT%H:%M:%S.%fZ"), - "_sum": {"spend": 0.05}, + "day": yesterday.date().isoformat(), + "spend": 0.05, }, { "api_key": "sk-test-key", "user": "test_user_1", "model": "gpt-4", - "startTime": yesterday.strftime("%Y-%m-%dT%H:%M:%S.%fZ"), - "_sum": {"spend": 0.10}, + "day": yesterday.date().isoformat(), + "spend": 0.10, }, ] @@ -3850,47 +3846,30 @@ async def test_view_spend_logs_with_date_range_summarized(client, monkeypatch): """ from datetime import datetime, timedelta, timezone - # This simulates the summarized data that Prisma's `group_by` would return. mock_summarized_response = [ { "api_key": "sk-test-key", "user": "test_user_1", "model": "gpt-4", - "startTime": (datetime.now(timezone.utc) - timedelta(days=1)).strftime( - "%Y-%m-%dT%H:%M:%S.%fZ" - ), - "_sum": {"spend": 0.15}, + "day": (datetime.now(timezone.utc) - timedelta(days=1)).date().isoformat(), + "spend": 0.15, } ] - # This mock class will replace the real Prisma client. class MockDB: - def __init__(self): - self.litellm_spendlogs = self - - async def group_by(self, *args, **kwargs): - # We assert that the `gte` and `lte` values are strings in ISO format. - # If they were datetime objects, this test would fail. - where_clause = kwargs.get("where", {}) - start_time_filter = where_clause.get("startTime", {}) - - assert "gte" in start_time_filter - assert "lte" in start_time_filter - assert isinstance(start_time_filter["gte"], str) - assert isinstance(start_time_filter["lte"], str) - assert "T" in start_time_filter["gte"] # Check for ISO format 'T' separator - - # If the assertions pass, return the mock response. + async def query_raw(self, sql_query, *params): + assert isinstance(params[0], str) + assert isinstance(params[1], str) + assert "T" in params[0] + assert "T" in params[1] return mock_summarized_response class MockPrismaClient: def __init__(self): self.db = MockDB() - # Apply the monkeypatch to replace the real prisma_client with our mock. monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MockPrismaClient()) - # Define a date range for the test. start_date = (datetime.now(timezone.utc) - timedelta(days=2)).strftime("%Y-%m-%d") end_date = datetime.now(timezone.utc).strftime("%Y-%m-%d") @@ -3898,8 +3877,6 @@ async def test_view_spend_logs_with_date_range_summarized(client, monkeypatch): user_role=LitellmUserRoles.PROXY_ADMIN ) try: - # Call the endpoint with both start and end dates. - # We don't need `summarize=true` as it's the default. response = client.get( "/spend/logs", params={ @@ -3909,11 +3886,9 @@ async def test_view_spend_logs_with_date_range_summarized(client, monkeypatch): headers={"Authorization": "Bearer sk-test"}, ) - # ASSERTIONS assert response.status_code == 200 data = response.json() - # Check that the response is not empty and has the summarized structure. assert isinstance(data, list) assert len(data) > 0 assert "startTime" in data[0] @@ -3924,6 +3899,183 @@ async def test_view_spend_logs_with_date_range_summarized(client, monkeypatch): app.dependency_overrides.pop(ps.user_api_key_auth, None) +@pytest.mark.asyncio +async def test_view_spend_logs_summarize_groups_by_day_in_sql(client, monkeypatch): + mock_rows = [ + { + "day": "2024-01-01", + "api_key": "hashed::sk-abc", + "user": "u1", + "model": "gpt-4", + "spend": 0.1, + }, + { + "day": "2024-01-01", + "api_key": "hashed::sk-abc", + "user": "u1", + "model": "gpt-4o", + "spend": 0.2, + }, + ] + + class MockDB: + def __init__(self): + self.captured_sql = None + self.captured_params = None + + async def query_raw(self, sql_query, *params): + self.captured_sql = sql_query + self.captured_params = params + return mock_rows + + class MockPrismaClient: + def __init__(self): + self.db = MockDB() + + def hash_token(self, token): + return "hashed::" + token + + mock_prisma_client = MockPrismaClient() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + response = client.get( + "/spend/logs", + params={ + "start_date": "2024-01-01", + "end_date": "2024-01-03", + "api_key": "sk-abc", + "request_id": "req-123", + "user_id": "u1", + }, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200 + data = response.json() + sql = mock_prisma_client.db.captured_sql + assert "date_trunc('day'" in sql + assert "GROUP BY" in sql + assert "find_many" not in sql + assert not hasattr(mock_prisma_client.db, "group_by") + assert mock_prisma_client.db.captured_params == ( + "2024-01-01T00:00:00+00:00", + "2024-01-03T00:00:00+00:00", + "hashed::sk-abc", + "req-123", + "u1", + ) + assert len(data) == 3 + assert data[0]["startTime"] == "2024-01-01" + assert data[0]["spend"] == pytest.approx(0.3) + assert data[0]["models"] == {"gpt-4": 0.1, "gpt-4o": 0.2} + assert data[0]["users"] == {"u1": pytest.approx(0.3)} + assert data[0]["hashed::sk-abc"] == pytest.approx(0.3) + assert data[1] == { + "startTime": "2024-01-02", + "spend": 0, + "users": {}, + "models": {}, + } + assert data[2] == { + "startTime": "2024-01-03", + "spend": 0, + "users": {}, + "models": {}, + } + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_view_spend_logs_summarize_empty_rows(client, monkeypatch): + class MockDB: + async def query_raw(self, sql_query, *params): + return [] + + class MockPrismaClient: + def __init__(self): + self.db = MockDB() + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MockPrismaClient()) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + response = client.get( + "/spend/logs", + params={"start_date": "2024-01-01", "end_date": "2024-01-01"}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200 + assert response.json() == [] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_view_spend_logs_summarize_unhashed_api_key_without_padding(client, monkeypatch): + mock_rows = [ + { + "day": "2024-01-01", + "api_key": "plain-key", + "user": "u1", + "model": "gpt-4", + "spend": 0.4, + } + ] + + class MockDB: + def __init__(self): + self.captured_params = None + + async def query_raw(self, sql_query, *params): + self.captured_params = params + return mock_rows + + class MockPrismaClient: + def __init__(self): + self.db = MockDB() + + mock_prisma_client = MockPrismaClient() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + response = client.get( + "/spend/logs", + params={ + "start_date": "2024-01-01", + "end_date": "2024-01-01", + "api_key": "plain-key", + }, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200 + data = response.json() + assert mock_prisma_client.db.captured_params == ( + "2024-01-01T00:00:00+00:00", + "2024-01-01T00:00:00+00:00", + "plain-key", + ) + assert data == [ + { + "startTime": "2024-01-01", + "spend": pytest.approx(0.4), + "plain-key": pytest.approx(0.4), + "users": {"u1": pytest.approx(0.4)}, + "models": {"gpt-4": pytest.approx(0.4)}, + } + ] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @pytest.mark.asyncio async def test_ui_view_spend_logs_with_error_code(client): """Test filtering spend logs by error code""" @@ -4832,13 +4984,14 @@ class _CaptureFilterDB: def __init__(self): self.litellm_spendlogs = self self.captured_where = None + self.captured_params = None async def find_many(self, *args, **kwargs): self.captured_where = kwargs.get("where") return [] - async def group_by(self, *args, **kwargs): - self.captured_where = kwargs.get("where") + async def query_raw(self, sql_query, *params): + self.captured_params = params return [] diff --git a/tests/test_litellm/rust_bridge/native_route_wheel_test.py b/tests/test_litellm/rust_bridge/native_route_wheel_test.py index a89a985952b..a7f50a82a99 100644 --- a/tests/test_litellm/rust_bridge/native_route_wheel_test.py +++ b/tests/test_litellm/rust_bridge/native_route_wheel_test.py @@ -181,14 +181,6 @@ def assert_success(route: str, response: object) -> None: raise AssertionError(f"{route} returned {actual!r}, expected {expected!r}") -def assert_traced_success(route: str, response: object) -> None: - if not isinstance(response, dict): - raise TypeError(f"{route} returned {type(response).__name__}, expected a traced dict") - assert_success(route, response["response"]) - expected_function: Final = "audio_transcription" if route == "transcription" else route - assert response["trace"][0] == {"function": expected_function, "depth": 0} - - def success_value(route: str, response: dict[object, object]) -> object: if route == "ocr": return response["pages"][0]["markdown"] @@ -213,7 +205,6 @@ def exercise_sync(native: object, api_base: str) -> None: for route in ("ocr", "transcription", "messages", "chat_completions"): function: Final = getattr(native, route) assert_success(route, function(**route_kwargs(route, api_base, "success"))) - assert_traced_success(route, function(**route_kwargs(route, api_base, "success"), trace=True)) try: function(**route_kwargs(route, api_base, "429")) except (RuntimeError, native.RustUpstreamError) as error: @@ -226,7 +217,6 @@ async def exercise_async(native: object, api_base: str) -> None: for route in ("ocr", "transcription", "messages", "chat_completions"): function: Final = getattr(native, f"a{route}") assert_success(route, await function(**route_kwargs(route, api_base, "success"))) - assert_traced_success(route, await function(**route_kwargs(route, api_base, "success"), trace=True)) try: await function(**route_kwargs(route, api_base, "429")) except (RuntimeError, native.RustUpstreamError) as error: @@ -251,6 +241,8 @@ async def exercise_async_concurrency(native: object, api_base: str) -> None: def exercise_routes(native_path: Path, api_base: str) -> object: native: Final = load_native(native_path) + if hasattr(native, "_trace"): + raise AssertionError("release wheel exposed trace-parity diagnostics") exercise_sync(native, api_base) asyncio.run(exercise_async(native, api_base)) asyncio.run(exercise_async_concurrency(native, api_base)) diff --git a/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py b/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py index c9ec22ab0df..4bc7c21d8d2 100644 --- a/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py +++ b/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py @@ -6,6 +6,7 @@ from unittest.mock import MagicMock, patch # Adds the grandparent directory to sys.path to allow importing project modules import pytest +from azure.core.exceptions import ClientAuthenticationError from litellm.secret_managers.get_azure_ad_token_provider import ( get_azure_ad_token_provider, @@ -16,6 +17,143 @@ from litellm.types.secret_managers.get_azure_ad_token_provider import ( ) +class TestDeploymentIdentityCredential: + @staticmethod + def _chain_for(credential_type): + with patch("azure.identity.get_bearer_token_provider", return_value=lambda: "token") as bearer: + get_azure_ad_token_provider( + azure_scope="https://storage.azure.com/.default", + azure_credential=credential_type, + ) + bearer.assert_called_once() + with bearer.call_args.args[0] as chain: + return {type(link).__name__ for link in chain.credentials} + + @staticmethod + def _managed_identity_client_ids(credential_type): + with patch("azure.identity.get_bearer_token_provider", return_value=lambda: "token") as bearer: + get_azure_ad_token_provider( + azure_scope="https://storage.azure.com/.default", + azure_credential=credential_type, + ) + with bearer.call_args.args[0] as chain: + return [ + (link._credential._settings or {}).get("client_id") + for link in chain.credentials + if type(link).__name__ == "ManagedIdentityCredential" + ] + + @patch.dict( + os.environ, + { + "AZURE_CLIENT_ID": "workload-identity-client-id", + "AZURE_TENANT_ID": "workload-identity-tenant-id", + "AZURE_FEDERATED_TOKEN_FILE": "/var/run/secrets/azure/tokens/azure-identity-token", + }, + clear=True, + ) + def test_deployment_identity_reaches_workload_and_managed_identity_only(self): + assert self._chain_for(AzureCredentialType.DeploymentIdentityCredential) == { + "WorkloadIdentityCredential", + "ManagedIdentityCredential", + } + + @patch.dict( + os.environ, + { + "AZURE_CLIENT_ID": "workload-identity-client-id", + "AZURE_TENANT_ID": "workload-identity-tenant-id", + "AZURE_FEDERATED_TOKEN_FILE": "/var/run/secrets/azure/tokens/azure-identity-token", + "AZURE_TOKEN_CREDENTIALS": "dev", + }, + clear=True, + ) + def test_deployment_identity_survives_a_developer_only_token_credentials_setting(self): + """AZURE_TOKEN_CREDENTIALS=dev asks the SDK for developer credentials only, which is every + credential this chain drops, so the deployment's own identity has to win over it""" + assert self._chain_for(AzureCredentialType.DeploymentIdentityCredential) == { + "WorkloadIdentityCredential", + "ManagedIdentityCredential", + } + + @patch.dict( + os.environ, + { + "AZURE_CLIENT_ID": "azure-openai-client-id", + "AZURE_CLIENT_SECRET": "azure-openai-client-secret", + "AZURE_TENANT_ID": "azure-openai-tenant-id", + }, + clear=True, + ) + def test_default_azure_credential_keeps_its_full_chain(self): + """Azure OpenAI callers pass DefaultAzureCredential and must be unaffected by the + narrowing that the storage callback asks for""" + full_chain = self._chain_for(AzureCredentialType.DefaultAzureCredential) + + assert "EnvironmentCredential" in full_chain + assert "AzureCliCredential" in full_chain + assert "EnvironmentCredential" not in self._chain_for( + AzureCredentialType.DeploymentIdentityCredential + ) + + @patch.dict( + os.environ, + { + "AZURE_CLIENT_ID": "azure-openai-client-id", + "AZURE_CLIENT_SECRET": "azure-openai-client-secret", + "AZURE_TENANT_ID": "azure-openai-tenant-id", + }, + clear=True, + ) + def test_deployment_identity_refuses_to_mint_a_token_for_a_configured_service_principal(self): + """A host carrying only an Azure OpenAI client secret must get no token at all, and the + refusal must name the identities that were actually tried""" + provider = get_azure_ad_token_provider( + azure_scope="https://storage.azure.com/.default", + azure_credential=AzureCredentialType.DeploymentIdentityCredential, + ) + + with pytest.raises(ClientAuthenticationError) as refusal: + provider() + + assert "ManagedIdentityCredential" in str(refusal.value) + assert "EnvironmentCredential" not in str(refusal.value) + assert "AzureCliCredential" not in str(refusal.value) + assert "azure-openai-client-secret" not in str(refusal.value) + + @patch.dict( + os.environ, + { + "AZURE_CLIENT_ID": "azure-openai-client-id", + "AZURE_CLIENT_SECRET": "azure-openai-client-secret", + "AZURE_TENANT_ID": "azure-openai-tenant-id", + }, + clear=True, + ) + def test_deployment_identity_still_reaches_a_system_assigned_managed_identity(self): + """AZURE_CLIENT_ID names one identity for the whole proxy, and pointing it at Azure OpenAI + must not hide the system assigned identity the host runs as""" + client_ids = self._managed_identity_client_ids(AzureCredentialType.DeploymentIdentityCredential) + + assert "azure-openai-client-id" in client_ids + assert None in client_ids + + @patch.dict( + os.environ, + { + "AZURE_CLIENT_ID": "user-assigned-identity-client-id", + "AZURE_TOKEN_CREDENTIALS": "dev", + }, + clear=True, + ) + def test_deployment_identity_keeps_the_user_assigned_identity_under_a_dev_only_setting(self): + """AZURE_TOKEN_CREDENTIALS=dev asks the SDK for developer credentials only, and the + identity a host actually runs as has to survive that""" + assert "user-assigned-identity-client-id" in self._managed_identity_client_ids( + AzureCredentialType.DeploymentIdentityCredential + ) + + class TestGetAzureAdTokenProvider: @patch.dict( os.environ, diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 228588d974f..5fc96bcfbb1 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -3855,7 +3855,7 @@ def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch): input_only_tokens = router._count_pre_call_check_tokens(messages=None, input=short_input) with_instructions_tokens = router._count_pre_call_check_tokens( - messages=None, input=short_input, instructions=long_instructions + messages=None, input=short_input, request_kwargs={"instructions": long_instructions} ) assert with_instructions_tokens > input_only_tokens @@ -3871,6 +3871,164 @@ def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch): ) +_OVERSIZED_TOOL_DESCRIPTION = "look up the answer in the knowledge base. " * 40 + + +@pytest.mark.parametrize( + "prompt_kwargs, tool", + [ + pytest.param( + {"messages": [{"role": "user", "content": "hi"}]}, + { + "type": "function", + "function": { + "name": "lookup", + "description": _OVERSIZED_TOOL_DESCRIPTION, + "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}, + }, + }, + id="chat_completions_tool", + ), + pytest.param( + {"input": "hi"}, + { + "type": "function", + "name": "lookup", + "description": _OVERSIZED_TOOL_DESCRIPTION, + "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}, + }, + id="responses_tool", + ), + pytest.param( + {"messages": [{"role": "user", "content": "hi"}]}, + { + "name": "lookup", + "description": _OVERSIZED_TOOL_DESCRIPTION, + "input_schema": {"type": "object", "properties": {"q": {"type": "string"}}}, + }, + id="anthropic_messages_tool", + ), + ], +) +def test_pre_call_checks_counts_tool_definition_tokens(monkeypatch, prompt_kwargs, tool): + """ + Tool definitions are sent to the model as prompt tokens but never appear in + `messages` or `input`. A request whose prompt alone fits the context window but + whose prompt plus `tools` exceeds it must be rejected before dispatch, for the + Chat Completions, Responses and Anthropic Messages tool shapes alike. + """ + router = litellm.Router( + model_list=[ + {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, + ], + enable_pre_call_checks=True, + ) + deployments = [ + {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, + ] + + prompt_only_tokens = router._count_pre_call_check_tokens( + messages=prompt_kwargs.get("messages"), input=prompt_kwargs.get("input") + ) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": prompt_only_tokens} + ) + + assert len(router._pre_call_checks(model="m", healthy_deployments=deployments, **prompt_kwargs)) == 1 + with pytest.raises(litellm.ContextWindowExceededError): + router._pre_call_checks( + model="m", + healthy_deployments=deployments, + request_kwargs={"tools": [tool]}, + **prompt_kwargs, + ) + + +@pytest.mark.parametrize( + "system", + [ + pytest.param("You are a meticulous assistant. " * 40, id="system_string"), + pytest.param( + [{"type": "text", "text": "You are a meticulous assistant. " * 40}], + id="system_blocks", + ), + ], +) +def test_pre_call_checks_counts_anthropic_system_tokens(monkeypatch, system): + """ + The Anthropic Messages API carries the system prompt as a top-level `system` field, + not as a message. Its tokens reach the model, so a request whose `messages` fit but + whose `messages` plus `system` exceed the context window must be rejected. + """ + router = litellm.Router( + model_list=[ + {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, + ], + enable_pre_call_checks=True, + ) + deployments = [ + {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, + ] + messages = [{"role": "user", "content": "hi"}] + + messages_only_tokens = router._count_pre_call_check_tokens(messages=messages, input=None) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": messages_only_tokens}) + + assert len(router._pre_call_checks(model="m", healthy_deployments=deployments, messages=messages)) == 1 + with pytest.raises(litellm.ContextWindowExceededError): + router._pre_call_checks( + model="m", + healthy_deployments=deployments, + messages=messages, + request_kwargs={"system": system}, + ) + + +@pytest.mark.asyncio +async def test_aanthropic_messages_enforces_context_window_with_system_and_tools(): + """ + End-to-end router regression for /v1/messages: a request whose only oversized + content lives in the top-level `system` field or in `tools` must trip the pre-call + context-window check instead of being dispatched (the deployment uses mock_response, + so reaching the provider handler would return a response rather than raise). + """ + router = litellm.Router( + model_list=[ + { + "model_name": "small-ctx", + "litellm_params": {"model": "anthropic/claude-3-5-haiku-20241022", "mock_response": "hi"}, + "model_info": {"max_input_tokens": 20}, + } + ], + enable_pre_call_checks=True, + ) + messages = [{"role": "user", "content": "hi"}] + + response = await router.aanthropic_messages(model="small-ctx", messages=messages, max_tokens=5) + assert response is not None + + with pytest.raises(litellm.ContextWindowExceededError): + await router.aanthropic_messages( + model="small-ctx", + messages=messages, + max_tokens=5, + system="You are a meticulous assistant. " * 40, + ) + with pytest.raises(litellm.ContextWindowExceededError): + await router.aanthropic_messages( + model="small-ctx", + messages=messages, + max_tokens=5, + tools=[ + { + "name": "lookup", + "description": _OVERSIZED_TOOL_DESCRIPTION, + "input_schema": {"type": "object", "properties": {"q": {"type": "string"}}}, + } + ], + ) + + def test_count_pre_call_check_tokens_across_api_surfaces(): """ _count_pre_call_check_tokens must count tokens from chat `messages`, a Responses @@ -4943,6 +5101,40 @@ def test_team_wildcard_credentials_not_usable_after_delete_deployment(): ) +def test_global_wildcard_pattern_router_evicts_stale_entry_on_upsert_and_delete(): + """ + Regression for #29064: upsert_deployment removed the old deployment from + model_list but left it in the global pattern_router, so wildcard requests + round-robined between the stale and the corrected deployment. + """ + from litellm.types.router import Deployment, LiteLLM_Params + + router = litellm.Router( + model_list=[ + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/openai/*", "api_key": "sk-old"}, + "model_info": {"id": "global-wildcard"}, + } + ] + ) + + router.upsert_deployment( + Deployment( + model_name="openai/*", + litellm_params=LiteLLM_Params(model="openai/*", api_key="sk-new"), + model_info={"id": "global-wildcard"}, + ) + ) + + matches = router.pattern_router.route("openai/gpt-5.2") + assert matches is not None + assert [m["litellm_params"]["api_key"] for m in matches] == ["sk-new"] + + router.delete_deployment(id="global-wildcard") + assert router.pattern_router.patterns == {} + + def test_pattern_match_router_remove_deployment(): """ remove_deployment must drop only the deployment with the given model id and diff --git a/tests/test_rust_python_harness.py b/tests/test_rust_python_harness.py index a2b9c8e2a76..b27d1c83597 100644 --- a/tests/test_rust_python_harness.py +++ b/tests/test_rust_python_harness.py @@ -1,206 +1,86 @@ from __future__ import annotations import importlib -import json -import os -import subprocess -import sys from pathlib import Path from typing import Final import pytest -catalog = importlib.import_module("tests.rust-python-harness.catalog") -cli = importlib.import_module("tests.rust-python-harness.cli") models = importlib.import_module("tests.rust-python-harness.shared.reporting.models") -runner = importlib.import_module("tests.rust-python-harness.shared.reporting.pytest_runner") +strategy_module = importlib.import_module("tests.rust-python-harness.shared.reporting.strategy") ui = importlib.import_module("tests.rust-python-harness.shared.reporting.ui") -ledger_module = importlib.import_module("tests.rust-python-harness.shared.parity.ledger") -mapping_validator = importlib.import_module( - "tests.rust-python-harness.strategies.unit_tests.mapping_validator" -) +mapping_validator = importlib.import_module("tests.rust-python-harness.strategies.unit_tests_mapping.mapping_validator") +mappings = importlib.import_module("tests.rust-python-harness.strategies.unit_tests_mapping.mappings") +ocr_mapping = importlib.import_module("tests.rust-python-harness.strategies.unit_tests_mapping.cases.ocr") +cli = importlib.import_module("tests.rust-python-harness.cli") -load_catalog = catalog.load_catalog -load_ledger = ledger_module.load_ledger -ledger_path_for = mapping_validator.ledger_path_for -REPO_ROOT = mapping_validator.REPO_ROOT -audit_ledger = mapping_validator.audit_ledger -build_function_report = mapping_validator.build_function_report -_pick_values = cli._pick_values -_coverage_pytest_args = cli._coverage_pytest_args -_select = cli._select -_validate_ledger = cli._validate_ledger +audit_mapping = mapping_validator.audit_mapping +UNIT_TEST_CONTRACTS = mappings.UNIT_TEST_CONTRACTS +OCR_CONTRACT = ocr_mapping.OCR_CONTRACT +REPO_ROOT = Path(__file__).resolve().parents[1] CaseResult = models.CaseResult Coverage = models.Coverage HarnessCase = models.HarnessCase HarnessRun = models.HarnessRun RunStatus = models.RunStatus -SDK_FUNCTIONS = models.SDK_FUNCTIONS -section_confidence = models.section_confidence -run_pytest = runner.run_pytest -runnable_selectors = runner.runnable_selectors -selector_matches_node = runner.selector_matches_node +ModuleCaseSpec = strategy_module.ModuleCaseSpec +NotImplementedCaseSpec = strategy_module.NotImplementedCaseSpec +SkippedCaseSpec = strategy_module.SkippedCaseSpec _format_duration = ui._format_duration -_rerun_command = ui._rerun_command _summary = ui._summary -def _case( - *, selectors: tuple[str, ...] = (), coverage: Coverage = Coverage.COMPLETE -) -> HarnessCase: +def _case(module: str = "tests.example") -> HarnessCase: return HarnessCase( strategy_id="example", strategy_label="Example", sdk_function="messages", - coverage=coverage, - selectors=selectors, + spec=ModuleCaseSpec(coverage=Coverage.COMPLETE, module=module), ) -def _manifest() -> dict[str, object]: - return { - "order": 1, - "id": "example", - "label": "Example strategy", - "description": "Example description", - "functions": { - function: {"coverage": "planned", "selectors": []} - for function in SDK_FUNCTIONS - }, - } - - -def test_should_load_the_four_harness_strategies_in_order() -> None: - strategies = load_catalog() - - assert [strategy.id for strategy in strategies] == [ - "e2e_parity", - "trace_parity", - "unit_tests", - "existing_e2e_test_sdk", - ] - assert all( - tuple(case.sdk_function for case in strategy.cases) == SDK_FUNCTIONS - for strategy in strategies - ) - - -def test_should_reject_a_manifest_missing_an_sdk_function(tmp_path: Path) -> None: - strategy_directory = tmp_path / "example" - strategy_directory.mkdir() - manifest = _manifest() - del manifest["functions"]["count_tokens"] # type: ignore[index] - (strategy_directory / "strategy.json").write_text( - json.dumps(manifest), encoding="utf-8" - ) - - with pytest.raises(ValueError, match="functions must exactly match"): - load_catalog(tmp_path) - - @pytest.mark.parametrize( - ("selector", "nodeid", "matches"), + "module", [ - ("tests/test_parity.py", "tests/test_parity.py::test_one", True), - ("tests/test_parity.py::test_one", "tests/test_parity.py::test_one", True), - ( - "tests/test_parity.py::test_one", - "tests/test_parity.py::test_one[value]", - True, - ), - ("tests/test_parity.py::test_one", "tests/test_parity.py::test_two", False), - ("tests/ocr_tests/", "tests/ocr_tests/test_ocr_mistral.py::test_one", True), - ("tests/ocr_tests/", "tests/other_tests/test_ocr_mistral.py::test_one", False), + "tests.rust-python-harness.strategies.e2e_parity.sdk.ocr.test_sdk_parity", + "tests.rust-python-harness.strategies.trace_parity.sdk.ocr.case", + "tests.rust-python-harness.strategies.trace_parity.sdk.messages.case", + "tests.rust-python-harness.strategies.trace_parity.sdk.chat_completions.case", + "tests.rust-python-harness.strategies.trace_parity.sdk.transcription.case", + "tests.rust-python-harness.strategies.trace_parity.gateway.messages.case", ], ) -def test_should_match_pytest_file_and_node_selectors( - selector: str, nodeid: str, matches: bool -) -> None: - assert selector_matches_node(selector, nodeid) is matches +def test_implemented_namespace_case_modules_remain_importable(module: str) -> None: + assert importlib.import_module(module) -def test_should_only_return_selectors_whose_files_exist(tmp_path: Path) -> None: - existing = tmp_path / "tests" / "test_parity.py" - existing.parent.mkdir() - existing.write_text("", encoding="utf-8") - case = _case( - selectors=("tests/test_parity.py", "tests/test_missing.py::test_missing") +def test_should_mark_not_implemented_and_skipped_cases_without_running() -> None: + not_implemented: Final = CaseResult( + case=HarnessCase( + strategy_id="example", + strategy_label="Example", + sdk_function="messages", + spec=NotImplementedCaseSpec(reason="No case is registered."), + ) + ) + skipped: Final = CaseResult( + case=HarnessCase( + strategy_id="example", + strategy_label="Example", + sdk_function="messages", + spec=SkippedCaseSpec(reason="The surface does not apply."), + ) ) - assert runnable_selectors((case,), tmp_path) == ("tests/test_parity.py",) + not_implemented.set_initial_status() + skipped.set_initial_status() - -def test_should_treat_an_existing_folder_selector_as_runnable(tmp_path: Path) -> None: - (tmp_path / "tests" / "ocr_tests").mkdir(parents=True) - case = _case(selectors=("tests/ocr_tests/",)) - - assert runnable_selectors((case,), tmp_path) == ("tests/ocr_tests/",) - - -def test_should_mark_planned_and_not_applicable_cases_without_running() -> None: - planned = CaseResult(case=_case(coverage=Coverage.PLANNED)) - not_applicable = CaseResult(case=_case(coverage=Coverage.NOT_APPLICABLE)) - - planned.set_initial_status() - not_applicable.set_initial_status() - - assert planned.status is RunStatus.PLANNED - assert not_applicable.status is RunStatus.NOT_APPLICABLE - - -def test_should_treat_an_all_planned_filtered_run_as_success(tmp_path: Path) -> None: - exit_code, run = run_pytest( - cases=(_case(coverage=Coverage.PLANNED),), - repo_root=tmp_path, - on_update=lambda _: None, - ) - - assert exit_code == 0 - assert next(iter(run.results.values())).status is RunStatus.PLANNED - - -@pytest.mark.parametrize("strategy_id", ("e2e_parity", "existing_e2e_test_sdk")) -def test_should_run_namespace_package_relative_imports(tmp_path: Path, strategy_id: str) -> None: - package: Final = tmp_path / "manual_suite" / "relative-tests" - package.mkdir(parents=True) - (package / "__init__.py").write_text("", encoding="utf-8") - (package / "values.py").write_text("ANSWER = 42\n", encoding="utf-8") - (package / "test_relative.py").write_text( - "from .values import ANSWER\n\ndef test_answer():\n assert ANSWER == 42\n", - encoding="utf-8", - ) - result: Final = subprocess.run( - ( - sys.executable, - "-c", - "import importlib\n" - "from pathlib import Path\n" - "cli = importlib.import_module('tests.rust-python-harness.cli')\n" - "models = importlib.import_module('tests.rust-python-harness.shared.reporting.models')\n" - f"case = models.HarnessCase(strategy_id={strategy_id!r}, strategy_label='Example', " - "sdk_function='ocr', coverage=models.Coverage.COMPLETE, " - "selectors=('manual_suite/relative-tests/',))\n" - f"code, run = cli._resolve_runner({strategy_id!r})((case,), Path.cwd(), lambda _: None)\n" - "assert code == 0, code\n" - "assert next(iter(run.results.values())).passed == 1\n", - ), - cwd=tmp_path, - env={ - **os.environ, - "PYTHONPATH": os.pathsep.join((str(tmp_path), str(Path(__file__).resolve().parents[1]))), - "PYTEST_DISABLE_PLUGIN_AUTOLOAD": "1", - }, - capture_output=True, - text=True, - timeout=30, - check=False, - ) - - assert result.returncode == 0, result.stdout + result.stderr + assert not_implemented.status is RunStatus.NOT_IMPLEMENTED + assert skipped.status is RunStatus.SKIPPED def test_should_finalize_a_fully_passing_case() -> None: - result = CaseResult(case=_case(selectors=("tests/test_parity.py",))) + result = CaseResult(case=_case()) result.set_initial_status() result.collected.update({"one", "two"}) result.completed.update({"one", "two"}) @@ -212,7 +92,7 @@ def test_should_finalize_a_fully_passing_case() -> None: def test_should_replace_a_pass_with_a_teardown_error() -> None: - result = CaseResult(case=_case(selectors=("tests/test_parity.py",))) + result = CaseResult(case=_case()) result.set_initial_status() result.collected.add("one") @@ -225,136 +105,40 @@ def test_should_replace_a_pass_with_a_teardown_error() -> None: assert result.duration == pytest.approx(0.3) -def test_should_filter_the_catalog_by_strategy_and_sdk_function() -> None: - strategies = load_catalog() - - cases = _select(strategies, {"e2e_parity"}, {"messages"}) - - assert len(cases) == 1 - assert cases[0].key == "e2e_parity:messages" - - -def test_should_reject_an_unknown_strategy() -> None: - with pytest.raises(ValueError, match="Unknown strategy"): - _select(load_catalog(), {"not-real"}, set()) - - -def test_should_pick_multiple_interactive_filters() -> None: - answers = iter(["nope", "1, 3"]) - - selected = _pick_values( - "Examples", - (("one", "One"), ("two", "Two"), ("three", "Three")), - input_fn=lambda _: next(answers), - ) - - assert selected == {"one", "three"} - - def test_should_format_developer_facing_run_context() -> None: - run = HarnessRun.from_cases((_case(selectors=("tests/test_parity.py",)),)) + run = HarnessRun.from_cases((_case(),)) result = next(iter(run.results.values())) result.collected.add("tests/test_parity.py::test_one") result.record("tests/test_parity.py::test_one", RunStatus.PASSED, 1.25) assert _summary(run) == (1, 0, 0, 0) assert _format_duration(1.25) == "1.2s" - assert _rerun_command("tests/test_parity.py::test_one") == ( - "poetry run pytest tests/test_parity.py::test_one -q -o consider_namespace_packages=true" + + +def test_should_leave_functions_without_mapping_contracts_unimplemented() -> None: + assert "messages" not in UNIT_TEST_CONTRACTS + + +def test_should_derive_ocr_mapping_status_from_live_tests() -> None: + report = audit_mapping(OCR_CONTRACT, repo_root=REPO_ROOT) + + assert report.is_valid, ( + f"Missing Python tests: {list(report.missing_python_tests)}\n" + f"Missing Rust tests: {list(report.missing_rust_tests)}\n" + f"Duplicate Python mappings: {list(report.duplicate_python_mappings)}\n" + f"Invalid mapping exclusions: {list(report.invalid_mapping_exclusions)}\n" + f"Invalid parity exclusions: {list(report.invalid_unit_parity_exclusions)}" ) - assert _rerun_command("tests/test_parity.py::test_one[value with spaces]") == ( - "poetry run pytest 'tests/test_parity.py::test_one[value with spaces]' -q -o consider_namespace_packages=true" + assert report.mapped_count == len(OCR_CONTRACT.mapping.mappings) + assert report.total_count == ( + report.mapped_count + len(report.excluded_python_tests) + len(report.unmapped_python_tests) ) -def test_should_build_python_coverage_reports_below_the_target_directory( - tmp_path: Path, -) -> None: - args = _coverage_pytest_args(tmp_path) - - assert tmp_path.is_dir() - assert "--cov=litellm" in args - assert "--cov-context=test" in args - assert f"--cov-report=json:{tmp_path / 'python.json'}" in args - assert f"--cov-report=xml:{tmp_path / 'python.xml'}" in args - assert f"--cov-report=html:{tmp_path / 'python-html'}" in args - - -def test_should_report_confidence_for_each_sdk_section() -> None: - strategies = load_catalog() - cases = tuple(case for strategy in strategies for case in strategy.cases) - run = HarnessRun.from_cases(cases) - passing = run.results["e2e_parity:responses"] - passing.collected.add("tests/test_parity.py::test_one") - passing.record("tests/test_parity.py::test_one", RunStatus.PASSED) - - scores = { - score.sdk_function: score for score in section_confidence(run, strategies) - } - - assert scores["responses"].verified_strategies == 1 - assert scores["responses"].required_strategies == 4 - assert scores["responses"].percentage == 25 - assert scores["responses"].level.value == "MEDIUM" - assert scores["count_tokens"].percentage == 0 - assert scores["count_tokens"].level.value == "LOW" - - - -def test_should_report_no_ledger_for_a_function_without_one() -> None: - report = build_function_report("messages", repo_root=REPO_ROOT) - - assert report.has_ledger is False - assert report.is_clean is True - - -def test_should_report_ocr_ledger_stats_and_a_clean_audit() -> None: - ledger = load_ledger(ledger_path_for("ocr")) - - report = build_function_report("ocr", repo_root=REPO_ROOT) - - assert report.has_ledger is True - assert report.ledger.mapped_count == ledger.mapped_count - assert report.ledger.total_count == ledger.total_count - assert report.is_clean is True - - -def test_should_scope_validate_ledger_to_the_requested_function( - capsys: pytest.CaptureFixture[str], -) -> None: - exit_code = _validate_ledger({"messages"}) - - captured = capsys.readouterr() - assert exit_code == 0 - assert "messages" in captured.out - assert "no ledger yet" in captured.out - assert "ocr" not in captured.out - - -@pytest.mark.parametrize("strategy_id", (None, "e2e_parity", "trace_parity", "unit_tests", "existing_e2e_test_sdk")) -def test_should_validate_chat_completions_ledger_from_each_runner( - strategy_id: str | None, capsys: pytest.CaptureFixture[str] -) -> None: - exit_code: Final = cli.main( - ("--validate-ledger", "--function", "chat_completions"), strategy_id=strategy_id - ) +def test_strategy_subcommand_accepts_function_filter(capsys: pytest.CaptureFixture[str]) -> None: + exit_code: Final = cli.main(["run", "unit_tests_mapping", "--function", "messages"]) captured: Final = capsys.readouterr() assert exit_code == 0 - assert "chat_completions" in captured.out - assert "no ledger yet" in captured.out - assert "ocr" not in captured.out - - -def test_should_have_every_python_and_rust_ocr_test_accounted_for_in_the_ledger() -> None: - ledger = load_ledger(ledger_path_for("ocr")) - - report = audit_ledger(ledger, repo_root=REPO_ROOT) - - assert report.is_clean, ( - "\nOCR test-parity ledger is out of sync with the live test files.\n" - f"Ledger references a Python test that no longer exists: {list(report.missing_python_tests)}\n" - f"Python test exists but is not tracked in the ledger: {list(report.stale_python_tests)}\n" - f"Ledger references a Rust test that no longer exists: {list(report.missing_rust_tests)}\n" - f"Rust test exists but is not tracked in the ledger: {list(report.stale_rust_tests)}\n" - ) + assert "- messages: not_implemented" in captured.out + assert "unit_tests_mapping:messages: not_implemented" not in captured.out diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 2f85128b4b6..ab1a793e09d 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -3,7 +3,7 @@ "limit": 22328 }, "LIT002": { - "limit": 26760 + "limit": 26750 }, "LIT003": { "limit": 261 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16470 + "limit": 16468 }, "LIT011": { - "limit": 5516 + "limit": 5514 }, "LIT012": { "limit": 4489 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/CacheFormField.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/CacheFormField.tsx index 24ec9043d21..2de078c8a33 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/CacheFormField.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/CacheFormField.tsx @@ -12,6 +12,7 @@ import { ComboboxList, } from "@/components/ui/combobox"; import { Input } from "@/components/ui/input"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Switch } from "@/components/ui/switch"; import { Textarea } from "@/components/ui/textarea"; import { CacheField } from "./cacheSettingsFields"; @@ -64,6 +65,36 @@ const CacheFormField: React.FC = ({ field, embeddingModels, /> ); } + if (field.type === "select") { + const options = field.options ?? []; + const { id, "aria-invalid": ariaInvalid, "aria-describedby": ariaDescribedBy, name, onBlur, disabled } = rest; + return ( + + ); + } if (field.type === "model-select") { const selected = embeddingModels.find((model) => model.value === value) ?? null; return ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsFields.ts b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsFields.ts index 04dd6ff6038..9ebb9c191ed 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsFields.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsFields.ts @@ -1,4 +1,17 @@ -export type CacheFieldType = "string" | "password" | "integer" | "float" | "boolean" | "list" | "model-select"; +export type CacheFieldType = + | "string" + | "password" + | "integer" + | "float" + | "boolean" + | "list" + | "model-select" + | "select"; + +export interface CacheFieldOption { + readonly value: string; + readonly label: string; +} export type RedisType = "node" | "cluster" | "sentinel" | "semantic"; @@ -18,6 +31,7 @@ export interface CacheField { readonly helpText: string; readonly redisType: RedisType | null; readonly defaultValue?: string | number | boolean; + readonly options?: readonly CacheFieldOption[]; readonly rules?: CacheFieldRule[]; // Credential field: never prefilled into the form, and dropped from the save // payload when left untouched so the redacted marker is never persisted. @@ -179,6 +193,20 @@ export const CACHE_FIELDS: readonly CacheField[] = [ helpText: "Embedding model for semantic cache", redisType: "semantic", }, + { + name: "semantic_cache_scope", + label: "Semantic Cache Scope", + type: "select", + section: "semantic", + helpText: + "Who can share a semantic cache hit. Key shares hits between all end users of a key/team/org. End user also isolates per end user; requests without an end user fall back to the key scope.", + redisType: "semantic", + defaultValue: "key", + options: [ + { value: "key", label: "Key (shared by all end users of the key/team/org)" }, + { value: "end_user", label: "End user (isolated per end user)" }, + ], + }, { name: "ssl", label: "SSL", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsUtils.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsUtils.test.ts index c530519ee06..c851f4aad96 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsUtils.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/cacheSettingsUtils.test.ts @@ -25,6 +25,7 @@ describe("buildInitialValues", () => { const values = buildInitialValues({}); expect(values.port).toBe("6379"); expect(values.similarity_threshold).toBe("0.8"); + expect(values.semantic_cache_scope).toBe("key"); expect(values.ssl).toBe(false); expect(values.db).toBe(""); }); @@ -75,6 +76,13 @@ describe("buildCachePayload", () => { expect(payload.similarity_threshold).toBe(0.9); }); + it("should send the semantic cache scope only for a semantic cache", () => { + const semantic = buildCachePayload("semantic", { semantic_cache_scope: "end_user" }, { forTesting: false }); + expect(semantic.semantic_cache_scope).toBe("end_user"); + const node = buildCachePayload("node", { semantic_cache_scope: "end_user" }, { forTesting: false }); + expect(node).not.toHaveProperty("semantic_cache_scope"); + }); + it("should keep type redis when testing a semantic cache so the test endpoint accepts it", () => { const payload = buildCachePayload("semantic", { similarity_threshold: 0.9 }, { forTesting: true }); expect(payload.type).toBe("redis"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/index.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/index.integration.test.tsx index 07cc73cddc4..918d8947151 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/index.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/index.integration.test.tsx @@ -138,6 +138,26 @@ describe("CacheSettings advanced settings round-trip", () => { expect(updateCacheSettingsCall.mock.calls[0][1]).not.toHaveProperty("redis_startup_nodes"); }); + it("saves the semantic cache scope picked from the select and shows the loaded value", async () => { + getCacheSettingsCall.mockResolvedValue({ + current_values: { redis_type: "semantic", host: "redis.internal", semantic_cache_scope: "key" }, + }); + const user = userEvent.setup(); + renderSettings(); + const trigger = await screen.findByLabelText("Semantic Cache Scope"); + expect(trigger).toHaveTextContent("Key (shared by all end users of the key/team/org)"); + + await user.click(trigger); + await user.click(await screen.findByRole("option", { name: "End user (isolated per end user)" })); + await save(user); + + await waitFor(() => expect(updateCacheSettingsCall).toHaveBeenCalledTimes(1)); + expect(updateCacheSettingsCall.mock.calls[0][1]).toMatchObject({ + type: "redis-semantic", + semantic_cache_scope: "end_user", + }); + }); + it("does not block the save on a malformed value inside a collapsed advanced section", async () => { getCacheSettingsCall.mockResolvedValue({ current_values: { host: "redis.internal" } }); const user = userEvent.setup(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/TagTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/TagTable.tsx index 076166ac827..fb4793ab340 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/TagTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/TagTable.tsx @@ -42,6 +42,7 @@ const TagTable: React.FC = ({ data, onEdit, onDelete, onSelectTag paginationMode="client" columns={columns} getRowId={(tag, index) => tag.name || String(index)} + fillHeight sortingMode="client" sorting={sorting} onSortingChange={setSorting} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx index a04afcd6d45..583492bf837 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx @@ -126,7 +126,7 @@ const TagManagement: React.FC = ({ accessToken, userID, userRole }) => }, [accessToken]); return ( -
+
{selectedTagId ? ( = ({ accessToken, userID, userRole }) => editTag={editTag} /> ) : ( -
+

Tag Management

@@ -162,23 +162,21 @@ const TagManagement: React.FC = ({ accessToken, userID, userRole }) =>

- -
-
- { - setSelectedTagId(tag.name); - setEditTag(true); - }} - onDelete={handleDelete} - onSelectTag={setSelectedTagId} - /> -
+
+ { + setSelectedTagId(tag.name); + setEditTag(true); + }} + onDelete={handleDelete} + onSelectTag={setSelectedTagId} + />
{/* Create Tag Modal */} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx index f74f444ce26..1745c710c51 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx @@ -137,8 +137,8 @@ const VectorStoreManagement: React.FC = ({ accessToken, userID />
) : ( -
-
+
+

Vector Store Management

diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx index 5f6f26bd16c..7f349fa29ba 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx @@ -400,7 +400,7 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, } return ( -
+
{publicPage == false ? (
{/* Header with Title, Description and URL */} diff --git a/ui/litellm-dashboard/src/components/PublicModelHubTableColumns.tsx b/ui/litellm-dashboard/src/components/PublicModelHubTableColumns.tsx index ab0ed976149..de987360791 100644 --- a/ui/litellm-dashboard/src/components/PublicModelHubTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/PublicModelHubTableColumns.tsx @@ -6,6 +6,7 @@ import { DataTableSortHeader } from "@/components/shared/DataTable"; import { CellTooltip, IdentityCell, StatusBadge, type StatusTone } from "@/components/shared/table_cells"; import { Badge } from "@/components/ui/badge"; import { getProviderLogoAndName } from "@/components/provider_info_helpers"; +import { PUBLIC_MODEL_HUB_SORTABLE_FIELDS } from "@/components/publicModelHub/publicModelHubFilters"; export interface ModelGroupInfo { model_group: string; @@ -163,154 +164,150 @@ interface PublicModelHubColumnsDeps { onModelClick: (model: ModelGroupInfo) => void; } -export const getPublicModelHubColumns = ({ onModelClick }: PublicModelHubColumnsDeps): ColumnDef[] => [ - { - id: "model_group", - accessorKey: "model_group", - meta: { title: "Model Name" }, - header: ({ column }) => , - size: 200, - enableSorting: true, - sortingFn: "alphanumeric", - cell: ({ row }) => ( - onModelClick(row.original)} - /> - ), - }, - { - id: "providers", - accessorKey: "providers", - meta: { title: "Providers", skeleton: "chips" }, - header: ({ column }) => , - size: 150, - enableSorting: true, - sortingFn: (rowA, rowB) => - (rowA.original.providers ?? []).join(", ").localeCompare((rowB.original.providers ?? []).join(", ")), - cell: ({ row }) => , - }, - { - id: "mode", - accessorKey: "mode", - meta: { title: "Mode" }, - header: ({ column }) => , - size: 110, - enableSorting: true, - sortingFn: "alphanumeric", - cell: ({ row }) => ( - - {getModeIcon(row.original.mode || "")} - {row.original.mode || "Chat"} - - ), - }, - { - id: "max_input_tokens", - accessorKey: "max_input_tokens", - meta: { title: "Max Input", numeric: true }, - header: ({ column }) => , - size: 100, - enableSorting: true, - cell: ({ row }) => {formatTokens(row.original.max_input_tokens)}, - }, - { - id: "max_output_tokens", - accessorKey: "max_output_tokens", - meta: { title: "Max Output", numeric: true }, - header: ({ column }) => , - size: 100, - enableSorting: true, - cell: ({ row }) => {formatTokens(row.original.max_output_tokens)}, - }, - { - id: "input_cost_per_token", - accessorKey: "input_cost_per_token", - meta: { title: "Input $/1M", numeric: true }, - header: ({ column }) => , - size: 110, - enableSorting: true, - cell: ({ row }) => ( - - {row.original.input_cost_per_token ? formatCost(row.original.input_cost_per_token) : "Free"} - - ), - }, - { - id: "output_cost_per_token", - accessorKey: "output_cost_per_token", - meta: { title: "Output $/1M", numeric: true }, - header: ({ column }) => , - size: 110, - enableSorting: true, - cell: ({ row }) => ( - - {row.original.output_cost_per_token ? formatCost(row.original.output_cost_per_token) : "Free"} - - ), - }, - { - id: "features", - meta: { title: "Features", skeleton: "chips" }, - header: "Features", - size: 140, - enableSorting: false, - cell: ({ row }) => { - const features = Object.entries(row.original) - .filter(([key, value]) => key.startsWith("supports_") && value === true) - .map(([key]) => formatCapabilityName(key)); - return ; - }, - }, - { - id: "health_status", - accessorKey: "health_status", - meta: { title: "Health Status", skeleton: "badge" }, - header: ({ column }) => , - size: 130, - enableSorting: true, - cell: ({ row }) => { - const model = row.original; - const responseTimeLabel = model.health_response_time - ? `Response Time: ${Number(model.health_response_time).toFixed(2)}ms` - : "N/A"; - const lastCheckedLabel = model.health_checked_at - ? `Last Checked: ${new Date(model.health_checked_at).toLocaleString()}` - : "N/A"; - return ( - -
{responseTimeLabel}
-
{lastCheckedLabel}
- - } - trigger={ - - - - } +export const getPublicModelHubColumns = ({ onModelClick }: PublicModelHubColumnsDeps): ColumnDef[] => { + const columns: ColumnDef[] = [ + { + id: "model_group", + accessorKey: "model_group", + meta: { title: "Model Name" }, + header: ({ column }) => , + size: 200, + sortingFn: "alphanumeric", + cell: ({ row }) => ( + onModelClick(row.original)} /> - ); + ), }, - }, - { - id: "rpm", - accessorKey: "rpm", - meta: { title: "Limits" }, - header: ({ column }) => , - size: 150, - enableSorting: true, - cell: ({ row }) => ( - {formatLimits(row.original.rpm, row.original.tpm)} - ), - }, -]; + { + id: "providers", + accessorKey: "providers", + meta: { title: "Providers", skeleton: "chips" }, + header: ({ column }) => , + size: 150, + sortingFn: (rowA, rowB) => + (rowA.original.providers ?? []).join(", ").localeCompare((rowB.original.providers ?? []).join(", ")), + cell: ({ row }) => , + }, + { + id: "mode", + accessorKey: "mode", + meta: { title: "Mode" }, + header: ({ column }) => , + size: 110, + sortingFn: "alphanumeric", + cell: ({ row }) => ( + + {getModeIcon(row.original.mode || "")} + {row.original.mode || "Chat"} + + ), + }, + { + id: "max_input_tokens", + accessorKey: "max_input_tokens", + meta: { title: "Max Input", numeric: true }, + header: ({ column }) => , + size: 100, + cell: ({ row }) => {formatTokens(row.original.max_input_tokens)}, + }, + { + id: "max_output_tokens", + accessorKey: "max_output_tokens", + meta: { title: "Max Output", numeric: true }, + header: ({ column }) => , + size: 100, + cell: ({ row }) => {formatTokens(row.original.max_output_tokens)}, + }, + { + id: "input_cost_per_token", + accessorKey: "input_cost_per_token", + meta: { title: "Input $/1M", numeric: true }, + header: ({ column }) => , + size: 110, + cell: ({ row }) => ( + + {row.original.input_cost_per_token ? formatCost(row.original.input_cost_per_token) : "Free"} + + ), + }, + { + id: "output_cost_per_token", + accessorKey: "output_cost_per_token", + meta: { title: "Output $/1M", numeric: true }, + header: ({ column }) => , + size: 110, + cell: ({ row }) => ( + + {row.original.output_cost_per_token ? formatCost(row.original.output_cost_per_token) : "Free"} + + ), + }, + { + id: "features", + meta: { title: "Features", skeleton: "chips" }, + header: "Features", + size: 140, + cell: ({ row }) => { + const features = Object.entries(row.original) + .filter(([key, value]) => key.startsWith("supports_") && value === true) + .map(([key]) => formatCapabilityName(key)); + return ; + }, + }, + { + id: "health_status", + accessorKey: "health_status", + meta: { title: "Health Status", skeleton: "badge" }, + header: ({ column }) => , + size: 130, + cell: ({ row }) => { + const model = row.original; + const responseTimeLabel = model.health_response_time + ? `Response Time: ${Number(model.health_response_time).toFixed(2)}ms` + : "N/A"; + const lastCheckedLabel = model.health_checked_at + ? `Last Checked: ${new Date(model.health_checked_at).toLocaleString()}` + : "N/A"; + return ( + +
{responseTimeLabel}
+
{lastCheckedLabel}
+ + } + trigger={ + + + + } + /> + ); + }, + }, + { + id: "rpm", + accessorKey: "rpm", + meta: { title: "Limits" }, + header: ({ column }) => , + size: 150, + cell: ({ row }) => ( + {formatLimits(row.original.rpm, row.original.tpm)} + ), + }, + ]; + return columns.map((column) => ({ + ...column, + enableSorting: PUBLIC_MODEL_HUB_SORTABLE_FIELDS.includes(String(column.id)), + })); +}; interface PublicAgentHubColumnsDeps { onAgentClick: (agent: AgentCard) => void; diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index ef58237a6aa..4be91f22339 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -559,6 +559,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser { key: "your-teams", label: "Your Teams", + className: "flex min-h-0 flex-1 flex-col", children: ( <> = ({ accessToken, userID, userRole, premiumUser { key: "available-teams", label: "Available Teams", + className: "min-h-0 flex-1 overflow-y-auto", children: , }, ...(isProxyAdminRole(userRole || "") @@ -615,6 +617,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser { key: "default-settings", label: "Default Team Settings", + className: "min-h-0 flex-1 overflow-y-auto", children: , }, ] @@ -622,7 +625,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser ]; return ( -
+
{selectedTeamId ? ( = ({ accessToken, userID, userRole, premiumUser premiumUser={premiumUser} /> ) : ( - + } title="Teams" @@ -674,7 +677,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser )} /> {tabItems.map((item) => ( - + {item.children} ))} diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx index 0b38d51e36a..3fb19e522f3 100644 --- a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx +++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx @@ -164,7 +164,7 @@ export function TeamsTable({ userRole, userID, onSelectTeam, onEditTeam, onDelet isLoading={isLoading} loadingMessage="Loading teams..." noDataMessage="No teams found" - maxBodyHeight="calc(75vh - 210px)" + fillHeight size="compact" toolbar={(table) => ( <> diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index ebedf57af45..f28ae27b6bc 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -256,7 +256,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { } return ( -
+
} title="Virtual Keys" @@ -283,7 +283,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { isLoading={isLoading} loadingMessage="Loading keys..." noDataMessage="No keys found" - maxBodyHeight="calc(75vh - 210px)" + fillHeight size="compact" toolbar={(table) => ( <> diff --git a/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx b/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx index 4d9e2122739..9022d424369 100644 --- a/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx @@ -1,26 +1,58 @@ import React from "react"; +import { Input } from "@/components/ui/input"; import { Switch } from "@/components/ui/switch"; import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; -import { DEFAULT_DEPLOYMENT_AFFINITY } from "./ComplexityRouterConfig"; +import { DEFAULT_DEPLOYMENT_AFFINITY, DEFAULT_SESSION_AFFINITY_TTL_SECONDS } from "./ComplexityRouterConfig"; export const AffinityControls: React.FC<{ value: ComplexityRouterConfigValue; onChange: (value: ComplexityRouterConfigValue) => void; -}> = ({ value, onChange }) => ( - <> -
- onChange({ ...value, deployment_affinity: deploymentAffinity })} - aria-label="Pin a session to one deployment per model group" - /> - Pin a session to one deployment per model group -
- - Keeps a session on the same deployment within a group, so provider prompt caches stay warm. Turn off to - load-balance every turn. - - -); +}> = ({ value, onChange }) => { + const [ttlDraft, setTtlDraft] = React.useState(null); + const commitTtl = (raw: string) => { + setTtlDraft(null); + if (raw.trim() === "") { + onChange({ ...value, session_affinity_ttl_seconds: undefined }); + return; + } + const parsed = Number(raw); + if (!Number.isFinite(parsed)) return; + onChange({ ...value, session_affinity_ttl_seconds: Math.max(1, Math.round(parsed)) }); + }; + + return ( + <> +
+ onChange({ ...value, deployment_affinity: deploymentAffinity })} + aria-label="Pin a session to one deployment per model group" + /> + Pin a session to one deployment per model group +
+ + Keeps a session on the same deployment within a group, so provider prompt caches stay warm. Turn off to + load-balance every turn. + +
+ + setTtlDraft(event.target.value)} + onBlur={(event) => commitTtl(event.target.value)} + /> + + Refreshes after every request that reuses a pin. Empty tracks the backend default of{" "} + {DEFAULT_SESSION_AFFINITY_TTL_SECONDS} seconds. + +
+ + ); +}; diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx index 02cb0543e17..bdca5205b2e 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -947,6 +947,46 @@ describe("ComplexityRouterConfig affinity panel", () => { expect(screen.getByRole("switch", { name: "Pin a session to one deployment per model group" })).not.toBeChecked(); }); + + it("writes an idle TTL on blur and keeps the partial input as a draft while typing", () => { + const onChange = vi.fn(); + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Affinity")); + + const ttl = screen.getByLabelText("How long a pin survives idle (seconds)"); + expect(ttl).toHaveAttribute("placeholder", "3600"); + fireEvent.change(ttl, { target: { value: "300" } }); + expect(onChange).not.toHaveBeenCalled(); + fireEvent.blur(ttl); + + expect(onChange).toHaveBeenCalledWith({ ...defaultValue, session_affinity_ttl_seconds: 300 }); + }); + + it("clearing the idle TTL returns the router to its backend default", () => { + const onChange = vi.fn(); + const value = { ...defaultValue, session_affinity_ttl_seconds: 300 }; + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Affinity")); + + const ttl = screen.getByLabelText("How long a pin survives idle (seconds)"); + expect(ttl).toHaveValue("300"); + fireEvent.change(ttl, { target: { value: "" } }); + fireEvent.blur(ttl); + + expect(onChange).toHaveBeenCalledWith({ ...value, session_affinity_ttl_seconds: undefined }); + }); + + it("clamps a non-positive idle TTL to the backend's minimum", () => { + const onChange = vi.fn(); + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Affinity")); + + const ttl = screen.getByLabelText("How long a pin survives idle (seconds)"); + fireEvent.change(ttl, { target: { value: "0" } }); + fireEvent.blur(ttl); + + expect(onChange).toHaveBeenCalledWith({ ...defaultValue, session_affinity_ttl_seconds: 1 }); + }); }); describe("ComplexityRouterConfig default model", () => { diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 791fd1fcfbb..06363830d64 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -58,6 +58,7 @@ export const DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE = 3; export const DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS = 8000; export const MIN_QUOTED_CONTEXT_TURN_CHARS = 120; export const DEFAULT_SESSION_AFFINITY = false; +export const DEFAULT_SESSION_AFFINITY_TTL_SECONDS = 3600; export const DEFAULT_DEPLOYMENT_AFFINITY = true; export type ClassificationMode = "every_request" | "user_turn"; @@ -411,6 +412,7 @@ export interface ComplexityRouterConfigValue { hybrid_boundary_margin?: number; classification_mode?: ClassificationMode; session_affinity?: boolean; + session_affinity_ttl_seconds?: number; modality_routing?: boolean; modality_pin_override?: boolean; deployment_affinity?: boolean; diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index 836b4c6ff3d..d2f6b10c3a6 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -19,7 +19,6 @@ vi.mock( "@/app/(dashboard)/hooks/autoRouter/useAutoRouterPresets", async () => await import("../../../tests/mocks/autoRouterPresets"), ); - const getAllPresets = (): AutoRouterPreset[] => BUNDLED_PRESETS; const getPresetByKey = (key: string): AutoRouterPreset | undefined => BUNDLED_PRESETS.find((p) => p.key === key); @@ -155,6 +154,58 @@ describe("AddAutoRouterTab", () => { expect(screen.getByText("Complexity Tier Configuration")).toBeInTheDocument(); }); + it("hides automatic setup when no available model is recommended", async () => { + mockFetchAvailableModels.mockResolvedValue([ + { model_group: "unknown-model-a", mode: "chat" }, + { model_group: "unknown-model-b", mode: "chat" }, + ]); + renderWithProviders(); + + openTemplateDropdown(); + await waitFor(() => expect(optionByLabel("Anthropic Family")).toHaveTextContent("Missing:")); + expect(screen.queryByTestId("configure-automatically-button")).not.toBeInTheDocument(); + }); + + it("mixes preferred tier models even when one complete preset is available", async () => { + const anthropicPreset = getPresetByKey("anthropic_family")!; + mockFetchAvailableModels.mockResolvedValue( + [...getRequiredModelsInPreset(anthropicPreset), "gpt-5.6-luna"].map((model_group) => ({ + model_group, + mode: "chat", + })), + ); + mockFetchAllModelDeployments.mockResolvedValue([]); + renderWithProviders(); + + const button = await screen.findByTestId("configure-automatically-button"); + await userEvent.click(button); + + expect( + screen.getByText( + /Simple: gpt-5.6-luna.*Medium: claude-sonnet-5.*Complex: claude-opus-5.*Reasoning: claude-opus-5/, + ), + ).toBeInTheDocument(); + expect(toast.success).not.toHaveBeenCalledWith(expect.stringContaining("Configured with")); + }); + + it("mixes available models from the preferred tier catalog when no complete template fits", async () => { + mockFetchAvailableModels.mockResolvedValue( + ["gpt-5.6-luna", "claude-sonnet-5", "gpt-5.6-sol"].map((model_group) => ({ + model_group, + mode: "chat", + })), + ); + mockFetchAllModelDeployments.mockResolvedValue([]); + renderWithProviders(); + + const button = await screen.findByTestId("configure-automatically-button"); + await userEvent.click(button); + + expect( + screen.getByText(/Simple: gpt-5.6-luna.*Medium: claude-sonnet-5.*Complex: gpt-5.6-sol.*Reasoning: gpt-5.6-sol/), + ).toBeInTheDocument(); + }); + // Nothing is filled in, so there is nothing to submit. The button reports that itself instead of // accepting a click and answering with a toast. it("offers no submit at all until every tier has a model", async () => { @@ -493,7 +544,7 @@ describe("AddAutoRouterTab", () => { }); }); - it("carries session affinity turned on through to the create payload", async () => { + it("carries session affinity turned on and its idle window through to the create payload", async () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); @@ -503,12 +554,17 @@ describe("AddAutoRouterTab", () => { expandDetailedConfiguration(); await user.click(screen.getByText("Advanced: Classification Method")); await user.click(await screen.findByRole("radio", { name: /Once per session/ })); + await user.click(screen.getByText("Advanced: Affinity")); + const ttl = await screen.findByLabelText("How long a pin survives idle (seconds)"); + fireEvent.change(ttl, { target: { value: "300" } }); + fireEvent.blur(ttl); await user.click(screen.getByRole("button", { name: /add auto router/i })); await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).toMatchObject({ session_affinity: true, + session_affinity_ttl_seconds: 300, }); }); diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index c335beade91..a548c2c6533 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -64,6 +64,7 @@ import { } from "@/lib/autorouter_presets"; import { useAutoRouterPresets } from "@/app/(dashboard)/hooks/autoRouter/useAutoRouterPresets"; import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import { buildAutomaticRouterConfig, buildPreferredTierModels } from "./auto_setup"; interface AddAutoRouterTabProps { handleOk: () => void; @@ -242,6 +243,7 @@ const AddAutoRouterTab: React.FC = ({ refetch: refetchPresets, } = useAutoRouterPresets(); const presets = presetsData ?? NO_PRESETS; + const automaticSetupLoading = modelsLoading || presetsPending; const presetsUnavailable = presetsError && presetsData === undefined; // react-query keeps the last successful list around when a later refetch fails, so isError alone // can't tell "never loaded" apart from "loaded, then a background refetch errored" - only the @@ -266,6 +268,14 @@ const AddAutoRouterTab: React.FC = ({ ), [modelInfo], ); + const preferredTierModels = React.useMemo( + () => buildPreferredTierModels(presets, availability), + [presets, availability], + ); + const automaticRouterConfig = React.useMemo( + () => buildAutomaticRouterConfig(modelInfo, deployments ?? [], preferredTierModels), + [modelInfo, deployments, preferredTierModels], + ); // A preset's models can only be trusted against a successfully loaded list. Selection and the // greyed-out state derive from this one function, so a preset that cannot be selected can never @@ -313,6 +323,14 @@ const AddAutoRouterTab: React.FC = ({ setEscalationKeywords(prefill.escalationKeywords); }; + const handleAutomaticSetup = () => { + if (automaticRouterConfig === null) return; + setSelectedPreset(undefined); + applyPrefill({ ...buildEmptyPrefill(), complexityRouterConfig: automaticRouterConfig }); + setDetailsExpanded(false); + toast.success("Automatic setup created", { description: tierConfigSummary(automaticRouterConfig) }); + }; + const handlePresetChange = (presetKey: string | undefined) => { if (!presetKey || presetKey === "custom") { setSelectedPreset(presetKey); @@ -388,6 +406,7 @@ const AddAutoRouterTab: React.FC = ({ reasoningOverrideMinScore: complexityRouterConfig.reasoning_override_min_score, enableContextWindowEscalation: complexityRouterConfig.enable_context_window_escalation, contextWindowEscalationBuffer: complexityRouterConfig.context_window_escalation_buffer, + sessionAffinityTtlSeconds: complexityRouterConfig.session_affinity_ttl_seconds, }; const submitRecommendedRouter = async (name: string) => { @@ -499,77 +518,92 @@ const AddAutoRouterTab: React.FC = ({
handleAutoRouterSubmit())} noValidate> - - {({ ref, ...field }) => } - -
- - + )} + - return ( - -
-
{preset.label}
-
{preset.description}
- {disabledHint &&
{disabledHint}
} - {matchedHint &&
{matchedHint}
} -
- - ); - })} - -
-
Custom Configuration
-
Define your auto router from scratch
-
- - - - {modelsUnverifiable && ( -
- Could not load available models.{" "} - -
- )} - {presetsPending &&
Loading templates...
} - {presetsUnavailable && ( -
- Could not load templates, so only Custom Configuration is shown.{" "} - -
+ {!automaticSetupLoading && automaticRouterConfig && ( + )} + +
+ + + {modelsUnverifiable && ( +
+ Could not load available models.{" "} + +
+ )} + {presetsPending &&
Loading templates...
} + {presetsUnavailable && ( +
+ Could not load templates, so only Custom Configuration is shown.{" "} + +
+ )} +
{requiresTeamScope && ( diff --git a/ui/litellm-dashboard/src/components/add_model/auto_setup.test.ts b/ui/litellm-dashboard/src/components/add_model/auto_setup.test.ts new file mode 100644 index 00000000000..c5784db4501 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/auto_setup.test.ts @@ -0,0 +1,196 @@ +import { describe, expect, it } from "vitest"; +import type { AutoRouterDeployment } from "@/app/(dashboard)/hooks/models/useModels"; +import { buildModelAvailability } from "@/lib/autorouter_presets"; +import { buildAutomaticRouterConfig, buildPreferredTierModels, type PreferredTierModels } from "./auto_setup"; + +const models = (...names: string[]) => names.map((model_group) => ({ model_group, mode: "chat" })); +const reasoningModel = (model_group: string, supported_reasoning_efforts: string[]) => ({ + model_group, + mode: "chat", + supports_reasoning: true, + supported_reasoning_efforts, +}); +const deployment = (model_name: string, model = model_name): AutoRouterDeployment => ({ + model_name, + litellm_params: { model }, +}); +const tierModels = (config: ReturnType) => + config && Object.values(config.tiers).map((tier) => (typeof tier === "string" ? tier : tier[0])); + +describe("buildPreferredTierModels", () => { + it("recognizes curated models that are not in a preset", () => { + const available = ["gpt-5.6-luna", "claude-sonnet-5", "grok-4.6", "deepseek-v4-pro"]; + const availability = buildModelAvailability(available, []); + const preferred = buildPreferredTierModels([], availability); + const expected: PreferredTierModels = { + SIMPLE: ["gpt-5.6-luna"], + MEDIUM: ["claude-sonnet-5"], + COMPLEX: ["deepseek-v4-pro", "grok-4.6"], + REASONING: ["deepseek-v4-pro", "grok-4.6"], + }; + + expect(preferred).toEqual(expected); + }); + + it("prefers the current model ladder over older preset entries", () => { + const availability = buildModelAvailability(["gpt-5.6-luna", "gpt-4o-mini"], []); + const preferred = buildPreferredTierModels( + [ + { + key: "old", + label: "Old", + description: "Old model", + complexity_router_config: { + tiers: { SIMPLE: ["gpt-4o-mini"], MEDIUM: [], COMPLEX: [], REASONING: [] }, + classifier_type: "heuristic_v2", + }, + }, + ], + availability, + ); + + expect(preferred.SIMPLE).toEqual(["gpt-5.6-luna", "gpt-4o-mini"]); + }); +}); + +describe("buildAutomaticRouterConfig", () => { + it("selects one preferred model for each tier", () => { + const preferred: PreferredTierModels = { + SIMPLE: ["simple"], + MEDIUM: ["medium"], + COMPLEX: ["complex"], + REASONING: ["reasoning"], + }; + + expect( + tierModels(buildAutomaticRouterConfig(models("simple", "medium", "complex", "reasoning"), [], preferred)), + ).toEqual(["simple", "medium", "complex", "reasoning"]); + }); + + it.each([ + { + provider: "OpenAI", + available: ["gpt-5.6-luna", "gpt-5.6-terra", "gpt-6-astra"], + expected: ["gpt-5.6-luna", "gpt-5.6-terra", "gpt-6-astra", "gpt-6-astra"], + supportedEfforts: ["low", "medium", "high", "xhigh", "max"], + effort: "max", + }, + { + provider: "Anthropic", + available: ["claude-haiku-4-5", "claude-sonnet-5", "claude-opus-5"], + expected: ["claude-haiku-4-5", "claude-sonnet-5", "claude-opus-5", "claude-opus-5"], + supportedEfforts: ["low", "medium", "high", "max"], + effort: "max", + }, + { + provider: "Google", + available: ["gemini-3.5-flash-lite", "gemini-3.8-flash", "gemini-3.1-pro-preview"], + expected: ["gemini-3.5-flash-lite", "gemini-3.8-flash", "gemini-3.1-pro-preview", "gemini-3.1-pro-preview"], + supportedEfforts: ["low", "medium", "high"], + effort: "high", + }, + { + provider: "DeepSeek", + available: ["deepseek-v4-flash", "deepseek-v4-pro"], + expected: ["deepseek-v4-flash", "deepseek-v4-flash", "deepseek-v4-pro", "deepseek-v4-pro"], + supportedEfforts: ["none", "high"], + effort: "high", + }, + { + provider: "xAI", + available: ["grok-4.6"], + expected: ["grok-4.6", "grok-4.6", "grok-4.6", "grok-4.6"], + supportedEfforts: ["low", "medium", "high", "xhigh"], + effort: "xhigh", + }, + ])( + "uses the current $provider ladder and strongest advertised reasoning effort", + ({ available, expected, supportedEfforts, effort }) => { + const availability = buildModelAvailability(available, []); + const preferred = buildPreferredTierModels([], availability); + const modelInfo = models(...available).map((model) => + model.model_group === expected[3] ? reasoningModel(model.model_group, supportedEfforts) : model, + ); + + const config = buildAutomaticRouterConfig(modelInfo, [], preferred); + + expect(tierModels(config)).toEqual(expected); + expect(config?.tier_model_params).toEqual({ + REASONING: { [expected[3]]: { reasoning_effort: effort } }, + }); + }, + ); + + it("never exceeds the selected model group's advertised reasoning efforts", () => { + const available = ["gpt-5.6-luna", "gpt-5.6-terra", "gpt-5.6-sol"]; + const availability = buildModelAvailability(available, []); + const preferred = buildPreferredTierModels([], availability); + const modelInfo = [ + ...models("gpt-5.6-luna", "gpt-5.6-terra"), + reasoningModel("gpt-5.6-sol", ["none", "low", "medium", "high", "xhigh"]), + ]; + + const config = buildAutomaticRouterConfig(modelInfo, [], preferred); + + expect(config?.tier_model_params).toEqual({ + REASONING: { "gpt-5.6-sol": { reasoning_effort: "xhigh" } }, + }); + }); + + it("leaves reasoning effort unset when the proxy does not report supported values", () => { + const available = ["grok-4.6"]; + const availability = buildModelAvailability(available, []); + const preferred = buildPreferredTierModels([], availability); + + const config = buildAutomaticRouterConfig(models(...available), [], preferred); + + expect(config?.tier_model_params).toBeUndefined(); + }); + + it("reuses the closest available tier when a tier has no match", () => { + const preferred: PreferredTierModels = { + SIMPLE: ["simple"], + MEDIUM: [], + COMPLEX: ["complex"], + REASONING: [], + }; + + expect(tierModels(buildAutomaticRouterConfig(models("simple", "complex"), [], preferred))).toEqual([ + "simple", + "simple", + "complex", + "complex", + ]); + }); + + it("returns null when none of the available models are recommended", () => { + const preferred: PreferredTierModels = { + SIMPLE: ["missing-simple"], + MEDIUM: ["missing-medium"], + COMPLEX: ["missing-complex"], + REASONING: ["missing-reasoning"], + }; + + expect(buildAutomaticRouterConfig(models("unknown-model"), [], preferred)).toBeNull(); + }); + + it("ignores non-chat models and existing auto routers", () => { + const preferred: PreferredTierModels = { + SIMPLE: ["gpt-4o-mini", "smart-router"], + MEDIUM: [], + COMPLEX: [], + REASONING: [], + }; + const available = [ + { model_group: "gpt-4o-mini", mode: "chat" }, + { model_group: "image-model", mode: "image_generation" }, + { model_group: "smart-router", mode: "chat" }, + ]; + + expect( + tierModels( + buildAutomaticRouterConfig(available, [deployment("smart-router", "auto_router/complexity_router")], preferred), + ), + ).toEqual(["gpt-4o-mini", "gpt-4o-mini", "gpt-4o-mini", "gpt-4o-mini"]); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/auto_setup.ts b/ui/litellm-dashboard/src/components/add_model/auto_setup.ts new file mode 100644 index 00000000000..cd7d587e2c2 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/auto_setup.ts @@ -0,0 +1,95 @@ +import { isAutoRouterDeployment, type AutoRouterDeployment } from "@/app/(dashboard)/hooks/models/useModels"; +import type { ModelGroup } from "@/components/llm_calls/fetch_models"; +import { resolveAvailableModel, type AutoRouterPreset, type ModelAvailability } from "@/lib/autorouter_presets"; +import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; + +const TIER_NAMES = ["SIMPLE", "MEDIUM", "COMPLEX", "REASONING"] as const; +type TierName = (typeof TIER_NAMES)[number]; +export type PreferredTierModels = Record; + +const REASONING_EFFORT_STRENGTH = ["max", "xhigh", "high", "medium", "low", "minimal", "none"] as const; + +const CURRENT_TIER_MODELS: PreferredTierModels = { + SIMPLE: ["gpt-5.6-luna", "claude-haiku-4-5", "gemini-3.5-flash-lite", "deepseek-v4-flash"], + MEDIUM: ["gpt-5.6-terra", "claude-sonnet-5", "gemini-3.8-flash", "deepseek-v4-flash"], + COMPLEX: ["gpt-6-astra", "gpt-5.6-sol", "claude-opus-5", "gemini-3.1-pro-preview", "deepseek-v4-pro", "grok-4.6"], + REASONING: ["gpt-6-astra", "gpt-5.6-sol", "claude-opus-5", "gemini-3.1-pro-preview", "deepseek-v4-pro", "grok-4.6"], +}; + +export const buildPreferredTierModels = ( + presets: AutoRouterPreset[], + availability: ModelAvailability, +): PreferredTierModels => + Object.fromEntries( + TIER_NAMES.map((tier) => [ + tier, + Array.from( + new Set( + [ + ...CURRENT_TIER_MODELS[tier], + ...presets.flatMap((preset) => preset.complexity_router_config.tiers[tier]), + ].flatMap((model) => { + const resolved = resolveAvailableModel(model, availability); + return resolved ? [resolved] : []; + }), + ), + ), + ]), + ) as PreferredTierModels; + +const selectPreferredTierModels = ( + preferredByTier: PreferredTierModels, + usableNames: ReadonlySet, +): [string, string, string, string] | null => { + const preferred = TIER_NAMES.map((tier) => preferredByTier[tier].find((name) => usableNames.has(name))); + const candidates = preferred.flatMap((model, tier) => (model ? [{ model, tier }] : [])); + if (candidates.length === 0) return null; + + const nearest = (tier: number): string => + [...candidates].sort( + (left, right) => Math.abs(left.tier - tier) - Math.abs(right.tier - tier) || left.tier - right.tier, + )[0].model; + return preferred.map((model, tier) => model ?? nearest(tier)) as [string, string, string, string]; +}; + +export const buildAutomaticRouterConfig = ( + models: ModelGroup[], + deployments: AutoRouterDeployment[], + preferredByTier: PreferredTierModels, +): ComplexityRouterConfigValue | null => { + const autoRouterNames: ReadonlySet = new Set( + deployments + .filter(isAutoRouterDeployment) + .flatMap((deployment) => (deployment.model_name ? [deployment.model_name] : [])), + ); + const names = Array.from( + new Set( + models + .filter((model) => model.mode === undefined || model.mode === "chat") + .map((model) => model.model_group) + .filter((name) => name && !name.startsWith("auto_router/") && !autoRouterNames.has(name)), + ), + ); + if (names.length === 0) return null; + const usableNames: ReadonlySet = new Set(names); + const selected = selectPreferredTierModels(preferredByTier, usableNames); + if (selected === null) return null; + + const supportedReasoningEfforts = models.find( + (model) => model.model_group === selected[3], + )?.supported_reasoning_efforts; + const reasoningEffort = REASONING_EFFORT_STRENGTH.find((effort) => supportedReasoningEfforts?.includes(effort)); + + return { + tiers: { + SIMPLE: [selected[0]], + MEDIUM: [selected[1]], + COMPLEX: [selected[2]], + REASONING: [selected[3]], + }, + classifier_type: "heuristic_v2", + ...(reasoningEffort && { + tier_model_params: { REASONING: { [selected[3]]: { reasoning_effort: reasoningEffort } } }, + }), + }; +}; diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index 53ab859baa0..9ee555f5dd2 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -73,6 +73,15 @@ describe("buildComplexityRouterConfig", () => { expect(config.context_window_escalation_buffer).toBe(0.9); }); + it("omits session_affinity_ttl_seconds when untouched, so the router tracks the backend default", () => { + expect(buildComplexityRouterConfig(baseParams)).not.toHaveProperty("session_affinity_ttl_seconds"); + }); + + it("emits an explicit session affinity idle window", () => { + const config = buildComplexityRouterConfig({ ...baseParams, sessionAffinityTtlSeconds: 300 }); + expect(config.session_affinity_ttl_seconds).toBe(300); + }); + it("trims escalation keywords and drops blank entries", () => { const config = buildComplexityRouterConfig({ ...baseParams, diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index b8af55e8c8a..956e593a234 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -138,6 +138,7 @@ export interface BuildComplexityRouterConfigParams { tierModelParams?: TierModelParamsByTier; enableContextWindowEscalation?: boolean; contextWindowEscalationBuffer?: number; + sessionAffinityTtlSeconds?: number; } /** @@ -175,6 +176,7 @@ export interface ComplexityRouterConfigPayload { hybrid_boundary_margin?: number; classification_mode: ClassificationMode; session_affinity: boolean; + session_affinity_ttl_seconds?: number; deployment_affinity: boolean; modality_routing: boolean; modality_pin_override: boolean; @@ -456,6 +458,7 @@ export const buildComplexityRouterConfig = ({ tierModelParams, enableContextWindowEscalation, contextWindowEscalationBuffer, + sessionAffinityTtlSeconds, }: BuildComplexityRouterConfigParams): ComplexityRouterConfigPayload => { const serializedTierModelConfigs = customTierSet ? serializeTierModelConfigs( @@ -522,6 +525,9 @@ export const buildComplexityRouterConfig = ({ ...(contextWindowEscalationBuffer !== undefined && { context_window_escalation_buffer: contextWindowEscalationBuffer, }), + ...(sessionAffinityTtlSeconds !== undefined && { + session_affinity_ttl_seconds: sessionAffinityTtlSeconds, + }), ...scorerKnobs, }; if (!customTierSet) return payload; diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts index f24f3033901..877b35199a4 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts @@ -257,6 +257,43 @@ describe("buildUpdatedComplexityRouterConfig session affinity", () => { }); }); +describe("buildUpdatedComplexityRouterConfig session affinity ttl", () => { + it("writes an edited idle window", () => { + const result = buildUpdatedComplexityRouterConfig(STORED, { ...FORM_VALUE, session_affinity_ttl_seconds: 300 }); + expect(result.session_affinity_ttl_seconds).toBe(300); + }); + + it("carries a stored idle window through an untouched open-and-save", () => { + const stored = { ...STORED, session_affinity_ttl_seconds: 900 }; + const result = buildUpdatedComplexityRouterConfig(stored, hydrateComplexityRouterConfig(stored, undefined)); + expect(result.session_affinity_ttl_seconds).toBe(900); + }); + + it("drops the key when the field is cleared, so the router goes back to tracking the backend default", () => { + const result = buildUpdatedComplexityRouterConfig( + { ...STORED, session_affinity_ttl_seconds: 900 }, + { ...FORM_VALUE, session_affinity_ttl_seconds: undefined }, + ); + expect(result).not.toHaveProperty("session_affinity_ttl_seconds"); + }); + + it("keeps the idle window on a custom tier set, whose deployment pin still uses it", () => { + const result = buildUpdatedComplexityRouterConfig(STORED, { + ...FORM_VALUE, + session_affinity_ttl_seconds: 300, + custom_tier_set: { + tiers: [ + { id: "a", name: "CASUAL", definition: "small talk", models: ["gpt-4o-mini"] }, + { id: "b", name: "AUDIT", definition: "security review", models: ["o1"] }, + ], + fallback_tier_id: "a", + }, + }); + expect(result.session_affinity).toBe(false); + expect(result.session_affinity_ttl_seconds).toBe(300); + }); +}); + describe("buildUpdatedComplexityRouterConfig modality pin override", () => { it("writes modality_pin_override explicitly both ways", () => { expect( @@ -529,6 +566,7 @@ describe("managed keys survive an untouched open-and-save", () => { classifier_fallback: "default_model", classification_mode: "user_turn", session_affinity: true, + session_affinity_ttl_seconds: 300, modality_routing: true, modality_pin_override: true, deployment_affinity: false, diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx index 6067e72e547..970bcaa545f 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx @@ -550,6 +550,49 @@ describe("EditAutoRouterModal deployment affinity", () => { expect(savedConfig().deployment_affinity).toBe(false); }); + it("preserves an idle TTL through an untouched save", async () => { + const user = userEvent.setup(); + renderWithStoredConfig({ ...STORED_CONFIG, session_affinity_ttl_seconds: 300 }); + + await user.click(await screen.findByText("Advanced: Affinity")); + expect(await screen.findByLabelText("How long a pin survives idle (seconds)")).toHaveValue("300"); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig().session_affinity_ttl_seconds).toBe(300); + }); + + it("persists an edited idle TTL", async () => { + const user = userEvent.setup(); + renderWithStoredConfig(STORED_CONFIG); + + await user.click(await screen.findByText("Advanced: Affinity")); + const ttl = await screen.findByLabelText("How long a pin survives idle (seconds)"); + fireEvent.change(ttl, { target: { value: "300" } }); + fireEvent.blur(ttl); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig().session_affinity_ttl_seconds).toBe(300); + }); + + it("removes the idle TTL when cleared", async () => { + const user = userEvent.setup(); + renderWithStoredConfig({ ...STORED_CONFIG, session_affinity_ttl_seconds: 300 }); + + await user.click(await screen.findByText("Advanced: Affinity")); + const ttl = await screen.findByLabelText("How long a pin survives idle (seconds)"); + fireEvent.change(ttl, { target: { value: "" } }); + fireEvent.blur(ttl); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig()).not.toHaveProperty("session_affinity_ttl_seconds"); + }); + // modality_pin_override is a managed key, so the modal rewrites it from form state on save. A // hydration gap would silently turn a stored override off on the next untouched save. it("shows a stored modality_pin_override=true as on and preserves it through an untouched save", async () => { diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index 54766df93ec..ea1e5cba6a3 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -104,6 +104,7 @@ export interface StoredComplexityRouterConfig { dimension_weights?: unknown; reasoning_override_min_score?: unknown; session_affinity?: unknown; + session_affinity_ttl_seconds?: unknown; modality_routing?: unknown; modality_pin_override?: unknown; deployment_affinity?: unknown; @@ -182,6 +183,11 @@ export const hydrateComplexityRouterConfig = ( reasoning_override_min_score: hydrateReasoningOverrideMinScore(parsedConfig.reasoning_override_min_score), session_affinity: typeof parsedConfig.session_affinity === "boolean" ? parsedConfig.session_affinity : DEFAULT_SESSION_AFFINITY, + session_affinity_ttl_seconds: + typeof parsedConfig.session_affinity_ttl_seconds === "number" && + Number.isFinite(parsedConfig.session_affinity_ttl_seconds) + ? parsedConfig.session_affinity_ttl_seconds + : undefined, modality_routing: typeof parsedConfig.modality_routing === "boolean" ? parsedConfig.modality_routing : false, modality_pin_override: typeof parsedConfig.modality_pin_override === "boolean" ? parsedConfig.modality_pin_override : false, @@ -224,6 +230,7 @@ export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ "hybrid_boundary_margin", "classification_mode", "session_affinity", + "session_affinity_ttl_seconds", "modality_routing", "modality_pin_override", "deployment_affinity", @@ -322,6 +329,7 @@ export const buildUpdatedComplexityRouterConfig = ( classifierContextIncludeAssistantTurns: value.classifier_context_include_assistant_turns, classifierFallback: value.classifier_fallback, sessionAffinity: value.session_affinity ?? DEFAULT_SESSION_AFFINITY, + sessionAffinityTtlSeconds: value.session_affinity_ttl_seconds, modalityRouting: value.modality_routing ?? false, modalityPinOverride: value.modality_pin_override ?? false, deploymentAffinity: value.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY, diff --git a/ui/litellm-dashboard/src/components/publicModelHub/publicModelHubFilters.test.ts b/ui/litellm-dashboard/src/components/publicModelHub/publicModelHubFilters.test.ts new file mode 100644 index 00000000000..d17f1ccae10 --- /dev/null +++ b/ui/litellm-dashboard/src/components/publicModelHub/publicModelHubFilters.test.ts @@ -0,0 +1,85 @@ +import type { ColumnFiltersState } from "@tanstack/react-table"; +import { describe, expect, it } from "vitest"; + +import { + FEATURE_FILTER_ID, + MODE_FILTER_ID, + PROVIDER_FILTER_ID, + featureLabel, + readFilterValues, + serializePublicModelHubFilters, + withFilterValue, +} from "./publicModelHubFilters"; + +describe("serializePublicModelHubFilters", () => { + it("sends each multi-select as the route's comma separated in filter", () => { + const filters: ColumnFiltersState = [ + { id: MODE_FILTER_ID, value: ["chat", "embedding"] }, + { id: PROVIDER_FILTER_ID, value: ["openai", "anthropic"] }, + { id: FEATURE_FILTER_ID, value: ["vision"] }, + ]; + + expect(serializePublicModelHubFilters(filters)).toEqual({ + "filter[mode][in]": "chat,embedding", + "filter[providers][in]": "openai,anthropic", + "filter[features][in]": "vision", + }); + }); + + it("omits blank filters rather than sending parameters the route rejects", () => { + const filters: ColumnFiltersState = [ + { id: MODE_FILTER_ID, value: [] }, + { id: PROVIDER_FILTER_ID, value: [] }, + ]; + + expect(serializePublicModelHubFilters(filters)).toEqual({}); + }); + + it("ignores filter ids the route does not declare", () => { + expect(serializePublicModelHubFilters([{ id: "health_status", value: ["healthy"] }])).toEqual({}); + }); +}); + +describe("readFilterValues", () => { + it("reads back the values of the filter it names", () => { + const filters: ColumnFiltersState = [ + { id: MODE_FILTER_ID, value: ["chat"] }, + { id: FEATURE_FILTER_ID, value: ["vision", "reasoning"] }, + ]; + + expect(readFilterValues(filters, FEATURE_FILTER_ID)).toEqual(["vision", "reasoning"]); + expect(readFilterValues(filters, PROVIDER_FILTER_ID)).toEqual([]); + }); +}); + +describe("withFilterValue", () => { + it("adds a filter that is not set yet", () => { + expect(withFilterValue([], PROVIDER_FILTER_ID, ["openai"])).toEqual([ + { id: PROVIDER_FILTER_ID, value: ["openai"] }, + ]); + }); + + it("replaces a filter instead of stacking a second one", () => { + const filters: ColumnFiltersState = [{ id: PROVIDER_FILTER_ID, value: ["openai"] }]; + + expect(withFilterValue(filters, PROVIDER_FILTER_ID, ["anthropic"])).toEqual([ + { id: PROVIDER_FILTER_ID, value: ["anthropic"] }, + ]); + }); + + it("drops a cleared filter and leaves the others alone", () => { + const filters: ColumnFiltersState = [ + { id: MODE_FILTER_ID, value: ["chat"] }, + { id: PROVIDER_FILTER_ID, value: ["openai"] }, + ]; + + expect(withFilterValue(filters, PROVIDER_FILTER_ID, [])).toEqual([{ id: MODE_FILTER_ID, value: ["chat"] }]); + }); +}); + +describe("featureLabel", () => { + it("renders a route feature the way the hub has always labelled it", () => { + expect(featureLabel("vision")).toBe("Vision"); + expect(featureLabel("parallel_function_calling")).toBe("Parallel Function Calling"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/publicModelHub/publicModelHubFilters.ts b/ui/litellm-dashboard/src/components/publicModelHub/publicModelHubFilters.ts new file mode 100644 index 00000000000..b217fdae955 --- /dev/null +++ b/ui/litellm-dashboard/src/components/publicModelHub/publicModelHubFilters.ts @@ -0,0 +1,60 @@ +import type { ColumnFilter, ColumnFiltersState } from "@tanstack/react-table"; + +export const MODE_FILTER_ID = "mode"; +export const PROVIDER_FILTER_ID = "providers"; +export const FEATURE_FILTER_ID = "features"; + +export const PUBLIC_MODEL_HUB_SORTABLE_FIELDS: readonly string[] = [ + "model_group", + "mode", + "providers", + "max_input_tokens", + "max_output_tokens", + "input_cost_per_token", + "output_cost_per_token", + "rpm", + "tpm", +]; + +type QueryEntry = readonly [string, string]; + +type FilterValue = string | string[]; + +const entries = (key: string, value: string): QueryEntry[] => (value === "" ? [] : [[key, value]]); + +const asStringArray = (value: unknown): string[] => + Array.isArray(value) ? value.filter((item): item is string => typeof item === "string") : []; + +const inFilter = (field: string, value: unknown): QueryEntry[] => + entries(`filter[${field}][in]`, asStringArray(value).join(",")); + +const filterParams = (filter: ColumnFilter): QueryEntry[] => { + switch (filter.id) { + case MODE_FILTER_ID: + case PROVIDER_FILTER_ID: + case FEATURE_FILTER_ID: + return inFilter(filter.id, filter.value); + default: + return []; + } +}; + +export const serializePublicModelHubFilters = (filters: ColumnFiltersState): Readonly> => + Object.fromEntries(filters.flatMap(filterParams)); + +export const readFilterValues = (filters: ColumnFiltersState, id: string): string[] => + asStringArray(filters.find((filter) => filter.id === id)?.value); + +const isEmpty = (value: FilterValue): boolean => (Array.isArray(value) ? value.length === 0 : value.trim() === ""); + +export const withFilterValue = (filters: ColumnFiltersState, id: string, value: FilterValue): ColumnFiltersState => { + const others = filters.filter((filter) => filter.id !== id); + return isEmpty(value) ? others : [...others, { id, value }]; +}; + +/** `supports_vision` reaches the route as `vision`; the hub has always shown it as "Vision". */ +export const featureLabel = (feature: string): string => + feature + .split("_") + .map((word) => word.charAt(0).toUpperCase() + word.slice(1)) + .join(" "); diff --git a/ui/litellm-dashboard/src/components/publicModelHub/usePublicModelHubFacets.ts b/ui/litellm-dashboard/src/components/publicModelHub/usePublicModelHubFacets.ts new file mode 100644 index 00000000000..d6b356ac8f9 --- /dev/null +++ b/ui/litellm-dashboard/src/components/publicModelHub/usePublicModelHubFacets.ts @@ -0,0 +1,47 @@ +"use client"; + +import { useQueries } from "@tanstack/react-query"; + +import { apiClient } from "@/components/networking"; +import type { components } from "@/lib/http/schema"; + +import { PUBLIC_MODEL_HUB_PATH } from "./usePublicModelHubList"; + +type FacetResponse = components["schemas"]["FacetListResponse"]; + +export const MODEL_HUB_FACETS = ["providers", "modes", "features"] as const; + +export type ModelHubFacet = (typeof MODEL_HUB_FACETS)[number]; + +/** The route caps a page at 100, which is far above the distinct providers, modes or features any proxy publishes. */ +const FACET_PAGE_SIZE = 100; + +export interface PublicModelHubFacets { + providers: string[]; + modes: string[]; + features: string[]; +} + +const fetchFacet = (facet: ModelHubFacet, signal: AbortSignal): Promise => + apiClient.get(`${PUBLIC_MODEL_HUB_PATH}/${facet}`, { + query: { page_size: FACET_PAGE_SIZE }, + signal, + }); + +/** + * The values each filter dropdown offers, read from the route rather than derived from a + * page of rows, which can only ever show the values that page happens to contain. + */ +export const usePublicModelHubFacets = (enabled: boolean): PublicModelHubFacets => { + const results = useQueries({ + queries: MODEL_HUB_FACETS.map((facet) => ({ + queryKey: ["publicModelHub", "facet", facet], + queryFn: ({ signal }: { signal: AbortSignal }) => fetchFacet(facet, signal), + enabled, + staleTime: Infinity, + })), + }); + + const [providers, modes, features] = results.map((result) => result.data?.data ?? []); + return { providers, modes, features }; +}; diff --git a/ui/litellm-dashboard/src/components/publicModelHub/usePublicModelHubList.ts b/ui/litellm-dashboard/src/components/publicModelHub/usePublicModelHubList.ts new file mode 100644 index 00000000000..2be66a010d1 --- /dev/null +++ b/ui/litellm-dashboard/src/components/publicModelHub/usePublicModelHubList.ts @@ -0,0 +1,83 @@ +"use client"; + +import type { SortingState } from "@tanstack/react-table"; +import { useCallback } from "react"; + +import { + useResourceList, + type ResourceListPage, + type ResourceListQuery, + type ResourceListResult, +} from "@/app/(dashboard)/hooks/common/useResourceList"; +import { apiClient } from "@/components/networking"; +import type { ModelGroupInfo } from "@/components/PublicModelHubTableColumns"; + +import { + FEATURE_FILTER_ID, + MODE_FILTER_ID, + PROVIDER_FILTER_ID, + readFilterValues, + serializePublicModelHubFilters, + withFilterValue, +} from "./publicModelHubFilters"; + +export const PUBLIC_MODEL_HUB_PATH = "/public/v1/model_hub"; +export const PUBLIC_MODEL_HUB_PAGE_SIZE = 50; + +const QUERY_KEY = ["publicModelHub", "list"] as const; +const DEFAULT_SORTING: SortingState = [{ id: "model_group", desc: false }]; + +export interface PublicModelHubListResult extends ResourceListResult { + providerValues: string[]; + onProvidersChange: (values: string[]) => void; + modeValues: string[]; + onModesChange: (values: string[]) => void; + featureValues: string[]; + onFeaturesChange: (values: string[]) => void; + hasActiveQuery: boolean; +} + +const fetchPage = async (query: ResourceListQuery, signal: AbortSignal): Promise> => { + try { + return await apiClient.get>(PUBLIC_MODEL_HUB_PATH, { query, signal }); + } catch (error) { + if (!signal.aborted) { + console.error("There was an error fetching the public model data", error); + } + throw error; + } +}; + +export const usePublicModelHubList = (enabled: boolean): PublicModelHubListResult => { + const listOptions = { + queryKey: QUERY_KEY, + fetchPage, + serializeFilters: serializePublicModelHubFilters, + defaultSorting: DEFAULT_SORTING, + defaultPageSize: PUBLIC_MODEL_HUB_PAGE_SIZE, + enabled, + }; + const list = useResourceList(listOptions); + + const { onColumnFiltersChange } = list; + + const setFilter = useCallback( + (id: string, values: string[]) => onColumnFiltersChange((previous) => withFilterValue(previous, id, values)), + [onColumnFiltersChange], + ); + + const onProvidersChange = useCallback((values: string[]) => setFilter(PROVIDER_FILTER_ID, values), [setFilter]); + const onModesChange = useCallback((values: string[]) => setFilter(MODE_FILTER_ID, values), [setFilter]); + const onFeaturesChange = useCallback((values: string[]) => setFilter(FEATURE_FILTER_ID, values), [setFilter]); + + return { + ...list, + providerValues: readFilterValues(list.columnFilters, PROVIDER_FILTER_ID), + onProvidersChange, + modeValues: readFilterValues(list.columnFilters, MODE_FILTER_ID), + onModesChange, + featureValues: readFilterValues(list.columnFilters, FEATURE_FILTER_ID), + onFeaturesChange, + hasActiveQuery: list.searchValue.trim() !== "" || list.columnFilters.length > 0, + }; +}; diff --git a/ui/litellm-dashboard/src/components/public_model_hub.test.tsx b/ui/litellm-dashboard/src/components/public_model_hub.test.tsx index fec46e98077..cb23dfd9bb6 100644 --- a/ui/litellm-dashboard/src/components/public_model_hub.test.tsx +++ b/ui/litellm-dashboard/src/components/public_model_hub.test.tsx @@ -1,8 +1,12 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { describe, it, expect, vi, beforeAll, beforeEach } from "vitest"; import { render, screen, waitFor, within, fireEvent } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; import { flexRender, getCoreRowModel, useReactTable } from "@tanstack/react-table"; import PublicModelHub from "./public_model_hub"; -import { getPublicMCPHubColumns, MCPServerData } from "./PublicModelHubTableColumns"; +import { getPublicMCPHubColumns, MCPServerData, ModelGroupInfo } from "./PublicModelHubTableColumns"; + +const { apiGetMock } = vi.hoisted(() => ({ apiGetMock: vi.fn() })); vi.mock("next/navigation", () => ({ useRouter: vi.fn(() => ({ @@ -16,6 +20,7 @@ vi.mock("./networking", async (importOriginal) => { const actual = await importOriginal(); return { ...actual, + apiClient: { ...actual.apiClient, get: apiGetMock }, modelHubPublicModelsCall: vi.fn().mockResolvedValue([]), getPublicModelHubInfo: vi.fn().mockResolvedValue({ docs_title: "LiteLLM Gateway", @@ -34,6 +39,68 @@ vi.mock("./navbar", () => ({ default: vi.fn(() =>
Navbar Component
), })); +const MODEL_HUB_PATH = "/public/v1/model_hub"; + +const FACET_VALUES: Record = { + [`${MODEL_HUB_PATH}/providers`]: ["anthropic", "openai"], + [`${MODEL_HUB_PATH}/modes`]: ["chat", "embedding"], + [`${MODEL_HUB_PATH}/features`]: ["function_calling", "vision"], +}; + +const MODEL_DEFAULTS = { + providers: ["openai"], + mode: "chat", + supports_function_calling: false, + supports_vision: false, + supports_parallel_function_calling: false, +}; + +const model = (overrides: Partial & { model_group: string }): ModelGroupInfo => ({ + ...MODEL_DEFAULTS, + ...overrides, +}); + +const DEFAULT_MODELS = [model({ model_group: "gpt-4" }), model({ model_group: "claude-3", providers: ["anthropic"] })]; + +const respondWith = (rows: ModelGroupInfo[], totalCount: number = rows.length, pageSize: number = 50) => + apiGetMock.mockImplementation((path: string) => { + const facet = FACET_VALUES[path]; + if (facet) { + return Promise.resolve({ + data: facet, + meta: { page: 1, page_size: 100, has_more: false }, + links: { self: path, prev: null, next: null }, + }); + } + return Promise.resolve({ + data: rows, + meta: { + total_count: totalCount, + page: 1, + page_size: pageSize, + total_pages: Math.max(Math.ceil(totalCount / pageSize), 1), + }, + links: { self: MODEL_HUB_PATH, first: MODEL_HUB_PATH, prev: null, next: null, last: MODEL_HUB_PATH }, + }); + }); + +type QueryRecord = Record; + +const modelCalls = () => apiGetMock.mock.calls.filter((call) => call[0] === MODEL_HUB_PATH); +const facetPaths = (): string[] => + apiGetMock.mock.calls.map((call) => String(call[0])).filter((path) => path.startsWith(`${MODEL_HUB_PATH}/`)); +const modelQueries = (): QueryRecord[] => modelCalls().map((call) => (call[1] as { query: QueryRecord }).query); +const lastModelQuery = (): QueryRecord => modelQueries()[modelQueries().length - 1]; + +const renderHub = () => { + const client = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } }); + return render( + + + , + ); +}; + beforeAll(() => { Object.defineProperty(window, "matchMedia", { writable: true, @@ -51,6 +118,8 @@ beforeAll(() => { }); beforeEach(() => { + vi.clearAllMocks(); + respondWith(DEFAULT_MODELS); Storage.prototype.getItem = vi.fn(() => "false"); Storage.prototype.setItem = vi.fn(); Object.defineProperty(window, "location", { @@ -64,58 +133,215 @@ beforeEach(() => { describe("PublicModelHub", () => { it("renders", () => { - const { container } = render(); + const { container } = renderHub(); expect(container).toBeInTheDocument(); }); + it("loads the first page of models from the paginated public endpoint", async () => { + renderHub(); + + expect(await screen.findByText("gpt-4")).toBeInTheDocument(); + expect(modelCalls()[0][0]).toBe(MODEL_HUB_PATH); + expect(modelQueries()[0]).toEqual({ page: 1, page_size: 50, sort: "model_group" }); + }); + + it("waits for the resolved proxy base url before asking for a page", async () => { + const networkingModule = await import("./networking"); + let publishConfig: () => void = () => {}; + vi.mocked(networkingModule.getUiConfig).mockReturnValueOnce( + new Promise((resolve) => { + publishConfig = () => resolve({} as Awaited>); + }), + ); + + renderHub(); + await new Promise((resolve) => setTimeout(resolve, 50)); + expect(modelCalls()).toHaveLength(0); + + publishConfig(); + + await waitFor(() => expect(modelCalls().length).toBeGreaterThan(0)); + }); + + it("stops calling the unpaginated public model hub route", async () => { + const networkingModule = await import("./networking"); + renderHub(); + + await waitFor(() => expect(apiGetMock).toHaveBeenCalled()); + expect(networkingModule.modelHubPublicModelsCall).not.toHaveBeenCalled(); + }); + + it("counts the whole catalogue from the response meta, not the rows on screen", async () => { + respondWith(DEFAULT_MODELS, 300); + renderHub(); + + await screen.findByText("gpt-4"); + expect(screen.getByTestId("pagination-range")).toHaveTextContent("of 300"); + expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 1 of 6"); + }); + + it("asks the server for the next page", async () => { + const user = userEvent.setup(); + respondWith(DEFAULT_MODELS, 300); + renderHub(); + await screen.findByText("gpt-4"); + + await user.click(screen.getByTestId("pagination-next")); + + await waitFor(() => expect(lastModelQuery().page).toBe(2)); + expect(lastModelQuery().page_size).toBe(50); + }); + + it("asks the server for a different page size", async () => { + const user = userEvent.setup(); + respondWith(DEFAULT_MODELS, 300); + renderHub(); + await screen.findByText("gpt-4"); + + await user.click(screen.getByTestId("pagination-page-size")); + await user.click(await screen.findByRole("option", { name: "25" })); + + await waitFor(() => expect(lastModelQuery().page_size).toBe(25)); + }); + + it("asks the server to sort, in the sort form the endpoint accepts", async () => { + const user = userEvent.setup(); + renderHub(); + await screen.findByText("gpt-4"); + + await user.click(screen.getByTestId("sort-header-model_group")); + await waitFor(() => expect(lastModelQuery().sort).toBe("-model_group")); + + await user.click(screen.getByTestId("sort-header-input_cost_per_token")); + await waitFor(() => expect(lastModelQuery().sort).toBe("-input_cost_per_token")); + }); + + it("renders the page in the order the server sent it, without re-sorting locally", async () => { + const user = userEvent.setup(); + respondWith([model({ model_group: "alpha-model" }), model({ model_group: "zeta-model" })], 300); + renderHub(); + await screen.findByText("alpha-model"); + + await user.click(screen.getByTestId("sort-header-model_group")); + await waitFor(() => expect(lastModelQuery().sort).toBe("-model_group")); + + const rendered = screen.getAllByText(/-model$/).map((cell) => cell.textContent); + expect(rendered).toEqual(["alpha-model", "zeta-model"]); + }); + + it("offers sorting on exactly the fields the endpoint accepts", async () => { + renderHub(); + await screen.findByText("gpt-4"); + + const sortable = screen + .getAllByTestId(/^sort-header-/) + .map((header) => header.getAttribute("data-testid")?.replace("sort-header-", "")); + + expect(sortable.sort()).toEqual([ + "input_cost_per_token", + "max_input_tokens", + "max_output_tokens", + "mode", + "model_group", + "output_cost_per_token", + "providers", + "rpm", + ]); + expect(screen.getByText("Health Status")).toBeInTheDocument(); + expect(screen.queryByTestId("sort-header-health_status")).not.toBeInTheDocument(); + }); + + it("searches on the server and returns to the first page", async () => { + const user = userEvent.setup(); + respondWith(DEFAULT_MODELS, 300); + renderHub(); + await screen.findByText("gpt-4"); + + await user.click(screen.getByTestId("pagination-next")); + await waitFor(() => expect(lastModelQuery().page).toBe(2)); + + await user.type(screen.getByPlaceholderText("Search model names..."), "claude"); + + await waitFor(() => expect(lastModelQuery().q).toBe("claude")); + expect(lastModelQuery().page).toBe(1); + }); + + it("filters by mode with the endpoint's in operator", async () => { + const user = userEvent.setup(); + renderHub(); + await screen.findByText("gpt-4"); + + await user.click(screen.getByPlaceholderText("Select modes")); + await user.click(await screen.findByRole("option", { name: "embedding" })); + + await waitFor(() => expect(lastModelQuery()["filter[mode][in]"]).toBe("embedding")); + }); + + it("filters by several providers at once, and returns to the first page", async () => { + const user = userEvent.setup(); + respondWith(DEFAULT_MODELS, 300); + renderHub(); + await screen.findByText("gpt-4"); + + await user.click(screen.getByTestId("pagination-next")); + await waitFor(() => expect(lastModelQuery().page).toBe(2)); + + await user.click(screen.getByPlaceholderText("Select providers")); + await user.click(await screen.findByRole("option", { name: /anthropic/i })); + await waitFor(() => expect(lastModelQuery()["filter[providers][in]"]).toBe("anthropic")); + expect(lastModelQuery().page).toBe(1); + + await user.click(await screen.findByRole("option", { name: /openai/i })); + + await waitFor(() => expect(lastModelQuery()["filter[providers][in]"]).toBe("anthropic,openai")); + }); + + it("filters by feature, which the table could not do while it paged", async () => { + const user = userEvent.setup(); + renderHub(); + await screen.findByText("gpt-4"); + + await user.click(screen.getByPlaceholderText("Select features")); + await user.click(await screen.findByRole("option", { name: "Vision" })); + + await waitFor(() => expect(lastModelQuery()["filter[features][in]"]).toBe("vision")); + }); + + it("offers the filter values the route reports, not the ones on the page", async () => { + respondWith([model({ model_group: "gpt-4" })], 1); + renderHub(); + await screen.findByText("gpt-4"); + + await waitFor(() => expect(facetPaths()).toContain(`${MODEL_HUB_PATH}/providers`)); + expect(facetPaths()).toEqual(expect.arrayContaining([`${MODEL_HUB_PATH}/modes`, `${MODEL_HUB_PATH}/features`])); + }); + it("displays health status correctly for models with health check information", async () => { - const mockModelsWithHealthChecks = [ + respondWith([ { + ...MODEL_DEFAULTS, model_group: "gpt-4", - providers: ["openai"], - mode: "chat", health_status: "healthy", health_response_time: 150.5, health_checked_at: "2024-01-15T10:30:00Z", - supports_function_calling: true, - supports_vision: false, - supports_parallel_function_calling: false, }, { + ...MODEL_DEFAULTS, model_group: "claude-3", providers: ["anthropic"], - mode: "chat", health_status: "unhealthy", health_response_time: 5000.0, health_checked_at: "2024-01-15T10:25:00Z", - supports_function_calling: true, - supports_vision: false, - supports_parallel_function_calling: false, }, - { - model_group: "gpt-3.5-turbo", - providers: ["openai"], - mode: "chat", - health_status: undefined, - health_response_time: undefined, - health_checked_at: undefined, - supports_function_calling: false, - supports_vision: false, - supports_parallel_function_calling: false, - }, - ]; + model({ model_group: "gpt-3.5-turbo" }), + ]); - const networkingModule = await import("./networking"); - vi.mocked(networkingModule.modelHubPublicModelsCall).mockResolvedValue(mockModelsWithHealthChecks); + renderHub(); - render(); - - // Wait for the component to load and render the table await waitFor(() => { expect(screen.getByText("gpt-4")).toBeInTheDocument(); }); - // Check the health status badge in each model's row await waitFor(() => { const gpt4Row = screen.getByText("gpt-4").closest("tr"); expect(gpt4Row).toBeInTheDocument(); @@ -134,19 +360,13 @@ describe("PublicModelHub", () => { expect(within(gpt35Row as HTMLElement).getByText("Unknown")).toBeInTheDocument(); }); }); - it("shows no models when the search has no matches (LIT-5230 regression)", async () => { - const networkingModule = await import("./networking"); - vi.mocked(networkingModule.modelHubPublicModelsCall).mockResolvedValue([ - { model_group: "gpt-4", providers: ["openai"], mode: "chat" }, - { model_group: "claude-3", providers: ["anthropic"], mode: "chat" }, - ]); - render(); + it("shows no models when the search has no matches (LIT-5230 regression)", async () => { + renderHub(); expect(await screen.findByText("gpt-4")).toBeInTheDocument(); - fireEvent.change(screen.getByPlaceholderText("Search model names... (smart search enabled)"), { - target: { value: "zzzz" }, - }); + respondWith([], 0); + fireEvent.change(screen.getByPlaceholderText("Search model names..."), { target: { value: "zzzz" } }); await waitFor(() => { expect(screen.queryByText("gpt-4")).not.toBeInTheDocument(); @@ -155,18 +375,23 @@ describe("PublicModelHub", () => { }); }); - it("handles non-array response gracefully (regression test for e.filter crash)", async () => { - const networkingModule = await import("./networking"); - // Mock the API to return an object (like an error response) instead of an array - vi.mocked(networkingModule.modelHubPublicModelsCall).mockResolvedValue({ - detail: "No models configured", - } as any); + it("reports the proxy as unavailable when the model page fails to load", async () => { + apiGetMock.mockRejectedValue(new Error("boom")); - render(); + renderHub(); + + expect(await screen.findByText(/Service unavailable/)).toBeInTheDocument(); + }); + + it("keeps the page usable when the response carries no rows", async () => { + respondWith([], 0); + + renderHub(); await waitFor(() => { expect(screen.getByTestId("navbar")).toBeInTheDocument(); expect(screen.getByText("Model Hub")).toBeInTheDocument(); + expect(screen.getByText("No models available")).toBeInTheDocument(); }); }); }); @@ -237,7 +462,7 @@ describe("public hub MCP details modal", () => { const networkingModule = await import("./networking"); vi.mocked(networkingModule.mcpHubPublicServersCall).mockResolvedValue([mockMcpServer]); - render(); + renderHub(); fireEvent.click(await screen.findByRole("tab", { name: /MCP Hub/i })); fireEvent.click(await screen.findByRole("button", { name: "exa_test" })); @@ -252,7 +477,7 @@ describe("public hub MCP details modal", () => { const networkingModule = await import("./networking"); vi.mocked(networkingModule.mcpHubPublicServersCall).mockResolvedValue([mockMcpServer]); - render(); + renderHub(); fireEvent.click(await screen.findByRole("tab", { name: /MCP Hub/i })); fireEvent.click(await screen.findByRole("button", { name: "exa_test" })); diff --git a/ui/litellm-dashboard/src/components/public_model_hub.tsx b/ui/litellm-dashboard/src/components/public_model_hub.tsx index 546c3b4e018..29a84175d75 100644 --- a/ui/litellm-dashboard/src/components/public_model_hub.tsx +++ b/ui/litellm-dashboard/src/components/public_model_hub.tsx @@ -21,6 +21,9 @@ import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/u import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import { MultiSelect } from "./shared/MultiSelect"; +import { featureLabel } from "./publicModelHub/publicModelHubFilters"; +import { usePublicModelHubFacets } from "./publicModelHub/usePublicModelHubFacets"; +import { usePublicModelHubList } from "./publicModelHub/usePublicModelHubList"; import { DataTable } from "./shared/DataTable"; import { toast } from "@/lib/toast"; import Navbar from "./navbar"; @@ -31,7 +34,6 @@ import { getPublicModelHubInfo, getUiConfig, mcpHubPublicServersCall, - modelHubPublicModelsCall, } from "./networking"; import { Plugin } from "./claude_code_plugins/types"; import SkillHubDashboard from "./AIHub/SkillHubDashboard"; @@ -68,25 +70,19 @@ function PublicHubEmptyState({ title, body }: { title: string; body: string }) { const PublicModelHub: React.FC = ({ accessToken, isEmbedded = false }) => { const anchor = useComboboxAnchor(); - const [modelHubData, setModelHubData] = useState(null); + const [proxyConfigured, setProxyConfigured] = useState(false); const [agentHubData, setAgentHubData] = useState(null); const [mcpHubData, setMcpHubData] = useState(null); const [pageTitle, setPageTitle] = useState("LiteLLM Gateway"); const [customDocsDescription, setCustomDocsDescription] = useState(null); const [litellmVersion, setLitellmVersion] = useState(""); const [usefulLinks, setUsefulLinks] = useState>({}); - const [loading, setLoading] = useState(true); const [agentLoading, setAgentLoading] = useState(true); const [mcpLoading, setMcpLoading] = useState(true); - const [searchTerm, setSearchTerm] = useState(""); const [agentSearchTerm, setAgentSearchTerm] = useState(""); const [mcpSearchTerm, setMcpSearchTerm] = useState(""); - const [selectedProviders, setSelectedProviders] = useState([]); - const [selectedModes, setSelectedModes] = useState([]); - const [selectedFeatures, setSelectedFeatures] = useState([]); const [selectedAgentSkills, setSelectedAgentSkills] = useState([]); const [selectedMcpTransports, setSelectedMcpTransports] = useState([]); - const [serviceStatus, setServiceStatus] = useState("I'm alive! ✓"); const [isModalVisible, setIsModalVisible] = useState(false); const [isAgentModalVisible, setIsAgentModalVisible] = useState(false); const [isMcpModalVisible, setIsMcpModalVisible] = useState(false); @@ -106,19 +102,7 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded console.error("Failed to get UI config:", error); // Continue anyway - might work with default proxyBaseUrl } - - const fetchPublicData = async () => { - try { - setLoading(true); - const _modelHubData = await modelHubPublicModelsCall(); - setModelHubData(Array.isArray(_modelHubData) ? _modelHubData : []); - } catch (error) { - console.error("There was an error fetching the public model data", error); - setServiceStatus("Service unavailable"); - } finally { - setLoading(false); - } - }; + setProxyConfigured(true); const fetchAgentData = async () => { try { @@ -166,7 +150,6 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded fetchPublicModelHubInfo(); - fetchPublicData(); fetchAgentData(); fetchMcpData(); fetchSkillData(); @@ -175,47 +158,6 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded initializeAndFetch(); }, []); - // Clear filters when filter values change to avoid confusion - useEffect(() => { - // This would clear selections if we had any selection functionality - // For now, it's just for consistency with the original component - }, [searchTerm, selectedProviders, selectedModes, selectedFeatures]); - - const getUniqueProviders = (data: ModelGroupInfo[]) => { - const providers = new Set(); - data.forEach((model) => { - (model.providers ?? []).forEach((provider) => providers.add(provider)); - }); - return Array.from(providers); - }; - - const getUniqueModes = (data: ModelGroupInfo[]) => { - const modes = new Set(); - data.forEach((model) => { - if (model.mode) modes.add(model.mode); - }); - return Array.from(modes); - }; - - const getUniqueFeatures = (data: ModelGroupInfo[]) => { - const features = new Set(); - data.forEach((model) => { - // Find all properties that start with 'supports_' and are true - Object.entries(model) - .filter(([key, value]) => key.startsWith("supports_") && value === true) - .forEach(([key]) => { - // Format the feature name (remove 'supports_' prefix and convert to title case) - const featureName = key - .replace(/^supports_/, "") - .split("_") - .map((word) => word.charAt(0).toUpperCase() + word.slice(1)) - .join(" "); - features.add(featureName); - }); - }); - return Array.from(features).sort(); - }; - const getUniqueAgentSkills = (data: AgentCard[]) => { const skills = new Set(); data.forEach((agent) => { @@ -234,39 +176,6 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded return Array.from(transports).sort(); }; - const filteredData = useMemo(() => { - if (!modelHubData || !Array.isArray(modelHubData)) return []; - - const searchResults = rankBySearchRelevance( - filterBySearchTerm(modelHubData, searchTerm, (model) => [model.model_group]), - searchTerm, - (model) => model.model_group, - ); - - // Apply other filters - return searchResults.filter((model) => { - const matchesProvider = - selectedProviders.length === 0 || selectedProviders.some((provider) => model.providers.includes(provider)); - const matchesMode = selectedModes.length === 0 || selectedModes.includes(model.mode || ""); - - // Check if model has any of the selected features - const matchesFeature = - selectedFeatures.length === 0 || - Object.entries(model) - .filter(([key, value]) => key.startsWith("supports_") && value === true) - .some(([key]) => { - const featureName = key - .replace(/^supports_/, "") - .split("_") - .map((word) => word.charAt(0).toUpperCase() + word.slice(1)) - .join(" "); - return selectedFeatures.includes(featureName); - }); - - return matchesProvider && matchesMode && matchesFeature; - }); - }, [modelHubData, searchTerm, selectedProviders, selectedModes, selectedFeatures]); - const filteredAgentData = useMemo(() => { if (!agentHubData || !Array.isArray(agentHubData)) return []; @@ -356,7 +265,14 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded return `$${(cost * 1_000_000).toFixed(4)}`; }; - const [modelSorting, setModelSorting] = useState([{ id: "model_group", desc: false }]); + const models = usePublicModelHubList(proxyConfigured); + const modelFacets = usePublicModelHubFacets(proxyConfigured); + const modeOptions = useMemo(() => modelFacets.modes.map((mode) => ({ label: mode, value: mode })), [modelFacets]); + const featureOptions = useMemo( + () => modelFacets.features.map((feature) => ({ label: featureLabel(feature), value: feature })), + [modelFacets], + ); + const serviceStatus = models.error ? "Service unavailable" : "I'm alive! ✓"; const [agentSorting, setAgentSorting] = useState([{ id: "name", desc: false }]); const [mcpSorting, setMcpSorting] = useState([{ id: "server_name", desc: false }]); @@ -367,22 +283,6 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded const hasAgents = Array.isArray(agentHubData) && agentHubData.length > 0; const hasMcpServers = Array.isArray(mcpHubData) && mcpHubData.length > 0; - const providerOptions = useMemo( - () => (Array.isArray(modelHubData) ? getUniqueProviders(modelHubData) : []), - [modelHubData], - ); - const modeOptions = useMemo( - () => - Array.isArray(modelHubData) ? getUniqueModes(modelHubData).map((mode) => ({ label: mode, value: mode })) : [], - [modelHubData], - ); - const featureOptions = useMemo( - () => - Array.isArray(modelHubData) - ? getUniqueFeatures(modelHubData).map((feature) => ({ label: feature, value: feature })) - : [], - [modelHubData], - ); const agentSkillOptions = useMemo( () => Array.isArray(agentHubData) @@ -495,9 +395,8 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded } /> - Smart search with relevance ranking - finds models containing your search terms, ranked by - relevance. Try searching 'xai grok-4', 'claude-4', 'gpt-4', or - 'sonnet' + Finds every published model whose name contains what you type, across all pages. Try + 'grok', 'claude', 'gpt-4', or 'sonnet'
@@ -505,9 +404,10 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded setSearchTerm(e.target.value)} + placeholder="Search model names..." + aria-label="Search model names" + value={models.searchValue} + onChange={(e) => models.onSearchChange(e.target.value)} className="border border-border rounded-lg pl-10 pr-4 py-2 w-full text-sm focus:outline-hidden focus:ring-2 focus:ring-ring focus:border-transparent bg-card" />
@@ -516,9 +416,9 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded

Provider:

setSelectedProviders(values)} + items={modelFacets.providers} + value={models.providerValues} + onValueChange={models.onProvidersChange} > } className="min-h-8 w-full py-1 text-sm"> @@ -567,8 +467,8 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded

Mode:

@@ -577,8 +477,8 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded

Features:

@@ -586,20 +486,23 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded
model.model_group || String(index)} - sortingMode="client" - sorting={modelSorting} - onSortingChange={setModelSorting} - isLoading={loading} + sortingMode="server" + sorting={models.sorting} + onSortingChange={models.onSortingChange} + paginationMode="server" + pagination={models.pagination} + onPaginationChange={models.onPaginationChange} + rowCount={models.rowCount} + isLoading={models.isLoading} loadingMessage="Loading models…" noDataMessage={ = ({ accessToken, isEmbedded } size="compact" /> - -
-

- Showing {filteredData.length} of {modelHubData?.length || 0} models -

-
{/* Agents Tab */} diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx index f502ea77127..149554a3ac3 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx @@ -613,8 +613,11 @@ describe("DataTable layout", () => { it("makes the header sticky and constrains body height when maxBodyHeight is set", () => { render(); - expect(screen.getByTestId("data-table-head")).toHaveClass("sticky"); - expect(screen.getByTestId("data-table-scroller")).toHaveStyle({ maxHeight: "240px" }); + const scroller = screen.getByTestId("data-table-scroller"); + expect(scroller).toHaveStyle({ maxHeight: "240px" }); + expect(scroller).toHaveClass("overflow-auto"); + expect(scroller).toHaveClass("[&_[data-slot=table-container]]:overflow-visible"); + expect(screen.getByTestId("data-table-head")).toHaveClass("sticky", "bg-background"); }); it("caps fillHeight at the parent's height instead of stretching to it, so a short table stays short", () => { diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx index 60267606951..17a5fe42d1c 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx @@ -59,18 +59,22 @@ const noop = () => {}; /** * Height-filling mode. The table still sizes to its rows; the parent's height is only a ceiling, so * a short table keeps its footer under the last row and a long one scrolls its rows instead of the - * page. `table-container` is the Table primitive's own overflow-x wrapper; left as a scroll box it - * captures the sticky header and the header scrolls away with the rows. And rows pass under that - * header, which the semi-transparent header row tint alone would not hide. + * page. */ const FILL_CLASSES = { outer: "flex max-h-full min-h-0 flex-col", frame: "flex min-h-0 flex-col", - body: "min-h-0 [&_[data-slot=table-container]]:overflow-visible", + body: "min-h-0", +} as const; + +const NO_FILL_CLASSES = { outer: "", frame: "", body: "" } as const; + +const STICKY_CLASSES = { + body: "[&_[data-slot=table-container]]:overflow-visible", header: "bg-background", } as const; -const NO_FILL_CLASSES = { outer: "", frame: "", body: "", header: "" } as const; +const NO_STICKY_CLASSES = { body: "", header: "" } as const; function columnDefId(column: ColumnDef): string | undefined { if ("id" in column && typeof column.id === "string") { @@ -533,6 +537,7 @@ export function DataTable(props: DataTableProps { @@ -593,13 +598,13 @@ export function DataTable(props: DataTableProps{toolbar(table)}
}
{table.getHeaderGroups().map((headerGroup) => ( diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx index aa9df4a0319..b0380255b95 100644 --- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx @@ -443,7 +443,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi }, []); return ( -
+
{selectedKey ? ( ) : ( -
+
( <> diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index 614b3b7af62..dcadc141103 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx @@ -216,24 +216,22 @@ const UserDashboard: React.FC = ({ const canCreateKey = userRole !== "Admin Viewer" && userRole !== "proxy_admin_viewer"; return ( -
-
- - ) : undefined - } - /> -
+
+ + ) : undefined + } + />
); }; diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx index e1ef3bb778d..c3d204e1ae3 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx @@ -85,6 +85,7 @@ export function RequestLogsTable({ data={data} columns={columns} getRowId={(row) => row.request_id} + fillHeight sortingMode="server" sorting={sorting} onSortingChange={onSortingChange} diff --git a/ui/litellm-dashboard/src/components/view_logs/index.tsx b/ui/litellm-dashboard/src/components/view_logs/index.tsx index f20285dffa0..aadf90fad6c 100644 --- a/ui/litellm-dashboard/src/components/view_logs/index.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/index.tsx @@ -27,6 +27,9 @@ const AUDIT_LOGS_TAB: LogsTab = { id: "audit logs", label: "Audit Logs" }; const DELETED_KEYS_TAB: LogsTab = { id: "deleted keys", label: "Deleted Keys" }; const DELETED_TEAMS_TAB: LogsTab = { id: "deleted teams", label: "Deleted Teams" }; +const tabContentClassName = (tabId: LogsTabId): string => + tabId === REQUEST_LOGS_TAB.id ? "flex min-h-0 flex-1 flex-col" : "min-h-0 flex-1 overflow-y-auto"; + export default function SpendLogsTable({ accessToken, token, userRole, userID, premiumUser }: SpendLogsTableProps) { const [activeTab, setActiveTab] = useState(REQUEST_LOGS_TAB.id); const canViewAuditLogs = useCan("viewAuditLogs"); @@ -78,8 +81,8 @@ export default function SpendLogsTable({ accessToken, token, userRole, userID, p }; return ( -
- setActiveTab(value as LogsTabId)}> +
+ setActiveTab(value as LogsTabId)} className="min-h-0 flex-1"> {tabs.map((tab) => ( @@ -88,7 +91,7 @@ export default function SpendLogsTable({ accessToken, token, userRole, userID, p ))} {tabs.map((tab) => ( - + {renderPanel(tab.id)} ))} diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts index 3dd9911c794..ffede7b2a6b 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts @@ -83,9 +83,19 @@ describe("autorouter_presets", () => { expect(config.tier_boundaries).toBeUndefined(); expect(config.token_thresholds).toBeUndefined(); expect(config.dimension_weights).toBeUndefined(); + expect(config.session_affinity_ttl_seconds).toBeUndefined(); } }); + it("carries a preset's session affinity idle window into the prefilled form state", () => { + const config = getPresetByKey("anthropic_family")!.complexity_router_config; + const prefill = buildPresetPrefill({ ...config, session_affinity_ttl_seconds: 300 }, groupsOnly([])); + expect(prefill.complexityRouterConfig.session_affinity_ttl_seconds).toBe(300); + expect( + buildPresetPrefill(config, groupsOnly([])).complexityRouterConfig.session_affinity_ttl_seconds, + ).toBeUndefined(); + }); + it("keeps the model-family presets on the heuristic classifier", () => { for (const key of ["anthropic_family", "gemini_family", "openai_family"]) { expect(getPresetByKey(key)!.complexity_router_config.classifier_type).toBe("heuristic"); diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.ts index ff482bd23b8..2aafbfcfcd9 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.ts @@ -152,7 +152,7 @@ export const deploymentRefsFromModelInfo = ( return row.model_name && underlyingModels.length > 0 ? [{ modelGroup: row.model_name, underlyingModels }] : []; }); -const resolveAvailableModel = (requiredModel: string, availability: ModelAvailability): string | undefined => { +export const resolveAvailableModel = (requiredModel: string, availability: ModelAvailability): string | undefined => { const { modelGroups, underlyingIndex } = availability; if (modelGroups.has(requiredModel)) return requiredModel; const normalized = normalizeModelName(requiredModel); @@ -284,6 +284,7 @@ export const buildPresetPrefill = ( classifier_context_include_assistant_turns: config.classifier_context_include_assistant_turns, classification_mode: config.classification_mode ?? DEFAULT_CLASSIFICATION_MODE, session_affinity: config.session_affinity ?? DEFAULT_SESSION_AFFINITY, + session_affinity_ttl_seconds: config.session_affinity_ttl_seconds, deployment_affinity: config.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY, modality_routing: config.modality_routing ?? false, modality_pin_override: config.modality_pin_override ?? false, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index c3d3439f83a..6e1612b8d34 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -12382,6 +12382,36 @@ export interface paths { patch?: never; trace?: never; }; + "/public/v1/model_hub/{facet}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Public Model Hub Facet + * @description The distinct providers, modes or features across the published model groups, for the + * Model Hub's filter dropdowns. No authentication. + * + * Carries the same filters and search as the list route, so a dropdown offers exactly + * the values the table can show: asking for providers under `filter[mode][in]=chat` + * lists only the providers that serve a chat model. + * + * Example curl: + * ``` + * curl --location --globoff 'http://0.0.0.0:4000/public/v1/model_hub/providers?filter[mode][in]=chat&page_size=50' + * ``` + */ + get: operations["public_model_hub_facet_public_v1_model_hub__facet__get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/queue/chat/completions": { parameters: { query?: never; @@ -25503,6 +25533,12 @@ export interface components { * @description Documents all the fields supported by `general_settings` in config.yaml */ ConfigGeneralSettings: { + /** + * Admission Queue Timeout Seconds + * @description maximum time a request waits for a worker slot + * @default 1 + */ + admission_queue_timeout_seconds: number; /** * Alert To Webhook Url * @description Mapping of alert type to webhook url. e.g. `alert_to_webhook_url: {'budget_alerts': 'https://nothooks.slack.com/services/T00000000/B00000000/XXXXXXXXXXXXXXXXXXXXXXXX'}` @@ -25711,11 +25747,21 @@ export interface components { * @description max file size in MB for /v1/files uploads, for any purpose, if a file is larger than this size it will be rejected before being forwarded to the provider */ max_file_size_mb?: number | null; + /** + * Max In Flight Requests Per Worker + * @description maximum concurrent requests handled by each worker + */ + max_in_flight_requests_per_worker?: number | null; /** * Max Parallel Requests * @description maximum parallel requests for each api key */ max_parallel_requests?: number | null; + /** + * Max Queued Requests Per Worker + * @description maximum requests waiting for a worker slot + */ + max_queued_requests_per_worker?: number | null; /** * Max Request Size Mb * @description max request size in MB, if a request is larger than this size it will be rejected @@ -55026,6 +55072,37 @@ export interface operations { }; }; }; + public_model_hub_facet_public_v1_model_hub__facet__get: { + parameters: { + query?: never; + header?: never; + path: { + facet: "providers" | "modes" | "features"; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["FacetListResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; async_queue_request_queue_chat_completions_post: { parameters: { query?: {