Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_mistral_voxtral_tts_speech

# Conflicts:
#	tests/test_litellm/test_router.py
This commit is contained in:
mateo-berri 2026-09-03 11:00:31 -07:00
commit 9f806be6e8
356 changed files with 69542 additions and 1840 deletions

View file

@ -4,6 +4,11 @@ description: >-
by a job nor listed here, so every entry below is a decision on the record.
test_paths:
- reason: >-
The Rust/Python parity harness is run manually through its local CLI. Recorded replay,
fixture generation, and harness checks are intentionally outside pull request CI
paths:
- tests/rust-python-harness
- reason: >-
What is left of the caching suite in tests/local_testing that runs nowhere. Every job that
globs that directory either deselects it (local_testing_part1 and part2 carry `-k "... and

View file

@ -6,6 +6,7 @@ import shutil
import subprocess
import tempfile
import time
from dataclasses import dataclass, replace
from pathlib import Path
from typing import Optional
@ -45,6 +46,38 @@ _MIGRATION_TS_RE = re.compile(r"^(\d{14})_")
_MIGRATION_DEADLOCK_MARKER = "deadlock detected"
MAX_MIGRATE_DEPLOY_ATTEMPTS = 4
@dataclass(frozen=True)
class _MigrateAttemptBudget:
"""Retries left, and the recoveries already run.
A recovery that lands something new costs nothing, so a database full of
objects `prisma db push` created works through them one per pass. Anything
that made no progress spends an attempt, so a stuck run still gives up.
"""
attempts_left: int
recoveries: frozenset[str] = frozenset()
@property
def exhausted(self) -> bool:
return self.attempts_left <= 0
@property
def attempt_number(self) -> int:
return MAX_MIGRATE_DEPLOY_ATTEMPTS - self.attempts_left + 1
def spend(self) -> "_MigrateAttemptBudget":
return replace(self, attempts_left=self.attempts_left - 1)
def after_recovery(self, recovery: str) -> "_MigrateAttemptBudget":
if recovery in self.recoveries:
return self.spend()
return replace(self, recoveries=self.recoveries | {recovery})
_SPEND_LOGS_ALTER_RE = re.compile(r'^ALTER\s+TABLE\s+"LiteLLM_SpendLogs"\s', re.IGNORECASE)
_SPEND_LOGS_ARTIFACT_DROP_RE = re.compile(
r'^DROP\s+TABLE\s+"LiteLLM_SpendLogs_[^"]*"', re.IGNORECASE
@ -716,6 +749,9 @@ class ProxyExtrasDBManager:
Ahead-of-HEAD state (DB has migrations newer than this build ships)
is logged as a warning, not a fatal error users whose DBs got into
weird shapes from the old thrashing should still be able to start.
The retry budget only counts attempts that made no progress: see
_MigrateAttemptBudget.
"""
schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma"
migrations_dir = ProxyExtrasDBManager._get_prisma_dir()
@ -749,8 +785,9 @@ class ProxyExtrasDBManager:
original_dir = os.getcwd()
os.chdir(migrations_dir)
deploy_timeout = prisma_migrate_deploy_timeout()
budget = _MigrateAttemptBudget(attempts_left=MAX_MIGRATE_DEPLOY_ATTEMPTS)
try:
for attempt in range(4):
while not budget.exhausted:
try:
result = subprocess.run(
[_get_prisma_command(), "migrate", "deploy"],
@ -767,168 +804,155 @@ class ProxyExtrasDBManager:
logger.warning(
"prisma migrate deploy attempt %s timed out after %ss, retrying. "
"Raise %s if this database needs longer to apply its pending migrations.",
attempt + 1,
budget.attempt_number,
deploy_timeout,
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
)
time.sleep(random.randrange(5, 15))
continue
next_budget = budget.spend()
except subprocess.CalledProcessError as e:
stderr = e.stderr or ""
next_budget = ProxyExtrasDBManager._budget_after_deploy_failure(
e, budget, schema_path
)
if "P3005" in stderr and "database schema is not empty" in stderr:
logger.info(
"Schema exists but no migrations ledger — creating baseline"
)
ProxyExtrasDBManager._create_baseline_migration(schema_path)
continue
if "P3009" in stderr:
migration_match = re.search(r"`(\d+_\S+?)`", stderr)
if (
migration_match
and ProxyExtrasDBManager._is_idempotent_error(stderr)
):
name = migration_match.group(1)
logger.info(
f"Migration {name} failed idempotently — marking applied and retrying"
)
try:
ProxyExtrasDBManager._roll_back_migration(name)
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
):
pass # may already be rolled-back
try:
ProxyExtrasDBManager._resolve_specific_migration(name)
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
) as resolve_err:
# We're already inside the outer
# `except CalledProcessError` handler —
# re-raising CalledProcessError from here
# would escape as itself, bypassing
# proxy_cli.py's `except RuntimeError`.
raise RuntimeError(
f"Failed to mark migration {name} as applied "
f"after idempotent recovery. Manual "
f"intervention may be required.\n\n"
f"Detail: {resolve_err}"
) from resolve_err
continue
if migration_match:
migration_name = migration_match.group(1)
ledger_logs = ProxyExtrasDBManager._failed_migration_logs(migration_name)
if ledger_logs is not None and (
ledger_logs == "" or _MIGRATION_DEADLOCK_MARKER in ledger_logs
):
logger.info(
"Migration %s failed in a concurrent migrate deploy "
"deadlock race, rolling its ledger row back and retrying",
migration_name,
)
ProxyExtrasDBManager._roll_back_migration_best_effort(migration_name)
time.sleep(random.randrange(5, 15))
continue
raise RuntimeError(
"Database migration failed and cannot be auto-recovered. "
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
) from e
if "P3018" in stderr:
if ProxyExtrasDBManager._is_permission_error(stderr):
raise RuntimeError(
"Database migration failed due to insufficient "
"permissions. Please grant the required privileges "
f"and retry.\n\nPrisma error:\n{stderr}"
) from e
migration_match = re.search(
r"Migration name: (\d+_\S+)", stderr
)
if (
migration_match
and ProxyExtrasDBManager._is_idempotent_error(stderr)
):
name = migration_match.group(1)
logger.info(
f"Migration {name} SQL hit idempotent error — marking applied and retrying"
)
try:
ProxyExtrasDBManager._roll_back_migration(name)
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
):
pass # may already be rolled-back
try:
ProxyExtrasDBManager._resolve_specific_migration(name)
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
) as resolve_err:
raise RuntimeError(
f"Failed to mark migration {name} as applied "
f"after idempotent recovery. Manual "
f"intervention may be required.\n\n"
f"Detail: {resolve_err}"
) from resolve_err
continue
if migration_match and _MIGRATION_DEADLOCK_MARKER in stderr:
logger.info(
"Migration %s deadlocked against a concurrent "
"migrate deploy, rolling its ledger row back "
"and retrying",
migration_match.group(1),
)
ProxyExtrasDBManager._roll_back_migration_best_effort(
migration_match.group(1)
)
time.sleep(random.randrange(5, 15))
continue
raise RuntimeError(
"Database migration failed and cannot be auto-recovered. "
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
) from e
if _MIGRATION_DEADLOCK_MARKER in stderr:
logger.info(
"prisma migrate deploy attempt %s deadlocked against "
"a concurrent migrate deploy, retrying",
attempt + 1,
)
time.sleep(random.randrange(5, 15))
continue
if "P1002" in stderr and "advisory lock" in stderr:
logger.info(
"prisma migrate deploy attempt %s timed out waiting for "
"the advisory lock a concurrent migrate deploy holds, retrying",
attempt + 1,
)
time.sleep(random.randrange(5, 15))
continue
raise RuntimeError(
"Database migration failed and cannot be auto-recovered. "
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
) from e
if next_budget.attempts_left < budget.attempts_left:
time.sleep(random.randrange(5, 15))
budget = next_budget # rebind-ok: the loop carries the budget from one migrate deploy pass to the next
raise RuntimeError(
"Database migration failed after 4 attempts (retry loop "
"exhausted by timeouts, deadlock retries, or repeated "
"idempotent-recovery continues). Check database connectivity, "
f"Database migration failed after {MAX_MIGRATE_DEPLOY_ATTEMPTS} "
"attempts that made no progress (timeouts, deadlock retries, or a "
"recovery that had already run once). Check database connectivity, "
"load, and _prisma_migrations ledger state, and raise "
f"{PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out."
)
finally:
os.chdir(original_dir)
@staticmethod
def _budget_after_deploy_failure(
error: subprocess.CalledProcessError,
budget: "_MigrateAttemptBudget",
schema_path: str,
) -> "_MigrateAttemptBudget":
"""Recover from one failed `prisma migrate deploy`, and price the pass.
Returns the budget the next pass runs under, or raises when the failure
is not one this resolver knows how to recover from.
"""
stderr = error.stderr or ""
if "P3005" in stderr and "database schema is not empty" in stderr:
logger.info("Schema exists but no migrations ledger — creating baseline")
if ProxyExtrasDBManager._create_baseline_migration(schema_path):
return budget.after_recovery("baseline")
return budget.spend()
if "P3009" in stderr:
migration_match = re.search(r"`(\d+_\S+?)`", stderr)
if migration_match and ProxyExtrasDBManager._is_idempotent_error(stderr):
name = migration_match.group(1)
logger.info(
f"Migration {name} failed idempotently — marking applied and retrying"
)
ProxyExtrasDBManager._mark_migration_applied(name)
return budget.after_recovery(f"resolved:{name}")
if migration_match:
migration_name = migration_match.group(1)
ledger_logs = ProxyExtrasDBManager._failed_migration_logs(migration_name)
if ledger_logs is not None and (
ledger_logs == "" or _MIGRATION_DEADLOCK_MARKER in ledger_logs
):
logger.info(
"Migration %s failed in a concurrent migrate deploy "
"deadlock race, rolling its ledger row back and retrying",
migration_name,
)
ProxyExtrasDBManager._roll_back_migration_best_effort(migration_name)
return budget.spend()
raise RuntimeError(
"Database migration failed and cannot be auto-recovered. "
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
) from error
if "P3018" in stderr:
if ProxyExtrasDBManager._is_permission_error(stderr):
raise RuntimeError(
"Database migration failed due to insufficient "
"permissions. Please grant the required privileges "
f"and retry.\n\nPrisma error:\n{stderr}"
) from error
migration_match = re.search(r"Migration name: (\d+_\S+)", stderr)
if migration_match and ProxyExtrasDBManager._is_idempotent_error(stderr):
name = migration_match.group(1)
logger.info(
f"Migration {name} SQL hit idempotent error — marking applied and retrying"
)
ProxyExtrasDBManager._mark_migration_applied(name)
return budget.after_recovery(f"resolved:{name}")
if migration_match and _MIGRATION_DEADLOCK_MARKER in stderr:
logger.info(
"Migration %s deadlocked against a concurrent "
"migrate deploy, rolling its ledger row back "
"and retrying",
migration_match.group(1),
)
ProxyExtrasDBManager._roll_back_migration_best_effort(
migration_match.group(1)
)
return budget.spend()
raise RuntimeError(
"Database migration failed and cannot be auto-recovered. "
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
) from error
if _MIGRATION_DEADLOCK_MARKER in stderr:
logger.info(
"prisma migrate deploy attempt %s deadlocked against "
"a concurrent migrate deploy, retrying",
budget.attempt_number,
)
return budget.spend()
if "P1002" in stderr and "advisory lock" in stderr:
logger.info(
"prisma migrate deploy attempt %s timed out waiting for "
"the advisory lock a concurrent migrate deploy holds, retrying",
budget.attempt_number,
)
return budget.spend()
raise RuntimeError(
"Database migration failed and cannot be auto-recovered. "
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
) from error
@staticmethod
def _mark_migration_applied(name: str) -> None:
"""Roll a failed ledger row back if it is still there, then mark it applied."""
try:
ProxyExtrasDBManager._roll_back_migration(name)
except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
pass # may already be rolled-back
try:
ProxyExtrasDBManager._resolve_specific_migration(name)
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
) as resolve_err:
# We're called from inside an `except CalledProcessError` handler —
# re-raising CalledProcessError from here would escape as itself,
# bypassing proxy_cli.py's `except RuntimeError`.
raise RuntimeError(
f"Failed to mark migration {name} as applied "
f"after idempotent recovery. Manual "
f"intervention may be required.\n\n"
f"Detail: {resolve_err}"
) from resolve_err
@staticmethod
def apply_replica_identity_full_if_requested() -> bool:
"""

View file

@ -27,7 +27,7 @@ rand = "0.8"
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] }
rstest = "0.26.1"
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
serde_json = { version = "1.0", features = ["float_roundtrip"] }
sha2 = "0.10"
subtle = "2"
thiserror = "2.0"

View file

@ -14,7 +14,7 @@ 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"];
const AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS: &[&str] = &["pages", "features"];
pub struct AzureAiOcrConfig;
pub struct AzureDocumentIntelligenceOcrConfig;
@ -192,6 +192,46 @@ fn normalize_pages_param(pages: &Value) -> Result<Option<String>, Error> {
}
}
fn feature_token_is_valid(token: &str) -> bool {
let Some((first, rest)) = token.as_bytes().split_first() else {
return false;
};
first.is_ascii_alphabetic() && rest.iter().all(u8::is_ascii_alphanumeric)
}
fn invalid_features_error(features: &Value) -> Error {
Error::InvalidRequest(format!(
"Invalid `features` for Azure Document Intelligence: {features:?}. Expected a list of feature names or a comma-separated string like 'keyValuePairs' or 'keyValuePairs,languages'."
))
}
fn normalize_features_param(features: &Value) -> Result<Option<String>, Error> {
let normalized = match features {
Value::String(value) => value
.split(',')
.map(str::trim)
.collect::<Vec<_>>()
.join(","),
Value::Array(values) if values.is_empty() => return Ok(None),
Value::Array(values) => values
.iter()
.map(Value::as_str)
.collect::<Option<Vec<_>>>()
.ok_or_else(|| invalid_features_error(features))?
.into_iter()
.map(str::trim)
.collect::<Vec<_>>()
.join(","),
_ => return Err(invalid_features_error(features)),
};
if normalized.split(',').all(feature_token_is_valid) {
Ok(Some(normalized))
} else {
Err(invalid_features_error(features))
}
}
pub fn complete_document_intelligence_url(
api_base: Option<&str>,
model: &str,
@ -213,6 +253,13 @@ pub fn complete_document_intelligence_url(
url.push_str(&normalized);
}
if let Some(features) = optional_params.get("features")
&& let Some(normalized) = normalize_features_param(features)?
{
url.push_str("&features=");
url.push_str(&normalized);
}
Ok(url)
}
@ -475,6 +522,103 @@ mod tests {
);
}
#[test]
fn document_intelligence_url_normalizes_features() {
let params = serde_json::Map::from_iter([(
"features".to_string(),
json!("keyValuePairs, languages"),
)]);
let url = complete_document_intelligence_url(
Some("https://example.cognitiveservices.azure.com"),
"prebuilt-layout",
&params,
&|_| None,
)
.expect("url builds");
assert_eq!(
url,
"https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout:analyze?api-version=2024-11-30&features=keyValuePairs,languages"
);
}
#[test]
fn document_intelligence_url_combines_pages_and_feature_list() {
let params = serde_json::Map::from_iter([
("pages".to_string(), json!([0, 1, 2])),
(
"features".to_string(),
json!([" keyValuePairs ", "languages"]),
),
]);
let url = complete_document_intelligence_url(
Some("https://example.cognitiveservices.azure.com"),
"prebuilt-layout",
&params,
&|_| None,
)
.expect("url builds");
assert_eq!(
url,
"https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout:analyze?api-version=2024-11-30&pages=1,2,3&features=keyValuePairs,languages"
);
}
#[test]
fn document_intelligence_url_omits_empty_feature_list() {
let params = serde_json::Map::from_iter([("features".to_string(), json!([]))]);
let url = complete_document_intelligence_url(
Some("https://example.cognitiveservices.azure.com"),
"prebuilt-layout",
&params,
&|_| None,
)
.expect("url builds");
assert_eq!(
url,
"https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout:analyze?api-version=2024-11-30"
);
}
#[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",
&params,
&|_| None,
)
.expect_err("invalid features must fail");
assert!(
matches!(error, Error::InvalidRequest(message) if message.contains("Invalid `features`")),
"features={features:?}"
);
}
}
#[test]
fn document_intelligence_maps_features() {
let params = Map::from_iter([
("features".to_string(), json!(["keyValuePairs"])),
("unsupported".to_string(), json!(true)),
]);
assert_eq!(
AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG.map_ocr_params(&params),
Map::from_iter([("features".to_string(), json!(["keyValuePairs"]))])
);
}
#[test]
fn document_intelligence_request_uses_base64_source_for_data_uri() {
let body = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG

View file

@ -59,3 +59,41 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
module.add("RustBridgeDeclined", py.get_type::<RustBridgeDeclined>())?;
module.add("RustUpstreamError", py.get_type::<RustUpstreamError>())
}
pub(crate) fn ocr_error_to_pyerr(err: Error) -> PyErr {
match err {
Error::MissingField("document_url" | "image_url") => {
PyValueError::new_err("Document URL is required")
}
Error::Http { status, body } => RustUpstreamError::new_err((status, body)),
other => core_error_to_pyerr(other),
}
}
#[cfg(test)]
mod ocr_error_tests {
use super::*;
#[test]
fn ocr_errors_preserve_python_validation_and_provider_details() {
Python::initialize();
Python::attach(|py| {
for field in ["document_url", "image_url"] {
let mapped = ocr_error_to_pyerr(Error::MissingField(field));
assert!(mapped.is_instance_of::<PyValueError>(py));
assert_eq!(mapped.value(py).to_string(), "Document URL is required");
}
let mapped = ocr_error_to_pyerr(Error::Http {
status: 429,
body: r#"{"message":"rate limited"}"#.to_string(),
});
assert!(mapped.is_instance_of::<RustUpstreamError>(py));
let args: (u16, String) = mapped
.value(py)
.getattr("args")
.and_then(|args| args.extract())
.expect("OCR failures retain status and unprefixed provider message");
assert_eq!(args, (429, r#"{"message":"rate limited"}"#.to_string()));
});
}
}

View file

@ -5,7 +5,7 @@ use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr};
use pyo3::prelude::*;
use serde_json::Value;
use crate::errors::core_error_to_pyerr;
use crate::errors::ocr_error_to_pyerr;
use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty};
fn prepare_ocr(
@ -69,5 +69,5 @@ bridge_route! {
timeout_seconds: Option<f64>,
},
prepare = prepare_ocr,
errors = core_error_to_pyerr,
errors = ocr_error_to_pyerr,
}

View file

@ -1554,6 +1554,22 @@ def with_prompt_cache_breakpoint(target: _MarkedT, marker: object) -> _MarkedT:
return cast(_MarkedT, marked) # cast-ok: same block shape as the input plus the marker key
LITELLM_INTERNAL_MESSAGE_FIELDS: Final = frozenset({"thinking_blocks", "reasoning_content", "provider_specific_fields"})
def strip_litellm_internal_message_fields(message: AllMessageValues) -> AllMessageValues:
"""Drop the fields litellm attaches to assistant messages (e.g. when translating Anthropic thinking
blocks) that OpenAI-compatible endpoints with strict schemas reject as extra inputs."""
if LITELLM_INTERNAL_MESSAGE_FIELDS.isdisjoint(message):
return message
return cast( # cast-ok: same TypedDict minus internal keys
AllMessageValues,
{ # mutable-ok: provider transforms mutate message dicts in place downstream
key: value for key, value in message.items() if key not in LITELLM_INTERNAL_MESSAGE_FIELDS
},
)
def filter_value_from_dict(dictionary: dict, key: str, depth: int = 0) -> Any:
"""
Filters a value from a dictionary

View file

@ -48,7 +48,9 @@ _BASE_SUFFIXES_TO_STRIP: Final = (
)
# Per Bedrock Mantle Responses API validation errors.
_BEDROCK_MANTLE_SUPPORTED_RESPONSE_TOOL_TYPES = frozenset({"function", "mcp", "custom", "namespace", "tool_search"})
_BEDROCK_MANTLE_SUPPORTED_RESPONSE_TOOL_TYPES: Final = frozenset(
{"function", "mcp", "custom", "namespace", "tool_search", "web_search"}
)
_BEDROCK_MANTLE_SUPPORTED_SERVICE_TIERS: Final = frozenset({"auto", "default"})

View file

@ -3,7 +3,7 @@ Translates from OpenAI's `/v1/chat/completions` to Databricks' `/chat/completion
"""
import os
from collections.abc import AsyncIterator, Coroutine, Iterator
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping
from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload
import httpx
@ -15,6 +15,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo
_should_convert_tool_call_to_json_mode,
)
from litellm.litellm_core_utils.prompt_templates.common_utils import (
strip_litellm_internal_message_fields,
strip_name_from_message,
)
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
@ -55,6 +56,14 @@ from ...openai_like.chat.transformation import OpenAILikeChatConfig
from ..common_utils import DatabricksBase, DatabricksException
def _is_bare_assistant_message(message_dict: Mapping[str, object]) -> bool:
"""Databricks rejects assistant messages with neither content nor tool calls, e.g. a replayed
thinking-only turn once its `thinking_blocks` are stripped."""
return message_dict.get("role") == "assistant" and not any(
message_dict.get(key) for key in ("content", "tool_calls", "function_call")
)
def _sanitize_empty_content(message_dict: dict[str, Any]) -> None:
"""
Remove or filter content so empty text blocks are not sent.
@ -423,6 +432,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
"""
Databricks does not support:
- 'name' in user message.
- litellm's internal `thinking_blocks` / `reasoning_content` on assistant messages.
"""
new_messages = []
for idx, message in enumerate(messages):
@ -431,10 +441,13 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
else:
_message = message
_message = strip_name_from_message(_message, allowed_name_roles=["user"])
_message = strip_litellm_internal_message_fields(_message)
# Move message-level cache_control into a content block when content is a string.
if "cache_control" in _message and isinstance(_message.get("content"), str):
_message = self._move_cache_control_into_string_content_block(_message)
_sanitize_empty_content(cast(dict[str, Any], _message))
if _is_bare_assistant_message(_message):
continue
new_messages.append(_message)
if "claude" not in model:

View file

@ -52933,7 +52933,8 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"supports_web_search": true
},
"bedrock_mantle/openai.gpt-5.6-terra": {
"input_cost_per_token": 2.2e-06,
@ -52966,7 +52967,8 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"supports_web_search": true
},
"bedrock_mantle/openai.gpt-5.6-cyber": {
"input_cost_per_token": 1.375e-05,
@ -53027,7 +53029,8 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"supports_web_search": true
},
"us.openai.gpt-5.6-sol": {
"input_cost_per_token": 4.4e-06,
@ -53213,7 +53216,8 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"supports_web_search": true
},
"bedrock_mantle/openai.gpt-5.4": {
"input_cost_per_token": 2.75e-06,
@ -53243,7 +53247,8 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"supports_web_search": true
},
"bedrock_mantle/google.gemma-4-31b": {
"input_cost_per_token": 1.4e-07,

View file

@ -3,7 +3,7 @@ from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, cast
from typing import TYPE_CHECKING, Final, Literal, cast
from fastapi import HTTPException
from starlette.datastructures import Headers
@ -305,6 +305,12 @@ def _admission_failure_fallback(
raise exc
@dataclass(frozen=True, slots=True)
class MCPServerAccess:
server_ids: tuple[str, ...]
scope: Literal["unscoped", "scoped", "unresolved"] = "unscoped"
@dataclass(frozen=True, slots=True)
class DcrBridgeTarget:
"""The single DCR-bridge server a request targets, paired with the exact name the caller
@ -1456,6 +1462,18 @@ class MCPRequestHandler:
*,
keyless_source: bool = False,
) -> list[str]:
access: Final = await MCPRequestHandler.get_mcp_server_access(
user_api_key_auth,
keyless_source=keyless_source,
)
return list(access.server_ids)
@staticmethod
async def get_mcp_server_access(
user_api_key_auth: UserAPIKeyAuth | None = None,
*,
keyless_source: bool = False,
) -> MCPServerAccess:
"""
Get list of allowed MCP servers for the given user/key based on permissions.
@ -1478,13 +1496,17 @@ class MCPRequestHandler:
"""
from litellm.proxy.proxy_server import general_settings
key_object_permission: Final = MCPRequestHandler._get_key_object_permission(user_api_key_auth)
try:
# A keyless admitted subject resolves per source BEFORE any single-source rule here. Ordering
# matters: the no_mcp_servers opt-out below reads the caller's own object_permission, so above
# this branch a user's own opt-out would wrongly zero their TEAMS' grants too (each source is
# independent; an opt-out silences only its own source, inside the recursive call).
if _is_mcp_admitted_user_subject(user_api_key_auth) and user_api_key_auth is not None:
return await MCPRequestHandler._resolve_admitted_subject_servers(user_api_key_auth)
return MCPServerAccess(
server_ids=tuple(await MCPRequestHandler._resolve_admitted_subject_servers(user_api_key_auth)),
)
# Get allowed servers from key and team
allowed_mcp_servers_for_key = await MCPRequestHandler._get_allowed_mcp_servers_for_key(user_api_key_auth)
@ -1492,7 +1514,7 @@ class MCPRequestHandler:
# The key explicitly opted out of every MCP server. This overrides
# team inheritance and additive grants (mirrors no-default-models).
if SpecialMCPServerNames.no_mcp_servers.value in allowed_mcp_servers_for_key:
return []
return MCPServerAccess(server_ids=(), scope="scoped")
allowed_mcp_servers_for_team = await MCPRequestHandler._get_allowed_mcp_servers_for_team(user_api_key_auth)
@ -1572,7 +1594,7 @@ class MCPRequestHandler:
"require_end_user_mcp_access_defined=True and end_user %s has no MCP permissions - blocking MCP access",
user_api_key_auth.end_user_id,
)
return []
return MCPServerAccess(server_ids=(), scope="scoped")
#########################################################
# Check agent permissions if agent_id is set on the key
@ -1601,14 +1623,22 @@ class MCPRequestHandler:
#########################################################
# Apply org-level ceiling if org_id is set
#########################################################
allowed_mcp_servers = await MCPRequestHandler._apply_primary_org_ceiling(
allowed_mcp_servers, org_restricts = await MCPRequestHandler._apply_primary_org_ceiling(
allowed_mcp_servers,
user_api_key_auth,
has_lower_level_mcp_restrictions,
keyless_source=keyless_source,
)
return list(set(allowed_mcp_servers))
declares_key_mcp_scope: Final = getattr(key_object_permission, "mcp_servers", None) is not None
return MCPServerAccess(
server_ids=tuple(set(allowed_mcp_servers)),
scope=(
"scoped"
if has_lower_level_mcp_restrictions or org_restricts or declares_key_mcp_scope
else "unscoped"
),
)
except Exception as e:
if isinstance(e, UnloadableEntitlementError):
# A ceiling we KNOW exists and cannot read. Denying is the only answer that does not
@ -1616,7 +1646,10 @@ class MCPRequestHandler:
verbose_logger.warning("Denying MCP access, entitlement unreadable: %s", e)
else:
verbose_logger.warning("Failed to get allowed MCP servers: %s", e)
return []
return MCPServerAccess(
server_ids=(),
scope="scoped" if getattr(key_object_permission, "mcp_servers", None) is not None else "unresolved",
)
@staticmethod
async def _apply_primary_org_ceiling(
@ -1624,7 +1657,7 @@ class MCPRequestHandler:
user_api_key_auth: UserAPIKeyAuth | None,
has_lower_level_mcp_restrictions: bool,
keyless_source: bool = False,
) -> list[str]:
) -> tuple[list[str], bool]:
"""Cap the resolved server list by this caller's org ceiling: an explicit org list intersects
lower-level restrictions (else becomes the ceiling); no org or an empty list leaves it unchanged.
@ -1638,7 +1671,7 @@ class MCPRequestHandler:
cannot be read raises out of ``_get_allowed_mcp_servers_for_org`` and never arrives here as
``None``, so key auth cannot silently shed a ceiling an operator did configure."""
if not (user_api_key_auth and user_api_key_auth.org_id):
return allowed_mcp_servers
return allowed_mcp_servers, False
allowed_mcp_servers_for_org: Final = await MCPRequestHandler._get_allowed_mcp_servers_for_org(user_api_key_auth)
if allowed_mcp_servers_for_org is None:
verbose_logger.warning(
@ -1646,9 +1679,9 @@ class MCPRequestHandler:
user_api_key_auth.org_id,
"denying (keyless admitted subject)" if keyless_source else "leaving uncapped (key auth)",
)
return [] if keyless_source else allowed_mcp_servers
return ([] if keyless_source else allowed_mcp_servers), False
if len(allowed_mcp_servers_for_org) == 0:
return allowed_mcp_servers
return allowed_mcp_servers, False
if has_lower_level_mcp_restrictions or keyless_source:
# Org can only cap lower-level restrictions. A keyless admitted source ALWAYS takes this
# arm: its model unions GRANTS, so an org list may only narrow a source, never become one.
@ -1657,7 +1690,7 @@ class MCPRequestHandler:
# No lower-level restrictions → org list becomes the ceiling.
capped = allowed_mcp_servers_for_org
verbose_logger.debug("Applied org ceiling filter. Final allowed servers: %s", capped)
return capped
return capped, True
@staticmethod
def _scoped_source_auth(

View file

@ -55,6 +55,7 @@ from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
MCPServerAccess,
_is_mcp_admitted_user_subject,
)
from litellm.proxy._experimental.mcp_server.elicitation_handler import (
@ -2958,7 +2959,13 @@ class MCPServerManager:
return None
return user_api_key_auth.mcp_session_resource_server_id
async def get_allowed_mcp_servers(self, user_api_key_auth: UserAPIKeyAuth | None = None) -> list[str]:
async def get_allowed_mcp_servers(
self,
user_api_key_auth: UserAPIKeyAuth | None = None,
*,
access: MCPServerAccess | None = None,
general_settings: Mapping[str, object] | None = None,
) -> list[str]:
"""
Get the allowed MCP Servers for the user.
@ -2967,6 +2974,9 @@ class MCPServerManager:
2. If admin and no object_permission, return all servers
3. Otherwise, use standard permission checks
"""
from litellm.proxy.proxy_server import general_settings as proxy_general_settings
resolved_general_settings: Final = proxy_general_settings if general_settings is None else general_settings
allow_all_server_ids: Final = self.get_allow_all_keys_server_ids()
# A keyless admitted subject is resolved per grant source, and channel decisions that are
@ -3007,11 +3017,16 @@ class MCPServerManager:
# whole registry, for keys AND admitted session subjects alike (one predicate owns the
# question). Seeded into the union rather than returned early so the session resource
# scope below still bounds a per-server envelope held by an admin.
combined_servers: Final = (
set(self.get_registry().keys())
if await MCPRequestHandler.admin_view_unscoped(user_api_key_auth)
else set(await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth))
admin_unscoped: Final = await MCPRequestHandler.admin_view_unscoped(user_api_key_auth)
resolved_access: Final = (
MCPServerAccess(server_ids=())
if admin_unscoped
else access or await MCPRequestHandler.get_mcp_server_access(user_api_key_auth)
)
resolved_server_ids: Final = (
set(self.get_registry().keys()) if admin_unscoped else set(resolved_access.server_ids)
)
combined_servers: Final = set(resolved_server_ids)
verbose_logger.debug("Allowed MCP Servers for user api key auth: %s", combined_servers)
combined_servers.update(
await self.operator_open_server_ids(
@ -3052,6 +3067,18 @@ class MCPServerManager:
]
combined_servers.update(delegate_server_ids)
restrict_allow_all: Final = (
resolved_general_settings.get("mcp_allow_all_keys_respects_mcp_scope", False)
and user_api_key_auth is not None
and user_api_key_auth.via_virtual_key
and resolved_access.scope != "unscoped"
)
if restrict_allow_all:
combined_servers.difference_update(
set(allow_all_server_ids)
- resolved_server_ids
- (set(submitted_server_ids) if resolved_access.scope != "unresolved" else set())
)
if len(combined_servers) == 0:
verbose_logger.debug("No allowed MCP Servers found for user api key auth.")
scope = MCPServerManager._admitted_session_resource_scope(user_api_key_auth)

View file

@ -1216,9 +1216,9 @@ class GenerateKeyRequest(KeyRequestBase):
organization_id: str | None = None
project_id: str | None = None
@field_validator("team_id", mode="before")
@field_validator("team_id", "organization_id", mode="before")
@classmethod
def treat_cleared_team_id_as_unset(cls, v: object) -> object:
def treat_cleared_id_as_unset(cls, v: object) -> object:
if v == "":
return None
return v
@ -1278,6 +1278,13 @@ class UpdateKeyRequest(KeyRequestBase):
rotation_interval: str | None = None
organization_id: str | None = None
@field_validator("organization_id", mode="before")
@classmethod
def treat_cleared_organization_id_as_unset(cls, v: object) -> object:
if v == "":
return None
return v
@model_validator(mode="after")
def validate_temp_budget(self) -> "UpdateKeyRequest":
if self.temp_budget_increase is not None or self.temp_budget_expiry is not None:
@ -1923,6 +1930,13 @@ class NewTeamRequest(TeamBase):
model_config = ConfigDict(protected_namespaces=())
@field_validator("team_id", mode="before")
@classmethod
def treat_blank_team_id_as_unset(cls, v: object) -> object:
if isinstance(v, str) and not v.strip():
return None
return v
class GlobalEndUsersSpend(LiteLLMPydanticObjectBase):
api_key: str | None = None

View file

@ -939,35 +939,32 @@ async def make_agent_public(
if agent is None:
raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found")
if litellm.public_agent_groups is None:
litellm.public_agent_groups = []
# handle duplicates
if not AGENT_REGISTRY.ids_for_agent(agent.agent_id).isdisjoint(litellm.public_agent_groups):
config: Final = await proxy_config.get_config()
current_public_agent_groups: Final = list(litellm.public_agent_groups or [])
if not AGENT_REGISTRY.ids_for_agent(agent.agent_id).isdisjoint(current_public_agent_groups):
raise HTTPException(
status_code=400,
detail=f"Agent with name {agent.agent_name} already in public agent groups",
)
litellm.public_agent_groups.append(agent.agent_id)
updated_public_agent_groups: Final = [*current_public_agent_groups, agent.agent_id]
# Load existing config
config: Final = await proxy_config.get_config()
# Update config with new settings
if "litellm_settings" not in config or config["litellm_settings"] is None:
config["litellm_settings"] = {}
config["litellm_settings"]["public_agent_groups"] = litellm.public_agent_groups
config["litellm_settings"]["public_agent_groups"] = updated_public_agent_groups
# Save the updated config
await proxy_config.save_config(new_config=config)
litellm.public_agent_groups = updated_public_agent_groups
verbose_proxy_logger.debug(
"Updated public agent groups to: %s by user: %s", litellm.public_agent_groups, user_api_key_dict.user_id
"Updated public agent groups to: %s by user: %s", updated_public_agent_groups, user_api_key_dict.user_id
)
return {
"message": "Successfully updated public agent groups",
"public_agent_groups": litellm.public_agent_groups,
"public_agent_groups": updated_public_agent_groups,
"updated_by": user_api_key_dict.user_id,
}
except HTTPException:

View file

@ -489,7 +489,7 @@ lite codex exec "summarize the repo"
Each command resolves your LiteLLM key (logging in via SSO when none is stored and you are at a terminal; otherwise it expects `LITELLM_PROXY_API_KEY` or `--api-key`), checks the key against the proxy so bad credentials fail immediately instead of deep inside the agent, exports the environment variables the agent reads, then replaces itself with the agent process.
The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol).
The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. It also gets `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` (again unless you already set it) so Claude Code v2.1.129+ fills its `/model` picker from the proxy's `/v1/models`; Claude Code only lists entries whose id contains `claude` or `anthropic`, and older versions ignore the variable. Export it as `0` to turn discovery off. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol).
Options (these belong to the wrapper, so put them before the agent's own flags):
@ -505,7 +505,7 @@ The credential is short-lived by design (default 24h, configurable via `LITELLM_
### Route Every Claude Code Session Through the Proxy
`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` when that key is missing, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it.
`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when those keys are missing, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it.
Two things need to already be true: you've run `lite login` (or `lite login --pkce`, whose key the helper renews on its own), since the apiKeyHelper depends on that stored token, and the proxy is already reachable, since `lite up` does not start one for you.
@ -529,7 +529,7 @@ Cursor is not supported: it has no equivalent file-based config to hot-patch thi
lite --base-url https://your-proxy.example.com login --config-claude
```
It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABLE_TOOL_SEARCH`, and `apiKeyHelper`, but persistently: there is no backup, nothing to restore, and no foreground process to keep alive. Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag.
It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABLE_TOOL_SEARCH`, `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY`, and `apiKeyHelper`, but persistently: there is no backup, nothing to restore, and no foreground process to keep alive. Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag.
Because the credential is reached through `apiKeyHelper` rather than copied into the file, a later `lite login` refreshes it with no further action: Claude Code re-runs the helper on every request and picks up whatever token the most recent login stored. Nothing secret is written to `settings.json`.

View file

@ -16,6 +16,8 @@ ANTHROPIC_AUTH_TOKEN_ENV: Final = "ANTHROPIC_AUTH_TOKEN"
ANTHROPIC_API_KEY_ENV: Final = "ANTHROPIC_API_KEY"
ENABLE_TOOL_SEARCH_ENV: Final = "ENABLE_TOOL_SEARCH"
ENABLE_TOOL_SEARCH_VALUE: Final = "true"
ENABLE_GATEWAY_MODEL_DISCOVERY_ENV: Final = "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"
ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE: Final = "1"
OPENAI_BASE_URL_ENV: Final = "OPENAI_BASE_URL"
OPENAI_API_KEY_ENV: Final = "OPENAI_API_KEY"
@ -67,7 +69,9 @@ def build_agent_env(
Anthropic key cannot win over the bearer token we set. ENABLE_TOOL_SEARCH
defaults to true because Claude Code turns tool search off when
ANTHROPIC_BASE_URL is not a first-party Anthropic host; a value already in
the environment is left alone.
the environment is left alone. CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY
defaults to 1 so Claude Code (v2.1.129+) fills its /model picker from the
proxy's /v1/models; likewise left alone when already set.
"""
env: Final = dict(base_env)
root: Final = base_url.rstrip("/")
@ -77,6 +81,8 @@ def build_agent_env(
env.pop(ANTHROPIC_API_KEY_ENV, None)
if ENABLE_TOOL_SEARCH_ENV not in env:
env[ENABLE_TOOL_SEARCH_ENV] = ENABLE_TOOL_SEARCH_VALUE
if ENABLE_GATEWAY_MODEL_DISCOVERY_ENV not in env:
env[ENABLE_GATEWAY_MODEL_DISCOVERY_ENV] = ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE
if PROFILE_OPENAI in profiles:
env[OPENAI_BASE_URL_ENV] = root + "/v1"
env[OPENAI_API_KEY_ENV] = api_key

View file

@ -26,6 +26,8 @@ ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL"
ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY"
ENABLE_TOOL_SEARCH_KEY: Final = "ENABLE_TOOL_SEARCH"
ENABLE_TOOL_SEARCH_VALUE: Final = "true"
ENABLE_GATEWAY_MODEL_DISCOVERY_KEY: Final = "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"
ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE: Final = "1"
CLAUDE_SETTINGS_PATH: Final = Path.home() / ".claude" / "settings.json"
BACKUP_PATH: Final = Path.home() / ".litellm" / "claude_settings_backup.json"
@ -77,13 +79,16 @@ def merge_claude_settings(
stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued
token (same reasoning as build_agent_env in agents.py). ENABLE_TOOL_SEARCH
defaults to true because Claude Code turns tool search off when
ANTHROPIC_BASE_URL is not a first-party Anthropic host; an existing value is
left alone. Every other key is preserved untouched.
ANTHROPIC_BASE_URL is not a first-party Anthropic host, and
CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY defaults to 1 so the /model picker
is filled from the proxy's /v1/models; existing values of both are left
alone. Every other key is preserved untouched.
"""
raw_env: Final = settings.get(ENV_KEY, {})
base_env: Final = raw_env if isinstance(raw_env, dict) else {}
env: Final = {
ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE,
ENABLE_GATEWAY_MODEL_DISCOVERY_KEY: ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE,
**{key: value for key, value in base_env.items() if key != ANTHROPIC_API_KEY_KEY},
ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"),
}
@ -156,6 +161,8 @@ __all__ = (
"AUTOROUTE_BACKUP_PATH",
"BACKUP_PATH",
"CLAUDE_SETTINGS_PATH",
"ENABLE_GATEWAY_MODEL_DISCOVERY_KEY",
"ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE",
"ENABLE_TOOL_SEARCH_KEY",
"ENABLE_TOOL_SEARCH_VALUE",
"ENV_KEY",

View file

@ -135,7 +135,7 @@ async def get_credentials(
]
return {"success": True, "credentials": masked_credentials}
except Exception as e:
return handle_exception_on_proxy(e)
raise handle_exception_on_proxy(e)
@router.get(
@ -239,13 +239,18 @@ async def delete_credential(
status_code=500,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
await CredentialsRepository(prisma_client).delete_by_name(credential_name)
deleted: Final = await CredentialsRepository(prisma_client).delete_by_name(credential_name)
if deleted is None:
raise HTTPException(
status_code=404,
detail="Credential not found. Got credential name: " + credential_name,
)
## DELETE FROM LITELLM ##
litellm.credential_list = [cred for cred in litellm.credential_list if cred.credential_name != credential_name]
return {"success": True, "message": "Credential deleted successfully"}
except Exception as e:
return handle_exception_on_proxy(e)
raise handle_exception_on_proxy(e)
def update_db_credential(

View file

@ -124,6 +124,42 @@ def add_missing_query_params(url: str, params: Mapping[str, str | int | float])
return urllib.parse.urlunsplit(parsed._replace(query=query))
LIBPQ_VERIFY_SSLMODES: Final[frozenset[str]] = frozenset({"verify-ca", "verify-full"})
def translate_libpq_ssl_params(url: str) -> str:
"""Rewrite libpq's certificate-verification params into Prisma's dialect.
Prisma's engine only knows ``sslmode=disable|prefer|require``, ``sslcert``
(the CA bundle) and ``sslaccept=strict``. It silently discards
``sslrootcert`` and downgrades ``sslmode=verify-ca`` / ``verify-full`` to
``prefer``, so a URL copied from libpq / RDS docs connects over TLS with no
certificate check at all. ``verify-ca`` and ``verify-full`` both become
``require`` (Prisma has no CA-only mode), ``sslrootcert`` becomes
``sslcert``, and either one turns on ``sslaccept=strict`` (chain and
hostname), matching libpq where a root cert makes ``require`` verify.
Prisma params the operator pinned themselves win; anything else is left
untouched.
"""
parsed: Final = urllib.parse.urlsplit(url)
pairs: Final = tuple(urllib.parse.parse_qsl(parsed.query, keep_blank_values=True))
keys: Final = frozenset(key for key, _ in pairs)
wants_verify: Final = any(key == "sslmode" and value in LIBPQ_VERIFY_SSLMODES for key, value in pairs)
if not wants_verify and "sslrootcert" not in keys:
return url
translated: Final = tuple(
("sslmode", "require") if key == "sslmode" and value in LIBPQ_VERIFY_SSLMODES else (key, value)
for key, value in pairs
if key != "sslrootcert"
)
root_cert: Final = tuple(
("sslcert", value) for key, value in pairs if key == "sslrootcert" and "sslcert" not in keys
)
strict: Final = () if "sslaccept" in keys else (("sslaccept", "strict"),)
query: Final = urllib.parse.urlencode(translated + root_cert + strict)
return urllib.parse.urlunsplit(parsed._replace(query=query))
def reader_shareable_params(params: Mapping[str, str | int | float]) -> Mapping[str, str | int | float]:
"""Return the subset of ``params`` the read replica is allowed to inherit."""
return MappingProxyType({key: value for key, value in params.items() if key in CONNECTION_PARAM_KEYS})
@ -403,6 +439,11 @@ class DatabaseURLSettings(BaseSettings):
self._raise_for_unsupported_scheme()
wrote_writer: Final = self.apply_writer_url_to_env()
for env_var in ("DATABASE_URL", "DIRECT_URL"):
url = os.environ.get(env_var)
if url:
os.environ[env_var] = translate_libpq_ssl_params(url)
# DATABASE_DISABLE_PREPARED_STATEMENTS maps to Prisma's `pgbouncer=true`
# URL param, same as the CLI's `database_disable_prepared_statements`
# config key. An explicit `pgbouncer` value already on the URL wins.
@ -418,7 +459,7 @@ class DatabaseURLSettings(BaseSettings):
reader_url: Final = self.build_reader_url() or self.database_url_read_replica
if reader_url is not None:
os.environ["DATABASE_URL_READ_REPLICA"] = add_missing_query_params(
reader_url,
translate_libpq_ssl_params(reader_url),
connection_params_from_url(os.environ.get("DATABASE_URL", "")),
)

View file

@ -68,6 +68,10 @@ from litellm.proxy.management_endpoints.team_endpoints import (
from litellm.proxy.management_endpoints.team_endpoints import (
update_team as _legacy_update_team,
)
from litellm.proxy.management_helpers.access_group_model_sync import (
sync_access_groups_for_deleted_model,
sync_access_groups_for_renamed_model,
)
from litellm.proxy.management_helpers.audit_logs import create_object_audit_log
from litellm.proxy.spend_tracking.ptu_feature_flag import (
PTU_COST_ATTRIBUTION_ENV_VAR,
@ -715,6 +719,7 @@ async def patch_model(
existing_params=db_model.litellm_params,
)
requested_model_name: Final = patch_data.model_name
# Handle team model updates with proper alias management
update_data: Final = await _update_team_model_in_db(
db_model=db_model,
@ -741,6 +746,20 @@ async def patch_model(
param=None,
)
stored_model_name: Final = update_data.get("model_name")
if (
stored_model_name is not None
and stored_model_name == requested_model_name
and stored_model_name != db_model.model_name
):
await sync_access_groups_for_renamed_model(
prisma_client=prisma_client,
model_id=model_id,
old_name=db_model.model_name,
new_name=stored_model_name,
llm_router=llm_router,
)
# Clear cache and reload models (uses config setting or defaults to preserving config models for DB updates)
live_before_reload: Final = live_model_ids_snapshot()
reload_outcome: Final = await clear_cache()
@ -1673,6 +1692,12 @@ async def delete_model(
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
)
await sync_access_groups_for_deleted_model(
prisma_client=prisma_client,
model_id=model_info.id,
model_name=model_params.model_name,
llm_router=llm_router,
)
## CREATE AUDIT LOG ##
asyncio.create_task(
@ -2027,25 +2052,36 @@ async def update_model(
model_params.litellm_params[k] = encrypted_value
### MERGE WITH EXISTING DATA ###
merged_dictionary: Final = {}
_mp: Final[dict[str, object]] = model_params.litellm_params.dict()
merged_dictionary: Final = {
key: _existing_litellm_params_dict[key] if value is None else value
for key, value in _mp.items()
if value is not None or _existing_litellm_params_dict.get(key) is not None
}
for key, value in _mp.items():
if value is not None:
merged_dictionary[key] = value
elif key in _existing_litellm_params_dict and _existing_litellm_params_dict[key] is not None:
merged_dictionary[key] = _existing_litellm_params_dict[key]
else:
pass
renamed_to: Final = (
model_params.model_name
if model_params.model_name not in (None, deployment.model_name)
and deployment.model_info.team_id is None
else None
)
_data: Final[dict[str, str]] = {
"litellm_params": json.dumps(merged_dictionary),
"updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME,
**({} if renamed_to is None else {"model_name": renamed_to}),
}
model_response: Final = await _proxy_model_table(prisma_client).update(
where={"model_id": _model_id},
data=_data,
)
if renamed_to is not None:
await sync_access_groups_for_renamed_model(
prisma_client=prisma_client,
model_id=_model_id,
old_name=deployment.model_name,
new_name=renamed_to,
llm_router=llm_router,
)
# Clear cache and reload models (uses config setting or defaults to preserving config models for DB updates)
live_before_reload: Final = live_model_ids_snapshot()

View file

@ -0,0 +1,119 @@
"""
Keep `litellm_accessgrouptable.access_model_names` pointing at deployment names that still exist.
Unified access groups store model names, not ids, so a deployment rename or delete that leaves
the arrays alone strands every group on a name nothing serves any more.
"""
from collections.abc import Sequence
from typing import Final, Protocol
from pydantic import BaseModel
from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient
from litellm.proxy.management_helpers.access_group_team_sync import invalidate_access_group_caches
from litellm.repositories.table_repositories import AccessGroupRepository
from litellm.router import Router
class _TouchedGroupRow(BaseModel):
access_group_id: str
class _DeploymentCountRow(BaseModel):
deployment_count: int
class _RawExecutor(Protocol):
async def query_raw(self, query: str, *args: str) -> Sequence[object]: ...
_BACKING_DEPLOYMENTS_SQL: Final = (
'SELECT COUNT(*)::int AS deployment_count FROM "LiteLLM_ProxyModelTable" WHERE "model_name" = $1'
)
_REPLACE_MODEL_NAME_SQL: Final = (
'UPDATE "LiteLLM_AccessGroupTable" '
'SET "access_model_names" = array_replace(array_remove("access_model_names", $2), $1, $2) '
'WHERE $1 = ANY("access_model_names") '
'RETURNING "access_group_id"'
)
_APPEND_MODEL_NAME_SQL: Final = (
'UPDATE "LiteLLM_AccessGroupTable" '
'SET "access_model_names" = array_append("access_model_names", $2) '
'WHERE $1 = ANY("access_model_names") AND NOT ($2 = ANY("access_model_names")) '
'RETURNING "access_group_id"'
)
_REMOVE_MODEL_NAME_SQL: Final = (
'UPDATE "LiteLLM_AccessGroupTable" '
'SET "access_model_names" = array_remove("access_model_names", $1) '
'WHERE $1 = ANY("access_model_names") '
'RETURNING "access_group_id"'
)
def _raw_executor(prisma_client: object) -> _RawExecutor:
db: Final = AccessGroupRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client
return WriterPinnedClient(db).db # pyright: ignore[reportAny, reportReturnType] # untyped Prisma client behind the pin
def _config_sourced_sibling(llm_router: Router, deployment_id: str, model_id: str) -> bool:
if deployment_id == model_id:
return False
deployment: Final = llm_router.get_deployment(model_id=deployment_id)
return deployment is not None and not deployment.model_info.db_model
def _served_by_a_config_deployment(llm_router: Router | None, model_name: str, model_id: str) -> bool:
if llm_router is None:
return False
return any(
_config_sourced_sibling(llm_router, deployment_id, model_id)
for deployment_id in llm_router.get_model_ids(model_name=model_name)
)
async def _still_backed(executor: _RawExecutor, llm_router: Router | None, model_name: str, model_id: str) -> bool:
if _served_by_a_config_deployment(llm_router, model_name, model_id):
return True
count_rows: Final = await executor.query_raw(_BACKING_DEPLOYMENTS_SQL, model_name)
return any(_DeploymentCountRow.model_validate(row).deployment_count > 0 for row in count_rows)
async def _rewrite_groups(executor: _RawExecutor, sql: str, *names: str) -> None:
touched_rows: Final = await executor.query_raw(sql, *names)
await invalidate_access_group_caches(
tuple(_TouchedGroupRow.model_validate(row).access_group_id for row in touched_rows)
)
async def sync_access_groups_for_renamed_model(
prisma_client: object,
*,
model_id: str,
old_name: str,
new_name: str,
llm_router: Router | None,
) -> None:
if old_name == new_name:
return
executor: Final = _raw_executor(prisma_client)
old_name_still_backed: Final = await _still_backed(executor, llm_router, old_name, model_id)
await _rewrite_groups(
executor, _APPEND_MODEL_NAME_SQL if old_name_still_backed else _REPLACE_MODEL_NAME_SQL, old_name, new_name
)
async def sync_access_groups_for_deleted_model(
prisma_client: object,
*,
model_id: str,
model_name: str,
llm_router: Router | None,
) -> None:
executor: Final = _raw_executor(prisma_client)
if await _still_backed(executor, llm_router, model_name, model_id):
return
await _rewrite_groups(executor, _REMOVE_MODEL_NAME_SQL, model_name)

View file

@ -1228,6 +1228,7 @@ def run_server(
add_missing_query_params,
idle_lifetime_params,
reader_shareable_params,
translate_libpq_ssl_params,
unsupported_db_scheme,
unsupported_db_scheme_message,
)
@ -1275,11 +1276,15 @@ def run_server(
writer_url,
connection_url_params,
)
os.environ["DATABASE_URL"] = add_missing_query_params(modified_url, lifetime_params)
os.environ["DATABASE_URL"] = translate_libpq_ssl_params(
add_missing_query_params(modified_url, lifetime_params)
)
if os.getenv("DIRECT_URL", None) is not None:
database_url = os.getenv("DIRECT_URL")
modified_url = append_query_params(database_url, connection_url_params)
os.environ["DIRECT_URL"] = add_missing_query_params(modified_url, lifetime_params)
os.environ["DIRECT_URL"] = translate_libpq_ssl_params(
add_missing_query_params(modified_url, lifetime_params)
)
# The reader pool is a real pool against the same configured cap, so it
# gets the allowlisted pool params. Schema-affecting ones, including any
# the operator smuggled in through database_extra_connection_params, stay
@ -1292,14 +1297,16 @@ def run_server(
db_statement_timeout,
db_lock_timeout,
)
os.environ["DATABASE_URL_READ_REPLICA"] = add_missing_query_params(
os.environ["DATABASE_URL_READ_REPLICA"] = translate_libpq_ssl_params(
add_missing_query_params(
_with_query_value(read_replica_url, "options", reader_options)
if reader_options
else read_replica_url,
reader_shareable_params(connection_url_params),
),
lifetime_params,
add_missing_query_params(
_with_query_value(read_replica_url, "options", reader_options)
if reader_options
else read_replica_url,
reader_shareable_params(connection_url_params),
),
lifetime_params,
)
)
subprocess.run(["prisma"], capture_output=True)
is_prisma_runnable = True

View file

@ -4,6 +4,7 @@ import json
import os
from collections.abc import Mapping, Sequence
from datetime import datetime, timedelta, timezone
from types import MappingProxyType
from typing import (
TYPE_CHECKING,
Annotated,
@ -12,6 +13,7 @@ from typing import (
Literal,
NamedTuple,
Protocol,
TypeAlias,
TypedDict,
TypeVar,
cast, # noqa: TID251 # prisma group_by returns untyped aggregate mappings
@ -19,6 +21,7 @@ from typing import (
import fastapi
from fastapi import APIRouter, Depends, HTTPException, Request, Response, status
from pydantic import TypeAdapter
from typing_extensions import ReadOnly
import litellm
@ -55,9 +58,21 @@ router: Final = APIRouter()
SPEND_LOGS_PAGINATION_COUNT_CAP: Final = 10000
_SESSION_GROUP_KEY_SQL: Final = "COALESCE(NULLIF(session_id, ''), request_id), api_key"
_SESSION_KEY_EXPR: Final = "COALESCE(NULLIF(session_id, ''), request_id)"
_SESSION_GROUP_KEY_SQL: Final = f"{_SESSION_KEY_EXPR}, api_key"
_MCP_CALL_TYPES_SQL: Final = "('call_mcp_tool', 'list_mcp_tools')"
_AGENT_CALL_TYPE_SQL: Final = "'asend_message'"
_SPEND_LOG_LIST_COLUMNS: Final = """
request_id, call_type, api_key, spend, total_tokens,
prompt_tokens, completion_tokens, "startTime", "endTime",
"completionStartTime", model, model_id, model_group,
custom_llm_provider, api_base, "user", metadata,
cache_hit, cache_key, request_tags, team_id,
organization_id, end_user, requester_ip_address,
session_id, status, mcp_namespaced_tool_name, agent_id,
COALESCE(request_duration_ms,
(EXTRACT(EPOCH FROM ("endTime" - "startTime")) * 1000)::INTEGER) AS request_duration_ms
"""
_INTERNAL_HEALTH_CHECK_API_KEYS: Final = (
LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
@ -158,6 +173,26 @@ class _SessionSpendRow(TypedDict):
session_cache_hit_count: ReadOnly[int]
session_llm_count: ReadOnly[int]
session_agent_count: ReadOnly[int]
session_models: ReadOnly[Sequence[str]]
_SESSION_MODELS_LIMIT: Final = 10
_SESSION_MODEL_NAME_MAX_LEN: Final = 256
class _SessionSpendStats(NamedTuple):
session_total_count: int
session_total_spend: float
mcp_tool_call_count: int
mcp_tool_call_spend: float
session_cache_hit_count: int
session_llm_count: int
session_agent_count: int
session_models: Sequence[str]
session_models_truncated: bool
_SessionSpendMap: TypeAlias = Mapping[tuple[str, str], _SessionSpendStats]
class _SpendSumAggregate(TypedDict, total=False):
@ -2281,6 +2316,13 @@ async def ui_view_spend_logs(
default=False,
description="Paginate over sessions instead of raw logs: one representative row per session, total counts sessions",
),
session_cursor: str | None = fastapi.Query(
default=None,
description=(
"Keyset cursor '<last_activity>|<api_key>|<session_key>' from a previous group_by_session page. "
"UI route only, honored when sorting by startTime"
),
),
):
"""
View spend logs with pagination support.
@ -2614,6 +2656,18 @@ async def ui_view_spend_logs(
sql_params.append(f"%{error_message}%")
p += 1
if group_by_session is True and not is_v2 and not is_request_id_lookup and sort_by == "startTime":
return await _ui_session_grouped_spend_logs(
prisma_client=prisma_client,
sql_conditions=sql_conditions,
sql_params=sql_params,
next_param_index=p,
page=page,
page_size=page_size,
sort_desc=order_direction != "asc",
session_cursor=session_cursor,
)
# Build the ORDER BY expression. ttft_ms is computed from
# completionStartTime - startTime; non-streaming rows (where
# completionStartTime is null or equals endTime) yield NULL, so we
@ -2655,19 +2709,11 @@ async def ui_view_spend_logs(
total_is_capped: Final = raw_total > SPEND_LOGS_PAGINATION_COUNT_CAP
total_records: Final = SPEND_LOGS_PAGINATION_COUNT_CAP if total_is_capped else raw_total
select_columns: Final = """request_id, call_type, api_key, spend, total_tokens,
prompt_tokens, completion_tokens, "startTime", "endTime",
"completionStartTime", model, model_id, model_group,
custom_llm_provider, api_base, "user", metadata,
cache_hit, cache_key, request_tags, team_id,
organization_id, end_user, requester_ip_address,
session_id, status, mcp_namespaced_tool_name, agent_id,
COALESCE(request_duration_ms, (EXTRACT(EPOCH FROM ("endTime" - "startTime")) * 1000)::INTEGER) AS request_duration_ms"""
sql_query: Final = (
f"""
SELECT * FROM (
SELECT DISTINCT ON ({_SESSION_GROUP_KEY_SQL})
{select_columns}
{_SPEND_LOG_LIST_COLUMNS}
FROM "LiteLLM_SpendLogs"
WHERE {joined_conditions}
ORDER BY {_SESSION_GROUP_KEY_SQL}, call_type IN {_MCP_CALL_TYPES_SQL}, "startTime" DESC
@ -2678,7 +2724,7 @@ async def ui_view_spend_logs(
if session_grouping
else f"""
SELECT
{select_columns}
{_SPEND_LOG_LIST_COLUMNS}
FROM "LiteLLM_SpendLogs"
WHERE {joined_conditions}
ORDER BY {_order_expr} {_sql_dir}{_nulls_clause}
@ -2711,6 +2757,162 @@ async def ui_view_spend_logs(
raise handle_exception_on_proxy(e)
class _SessionPageRow(TypedDict):
session_key: ReadOnly[str]
api_key: ReadOnly[str]
last_activity: ReadOnly[str]
def _parse_session_cursor(session_cursor: str | None) -> tuple[str, str, str] | None:
if session_cursor is None or session_cursor.count("|") < 2:
return None
last_activity, _, rest = session_cursor.partition("|")
api_key, _, session_key = rest.partition("|")
if not last_activity or not session_key:
return None
return (last_activity, session_key, api_key)
async def _fetch_session_representatives(
prisma_client: "PrismaClient",
where_clause: str,
sql_params: Sequence[object],
next_param_index: int,
session_keys: Sequence[tuple[str, str]],
) -> list[dict[str, object]]: # mutable-ok: _build_ui_spend_logs_response writes session counts onto each row
"""Fetch the newest non-MCP row of each ``(session_key, api_key)`` session, in ``session_keys`` order."""
rep_query: Final = f"""
SELECT * FROM (
SELECT DISTINCT ON ({_SESSION_GROUP_KEY_SQL})
{_SPEND_LOG_LIST_COLUMNS}
FROM "LiteLLM_SpendLogs"
WHERE {where_clause}
AND ({_SESSION_GROUP_KEY_SQL}) IN (
SELECT * FROM unnest(${next_param_index}::text[], ${next_param_index + 1}::text[])
)
ORDER BY {_SESSION_GROUP_KEY_SQL}, call_type IN {_MCP_CALL_TYPES_SQL}, "startTime" DESC
) AS session_representatives
"""
rep_rows: Final[Sequence[dict[str, object]]] = await _query_raw( # mutable-ok: rows are enriched in place
prisma_client,
rep_query,
*sql_params,
[session_key for session_key, _ in session_keys], # mutable-ok: prisma serializes array params from a list
[api_key for _, api_key in session_keys], # mutable-ok: prisma serializes array params from a list
)
rep_by_key: Final[Mapping[tuple[str, str], dict[str, object]]] = MappingProxyType( # mutable-ok: same rows
{(str(row["session_id"] or row["request_id"]), str(row["api_key"])): row for row in rep_rows}
)
return [rep_by_key[key] for key in session_keys if key in rep_by_key] # mutable-ok: rows are enriched in place
async def _ui_session_grouped_spend_logs(
prisma_client: "PrismaClient",
sql_conditions: Sequence[str],
sql_params: Sequence[object],
next_param_index: int,
page: int,
page_size: int,
sort_desc: bool,
session_cursor: str | None,
) -> Mapping[str, object]:
"""
One row per session, keyset-paginated by session last activity.
Sessions are derived on the fly from ``LiteLLM_SpendLogs`` (no extra
table): rows sharing a ``session_id`` and ``api_key`` form a session, rows
without a session id are singletons keyed by ``request_id``. A page is the
next ``page_size`` sessions ordered by ``(MAX(startTime), session_key,
api_key)``, resumed from the ``session_cursor`` keyset
``'<last_activity>|<api_key>|<session_key>'`` instead of an OFFSET, so
page depth does not degrade the query plan. Each session is represented
by its newest non-MCP row, enriched by ``_build_ui_spend_logs_response``
exactly like the flat listing, and the response carries
``next_session_cursor`` / ``has_more`` while ``total`` counts sessions
(capped like the flat total).
"""
where_clause: Final = " AND ".join(sql_conditions) if sql_conditions else "TRUE"
cmp_op: Final = "<" if sort_desc else ">"
direction: Final = "DESC" if sort_desc else "ASC"
cursor: Final = _parse_session_cursor(session_cursor)
having_clause: Final = (
f'HAVING (MAX("startTime"), {_SESSION_GROUP_KEY_SQL}) {cmp_op} '
f"(${next_param_index}::timestamp, ${next_param_index + 1}, ${next_param_index + 2})"
if cursor
else ""
)
cursor_params: Final[tuple[object, ...]] = cursor if cursor else ()
limit_index: Final = next_param_index + len(cursor_params)
page_query: Final = f"""
SELECT {_SESSION_KEY_EXPR} AS session_key,
api_key,
MAX("startTime")::text AS last_activity
FROM "LiteLLM_SpendLogs"
WHERE {where_clause}
GROUP BY {_SESSION_GROUP_KEY_SQL}
{having_clause}
ORDER BY MAX("startTime") {direction}, {_SESSION_KEY_EXPR} {direction}, api_key {direction}
LIMIT ${limit_index}
"""
page_rows: Final[Sequence[_SessionPageRow]] = await _query_raw(
prisma_client, page_query, *sql_params, *cursor_params, page_size + 1
)
has_more: Final = len(page_rows) > page_size
visible_rows: Final = page_rows[:page_size]
next_cursor: Final = (
f"{visible_rows[-1]['last_activity']}|{visible_rows[-1]['api_key']}|{visible_rows[-1]['session_key']}"
if has_more and visible_rows
else None
)
count_query: Final = f"""
SELECT COUNT(*) AS total_count
FROM (
SELECT 1
FROM "LiteLLM_SpendLogs"
WHERE {where_clause}
GROUP BY {_SESSION_GROUP_KEY_SQL}
LIMIT ${next_param_index}
) AS bounded_sessions
"""
count_rows: Final[Sequence[_SpendLogsCountRow]] = await _query_raw(
prisma_client, count_query, *sql_params, SPEND_LOGS_PAGINATION_COUNT_CAP + 1
)
raw_total: Final = int(count_rows[0]["total_count"]) if count_rows else 0
total_is_capped: Final = raw_total > SPEND_LOGS_PAGINATION_COUNT_CAP
total_records: Final = SPEND_LOGS_PAGINATION_COUNT_CAP if total_is_capped else raw_total
session_keys: Final = tuple((row["session_key"], row["api_key"]) for row in visible_rows)
data: Final[list[dict[str, object]]] = ( # mutable-ok: _build_ui_spend_logs_response writes onto each row
await _fetch_session_representatives(
prisma_client=prisma_client,
where_clause=where_clause,
sql_params=sql_params,
next_param_index=next_param_index,
session_keys=session_keys,
)
if session_keys
else [] # mutable-ok: downstream enrichment mutates rows in place
)
_hydrate_spend_log_metadata(data)
total_pages: Final = (total_records + page_size - 1) // page_size
response: Final[Mapping[str, object]] = await _build_ui_spend_logs_response(
prisma_client,
data,
total_records,
page,
page_size,
total_pages,
enrich_session_counts=True,
total_is_capped=total_is_capped,
)
return {**response, "next_session_cursor": next_cursor, "has_more": has_more} # mutable-ok: FastAPI response body
class RequestResponsePayload(NamedTuple):
messages: str | list | dict | None
response: str | list | dict | None
@ -4080,7 +4282,7 @@ async def _build_ui_spend_logs_response(
total_pages: int,
enrich_session_counts: bool = True,
total_is_capped: bool = False,
) -> dict:
) -> dict[str, object]:
"""
Build the paginated response for the UI spend-logs endpoint.
@ -4121,7 +4323,7 @@ async def _build_ui_spend_logs_response(
}
)
session_spend_map: dict[tuple[str, str], dict[str, int | float]] = {}
session_spend_map: _SessionSpendMap = {}
if enrich_session_counts and session_ids:
from prisma.errors import PrismaError
@ -4139,40 +4341,60 @@ async def _build_ui_spend_logs_response(
rows: Final[Sequence[_SessionSpendRow]] = await _query_raw(
prisma_client,
f"""
SELECT session_id, api_key,
COUNT(*)::int AS session_total_count,
COALESCE(SUM(spend), 0)::double precision AS session_total_spend,
COUNT(*) FILTER (
WHERE call_type IN {_MCP_CALL_TYPES_SQL}
)::int AS mcp_tool_call_count,
COALESCE(SUM(spend) FILTER (
WHERE call_type IN {_MCP_CALL_TYPES_SQL}
), 0)::double precision AS mcp_tool_call_spend,
COUNT(*) FILTER (WHERE LOWER(cache_hit) = 'true')::int AS session_cache_hit_count,
COUNT(*) FILTER (
WHERE call_type NOT IN {_MCP_CALL_TYPES_SQL} AND call_type != {_AGENT_CALL_TYPE_SQL}
)::int AS session_llm_count,
COUNT(*) FILTER (WHERE call_type = {_AGENT_CALL_TYPE_SQL})::int AS session_agent_count
FROM "LiteLLM_SpendLogs"
WHERE session_id = ANY($1::text[])
AND api_key = ANY($2::text[])
GROUP BY session_id, api_key
SELECT s.*, COALESCE(m.session_models, ARRAY[]::text[]) AS session_models
FROM (
SELECT session_id, api_key,
COUNT(*)::int AS session_total_count,
COALESCE(SUM(spend), 0)::double precision AS session_total_spend,
COUNT(*) FILTER (
WHERE call_type IN {_MCP_CALL_TYPES_SQL}
)::int AS mcp_tool_call_count,
COALESCE(SUM(spend) FILTER (
WHERE call_type IN {_MCP_CALL_TYPES_SQL}
), 0)::double precision AS mcp_tool_call_spend,
COUNT(*) FILTER (WHERE LOWER(cache_hit) = 'true')::int AS session_cache_hit_count,
COUNT(*) FILTER (
WHERE call_type NOT IN {_MCP_CALL_TYPES_SQL} AND call_type != {_AGENT_CALL_TYPE_SQL}
)::int AS session_llm_count,
COUNT(*) FILTER (WHERE call_type = {_AGENT_CALL_TYPE_SQL})::int AS session_agent_count
FROM "LiteLLM_SpendLogs"
WHERE session_id = ANY($1::text[])
AND api_key = ANY($2::text[])
GROUP BY session_id, api_key
) s
LEFT JOIN LATERAL (
SELECT ARRAY_AGG(d.model ORDER BY d.model) AS session_models
FROM (
SELECT DISTINCT LEFT(model, $3::int) AS model
FROM "LiteLLM_SpendLogs"
WHERE session_id = s.session_id
AND api_key = s.api_key
AND model IS NOT NULL AND model <> ''
ORDER BY 1
LIMIT $4::int
) d
) m ON TRUE
""",
session_ids,
authorized_api_keys,
_SESSION_MODEL_NAME_MAX_LEN,
_SESSION_MODELS_LIMIT + 1,
)
session_spend_map = {
(row["session_id"], row["api_key"]): {
"session_total_count": int(row.get("session_total_count") or 0),
"session_total_spend": float(row.get("session_total_spend") or 0.0),
"mcp_tool_call_count": int(row.get("mcp_tool_call_count") or 0),
"mcp_tool_call_spend": float(row.get("mcp_tool_call_spend") or 0.0),
"session_cache_hit_count": int(row.get("session_cache_hit_count") or 0),
"session_llm_count": int(row.get("session_llm_count") or 0),
"session_agent_count": int(row.get("session_agent_count") or 0),
}
(row["session_id"], row["api_key"]): _SessionSpendStats(
session_total_count=int(row.get("session_total_count") or 0),
session_total_spend=float(row.get("session_total_spend") or 0.0),
mcp_tool_call_count=int(row.get("mcp_tool_call_count") or 0),
mcp_tool_call_spend=float(row.get("mcp_tool_call_spend") or 0.0),
session_cache_hit_count=int(row.get("session_cache_hit_count") or 0),
session_llm_count=int(row.get("session_llm_count") or 0),
session_agent_count=int(row.get("session_agent_count") or 0),
session_models=models[:_SESSION_MODELS_LIMIT],
session_models_truncated=len(models) > _SESSION_MODELS_LIMIT,
)
for row in rows
if row.get("session_id") and row.get("api_key") is not None
for models in (TypeAdapter(list[str]).validate_python(row.get("session_models") or ()),)
}
except PrismaError:
verbose_proxy_logger.debug(
@ -4187,15 +4409,17 @@ async def _build_ui_spend_logs_response(
sid = row_dict.get("session_id")
row_api_key = row_dict.get("api_key")
session_stats = session_spend_map.get((sid, row_api_key)) if sid and row_api_key is not None else None
row_dict["session_total_count"] = int(session_stats["session_total_count"]) if session_stats else 1
row_dict["session_total_count"] = session_stats.session_total_count if session_stats else 1
if session_stats:
row_dict["session_total_spend"] = session_stats["session_total_spend"]
if session_stats["mcp_tool_call_count"]:
row_dict["mcp_tool_call_count"] = session_stats["mcp_tool_call_count"]
row_dict["mcp_tool_call_spend"] = session_stats["mcp_tool_call_spend"]
row_dict["session_cache_hit_count"] = session_stats["session_cache_hit_count"]
row_dict["session_llm_count"] = session_stats["session_llm_count"]
row_dict["session_agent_count"] = session_stats["session_agent_count"]
row_dict["session_total_spend"] = session_stats.session_total_spend
if session_stats.mcp_tool_call_count:
row_dict["mcp_tool_call_count"] = session_stats.mcp_tool_call_count
row_dict["mcp_tool_call_spend"] = session_stats.mcp_tool_call_spend
row_dict["session_cache_hit_count"] = session_stats.session_cache_hit_count
row_dict["session_llm_count"] = session_stats.session_llm_count
row_dict["session_agent_count"] = session_stats.session_agent_count
row_dict["session_models"] = session_stats.session_models
row_dict["session_models_truncated"] = session_stats.session_models_truncated
enriched.append(row_dict)
response_data: list = enriched
else:

View file

@ -50,6 +50,7 @@ from litellm.constants import (
DEFAULT_HEALTH_CHECK_INTERVAL,
DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER,
DEFAULT_MAX_LRU_CACHE_SIZE,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
RUNTIME_UPDATABLE_ROUTER_SETTINGS,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
)
@ -231,6 +232,7 @@ from litellm.types.router import (
)
from litellm.types.services import ServiceTypes
from litellm.types.utils import (
AUTOROUTER_CLASSIFIER_CALL_ORIGIN,
PROMPT_QUOTING_ROUTING_DECISION_FIELDS,
CustomPricingLiteLLMParams,
GenericBudgetConfigType,
@ -2168,6 +2170,76 @@ class Router:
verbose_router_logger.debug("Error occurred while printing deployment - %s", e)
raise e
@staticmethod
def _deployment_params_with_request_reasoning_override(
deployment_params: Mapping[str, object], request_kwargs: Mapping[str, object]
) -> dict[str, object]: # mutable-ok: litellm's request pipeline consumes a mutable kwargs mapping
"""Return deployment params whose equivalent effort controls cannot outrank a request override.
Providers expose the same setting through several native carriers. A request-level
``reasoning_effort`` is the portable override, so a deployment's ``thinking`` or nested
``*.effort`` must not remain beside it and either win or trigger a conflicting-params 400.
Every changed mapping is copied so the Router's shared deployment config stays immutable.
"""
sanitized: Final = dict(deployment_params) # mutable-ok: request-local copy protects shared Router state
if request_kwargs.get("reasoning_effort") is None:
return sanitized
sanitized.pop("thinking", None)
Router._pop_effort_from_nested_carrier(sanitized, "output_config")
Router._pop_effort_from_nested_carrier(sanitized, "reasoning")
extra_body: Final = sanitized.get("extra_body")
if isinstance(extra_body, Mapping):
sanitized_extra_body: Final = dict(extra_body) # mutable-ok: request-local nested copy
sanitized_extra_body.pop("reasoning_effort", None)
sanitized_extra_body.pop("thinking", None)
Router._pop_effort_from_nested_carrier(sanitized_extra_body, "output_config")
Router._pop_effort_from_nested_carrier(sanitized_extra_body, "reasoning")
if sanitized_extra_body:
sanitized["extra_body"] = sanitized_extra_body
else:
sanitized.pop("extra_body", None)
return sanitized
@staticmethod
def _is_classifier_internal_call(kwargs: Mapping[str, object]) -> bool:
metadata: Final = kwargs.get("metadata")
litellm_metadata: Final = kwargs.get("litellm_metadata")
return any(
isinstance(candidate, Mapping)
and candidate.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == AUTOROUTER_CLASSIFIER_CALL_ORIGIN
for candidate in (metadata, litellm_metadata)
)
def _drop_unsupported_classifier_reasoning_effort(
self,
deployment: DeploymentTypedDict,
model: str,
kwargs: dict[str, object], # mutable-ok: fallback must update the active request and its log body together
) -> None:
"""Let a classifier fallback without reasoning support remain a usable fallback.
The dashboard only offers explicitly advertised levels, but an existing config can outlive
a model change and fallbacks can target a different group. Unknown capability fails open;
only a provider that explicitly rejects the parameter has it removed.
"""
if kwargs.get("reasoning_effort") is None or not self._is_classifier_internal_call(kwargs):
return
if self._deployment_accepts_param(deployment, model, "reasoning_effort"):
return
verbose_router_logger.warning(
"litellm.router.py: dropping classifier reasoning_effort for model=%s because the selected deployment does not support it",
model,
)
kwargs.pop("reasoning_effort", None)
proxy_server_request: Final = kwargs.get("proxy_server_request")
if not isinstance(proxy_server_request, dict):
return
body: Final = proxy_server_request.get("body")
if isinstance(body, dict):
body.pop("reasoning_effort", None)
### COMPLETION, EMBEDDING, IMG GENERATION FUNCTIONS
def completion(self, model: str, messages: list[dict[str, str]], **kwargs) -> ModelResponse | CustomStreamWrapper:
@ -2203,9 +2275,16 @@ class Router:
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
self._drop_unsupported_classifier_reasoning_effort(
deployment=cast(DeploymentTypedDict, deployment), # cast-ok: selection returns a router deployment
model=model,
kwargs=kwargs,
)
# Check for silent model experiment
# Make a local copy of litellm_params to avoid mutating the Router's state
litellm_params: Final = deployment["litellm_params"].copy()
litellm_params: Final = self._deployment_params_with_request_reasoning_override(
deployment["litellm_params"], kwargs
)
silent_model: Final = litellm_params.pop("silent_model", None)
if silent_model is not None:
@ -3216,6 +3295,11 @@ class Router:
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
self._drop_unsupported_classifier_reasoning_effort(
deployment=cast(DeploymentTypedDict, deployment), # cast-ok: selection returns a router deployment
model=model,
kwargs=kwargs,
)
_timeout_debug_deployment_dict = deployment
end_time: Final = time.time()
@ -3237,7 +3321,9 @@ class Router:
# Check for silent model experiment
# Make a local copy of litellm_params to avoid mutating the Router's state
litellm_params: Final = deployment["litellm_params"].copy()
litellm_params: Final = self._deployment_params_with_request_reasoning_override(
deployment["litellm_params"], kwargs
)
silent_model: Final = litellm_params.pop("silent_model", None)
if silent_model is not None:
@ -10255,6 +10341,8 @@ class Router:
total_itpm: int | None = None
total_otpm: int | None = None
configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None
reasoning_efforts_initialized = False
reasoning_efforts_unknown = False
model_list: Final = self.get_model_list(model_name=model_group)
if model_list is None:
return None
@ -10441,10 +10529,23 @@ class Router:
if model_info.get("rpm", None) is not None and _deployment_rpm is None:
_deployment_rpm = model_info.get("rpm")
model_group_info.supported_reasoning_efforts = intersect_supported_reasoning_efforts(
model_group_info.supported_reasoning_efforts,
resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=deployment_is_mapped),
deployment_reasoning_efforts = (
resolve_supported_reasoning_efforts( # rebind-ok: recalculated per deployment
model_info, deployment_is_mapped=deployment_is_mapped
)
)
if deployment_reasoning_efforts is None:
reasoning_efforts_unknown = True
model_group_info.supported_reasoning_efforts = None
elif not reasoning_efforts_initialized:
reasoning_efforts_initialized = True
if not reasoning_efforts_unknown:
model_group_info.supported_reasoning_efforts = deployment_reasoning_efforts
elif not reasoning_efforts_unknown:
model_group_info.supported_reasoning_efforts = intersect_supported_reasoning_efforts(
model_group_info.supported_reasoning_efforts,
deployment_reasoning_efforts,
)
if _deployment_tpm is not None:
if total_tpm is None:
@ -12308,10 +12409,12 @@ class Router:
@staticmethod
def _pop_effort_from_nested_carrier(request_kwargs: dict[str, object], carrier: str) -> None:
nested: Final = request_kwargs.get(carrier)
if not isinstance(nested, dict):
if not isinstance(nested, Mapping):
return
nested.pop("effort", None)
if not nested:
sanitized: Final = {key: value for key, value in nested.items() if key != "effort"}
if sanitized:
request_kwargs[carrier] = sanitized # rebind-ok: copy-on-write, so a shared nested carrier is never edited
else:
request_kwargs.pop(carrier, None)
@staticmethod

View file

@ -255,7 +255,8 @@ model_list:
classifier_type: heuristic_first
heuristic_first_max_tier: SIMPLE
classifier_llm_config:
model: gpt-4o-mini
model: gpt-5-mini
reasoning_effort: low
tiers:
SIMPLE: gpt-4o-mini
MEDIUM: gpt-4o
@ -263,6 +264,10 @@ model_list:
REASONING: o1-preview
```
`classifier_llm_config.reasoning_effort` applies only to the internal classifier call. Omit it to
keep the classifier deployment or provider default, or set a supported value such as `none` or
`low` to override that call.
A request short-circuits, meaning it routes on the scorer's own tier with no classifier call, when
two things hold: the scorer landed at or below `heuristic_first_max_tier`, and it produced at least
one signal. Everything else goes to the classifier, which then decides as it normally would.

View file

@ -28,6 +28,7 @@ from pydantic import BaseModel, create_model
from litellm._logging import verbose_router_logger
from litellm.constants import (
EMPTY_MAPPING,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
RETURN_RAW_MODEL_NAME_METADATA_KEY,
SESSION_ID_GENERATED_METADATA_KEY,
)
@ -42,6 +43,7 @@ from litellm.router_strategy.complexity_router.tier_predictor import (
TierSuccessPredictor,
resolve_tier_artifact,
)
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import (
AUTOROUTER_CLASSIFIER_CALL_ORIGIN,
ModelResponse,
@ -1668,20 +1670,27 @@ class ComplexityRouter(CustomLogger):
)
request_metadata = (request_kwargs or {}).get("litellm_metadata") or (request_kwargs or {}).get("metadata")
metadata: Final = forwarded_internal_call_metadata(request_metadata, AUTOROUTER_CLASSIFIER_CALL_ORIGIN)
metadata: Final = { # mutable-ok: SDK metadata kwarg is enriched by the request pipeline
**forwarded_internal_call_metadata(request_metadata, AUTOROUTER_CLASSIFIER_CALL_ORIGIN),
INTERNAL_CALL_ORIGIN_METADATA_KEY: AUTOROUTER_CLASSIFIER_CALL_ORIGIN,
}
turn_off_message_logging: Final = _effective_turn_off_message_logging(request_kwargs)
messages_for_call: Final = [
messages_for_call: Final[list[AllMessageValues]] = [ # mutable-ok: SDK request payload list is built once
{"role": "system", "content": classifier_system_prompt},
{"role": "user", "content": user_payload},
]
response_format: Final = classifier_response_format
classifier_call_params: Mapping[str, str] = EMPTY_MAPPING
if llm_config.reasoning_effort is not None:
classifier_call_params = MappingProxyType({"reasoning_effort": llm_config.reasoning_effort})
proxy_server_request: Final = {
"body": {
"model": llm_config.model,
"messages": messages_for_call,
"response_format": response_format,
**classifier_call_params,
}
}
@ -1693,6 +1702,7 @@ class ComplexityRouter(CustomLogger):
metadata=metadata,
proxy_server_request=proxy_server_request,
turn_off_message_logging=turn_off_message_logging,
**classifier_call_params,
**_parent_session_kwargs(request_kwargs),
)
content: Final = response.choices[0].message.content

View file

@ -12,6 +12,7 @@ from typing import Annotated, Final, Literal
from pydantic import BaseModel, ConfigDict, Field, SkipValidation, field_serializer, field_validator, model_validator
from litellm.types.llms.openai import REASONING_EFFORT
from litellm.types.router import AdaptiveRouterWeights, ClassifierPlugin, RoutingPlugin
from .tier_predictor import TrainedTierArtifact
@ -432,6 +433,13 @@ class ClassifierLLMConfig(BaseModel):
model: str = Field(
description="Model name (from the router's model_list) to call for classification",
)
reasoning_effort: REASONING_EFFORT | None = Field(
default=None,
description=(
"Reasoning effort override for classifier calls. Leave unset to use "
"the classifier deployment or provider default."
),
)
timeout_ms: int = Field(
default=3000,
description="Timeout budget for the classification call, in milliseconds",

View file

@ -52933,7 +52933,8 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"supports_web_search": true
},
"bedrock_mantle/openai.gpt-5.6-terra": {
"input_cost_per_token": 2.2e-06,
@ -52966,7 +52967,8 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"supports_web_search": true
},
"bedrock_mantle/openai.gpt-5.6-cyber": {
"input_cost_per_token": 1.375e-05,
@ -53027,7 +53029,8 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"supports_web_search": true
},
"us.openai.gpt-5.6-sol": {
"input_cost_per_token": 4.4e-06,
@ -53213,7 +53216,8 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"supports_web_search": true
},
"bedrock_mantle/openai.gpt-5.4": {
"input_cost_per_token": 2.75e-06,
@ -53243,7 +53247,8 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"supports_web_search": true
},
"bedrock_mantle/google.gemma-4-31b": {
"input_cost_per_token": 1.4e-07,

View file

@ -174,6 +174,8 @@ litellm-proxy = "litellm.proxy.client.cli:cli"
[dependency-groups]
dev = [
"diff-cover==9.7.2",
"hypothesis==6.165.10",
"reportlab==5.0.1",
"basedpyright==1.39.7",
"keyring==25.7.0",
"pytest==9.0.3",

View file

@ -207,13 +207,16 @@ async def completions(request: Request) -> Response:
async def embeddings(request: Request) -> Response:
body = await _parse_body(request)
model = _requested_model(body)
if model == _SLOW_MODEL:
await asyncio.sleep(_SLOW_RESPONSE_SECONDS)
raw_input = body.get("input", "")
count = len(raw_input) if isinstance(raw_input, list) else 1
return JSONResponse(
{
"object": "list",
"data": [{"object": "embedding", "index": i, "embedding": [0.0] * 1536} for i in range(max(count, 1))],
"model": _requested_model(body),
"model": model,
"usage": {"prompt_tokens": 5, "total_tokens": 5},
}
)

View file

@ -172,6 +172,7 @@ pylint: >=3.3.9 # GPLv2 license
langchain-mcp-adapters: >=0.2.1 # MIT License
langgraph: >=1.0.10 # MIT License
langgraph-prebuilt: >=1.0.8 # MIT License - https://github.com/langchain-ai/langgraph/blob/main/LICENSE
hypothesis: >=6.165.10 # MPL 2.0 license
pytest-rerunfailures: >=15.1 # MPL 2.0 license
pytest-recording: >=0.13.4 # MIT license
expression: >=5.6.0 # MIT License - https://github.com/cognitedata/Expression/blob/main/LICENSE

View file

@ -23,10 +23,12 @@ test.describe("AI Hub (internal admin view)", () => {
await expect(modal.getByText(/Select All \(\d+\)/)).toBeVisible({ timeout: 5_000 });
// Step 1: pick the seeded models via "Select All"
await modal.getByText(/Select All/i).click();
await modal.getByRole("checkbox", { name: /Select All/ }).check();
// Move to confirm step
await modal.getByRole("button", { name: "Next" }).click();
const next = modal.getByRole("button", { name: "Next" });
await expect(next).toBeEnabled();
await next.click();
await expect(modal.getByText("Confirm Making Models Public")).toBeVisible({ timeout: 5_000 });
// Submit

View file

@ -703,3 +703,170 @@ class TestSpendLogsPartitionDetectionMissingPsycopg:
assert any(
"psycopg is not installed" in record.message for record in caplog.records
)
_ATTEMPT_BUDGET = 4
_P3005_STDERR = """Error: P3005
The database schema is not empty. Read more about how to baseline an existing production database: https://pris.ly/d/migrate-baseline
"""
def _p3018_stderr(migration_name):
return f"""Error: P3018
A migration failed to apply. New migrations cannot be applied before the error is recovered from.
Migration name: {migration_name}
Database error code: 42P07
Database error:
ERROR: relation "SomeTable" already exists
"""
class _MigrateDeployHarness:
"""Drives _setup_database_v2 with a scripted sequence of
`prisma migrate deploy` outcomes, with every recovery command faked out so
nothing touches a database or the packaged migrations directory."""
def __init__(self, monkeypatch, tmp_path, outcomes, repeat_last=False):
import subprocess as subprocess_module
import litellm_proxy_extras.utils as utils_module
self.deploy_calls = []
self.resolved = []
self.baselines = 0
self._outcomes = list(outcomes)
self._repeat_last = repeat_last
self._subprocess_module = subprocess_module
monkeypatch.delenv("DATABASE_URL", raising=False)
monkeypatch.setattr(
ProxyExtrasDBManager, "_get_prisma_dir", staticmethod(lambda: str(tmp_path))
)
monkeypatch.setattr(
ProxyExtrasDBManager,
"_create_baseline_migration",
staticmethod(self._fake_baseline),
)
monkeypatch.setattr(
ProxyExtrasDBManager,
"_roll_back_migration",
staticmethod(lambda name: None),
)
monkeypatch.setattr(
ProxyExtrasDBManager,
"_resolve_specific_migration",
staticmethod(self.resolved.append),
)
monkeypatch.setattr(utils_module.subprocess, "run", self._fake_run)
monkeypatch.setattr(utils_module.time, "sleep", lambda seconds: None)
self.baseline_succeeds = True
def _fake_baseline(self, *args, **kwargs):
self.baselines += 1
return self.baseline_succeeds
def _next_outcome(self):
if self._outcomes:
if self._repeat_last and len(self._outcomes) == 1:
return self._outcomes[0]
return self._outcomes.pop(0)
raise AssertionError("prisma migrate deploy called more times than scripted")
def _fake_run(self, cmd, **kwargs):
assert cmd[1:] == ["migrate", "deploy"], f"unexpected prisma command: {cmd}"
self.deploy_calls.append(cmd)
outcome = self._next_outcome()
if outcome == "ok":
return _FakeCompleted()
if outcome == "timeout":
raise self._subprocess_module.TimeoutExpired(cmd, 1)
raise self._subprocess_module.CalledProcessError(1, cmd, stderr=outcome)
def run(self):
return ProxyExtrasDBManager._setup_database_v2(use_migrate=True)
class TestMigrateDeployAttemptAccounting:
"""A `prisma db push` database has a full schema and no ledger, so the v2
resolver baselines it and then works through every migration whose objects
already exist. Those recoveries make progress, so they must not spend the
retry budget, which is there to stop a run that is getting nowhere."""
def test_a_push_created_database_finishes_bootstrapping(
self, monkeypatch, tmp_path
):
already_there = [
"20250329084805_new_cron_job_table",
"20250806095134_rename_alias_to_server_name_mcp_table",
"20260224203854_add_agent_object_permissions_table",
"20260301120000_fourth_table",
"20260302120000_fifth_table",
"20260303120000_sixth_table",
]
harness = _MigrateDeployHarness(
monkeypatch,
tmp_path,
[_P3005_STDERR]
+ [_p3018_stderr(name) for name in already_there]
+ ["ok"],
)
assert harness.run() is True
assert harness.baselines == 1
assert harness.resolved == already_there
assert len(harness.deploy_calls) == len(already_there) + 2
def test_repeated_recovery_of_one_migration_still_gives_up(
self, monkeypatch, tmp_path
):
harness = _MigrateDeployHarness(
monkeypatch,
tmp_path,
[_p3018_stderr("20250329084805_new_cron_job_table")],
repeat_last=True,
)
with pytest.raises(RuntimeError):
harness.run()
assert len(harness.deploy_calls) <= _ATTEMPT_BUDGET + 1
def test_timeouts_still_spend_the_budget(self, monkeypatch, tmp_path):
harness = _MigrateDeployHarness(
monkeypatch, tmp_path, ["timeout"], repeat_last=True
)
with pytest.raises(RuntimeError):
harness.run()
assert len(harness.deploy_calls) == _ATTEMPT_BUDGET
def test_a_baseline_that_never_lands_stops_after_the_budget(
self, monkeypatch, tmp_path
):
harness = _MigrateDeployHarness(
monkeypatch, tmp_path, [_P3005_STDERR], repeat_last=True
)
harness.baseline_succeeds = False
with pytest.raises(RuntimeError):
harness.run()
assert len(harness.deploy_calls) == _ATTEMPT_BUDGET
def test_an_unrecoverable_error_is_not_retried(self, monkeypatch, tmp_path):
harness = _MigrateDeployHarness(
monkeypatch,
tmp_path,
["Error: P3018\n\nMigration name: 20260101000000_x\n\nERROR: syntax error at or near \"SLECT\"\n"],
repeat_last=True,
)
with pytest.raises(RuntimeError):
harness.run()
assert len(harness.deploy_calls) == 1
assert harness.resolved == []

View file

@ -15,6 +15,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import litellm
from litellm import completion, completion_cost, embedding
from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE
litellm.set_verbose = False
@ -269,11 +270,14 @@ def test_openai_azure_embedding_timeouts():
def test_openai_embedding_timeouts():
try:
response = embedding(
model="text-embedding-ada-002",
model="openai/slow-endpoint",
input=["good morning from litellm"],
timeout=0.00001,
api_base=FAKE_OPENAI_API_BASE,
api_key="fake-key",
timeout=0.5,
)
print(response)
pytest.fail("Expected timeout error, the request returned instead")
except openai.APITimeoutError:
print("Good job got OpenAI timeout error!")
pass

View file

@ -1552,8 +1552,9 @@ def test_router_timeout():
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": "os.environ/OPENAI_API_KEY",
"model": "openai/slow-endpoint",
"api_base": FAKE_OPENAI_API_BASE,
"api_key": "fake-key",
},
}
]
@ -1562,7 +1563,7 @@ def test_router_timeout():
start_time = time.time()
try:
res = router.completion(
model="gpt-3.5-turbo", messages=messages, timeout=0.0001
model="gpt-3.5-turbo", messages=messages, timeout=0.5
)
print(res)
pytest.fail("this should have timed out")

View file

@ -12,6 +12,7 @@ import openai
import pytest
import litellm
from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE
@pytest.mark.parametrize(
@ -216,13 +217,16 @@ def test_timeout_streaming():
litellm.set_verbose = False
try:
response = litellm.completion(
model="gpt-3.5-turbo",
model="openai/slow-endpoint",
messages=[{"role": "user", "content": "hello, write a 20 pg essay"}],
timeout=0.0001,
api_base=FAKE_OPENAI_API_BASE,
api_key="fake-key",
timeout=0.5,
stream=True,
)
for chunk in response:
print(chunk)
pytest.fail("Did not raise error `openai.APITimeoutError`. The stream completed instead")
except openai.APITimeoutError as e:
print(
"Passed: Raised correct exception. Got openai.APITimeoutError\nGood Job", e

View file

@ -2815,6 +2815,7 @@ async def test_mcp_server_manager_with_access_groups_integration():
"""Integration test for MCPServerManager with access group filtering"""
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
MCPServerAccess,
)
from litellm.proxy._types import UserAPIKeyAuth
@ -2848,11 +2849,11 @@ async def test_mcp_server_manager_with_access_groups_integration():
)
# Mock the permission lookup to return staff access group
with patch.object(MCPRequestHandler, "get_allowed_mcp_servers") as mock_get_allowed:
mock_get_allowed.return_value = [
"staff-server-id",
"ops-server-id",
] # User has access to staff and ops
with patch.object(MCPRequestHandler, "get_mcp_server_access") as mock_get_allowed: # test-quality-ok: manager resolver seam
mock_get_allowed.return_value = MCPServerAccess(
server_ids=("staff-server-id", "ops-server-id"),
scope="scoped",
)
allowed_servers = await test_manager.get_allowed_mcp_servers(user_auth)
@ -2901,6 +2902,7 @@ async def test_get_allowed_mcp_servers_returns_empty_for_non_admin_without_permi
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
MCPServerAccess,
)
test_manager = MCPServerManager()
@ -2923,9 +2925,9 @@ async def test_get_allowed_mcp_servers_returns_empty_for_non_admin_without_permi
)
with patch.object(
MCPRequestHandler, "get_allowed_mcp_servers", new_callable=AsyncMock
MCPRequestHandler, "get_mcp_server_access", new_callable=AsyncMock
) as mock_permission_lookup:
mock_permission_lookup.return_value = []
mock_permission_lookup.return_value = MCPServerAccess(server_ids=())
allowed_servers = await test_manager.get_allowed_mcp_servers(user_auth)
assert allowed_servers == []

View file

@ -1,148 +1,105 @@
# Rust ↔ Python SDK parity harness
# Rust/Python migration harness
This folder is the operator-facing harness for the Rust migration test plan. It runs pytest normally, listens to test events in-process, and redraws a live matrix grouped by testing strategy and SDK-level function.
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
The matrix always has these SDK columns:
- `ocr / aocr`
- `messages / amessages`
- `responses / aresponses`
- `count_tokens`
- `chat_completions / acompletion`
- `transcription / atranscription`
The harness has four deliberately broad test-strategy folders:
| Strategy | Folder |
| --- | --- |
| Public SDK parity over generated and recorded inputs | [`e2e_fuzz_tests/`](e2e_fuzz_tests/) |
| Focused tests of Rust-owned behavior | [`unit_tests_rust/`](unit_tests_rust/) |
| Isolated transform and Python-to-Rust helper coverage | [`validate_sub_methods/`](validate_sub_methods/) |
| Already-existing live-API SDK tests | [`existing_e2e_test_sdk/`](existing_e2e_test_sdk/) |
## Run it
From the repository root:
```bash
poetry run python -m tests.rust-python-harness
```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/
```
The default runs every configured test once and updates all matching cells in real time. Narrow a run by strategy, SDK function, or both:
## Run locally
```bash
poetry run python -m tests.rust-python-harness --strategy e2e_fuzz_tests
poetry run python -m tests.rust-python-harness --function messages
poetry run python -m tests.rust-python-harness --strategy validate_sub_methods --function ocr
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
```
For a guided run, use the interactive picker. It asks which strategy rows and SDK
function columns to include, then hands the terminal to the live dashboard. It never
captures keys while tests are running, so Ctrl-C and pytest debugging remain safe.
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
```bash
poetry run python -m tests.rust-python-harness --interactive
```
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
Useful operator options:
## Strategy responsibilities
```bash
# Inspect coverage and pytest selectors without running anything.
poetry run python -m tests.rust-python-harness --list
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
# Stable line-oriented output for CI logs or redirected output.
poetry run python -m tests.rust-python-harness --plain
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
# Measure Python reference lines exercised by this parity run and build an HTML heatmap.
poetry run python -m tests.rust-python-harness --coverage
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
# Forward pytest options. Use the equals form when the value begins with a dash.
poetry run python -m tests.rust-python-harness --pytest-arg=-x
```
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
The process returns pytest's exit code. A configured selector that collects no test is also a failure. A planned cell has no selector yet and does not fail the run.
## Configure cases
The dashboard adapts to narrow terminals, shows elapsed time and unique-test progress,
and prints the three slowest tests when the run ends. Each failure includes a focused
`poetry run pytest ... -q` command. Redirected output and CI automatically use the
line-oriented plain renderer; `--plain` lets you opt into it locally.
The final screen includes a confidence score for every SDK section. It is the direct
ratio of required strategy rows with passing evidence, such as `1/3 = 33%`; High means
all required strategies passed, Medium means some passed, and Low means none passed.
This behavioral score is intentionally shown separately from Python and Rust LOC.
Coverage reports are written outside the three strategy folders at
`target/rust-python-harness/`. Open `python-html/index.html` to inspect executed and
missing Python lines; `python.json` and `python.xml` are available for automation.
Coverage is finalized after pytest exits, because worker processes must flush their
data first.
## Port coverage and confidence
Treat these as separate signals instead of one ambiguous coverage percentage:
| Signal | Tool | What it proves |
| --- | --- | --- |
| Python reference LOC | `coverage.py` / `pytest-cov` via `--coverage` | The mapped Python behavior ran |
| Rust port LOC | `cargo-llvm-cov` | The mapped Rust implementation ran |
| Parity contracts | This harness matrix | Python and Rust had the same observable behavior |
`validate_sub_methods/` owns the future source-section inventory that maps a stable
Python qualified symbol to its Rust symbol. That inventory is the denominator for
per-function rollups; raw coverage for the entire LiteLLM repository would obscure
the port's real gaps. `unit_tests_rust/` owns direct `cargo-llvm-cov` runs, while
`e2e_fuzz_tests/` owns behavioral parity and fuzz-case counts. Keep Python, Rust, and
parity percentages visible side by side and label section confidence High only when
the mapped implementation exists, every required strategy passes, and both sides meet
their LOC thresholds. Generated Rust LCOV/HTML and the combined index also belong in
`target/rust-python-harness/`, not in a fourth strategy folder.
## Read the matrix
| Mark | Meaning |
| --- | --- |
| `✓` | All collected tests passed |
| `✗` | At least one test failed |
| `!` | Test setup or teardown failed |
| `↷` | All collected tests skipped |
| `?` | A configured selector did not collect a test |
| `—` | Strategy is planned but has no test yet |
| `n/a` | Strategy does not apply to this SDK function |
| `◐` | The configured tests cover only part of the TDD's parity contract |
The initial end-to-end entries deliberately show `◐`: the repository has Rust bridge tests for OCR, Messages, and Responses websocket plumbing, but those are not yet frozen-Python-oracle comparisons. The remaining TDD cells stay visible as planned work instead of disappearing from a green summary.
## Attach parity tests
Each of the four folders contains a concise `README.md` and a `strategy.json`. Add a pytest file or node ID to the appropriate SDK function's `selectors` list:
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": "complete",
"selectors": [
"tests/rust-python-harness/validate_sub_methods/test_messages.py"
]
"coverage": "partial",
"selectors": ["tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py"]
}
```
Selectors use the same syntax as pytest. A file selector aggregates every test in the file; a node selector can target one test or parametrized family; a selector ending in `/` aggregates every test in that folder, recursively. The runner deduplicates selectors, so one test may intentionally prove more than one cell without executing twice.
Selectors use pytest file or node syntax. A selector ending in `/` includes tests recursively from that directory
Use these coverage values:
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
- `complete`: implements the full strategy contract for that SDK function.
- `partial`: useful coverage exists, but the TDD contract is not fully proven.
- `planned`: no runnable parity test exists yet.
- `not_applicable`: the strategy cannot apply, such as streaming for OCR.
Unit cases use `unit_suite` instead of `selectors`, pointing to a repository-relative JSON file with this shape:
Keep comparison mechanics in shared harness modules and provider/function facts in the owning strategy folder. A Python/Rust mismatch is a test failure; do not normalize away observable return types, exception classes, private response fields, chunk ordering, or callback payload differences merely to make a cell green.
```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"}]
}
```
## Architecture
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
- `catalog.py` validates and loads every strategy manifest.
- `models.py` owns typed strategy, case, coverage, and run-state models.
- `runner.py` maps live pytest events back to one or more matrix cells.
- `ui.py` renders the interactive Rich dashboard and a dependency-free plain fallback.
- `cli.py` handles filtering and preserves pytest exit semantics.
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
The harness is driven from Python, matching the SDK surface and existing test tooling. Rust remains responsible for the implementation under comparison; the harness does not move provider semantics into the PyO3 bridge.
## 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

View file

@ -2,92 +2,74 @@ from __future__ import annotations
import json
from pathlib import Path
from typing import Any
from typing import Final
from .models import Coverage, HarnessCase, SDK_FUNCTIONS, Strategy
from pydantic import BaseModel, ConfigDict, ValidationError
STRATEGIES_ROOT = Path(__file__).parent
from .shared.reporting.models import Coverage, HarnessCase, SDK_FUNCTIONS, Strategy
STRATEGIES_ROOT: Final = Path(__file__).parent / "strategies"
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
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:
with source.open(encoding="utf-8") as stream:
data = json.load(stream)
strategy_id = _require_string(data.get("id"), "id", source)
label = _require_string(data.get("label"), "label", source)
description = _require_string(data.get("description"), "description", source)
order = data.get("order")
if not isinstance(order, int):
raise ValueError(f"{source}: order must be an integer")
function_data = data.get("functions")
if not isinstance(function_data, dict):
raise ValueError(f"{source}: functions must be an object")
missing = set(SDK_FUNCTIONS) - set(function_data)
extra = set(function_data) - set(SDK_FUNCTIONS)
if missing or extra:
raise ValueError(
f"{source}: functions must exactly match {SDK_FUNCTIONS}; missing={missing}, extra={extra}"
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,
)
cases: list[HarnessCase] = []
for sdk_function in SDK_FUNCTIONS:
case_data = function_data[sdk_function]
if not isinstance(case_data, dict):
raise ValueError(f"{source}: functions.{sdk_function} must be an object")
try:
coverage = Coverage(case_data.get("coverage"))
except ValueError as exc:
raise ValueError(f"{source}: invalid coverage for {sdk_function}") from exc
selectors = case_data.get("selectors", [])
if not isinstance(selectors, list) or not all(
isinstance(item, str) and item for item in selectors
):
raise ValueError(
f"{source}: selectors for {sdk_function} must be a list of strings"
)
if coverage is Coverage.NOT_APPLICABLE and selectors:
raise ValueError(
f"{source}: not_applicable case {sdk_function} cannot have selectors"
)
cases.append(
HarnessCase(
strategy_id=strategy_id,
strategy_label=label,
sdk_function=sdk_function,
coverage=coverage,
selectors=tuple(selectors),
note=str(case_data.get("note", "")),
)
)
return Strategy(
order=order,
id=strategy_id,
label=label,
description=description,
directory=source.parent,
cases=tuple(cases),
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 = sorted(root.glob("*/strategy.json"))
sources: Final = tuple(sorted(root.glob("*/strategy.json")))
if not sources:
raise ValueError(f"No strategy manifests found below {root}")
strategies = tuple(
sorted(
(_load_strategy(source) for source in sources),
key=lambda strategy: strategy.order,
)
)
ids = [strategy.id for strategy in strategies]
if len(ids) != len(set(ids)):
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

View file

@ -6,10 +6,14 @@ from collections.abc import Sequence
from pathlib import Path
from .catalog import load_catalog
from .models import SDK_FUNCTIONS, HarnessCase, Strategy
from .runner import run_pytest
from .ui import make_dashboard
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"
@ -44,6 +48,7 @@ def _parser() -> argparse.ArgumentParser:
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",
@ -135,9 +140,9 @@ def _print_catalog(strategies: Sequence[Strategy]) -> None:
print(f"{strategy.id:20} {strategy.label}")
for case in strategy.cases:
selectors = (
", ".join(case.selectors) if case.selectors else "no test configured"
", ".join(case.selectors) if case.selectors else case.unit_suite or "no test configured"
)
print(f" {case.sdk_function:12} {case.coverage.value:14} {selectors}")
print(f" {case.surface}/{case.sdk_function:12} {case.coverage.value:14} {selectors}")
def _print_function_report(report: FunctionReport) -> None:
@ -172,7 +177,21 @@ def _validate_ledger(sdk_functions: set[str]) -> int:
return 0 if all(report.is_clean for report in reports) else 1
def main(argv: Sequence[str] | None = None) -> int:
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(
@ -181,7 +200,8 @@ def main(argv: Sequence[str] | None = None) -> int:
)
if args.validate_ledger:
return _validate_ledger(set(args.sdk_functions))
strategies = load_catalog()
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
@ -194,7 +214,8 @@ def main(argv: Sequence[str] | None = None) -> int:
sdk_functions = sdk_functions or picked_functions
try:
cases = _select(strategies, strategy_ids, sdk_functions)
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}
@ -210,11 +231,12 @@ def main(argv: Sequence[str] | None = None) -> int:
if args.coverage:
pytest_args.extend(_coverage_pytest_args())
with dashboard:
exit_code, run = run_pytest(
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():

View file

@ -1,3 +0,0 @@
# End-to-end fuzz tests
Runs the same SDK call through the Python and Rust paths using generated inputs and recorded provider responses. It compares public results, streams, callbacks, and exceptions to catch behavior differences a unit test can miss.

View file

@ -1,14 +0,0 @@
{
"order": 10,
"id": "e2e_fuzz_tests",
"label": "End-to-end fuzz tests",
"description": "Compare observable Python and Rust SDK behavior over generated and recorded inputs.",
"functions": {
"ocr": {"coverage": "partial", "selectors": ["tests/test_litellm/ocr/test_rust_bridge.py"], "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added."},
"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."}
}
}

View file

@ -0,0 +1,91 @@
# Implementation parity testing through the SDK interface
> Given the same SDK call and identical provider behavior, do two implementations expose the same SDK contract?
## What the harness compares
- A fixture contains a LiteLLM SDK input and a recorded upstream provider response
- The same LiteLLM input is transformed by isolated baseline and candidate implementations
- The resulting provider requests must match in method, path, headers, and body, excluding runtime-specific HTTP metadata
- The recorded provider response is then replayed unchanged to both workers
- The harness compares the values returned through the Python SDK interface
- Non-streaming responses are compared directly, including their concrete return type and public model fields
- Streaming responses are consumed and compared chunk by chunk, including wrapper type, chunk type and order, termination, and public exception behavior
- Failed SDK calls are compared by exception class, stable message, status, code, model, provider, and parameter fields
- Traceback paths and line numbers are excluded because they are runtime-specific
- Route-specific comparators and chunk normalizers handle differences in each public SDK contract
## Process isolation
- SDK object and stream parity runs both implementations sequentially in the same process so tests can retain returned objects
- Every test saves and restores the original bridge state
- A small subprocess smoke test verifies environment-based startup configuration and detects fallback to the Python HTTP implementation
## Streaming execution
The invocation callback passed to `run_in_process` must consume the stream before returning its `StreamOutcome`.
Use `consume_sync_stream` inside that callback, or await `consume_async_stream` inside the callback passed to
`run_in_process_async`. Provider requests are collected only after the callback completes. Streaming is explicit:
an iterable return value alone does not select stream consumption
The consumers retain the wrapper type, iteration capabilities, chunk types and order, and any partial output before
an error. Errors retain their creation or iteration phase and the full public `SDKError` fields, with traceback text
removed. `capture_sync_stream` and `capture_async_stream` consume through the same helpers and then serialize the
outcome for subprocess reports. A serialization failure raises as a harness failure rather than becoming an SDK error
Response models and stream chunks share a recursive comparator. It compares concrete model, container, and scalar
types, public fields and extras, and exact values while ignoring Pydantic private attributes at every nesting level.
An API may supply an explicit chunk normalizer for its public contract
Shared tests exercise a local SSE provider through recording, VCR cassette storage, replay, and typed event comparison
in sync and async modes. They cover fragmented events, split UTF-8 characters, CRLF framing, coalesced events, and
application errors within a normally completed HTTP stream. HTTP byte boundaries and decoded SDK event boundaries
are checked separately
OCR remains the only integrated LiteLLM route. These tests validate shared streaming machinery, not another route's
SDK parity. Connection interruption, early cancellation, and lifecycle timeout enforcement remain outside this coverage
## Hypothesis and property-based testing
- Hypothesis is a Python library for property-based testing
- Example-based tests use inputs selected by the test author
- Property-based tests define strategies for valid inputs and properties that must hold for every generated example
- Hypothesis generates combinations from those strategies and normally shrinks a failing example to a smaller reproducible case
- In this harness, Hypothesis is used only during fixture generation to expand the LiteLLM input corpus
- Each API owns the strategies that vary its supported inputs
- Fixture generation is deterministic, and each generated input is recorded with the raw provider response it received
- The parity tests use committed fixtures and do not call the provider or generate new Hypothesis examples
- Provider responses are replayed unchanged, so the parity test does not fuzz or validate provider behavior
- Because Hypothesis does not run the parity assertion directly, parity failures are not automatically shrunk
## API-owned fixtures
The shared package owns recording, replay, persistence, execution, comparison, and route-neutral media constructors.
Each API package owns its input models, explicit strategies, provider targets, route-specific assets, fixture directory,
and regeneration command. See the API package documentation for its configured contracts and recording command
## VCR cassettes
Fixtures use VCR's YAML `version: 1` format with ordered request/response `interactions`. VCR handles text and binary
body serialization. Each cassette also contains `recorded_at`, `ttl_seconds: 0` (committed fixtures never expire), and
`x-litellm` metadata holding the SDK input and request provenance. Streaming responses carry
`x-litellm-chunk-lengths` so local replay preserves the original byte boundaries
The recording server captures requests before forwarding their responses. Saved requests use the stable
`http://parity-provider.invalid` origin and strip authentication headers and credential query parameters. The upstream
request keeps its credentials. Provider response bytes and non-success statuses are preserved
Standard VCR can load these files and replay their interactions. Parity tests keep using the local HTTP server because
Rust HTTP calls do not pass through VCR's Python patches. The harness still compares the two implementations' requests
against each other; the saved request is available for inspection and VCR playback, not a new parity assertion
Refresh parity cassettes through the API's recording command. Generic VCR writers do not preserve the SDK metadata
Legacy JSON fixtures remain readable. Migrated cassettes mark reconstructed requests as `python_replay`; fresh
recordings use `recorded`. The metadata extensions follow the filesystem cassette layout proposed in
[PR #39338](https://github.com/BerriAI/litellm/pull/39338), without depending on its unmerged persistence backend
## References
- [Hypothesis documentation](https://hypothesis.readthedocs.io/en/latest/)
- [Hypothesis documentation source](https://github.com/HypothesisWorks/hypothesis/tree/master/hypothesis/docs)

View file

@ -0,0 +1,3 @@
import pytest
pytest.register_assert_rewrite("tests.rust-python-harness.shared.parity.compare")

View file

@ -0,0 +1,80 @@
from __future__ import annotations
from collections.abc import Mapping, Sequence
from typing import Final, cast
from pydantic import BaseModel
from .models import CapturedRequest, Execution
def validate_harness(baseline: Execution, candidate: Execution, baseline_user_agent: str) -> None:
for request in baseline.requests:
if request.user_agent != baseline_user_agent:
raise AssertionError(
f"baseline provider request did not carry sentinel user-agent {baseline_user_agent!r}: "
f"{request.user_agent!r}"
)
for request in candidate.requests:
if request.user_agent == baseline_user_agent:
raise AssertionError("candidate route fell back to the baseline HTTP implementation")
def _request_after_transformation(request: CapturedRequest) -> CapturedRequest:
return request.model_copy(update={"user_agent": None})
def assert_request_parity(baseline: tuple[CapturedRequest, ...], candidate: tuple[CapturedRequest, ...]) -> None:
baseline_requests: Final = tuple(_request_after_transformation(request) for request in baseline)
candidate_requests: Final = tuple(_request_after_transformation(request) for request in candidate)
assert_value_parity(baseline_requests, candidate_requests)
def _public_model_values(model: BaseModel) -> dict[str, object]:
fields: Final = (*type(model).model_fields, *type(model).model_computed_fields)
extras: Final = cast(Mapping[str, object], model.model_extra or {})
return {
**{name: cast(object, getattr(model, name)) for name in fields if not name.startswith("_")},
**{name: value for name, value in extras.items() if not name.startswith("_")},
}
def assert_model_parity(baseline: BaseModel, candidate: BaseModel) -> None:
assert_value_parity(baseline, candidate)
def assert_value_parity(baseline: object, candidate: object, *, path: str = "$") -> None:
assert type(baseline) is type(candidate), f"type mismatch at {path}: {type(baseline)} != {type(candidate)}"
if isinstance(baseline, BaseModel) and isinstance(candidate, BaseModel):
assert_value_parity(_public_model_values(baseline), _public_model_values(candidate), path=path)
return
if isinstance(baseline, Mapping) and isinstance(candidate, Mapping):
baseline_mapping: Final = cast(Mapping[object, object], baseline)
candidate_mapping: Final = cast(Mapping[object, object], candidate)
assert frozenset((type(key), key) for key in baseline_mapping) == frozenset(
(type(key), key) for key in candidate_mapping
), f"mapping keys differ at {path}"
for key in baseline_mapping:
assert_value_parity(baseline_mapping[key], candidate_mapping[key], path=f"{path}.{key}")
return
if (
isinstance(baseline, Sequence)
and not isinstance(baseline, (str, bytes))
and isinstance(candidate, Sequence)
and not isinstance(candidate, (str, bytes))
):
baseline_sequence: Final = cast(Sequence[object], baseline)
candidate_sequence: Final = cast(Sequence[object], candidate)
assert len(baseline_sequence) == len(candidate_sequence), f"sequence lengths differ at {path}"
for index, (baseline_item, candidate_item) in enumerate(
zip(baseline_sequence, candidate_sequence, strict=True)
):
assert_value_parity(baseline_item, candidate_item, path=f"{path}[{index}]")
return
assert baseline == candidate, f"value mismatch at {path}: {baseline!r} != {candidate!r}"
def assert_parity(baseline: Execution, candidate: Execution, baseline_user_agent: str) -> None:
validate_harness(baseline, candidate, baseline_user_agent)
assert_request_parity(baseline.requests, candidate.requests)
assert_value_parity(baseline.report, candidate.report)

View file

@ -0,0 +1,64 @@
from __future__ import annotations
from collections.abc import Mapping
from typing import ClassVar, Final, Generic, Literal, TypeVar, cast
from pydantic import BaseModel, ConfigDict, Field, JsonValue, model_validator
from .recorded_http import RecordedResponse
JsonObject = dict[str, JsonValue]
class FixtureModel(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid", populate_by_name=True, serialize_by_alias=True)
class SdkInputBase(FixtureModel):
fixture_only_fields: ClassVar[tuple[str, ...]] = ()
def as_sdk_kwargs(self) -> dict[str, object]:
return cast(
dict[str, object],
self.model_dump(
mode="python",
exclude_unset=True,
exclude=set(self.fixture_only_fields),
),
)
def canonical_input(self) -> dict[str, object]:
dumped: Final = cast(dict[str, object], self.model_dump(mode="json", exclude_unset=True))
fixture_fields: Final = {field: getattr(self, field) for field in self.fixture_only_fields}
return {**fixture_fields, **dumped}
class JsonSchemaDefinition(FixtureModel):
name: str
description: str | None = None
schema_definition: JsonObject = Field(alias="schema")
strict: bool = False
class JsonSchemaResponseFormat(FixtureModel):
type: Literal["json_schema"]
json_schema: JsonSchemaDefinition
InputT = TypeVar("InputT", bound=SdkInputBase)
class ParityCase(FixtureModel, Generic[InputT]):
litellm_input: InputT
provider_responses: tuple[RecordedResponse, ...]
@model_validator(mode="before")
@classmethod
def load_legacy_single_response(cls, value: object) -> object:
if not isinstance(value, Mapping):
return value
migrated: Final = dict(cast(Mapping[str, object], value))
provider_response: Final = migrated.pop("provider_response", None)
if "provider_responses" not in migrated and provider_response is not None:
migrated["provider_responses"] = (provider_response,)
return migrated

View file

@ -0,0 +1 @@
from __future__ import annotations

View file

@ -0,0 +1,147 @@
from __future__ import annotations
from collections.abc import Mapping
from datetime import datetime
from itertools import accumulate
from typing import Final, Literal
from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, TypeAdapter
from vcr.serialize import serialize
from vcr.serializers import yamlserializer
from .recording import RecordedInteraction
from ..recorded_http import (
HttpHeader,
RecordedHttpResponse,
RecordedHttpStreamResponse,
RecordedResponse,
RecordedStreamChunk,
)
_OBJECT: Final = TypeAdapter(dict[str, object])
class _CassetteModel(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid", populate_by_name=True)
class _Body(_CassetteModel):
string: str | bytes
def as_bytes(self) -> bytes:
return self.string.encode("utf-8") if isinstance(self.string, str) else self.string
class _Status(_CassetteModel):
code: int
message: str
class _Request(_CassetteModel):
method: str
uri: str
body: str | bytes | None
headers: dict[str, tuple[str, ...]]
class _Response(_CassetteModel):
status: _Status
headers: dict[str, tuple[str, ...]]
body: _Body
chunk_lengths: tuple[int, ...] | None = Field(default=None, alias="x-litellm-chunk-lengths")
def recorded_response(self) -> RecordedResponse:
headers: Final = tuple(
HttpHeader(name=name, value=value) for name, values in self.headers.items() for value in values
)
body: Final = self.body.as_bytes()
if self.chunk_lengths is None:
return RecordedHttpResponse.from_bytes(self.status.code, headers, body)
if any(length < 0 for length in self.chunk_lengths) or sum(self.chunk_lengths) != len(body):
raise ValueError("cassette stream chunk lengths do not match the response body")
offsets: Final = tuple(accumulate(self.chunk_lengths, initial=0))
return RecordedHttpStreamResponse(
kind="http_stream",
status_code=self.status.code,
headers=headers,
chunks=tuple(RecordedStreamChunk.from_bytes(body[start:end]) for start, end in zip(offsets, offsets[1:])),
)
class _Interaction(_CassetteModel):
request: _Request
response: _Response
class _ParityMetadata(_CassetteModel):
schema_version: Literal[1]
request_source: Literal["recorded", "python_replay"]
case: dict[str, object]
class ParityCassette(_CassetteModel):
version: Literal[1]
recorded_at: AwareDatetime
ttl_seconds: Literal[0]
interactions: tuple[_Interaction, ...]
parity: _ParityMetadata = Field(alias="x-litellm")
def case_data(self) -> dict[str, object]:
return {
**self.parity.case,
"provider_responses": tuple(item.response.recorded_response() for item in self.interactions),
}
def _response_dict(response: RecordedResponse) -> dict[str, object]:
headers: Final = {
name: [header.value for header in response.headers if header.name == name]
for name in dict.fromkeys(header.name for header in response.headers)
}
chunks: Final = (
tuple(chunk.data_bytes() for chunk in response.chunks)
if isinstance(response, RecordedHttpStreamResponse)
else None
)
body: Final = response.body_bytes() if isinstance(response, RecordedHttpResponse) else b"".join(chunks or ())
return {
"status": {"code": response.status_code, "message": ""},
"headers": headers,
"body": {"string": body},
**({"x-litellm-chunk-lengths": list(map(len, chunks))} if chunks is not None else {}),
}
def serialize_cassette(
case: Mapping[str, object],
interactions: tuple[RecordedInteraction, ...],
recorded_at: datetime,
request_source: Literal["recorded", "python_replay"],
) -> str:
normalized: Final = _OBJECT.validate_python(
yamlserializer.deserialize(
serialize(
{
"requests": [item.request for item in interactions],
"responses": [_response_dict(item.response) for item in interactions],
},
yamlserializer,
)
)
)
payload: Final = {
**normalized,
"recorded_at": recorded_at.isoformat(),
"ttl_seconds": 0,
"x-litellm": {
"schema_version": 1,
"request_source": request_source,
"case": {key: value for key, value in case.items() if key != "provider_responses"},
},
}
ParityCassette.model_validate(payload).case_data()
return str(yamlserializer.serialize(payload))
def deserialize_cassette(contents: str) -> ParityCassette:
return ParityCassette.model_validate(yamlserializer.deserialize(contents))

View file

@ -0,0 +1,34 @@
from __future__ import annotations
import argparse
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Final, cast
@dataclass(frozen=True, slots=True)
class RecordingArgs:
concurrency: int
examples: int
fixture_dir: Path | None
def _positive_int(value: str) -> int:
parsed: Final = int(value)
if parsed < 1:
raise argparse.ArgumentTypeError("must be at least 1")
return parsed
def parse_recording_args(argv: Sequence[str] | None = None) -> RecordingArgs:
parser: Final = argparse.ArgumentParser()
parser.add_argument("--concurrency", type=_positive_int, default=2)
parser.add_argument("--examples", type=_positive_int, default=4)
parser.add_argument("--fixture-dir", type=Path)
namespace: Final = parser.parse_args(argv)
return RecordingArgs(
concurrency=cast(int, namespace.concurrency),
examples=cast(int, namespace.examples),
fixture_dir=cast(Path | None, namespace.fixture_dir),
)

View file

@ -0,0 +1,22 @@
from __future__ import annotations
import queue
from typing import Final, TypeVar
from hypothesis import given, settings
from hypothesis.strategies import SearchStrategy
InputT = TypeVar("InputT")
def generate_case_inputs(strategy: SearchStrategy[InputT], examples: int) -> tuple[InputT, ...]:
generated: Final[queue.SimpleQueue[InputT | None]] = queue.SimpleQueue()
@settings(max_examples=examples, deadline=None, derandomize=True)
@given(case_input=strategy)
def generate_case(case_input: InputT) -> None:
generated.put(case_input)
generate_case()
generated.put(None)
return tuple(iter(generated.get, None))

View file

@ -0,0 +1,225 @@
from __future__ import annotations
import base64
from functools import cache
from io import BytesIO
from typing import Final
from urllib.parse import quote
from PIL import Image, ImageDraw
from reportlab.graphics.barcode import code128 # pyright: ignore[reportMissingTypeStubs] # ReportLab has no stubs
from reportlab.lib import colors # pyright: ignore[reportMissingTypeStubs] # ReportLab has no stubs
from reportlab.lib.pagesizes import letter # pyright: ignore[reportMissingTypeStubs] # ReportLab has no stubs
from reportlab.lib.utils import ImageReader # pyright: ignore[reportMissingTypeStubs] # ReportLab has no stubs
from reportlab.pdfgen import canvas # pyright: ignore[reportMissingTypeStubs] # ReportLab has no stubs
def dummy_image_url(text: str, font_size: int, width: int = 800, height: int = 300) -> str:
return f"https://dummyjson.com/image/{width}x{height}/ffffff/000000?text={quote(text)}&fontSize={font_size}"
_GLYPHS: Final = {
"D": ("11110", "10001", "10001", "10001", "10001", "10001", "11110"),
"O": ("01110", "10001", "10001", "10001", "10001", "10001", "01110"),
"C": ("01111", "10000", "10000", "10000", "10000", "10000", "01111"),
"1": ("00100", "01100", "00100", "00100", "00100", "00100", "01110"),
"2": ("01110", "10001", "00001", "00010", "00100", "01000", "11111"),
"3": ("11110", "00001", "00001", "01110", "00001", "00001", "11110"),
}
@cache
def structured_image_bytes() -> bytes:
image: Final = Image.new("RGB", (320, 80), "white")
draw: Final = ImageDraw.Draw(image)
scale: Final = 8
cursor_x = 24
for character in "DOC 123":
if character == " ":
cursor_x += scale * 3
continue
for glyph_y, row in enumerate(_GLYPHS[character]):
for glyph_x, filled in enumerate(row):
if filled == "1":
x = cursor_x + glyph_x * scale
y = 12 + glyph_y * scale
draw.rectangle((x, y, x + scale - 1, y + scale - 1), fill="black")
cursor_x += scale * 6
output: Final = BytesIO()
image.save(output, format="PNG")
return output.getvalue()
@cache
def structured_image_data_uri() -> str:
encoded: Final = base64.b64encode(structured_image_bytes()).decode("ascii")
return f"data:image/png;base64,{encoded}"
def _draw_header(pdf: canvas.Canvas, title: str, page_number: int) -> None:
pdf.setFillColor(colors.black)
pdf.setFont("Helvetica", 11)
pdf.drawString(45, 770, "Quarterly Operations Report")
pdf.setFont("Helvetica-Bold", 16)
pdf.drawString(45, 745, title)
pdf.setFont("Helvetica", 9)
pdf.drawString(45, 30, f"Confidential | Page {page_number} of 5")
def _draw_body(pdf: canvas.Canvas, page_number: int) -> None:
pdf.setFont("Helvetica", 10)
for line_number in range(1, 9):
pdf.drawString(
45,
500 - (line_number * 28),
f"Section {page_number}.{line_number}: Invoice totals, regional revenue, and reconciliation notes.",
)
def _diagram_image(width: int, height: int, accent: tuple[int, int, int]) -> Image.Image:
image: Final = Image.new("RGB", (width, height), (242, 246, 252))
draw: Final = ImageDraw.Draw(image)
for coordinate in range(0, max(width, height), 40):
draw.line((coordinate, 0, coordinate, height), fill=(32, 32, 32), width=3)
draw.line((0, coordinate, width, coordinate), fill=(32, 32, 32), width=3)
draw.line((0, 0, width, height), fill=accent, width=8)
draw.line((width, 0, 0, height), fill=accent, width=8)
draw.rectangle((width // 4, height // 4, width * 3 // 4, height * 3 // 4), outline=accent, width=6)
return image
def _draw_embedded_images(pdf: canvas.Canvas) -> None:
images: Final = (
(_diagram_image(320, 320, (51, 115, 217)), 455, 655, 70, 70),
(_diagram_image(360, 320, (38, 151, 92)), 455, 565, 70, 62),
(_diagram_image(120, 120, (219, 68, 55)), 455, 500, 45, 45),
)
for image, x, y, width, height in images:
pdf.drawImage( # pyright: ignore[reportUnknownMemberType] # ReportLab has no stubs
ImageReader(image), x, y, width=width, height=height, mask="auto"
)
def _draw_table_page(pdf: canvas.Canvas) -> None:
columns: Final = (45, 245, 405, 565)
tables: Final = (
(
(730, 695, 660, 625),
(
("Item", "Quantity", "Amount", 707),
("Document analysis", "2", "120.00", 672),
("Document verification", "1", "80.00", 637),
),
),
(
(600, 565, 530, 495),
(
("Item continued", "Quantity", "Amount", 577),
("Fixture validation", "3", "45.00", 542),
("Provider review", "1", "25.00", 507),
),
),
)
for rows, values in tables:
for x in columns:
pdf.line(x, rows[-1], x, rows[0])
for y in rows:
pdf.line(45, y, 565, y)
for item, quantity, amount, y in values:
pdf.drawString(55, y, item)
pdf.drawString(255, y, quantity)
pdf.drawString(415, y, amount)
def _draw_chart_page(pdf: canvas.Canvas) -> None:
bars: Final = ((70, 70), (170, 115), (270, 90), (370, 130))
pdf.setFillColor(colors.HexColor("#3373D9"))
for x, height in bars:
pdf.rect(x, 610, 65, height, fill=1, stroke=0)
pdf.setFillColor(colors.black)
for quarter, x in zip(("Q1", "Q2", "Q3", "Q4"), (90, 190, 290, 390), strict=True):
pdf.drawString(x, 590, quarter)
pdf.drawString(45, 550, "Formula: gross margin = (revenue - cost) / revenue")
_draw_embedded_images(pdf)
def _draw_metadata_page(pdf: canvas.Canvas) -> None:
pdf.setFont("Helvetica", 12)
pdf.drawString(45, 700, "Invoice Number: INV-2048")
pdf.drawString(45, 675, "Purchase Order: PO-4096")
pdf.setFillColor(colors.HexColor("#F2E65A"))
pdf.rect(40, 555, 500, 24, fill=1, stroke=0)
pdf.setFillColor(colors.black)
pdf.drawString(45, 560, "Highlighted total requiring review")
pdf.drawString(45, 530, "Reviewer comment: verify the highlighted total before approval")
pdf.setFillColor(colors.red)
pdf.drawString(45, 495, "Revised total: 245.00")
pdf.line(45, 501, 150, 501)
pdf.setFillColor(colors.black)
pdf.linkURL( # pyright: ignore[reportUnknownMemberType] # ReportLab has no stubs
"https://example.com/invoices/INV-2048", (45, 575, 300, 590), relative=0
)
pdf.highlightAnnotation( # pyright: ignore[reportUnknownMemberType] # ReportLab has no stubs
"Total highlighted for review",
Rect=(40, 555, 540, 579),
QuadPoints=(40, 579, 540, 579, 40, 555, 540, 555),
)
pdf.textAnnotation( # pyright: ignore[reportUnknownMemberType] # ReportLab has no stubs
"Verify the highlighted total", Rect=(520, 525, 540, 545)
)
pdf.drawString(45, 575, "https://example.com/invoices/INV-2048")
barcode: Final = code128.Code128("5901234123457", barHeight=70, barWidth=1.2)
barcode.drawOn(pdf, 90, 130)
def _draw_signature_page(pdf: canvas.Canvas) -> None:
pdf.saveState()
pdf.setFillColor(colors.lightgrey)
pdf.setFont("Helvetica-Bold", 54)
pdf.translate(110, 390)
pdf.rotate(25)
pdf.drawString(0, 0, "DRAFT")
pdf.restoreState()
pdf.setFillColor(colors.black)
pdf.setFont("Helvetica", 12)
pdf.drawString(45, 635, "Approved by: Jordan Lee")
pdf.line(45, 610, 310, 610)
pdf.bezier(55, 595, 75, 625, 112, 602, 155, 600)
pdf.drawString(45, 580, "Signature")
def _draw_appendix_page(pdf: canvas.Canvas) -> None:
pdf.setFont("Helvetica-Bold", 14)
pdf.drawString(45, 700, "1. Scope")
pdf.drawString(45, 650, "2. Findings")
pdf.drawString(45, 600, "3. Recommendations")
@cache
def structured_pdf_bytes() -> bytes:
output: Final = BytesIO()
pdf: Final = canvas.Canvas(output, pagesize=letter, pageCompression=0, invariant=1)
pdf.setTitle("Quarterly Operations Report")
pdf.setAuthor("LiteLLM parity fixture generator")
pdf.setSubject("Semantic document coverage for tables, figures, annotations, and metadata")
pdf.setKeywords("document, invoice, table, figure, annotation")
pages: Final = (
("Invoice Summary and Line Items", _draw_table_page),
("Revenue Chart and Formula Review", _draw_chart_page),
("Key Values, Link, Highlight, and Comment", _draw_metadata_page),
("Approval Signature and Watermark", _draw_signature_page),
("Appendix with Section Boundaries", _draw_appendix_page),
)
for page_number, (title, draw_page) in enumerate(pages, start=1):
_draw_header(pdf, title, page_number)
draw_page(pdf)
_draw_body(pdf, page_number)
pdf.showPage()
pdf.save()
return output.getvalue()
@cache
def structured_pdf_data_uri() -> str:
encoded: Final = base64.b64encode(structured_pdf_bytes()).decode("ascii")
return f"data:application/pdf;base64,{encoded}"

View file

@ -0,0 +1,201 @@
from __future__ import annotations
import logging
from concurrent.futures import Future, ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field
from pathlib import Path
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 (
FixtureInput,
canonical_json,
fixture_cache_key,
fixture_id,
fixture_path,
load_fixture,
save_fixture,
)
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]):
def execute(self, provider_url: str, case_input: InputT_contra) -> None: ...
@dataclass(frozen=True, slots=True)
class RecordingTarget(Generic[InputT]):
name: str
upstream: UpstreamEndpoint
strategy: SearchStrategy[InputT]
invocation: RecordingInvocation[InputT] = field(repr=False)
required_inputs: tuple[InputT, ...] = ()
@dataclass(frozen=True, slots=True)
class RecordingJob(Generic[InputT]):
target_name: str
directory: Path
upstream: UpstreamEndpoint
case_input: InputT
invocation: RecordingInvocation[InputT] = field(repr=False)
@property
def case_id(self) -> str:
return fixture_id(self.case_input, self.target_name)
@dataclass(frozen=True, slots=True)
class RecordedFixture:
target_name: str
case_id: str
path: Path
kind: Literal["recorded"] = field(default="recorded", init=False)
@dataclass(frozen=True, slots=True)
class CachedFixture:
target_name: str
case_id: str
path: Path
kind: Literal["cached"] = field(default="cached", init=False)
@dataclass(frozen=True, slots=True)
class FailedFixture:
target_name: str
case_id: str
error: Exception = field(repr=False)
kind: Literal["failed"] = field(default="failed", init=False)
RecordingOutcome = RecordedFixture | CachedFixture | FailedFixture
@dataclass(frozen=True, slots=True)
class RecordingSummary:
recorded: tuple[RecordedFixture, ...]
cached: tuple[CachedFixture, ...]
failed: tuple[FailedFixture, ...]
@property
def exit_code(self) -> int:
return 1 if self.failed else 0
def _unique_inputs(target: RecordingTarget[InputT], examples: int) -> tuple[InputT, ...]:
generated_inputs: Final = generate_case_inputs(target.strategy, examples)
case_inputs: Final = (*target.required_inputs, *generated_inputs)
return tuple({canonical_json(fixture_cache_key(case_input)): case_input for case_input in case_inputs}.values())
def build_recording_jobs(
targets: tuple[RecordingTarget[InputT], ...],
root: Path,
examples: int,
) -> tuple[RecordingJob[InputT], ...]:
if examples < 1:
raise ValueError("examples must be at least 1")
return tuple(
RecordingJob(
target_name=target.name,
directory=root / target.name,
upstream=target.upstream,
case_input=case_input,
invocation=target.invocation,
)
for target in targets
for case_input in _unique_inputs(target, examples)
)
def _record_job(job: RecordingJob[InputT], case_type: type[CaseT]) -> RecordedFixture | CachedFixture:
cached: Final = load_fixture(job.directory, job.case_input, case_type)
if cached is not None:
path: Final = fixture_path(job.directory, job.case_input)
return CachedFixture(
target_name=job.target_name,
case_id=job.case_id,
path=path if path.is_file() else path.with_suffix(".json"),
)
interactions: Final = record_upstream_interactions(
job.upstream,
job.case_input,
job.invocation.execute,
)
status: Final = interactions[-1].response.status_code
if status in {408, 429} or status >= 500:
raise RuntimeError(f"Upstream returned transient HTTP {status}; rerun recording to retry")
case: Final = case_type.model_validate(
{
"litellm_input": job.case_input,
"provider_responses": tuple(item.response for item in interactions),
}
)
saved_path: Final = save_fixture(job.directory, job.case_input, case, interactions)
return RecordedFixture(target_name=job.target_name, case_id=job.case_id, path=saved_path)
def _completed_outcome(
completed: int,
total: int,
job: RecordingJob[InputT],
future: Future[RecordedFixture | CachedFixture],
) -> RecordingOutcome:
try:
outcome: Final = future.result()
except Exception as error:
failed: Final = FailedFixture(target_name=job.target_name, case_id=job.case_id, error=error)
LOGGER.error(
"[%d/%d] failed %s %s: %s",
completed,
total,
failed.target_name,
failed.case_id,
type(error).__name__,
)
return failed
LOGGER.info("[%d/%d] %s %s %s", completed, total, outcome.kind, outcome.target_name, outcome.case_id)
return outcome
def record_fixtures(
targets: tuple[RecordingTarget[InputT], ...],
root: Path,
examples: int,
concurrency: int,
case_type: type[CaseT],
) -> RecordingSummary:
if concurrency < 1:
raise ValueError("concurrency must be at least 1")
jobs: Final = build_recording_jobs(targets, root, examples)
total: Final = len(jobs)
LOGGER.info("Recording %d fixtures across %d targets with concurrency %d", total, len(targets), concurrency)
with ThreadPoolExecutor(max_workers=concurrency) as executor:
future_jobs: Final = MappingProxyType({executor.submit(_record_job, job, case_type): job for job in jobs})
outcomes: Final = tuple(
_completed_outcome(completed, total, future_jobs[future], future)
for completed, future in enumerate(as_completed(future_jobs), start=1)
)
summary: Final = RecordingSummary(
recorded=tuple(outcome for outcome in outcomes if isinstance(outcome, RecordedFixture)),
cached=tuple(outcome for outcome in outcomes if isinstance(outcome, CachedFixture)),
failed=tuple(outcome for outcome in outcomes if isinstance(outcome, FailedFixture)),
)
LOGGER.info(
"Finished %d fixtures: %d recorded, %d cached, %d failed",
total,
len(summary.recorded),
len(summary.cached),
len(summary.failed),
)
return summary

View file

@ -0,0 +1,66 @@
from __future__ import annotations
import os
from collections.abc import Callable
from pathlib import Path
from typing import Final, TypeVar
import pytest
from pydantic import BaseModel, ValidationError
from .store import recorded_fixtures
CaseT = TypeVar("CaseT", bound=BaseModel)
def parametrize_recorded_fixtures(
metafunc: pytest.Metafunc,
*,
fixture_name: str,
case_type: type[CaseT],
env_var: str,
default_directory: Path,
regeneration_command: str,
id_builder: Callable[[CaseT], str],
marks_builder: Callable[[CaseT], tuple[pytest.MarkDecorator, ...]] | None = None,
) -> None:
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:
fixtures: Final = recorded_fixtures(directory, case_type)
except (ValidationError, ValueError) as error:
raise pytest.UsageError(
f"Invalid parity fixture bundle at {directory}. "
"Each fixture must use the current versioned envelope. "
f"Record fresh fixtures in an empty directory with: `{regeneration_command}`. "
f"Validation details: {error}"
) from error
if fixtures:
metafunc.parametrize(
fixture_name,
tuple(
pytest.param(
fixture,
id=id_builder(fixture),
marks=marks_builder(fixture) if marks_builder is not None else (),
)
for fixture in fixtures
),
)
return
if configured is not None:
raise pytest.UsageError(f"no recorded fixtures in {directory}")
metafunc.parametrize(
fixture_name,
(
pytest.param(
None,
marks=pytest.mark.skip(reason=f"no recorded fixtures in {directory}"),
id="no-recorded-fixtures",
),
),
)

View file

@ -0,0 +1,267 @@
from __future__ import annotations
import queue
import threading
from collections.abc import Callable, Generator, Iterable
from contextlib import contextmanager
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 (
dropped_request_headers,
dropped_response_headers,
is_streaming_response,
)
from ..recorded_http import (
HttpHeader,
RecordedHttpResponse,
RecordedHttpStreamResponse,
RecordedResponse,
RecordedStreamChunk,
)
_PARITY_PROVIDER_HOST: Final = "parity-provider.invalid"
_SECRET_HEADERS: Final = frozenset(
{
"authorization",
"proxy-authorization",
"cookie",
"x-api-key",
"api-key",
"anthropic-api-key",
"openai-api-key",
"azure-api-key",
"x-goog-api-key",
"ocp-apim-subscription-key",
"x-amz-security-token",
}
)
InputT = TypeVar("InputT")
@dataclass(frozen=True, slots=True)
class UpstreamEndpoint:
base_url: str
@dataclass(frozen=True, slots=True)
class RecordedInteraction:
request: Request
response: RecordedResponse
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))
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
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)
except queue.Empty as error:
raise RuntimeError("successful SDK call did not produce a recorded response") from error
remaining: Final = tuple(self.interactions.get_nowait() for _ in range(self.interactions.qsize()))
return (first, *remaining)
class _RecordingHandler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
def do_POST(self) -> None:
self._forward()
def do_GET(self) -> None:
self._forward()
def do_PUT(self) -> None:
self._forward()
def do_PATCH(self) -> None:
self._forward()
def do_DELETE(self) -> None:
self._forward()
def _forward(self) -> None:
provider: Final = self.server
assert isinstance(provider, _RecordingProvider)
length: Final = int(self.headers.get("content-length") or "0")
request_body: Final = self.rfile.read(length) if length else b""
raw_headers: Final = tuple(self.headers.raw_items())
excluded: Final = dropped_request_headers(raw_headers)
forwarded_headers: Final = tuple((name, value) for name, value in raw_headers if name.lower() not in excluded)
upstream_url: Final = f"{provider.spec.base_url.rstrip('/')}{self.path}"
try:
with httpx.stream(
self.command,
upstream_url,
headers=forwarded_headers,
content=request_body,
timeout=120,
) as upstream:
headers: Final = _end_to_end_headers(upstream.headers)
recorded_response: Final = self._record_upstream_response(upstream, headers)
except httpx.HTTPError as error:
self._send_response(502, (), str(error).encode("utf-8"))
return
recorded_request: Final = remove_query_parameters(
Request(
self.command,
f"http://{_PARITY_PROVIDER_HOST}{self.path}",
request_body,
{name: value for name, value in forwarded_headers if name.lower() not in _SECRET_HEADERS},
),
("api_key", "api-key", "key", "access_token", "subscription-key"),
)
provider.interactions.put(RecordedInteraction(recorded_request, recorded_response))
if isinstance(recorded_response, RecordedHttpResponse):
self._send_response(
recorded_response.status_code, recorded_response.headers, recorded_response.body_bytes()
)
def _record_upstream_response(
self,
upstream: httpx.Response,
headers: tuple[HttpHeader, ...],
) -> RecordedResponse:
content_type: Final = cast(str, upstream.headers.get("content-type", ""))
if is_streaming_response(content_type):
return self._record_stream(upstream, headers)
response_body: Final = b"".join(upstream.iter_bytes())
return RecordedHttpResponse.from_bytes(
status_code=upstream.status_code,
headers=headers,
body=response_body,
)
def _record_stream(
self,
upstream: httpx.Response,
headers: tuple[HttpHeader, ...],
) -> RecordedHttpStreamResponse:
self.send_response_only(upstream.status_code)
provider: Final = self.server
assert isinstance(provider, _RecordingProvider)
for header in headers:
self.send_header(header.name, local_response_header(header.name, header.value, provider.url))
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()
return RecordedHttpStreamResponse(
kind="http_stream",
status_code=upstream.status_code,
headers=headers,
chunks=chunks,
)
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()
yield RecordedStreamChunk.from_bytes(chunk)
def _send_response(self, status_code: int, headers: tuple[HttpHeader, ...], body: bytes) -> None:
self.send_response_only(status_code)
provider: Final = self.server
assert isinstance(provider, _RecordingProvider)
for header in headers:
self.send_header(header.name, local_response_header(header.name, header.value, provider.url))
self.send_header("content-length", str(len(body)))
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 _invoke_and_take_interactions(
recorder: _RecordingProvider,
case_input: InputT,
sdk_call: Callable[[str, InputT], object],
) -> tuple[RecordedInteraction, ...]:
try:
sdk_call(recorder.url, case_input)
except Exception as invocation_error:
try:
return recorder.take_interactions()
except RuntimeError:
raise invocation_error
return recorder.take_interactions()
def record_upstream_interactions(
spec: UpstreamEndpoint,
case_input: InputT,
sdk_call: Callable[[str, InputT], object],
) -> tuple[RecordedInteraction, ...]:
with _recording_provider(spec) as recorder:
return _invoke_and_take_interactions(recorder, case_input, sdk_call)
def record_upstream_responses(
spec: UpstreamEndpoint,
case_input: InputT,
sdk_call: Callable[[str, InputT], object],
) -> tuple[RecordedResponse, ...]:
return tuple(item.response for item in record_upstream_interactions(spec, case_input, sdk_call))

View file

@ -0,0 +1,126 @@
from __future__ import annotations
import hashlib
import json
import tempfile
from collections.abc import Mapping
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 .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]: ...
CaseT = TypeVar("CaseT", bound=BaseModel)
class FixtureEnvelope(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
schema_version: int
recorded_at: AwareDatetime
case: dict[str, object]
def canonical_json(value: Mapping[str, object]) -> str:
return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True)
def fixture_cache_key(case_input: FixtureInput) -> dict[str, object]:
return case_input.canonical_input()
def fixture_path(directory: Path, case_input: FixtureInput) -> Path:
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"
def load_fixture(directory: Path, case_input: FixtureInput, case_type: type[CaseT]) -> CaseT | None:
path: Final = fixture_path(directory, case_input)
if path.is_file():
return read_fixture(path, case_type)
legacy_path: Final = path.with_suffix(".json")
if not legacy_path.is_file():
return None
return read_fixture(legacy_path, case_type)
def save_fixture(
directory: Path,
case_input: FixtureInput,
case: BaseModel,
interactions: tuple[RecordedInteraction, ...],
*,
recorded_at: datetime | None = None,
request_source: Literal["recorded", "python_replay"] = "recorded",
) -> Path:
directory.mkdir(parents=True, exist_ok=True)
path: Final = fixture_path(directory, case_input)
serialized: Final = serialize_cassette(
cast(dict[str, object], case.model_dump(mode="json", exclude_unset=True)),
interactions,
recorded_at or datetime.now(timezone.utc),
request_source,
)
with tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", dir=directory, delete=False) as temporary:
temporary_path: Final = Path(temporary.name)
try:
temporary.write(serialized)
temporary.close()
temporary_path.replace(path)
finally:
temporary_path.unlink(missing_ok=True)
return path
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)
try:
cassette: Final = deserialize_cassette(contents)
return case_type.model_validate(cassette.case_data())
except ValueError as error:
raise ValueError(f"invalid parity cassette {path}") from error
def _load_fixture(raw_fixture: dict[str, object], path: Path, case_type: type[CaseT]) -> CaseT:
schema_version: Final = raw_fixture.get("schema_version")
if schema_version != FIXTURE_SCHEMA_VERSION:
raise ValueError(
f"fixture {path} has schema_version {schema_version!r}, expected {FIXTURE_SCHEMA_VERSION}; "
"delete it and regenerate the fixture bundle"
)
try:
envelope: Final = FixtureEnvelope.model_validate(raw_fixture)
return case_type.model_validate(envelope.case)
except ValidationError as error:
raise ValueError(f"invalid parity fixture {path} ({len(error.errors())} validation errors)") from error
def recorded_fixtures(directory: Path, case_type: type[CaseT]) -> tuple[CaseT, ...]:
if not directory.is_dir():
return ()
paths: Final = tuple(sorted((*directory.rglob("*.yaml"), *directory.rglob("*.json"))))
return tuple(read_fixture(path, case_type) for path in paths)
def fixture_directory(configured: Path | None, env_value: str | None, default: Path) -> Path:
return (configured or Path(env_value or 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}"

View file

@ -0,0 +1,95 @@
from __future__ import annotations
from datetime import datetime, timezone
from pathlib import Path
from typing import Final
import httpx
import pytest
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,
RecordedHttpStreamResponse,
RecordedResponse,
RecordedStreamChunk,
)
from ..replay import replay_server
_URI: Final = "http://parity-provider.invalid/operation?api-version=1"
class _Input(SdkInputBase):
model: str = "fixture-model"
@pytest.mark.parametrize("body", (b'{"text":"caf\xc3\xa9"}', b"\x00\xff\x80", b""))
def test_cassette_replays_repeated_requests_with_vcr_and_preserves_bytes(tmp_path: Path, body: bytes) -> None:
sdk_input: Final = _Input()
responses: Final = tuple(
RecordedHttpResponse.from_bytes(
status,
(HttpHeader(name="content-type", value="application/octet-stream"),),
body,
)
for status in (200, 429)
)
case: Final = ParityCase[_Input](litellm_input=sdk_input, provider_responses=responses)
interactions: Final = tuple(
RecordedInteraction(Request("POST", _URI, b"\xffrequest", {}), response) for response in responses
)
timestamp: Final = datetime(2020, 1, 1, tzinfo=timezone.utc)
path: Final = save_fixture(tmp_path, sdk_input, case, interactions, recorded_at=timestamp)
assert load_fixture(tmp_path, sdk_input, ParityCase[_Input]) == case
assert deserialize_cassette(path.read_text()).recorded_at == timestamp
with VCR().use_cassette(str(path), record_mode="none", match_on=("method", "uri", "body")) as cassette:
for status in (200, 429):
replayed: Final = httpx.post(_URI, content=b"\xffrequest")
assert replayed.status_code == status
assert replayed.content == body
assert cassette.all_played
def test_stream_cassette_preserves_chunk_boundaries_through_local_replay(tmp_path: Path) -> None:
sdk_input: Final = _Input()
chunks: Final = (b"data: caf\xc3", b"\xa9\n\n", b"data: [DONE]\n\n")
response: Final = RecordedHttpStreamResponse(
kind="http_stream",
status_code=200,
headers=(HttpHeader(name="content-type", value="text/event-stream"),),
chunks=tuple(RecordedStreamChunk.from_bytes(chunk) for chunk in chunks),
)
case: Final = ParityCase[_Input](litellm_input=sdk_input, provider_responses=(response,))
path: Final = save_fixture(
tmp_path, sdk_input, case, (RecordedInteraction(Request("POST", _URI, b"{}", {}), response),)
)
loaded: Final = load_fixture(tmp_path, sdk_input, ParityCase[_Input])
assert loaded == case
with replay_server() as server:
server.enqueue_response(loaded.provider_responses[0])
with httpx.stream("POST", f"{server.url}/operation", content=b"{}") as replayed:
assert tuple(replayed.iter_raw()) == chunks
server.take_requests(1)
path.write_text(path.read_text().replace("- 10\n", "- 999\n"))
with pytest.raises(ValueError, match="invalid parity cassette"):
load_fixture(tmp_path, sdk_input, ParityCase[_Input])
def test_cassette_preserves_duplicate_response_headers(tmp_path: Path) -> None:
sdk_input: Final = _Input()
response: Final[RecordedResponse] = RecordedHttpResponse.from_bytes(
200,
(HttpHeader(name="x-test", value="first"), HttpHeader(name="x-test", value="second")),
b"{}",
)
case: Final = ParityCase[_Input](litellm_input=sdk_input, provider_responses=(response,))
save_fixture(tmp_path, sdk_input, case, (RecordedInteraction(Request("POST", _URI, b"", {}), response),))
assert load_fixture(tmp_path, sdk_input, ParityCase[_Input]) == case

View file

@ -0,0 +1,20 @@
from __future__ import annotations
from typing import Final
from hypothesis import strategies as st
from pydantic import BaseModel, ConfigDict
from .inputs import generate_case_inputs
class _Input(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
identifier: str
def test_generate_case_inputs_is_deterministic() -> None:
strategy: Final = st.builds(_Input, identifier=st.integers().map(str))
assert generate_case_inputs(strategy, examples=4) == generate_case_inputs(strategy, examples=4)

View file

@ -0,0 +1,27 @@
from __future__ import annotations
import base64
from io import BytesIO
from typing import Final, cast
from PIL import Image
from .media import dummy_image_url, structured_image_bytes, structured_image_data_uri
def test_dummy_image_url_encodes_text_and_dimensions() -> None:
assert dummy_image_url("invoice 123", 24, width=320, height=80) == (
"https://dummyjson.com/image/320x80/ffffff/000000?text=invoice%20123&fontSize=24"
)
def test_structured_image_is_local_content_bearing_png() -> None:
png: Final = structured_image_bytes()
encoded: Final = structured_image_data_uri().partition(",")[2]
image: Final = Image.open(BytesIO(png))
colors: Final = cast(list[tuple[int, tuple[int, int, int]]], image.getcolors(maxcolors=2))
assert png.startswith(b"\x89PNG\r\n\x1a\n")
assert base64.b64decode(encoded, validate=True) == png
assert image.size == (320, 80)
assert {color for _, color in colors} == {(0, 0, 0), (255, 255, 255)}

View file

@ -0,0 +1,217 @@
from __future__ import annotations
import logging
import threading
from collections.abc import Generator
from contextlib import contextmanager
from dataclasses import dataclass
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Final, Literal
import httpx
import pytest
from hypothesis import strategies as st
from pydantic import BaseModel, ConfigDict
from .pipeline import (
RecordingInvocation,
RecordingTarget,
build_recording_jobs,
record_fixtures,
)
from .recording import UpstreamEndpoint
from .store import fixture_path
from ..recorded_http import RecordedResponse
class _FixtureInput(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
identifier: str
def canonical_input(self) -> dict[str, object]:
return {"identifier": self.identifier}
class _ParityCase(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
litellm_input: _FixtureInput
provider_responses: tuple[RecordedResponse, ...]
class _Upstream(ThreadingHTTPServer):
daemon_threads = True
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"
def do_POST(self) -> None:
length: Final = int(self.headers.get("content-length") or "0")
self.rfile.read(length)
body: Final = b"{}"
server: Final = self.server
assert isinstance(server, _Upstream)
self.send_response(server.response_status)
self.send_header("content-type", "application/json")
self.send_header("content-length", str(len(body)))
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)
@dataclass(frozen=True, slots=True)
class _OrderedInvocation:
order: Literal["slow", "fast"]
slow_started: threading.Event
fast_finished: threading.Event
def execute(self, provider_url: str, case_input: _FixtureInput) -> None:
if self.order == "slow":
self.slow_started.set()
if not self.fast_finished.wait(timeout=2):
raise TimeoutError("fast recording did not finish")
else:
if not self.slow_started.wait(timeout=2):
raise TimeoutError("slow recording did not start")
response: Final = httpx.post(f"{provider_url}/record", json={"id": case_input.identifier}, timeout=5)
response.raise_for_status()
if self.order == "fast":
self.fast_finished.set()
@dataclass(frozen=True, slots=True)
class _Invocation:
def execute(self, provider_url: str, case_input: _FixtureInput) -> None:
response: Final = httpx.post(f"{provider_url}/record", json={"id": case_input.identifier}, timeout=5)
response.raise_for_status()
def _target(
name: str,
upstream_url: str,
case_input: _FixtureInput,
invocation: RecordingInvocation[_FixtureInput],
) -> RecordingTarget[_FixtureInput]:
return RecordingTarget(
name=name,
upstream=UpstreamEndpoint(base_url=upstream_url),
strategy=st.just(case_input),
invocation=invocation,
required_inputs=(case_input,),
)
def test_build_jobs_keeps_required_inputs_before_generated_inputs_and_deduplicates(tmp_path: Path) -> None:
required: Final = _FixtureInput(identifier="required")
generated: Final = _FixtureInput(identifier="generated")
target: Final = RecordingTarget(
name="ordered",
upstream=UpstreamEndpoint(base_url="https://provider.invalid"),
strategy=st.just(generated),
invocation=_Invocation(),
required_inputs=(required, required),
)
jobs: Final = build_recording_jobs((target,), tmp_path, examples=1)
assert tuple(job.case_input.identifier for job in jobs) == ("required", "generated")
def test_progress_follows_completion_order(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
slow_started: Final = threading.Event()
fast_finished: Final = threading.Event()
with _upstream() as upstream:
targets: Final = (
_target(
"slow",
upstream.url,
_FixtureInput(identifier="slow"),
_OrderedInvocation("slow", slow_started, fast_finished),
),
_target(
"fast",
upstream.url,
_FixtureInput(identifier="fast"),
_OrderedInvocation("fast", slow_started, fast_finished),
),
)
with caplog.at_level(logging.INFO, logger="tests.rust-python-harness.shared.parity.fixtures.pipeline"):
summary: Final = record_fixtures(targets, tmp_path, 1, 2, _ParityCase)
progress: Final = tuple(record.message for record in caplog.records if record.message.startswith("["))
assert len(summary.recorded) == 2
assert summary.exit_code == 0
assert "recorded fast" in progress[0]
assert "recorded slow" in progress[1]
assert caplog.records[0].message == "Recording 2 fixtures across 2 targets with concurrency 2"
assert caplog.records[-1].message == "Finished 2 fixtures: 2 recorded, 0 cached, 0 failed"
def test_failure_does_not_stop_independent_recordings(tmp_path: Path) -> None:
stale_input: Final = _FixtureInput(identifier="stale")
stale_directory: Final = tmp_path / "stale"
stale_directory.mkdir()
fixture_path(stale_directory, stale_input).with_suffix(".json").write_text(
'{"schema_version": 0}\n', encoding="utf-8"
)
with _upstream() as upstream:
targets: Final = (
_target("stale", upstream.url, stale_input, _Invocation()),
_target("valid", upstream.url, _FixtureInput(identifier="valid"), _Invocation()),
)
summary: Final = record_fixtures(targets, tmp_path, 1, 2, _ParityCase)
assert len(summary.recorded) == 1
assert summary.recorded[0].target_name == "valid"
assert len(summary.failed) == 1
assert summary.failed[0].target_name == "stale"
assert summary.exit_code == 1
@pytest.mark.parametrize("status", (408, 429, 500, 503))
def test_transient_response_is_not_cached_and_can_be_retried(tmp_path: Path, status: int) -> None:
case_input: Final = _FixtureInput(identifier="retry")
with _upstream(status) as upstream:
target: Final = _target("retry", upstream.url, case_input, _Invocation())
failed: Final = record_fixtures((target,), tmp_path, 1, 1, _ParityCase)
assert failed.exit_code == 1
assert not fixture_path(tmp_path / "retry", case_input).exists()
with _upstream() as healthy_upstream:
healthy_target: Final = _target("retry", healthy_upstream.url, case_input, _Invocation())
retried: Final = record_fixtures((healthy_target,), tmp_path, 1, 1, _ParityCase)
assert retried.exit_code == 0
assert len(retried.recorded) == 1
def test_provider_rejected_response_can_be_recorded(tmp_path: Path) -> None:
with _upstream(400) as upstream:
target: Final = _target("rejected", upstream.url, _FixtureInput(identifier="invalid"), _Invocation())
summary: Final = record_fixtures((target,), tmp_path, 1, 1, _ParityCase)
assert summary.exit_code == 0
assert len(summary.recorded) == 1

View file

@ -0,0 +1,574 @@
from __future__ import annotations
import asyncio
import queue
import threading
from collections.abc import AsyncIterator, Callable, Generator, Iterator
from contextlib import contextmanager
from dataclasses import dataclass
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Final, Literal
import httpx
import pytest
from hypothesis import strategies as st
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 ..recorded_http import (
HttpHeader,
RecordedHttpStreamResponse,
RecordedResponse,
RecordedStreamChunk,
)
from ..replay import ReplayServer, replay_server
from ..stream import (
StreamCompleted,
StreamFailed,
StreamOutcome,
assert_stream_parity,
consume_async_stream,
consume_sync_stream,
)
_SSE_CHUNKS: Final = (
b'data: {"choices":[{"delta":{"content":"hello"}}]}\n\n',
b'data: {"choices":[{"delta":{"content":" world"}}]}\n\n',
b"data: [DONE]\n\n",
)
class _StreamEvent(BaseModel):
kind: Literal["delta", "done", "error"]
value: str
class _StreamApplicationError(Exception):
status_code: Final = 400
code: Final = "invalid_input"
type: Final = "validation_error"
param: Final = "input"
model: Final = "fixture-model"
llm_provider: Final = "fixture-provider"
def _stream_event(data: str) -> _StreamEvent:
event: Final = _StreamEvent.model_validate_json(data)
if event.kind == "error":
raise _StreamApplicationError(event.value)
return event
def _event_chunks(failed: bool) -> tuple[bytes, ...]:
terminal: Final = (
b'event: error\r\ndata: {"kind":"error","value":"invalid input"}\r\n\r\n'
if failed
else b'event: done\r\ndata: {"kind":"done","value":""}\r\n\r\n'
)
return (
b'event: delta\r\ndata: {"kind":"delta",\r\ndata: "value":"caf\xc3',
b'\xa9"}\r\n',
b'\r\nevent: delta\r\ndata: {"kind":"delta","value":"second"}\r\n\r\n' + terminal,
)
def _sync_events(api_base: str, case_input: _FixtureInput) -> Iterator[_StreamEvent]:
with httpx.stream("POST", f"{api_base}/stream", json={"id": case_input.identifier}, timeout=5) as response:
response.raise_for_status()
for event in SSEDecoder().iter_bytes(response.iter_bytes()):
yield _stream_event(event.data)
async def _async_events(api_base: str, case_input: _FixtureInput) -> AsyncIterator[_StreamEvent]:
async with httpx.AsyncClient(timeout=5) as client:
async with client.stream("POST", f"{api_base}/stream", json={"id": case_input.identifier}) as response:
response.raise_for_status()
async for event in SSEDecoder().aiter_bytes(response.aiter_bytes()):
yield _stream_event(event.data)
async def _consume_async_events(api_base: str, case_input: _FixtureInput) -> StreamOutcome:
async def create() -> AsyncIterator[_StreamEvent]:
return _async_events(api_base, case_input)
return await consume_async_stream(create)
async def _replay_events(
mode: Literal["sync", "async"],
provider: ReplayServer,
response: RecordedHttpStreamResponse,
case_input: _FixtureInput,
) -> InProcessExecution[StreamOutcome]:
if mode == "sync":
return run_in_process(
provider, (response,), lambda url: consume_sync_stream(lambda: _sync_events(url, case_input))
)
return await run_in_process_async(provider, (response,), lambda url: _consume_async_events(url, case_input))
class _FixtureInput(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
identifier: str
def canonical_input(self) -> dict[str, object]:
return {"identifier": self.identifier}
class _ParityCase(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
litellm_input: _FixtureInput
provider_responses: tuple[RecordedResponse, ...]
@dataclass(frozen=True, slots=True)
class _Invocation:
sdk_call: Callable[[str, _FixtureInput], object]
def execute(self, provider_url: str, case_input: _FixtureInput) -> None:
self.sdk_call(provider_url, case_input)
class _ControlledUpstream(ThreadingHTTPServer):
daemon_threads = True
def __init__(self, stream_chunks: tuple[bytes, ...]) -> None:
super().__init__(("127.0.0.1", 0), _ControlledUpstreamHandler)
self.stream_chunks: Final = stream_chunks
self.lock: Final = threading.Lock()
self.two_requests_started: Final = threading.Event()
self.active_requests: int = 0
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
self.request_count += 1
self.max_active_requests = max(self.max_active_requests, self.active_requests)
if self.active_requests == 2:
self.two_requests_started.set()
self.two_requests_started.wait(timeout=2)
def end_tracked_request(self) -> None:
with self.lock:
self.active_requests -= 1
class _ControlledUpstreamHandler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
def do_POST(self) -> None:
upstream: Final = self.server
assert isinstance(upstream, _ControlledUpstream)
length: Final = int(self.headers.get("content-length") or "0")
self.rfile.read(length)
if self.path == "/credentials?api_key=query-secret&api-version=1":
authorized: Final = self.headers.get("authorization") == "Bearer header-secret"
self._send_json(200 if authorized else 401, b"{}")
return
if self.path == "/upload":
self._send_json(200, b'{"file_id":"fixture://document.pdf"}')
return
if self.path == "/parse":
self._send_json(200, b'{"result":{"chunks":[]}}')
return
if self.path == "/analyze":
self.send_response(202)
self.send_header("operation-location", f"{upstream.url}/results/1")
self.send_header("content-length", "0")
self.end_headers()
return
if self.path in {"/v1/chat/completions", "/stream"}:
with upstream.lock:
upstream.request_count += 1
self.send_response(200)
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()
return
if self.path == "/error":
self._send_json(429, b'{"error":{"message":"rate limited"}}')
return
upstream.start_request()
try:
body: Final = b"{}"
self.send_response(200)
self.send_header("content-type", "application/json")
self.send_header("set-cookie", "session=must-not-be-recorded")
self.send_header("content-length", str(len(body)))
self.end_headers()
self.wfile.write(body)
finally:
upstream.end_tracked_request()
def do_GET(self) -> None:
if self.path == "/results/1":
self._send_json(200, b'{"status":"succeeded","analyzeResult":{"pages":[]}}')
return
self.send_error(404)
def do_PUT(self) -> None:
self.do_POST()
def do_PATCH(self) -> None:
self.do_POST()
def do_DELETE(self) -> None:
self.do_POST()
def _send_json(self, status: int, body: bytes) -> None:
self.send_response(status)
self.send_header("content-type", "application/json")
self.send_header("content-length", str(len(body)))
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 _case(identifier: str) -> _FixtureInput:
return _FixtureInput(identifier=identifier)
def _sdk_call(api_base: str, case_input: _FixtureInput) -> object:
return httpx.post(f"{api_base}/v1/operation", content=b"{}", timeout=5)
def _stream_sdk_call(api_base: str, case_input: _FixtureInput) -> object:
return httpx.post(f"{api_base}/v1/chat/completions", content=b"{}", timeout=5)
def _error_sdk_call(api_base: str, case_input: _FixtureInput) -> object:
response: Final = httpx.post(f"{api_base}/error", content=b"{}", timeout=5)
response.raise_for_status()
return response
def _method_sdk_call(method: str) -> Callable[[str, _FixtureInput], object]:
def call(api_base: str, case_input: _FixtureInput) -> object:
return httpx.request(method, f"{api_base}/method", json={"id": case_input.identifier}, timeout=5)
return call
def _multi_sdk_call(api_base: str, case_input: _FixtureInput) -> object:
upload: Final = httpx.post(f"{api_base}/upload", json={"document": case_input.identifier}, timeout=5)
upload.raise_for_status()
parsed: Final = httpx.post(f"{api_base}/parse", json={"input": upload.json()["file_id"]}, timeout=5)
parsed.raise_for_status()
return parsed
def _polling_sdk_call(api_base: str, case_input: _FixtureInput) -> object:
started: Final = httpx.post(f"{api_base}/analyze", json={"document": case_input.identifier}, timeout=5)
operation_location: Final = started.headers["operation-location"]
completed: Final = httpx.get(operation_location, timeout=5)
completed.raise_for_status()
return completed
def test_recording_deduplicates_per_target_and_caps_global_concurrency(tmp_path: Path) -> None:
shared_input: Final = _case("shared")
with _controlled_upstream() as upstream:
spec: Final = UpstreamEndpoint(base_url=upstream.url)
targets: Final = (
RecordingTarget(
name="first",
upstream=spec,
strategy=st.just(shared_input),
invocation=_Invocation(_sdk_call),
required_inputs=(shared_input, shared_input),
),
RecordingTarget(
name="second",
upstream=spec,
strategy=st.just(shared_input),
invocation=_Invocation(_sdk_call),
required_inputs=(shared_input,),
),
)
summary: Final = record_fixtures(targets, tmp_path, examples=1, concurrency=2, case_type=_ParityCase)
assert len(summary.recorded) == 2
assert {result.target_name for result in summary.recorded} == {"first", "second"}
assert summary.cached == ()
assert summary.failed == ()
assert upstream.request_count == 2
assert upstream.max_active_requests == 2
assert len(recorded_fixtures(tmp_path, _ParityCase)) == 2
for path in tmp_path.rglob("*.yaml"):
contents = path.read_text(encoding="utf-8")
assert f"schema_version: {FIXTURE_SCHEMA_VERSION}" in contents
assert "recorded_at:" in contents
def test_pipeline_rejects_stale_fixture_before_provider_call(tmp_path: Path) -> None:
case_input: Final = _case("stale")
directory: Final = tmp_path / "stale-target"
directory.mkdir()
path: Final = fixture_path(directory, case_input).with_suffix(".json")
path.write_text('{"schema_version": 0}\n', encoding="utf-8")
target: Final = RecordingTarget(
name="stale-target",
upstream=UpstreamEndpoint(base_url="http://127.0.0.1:1"),
strategy=st.just(case_input),
invocation=_Invocation(_sdk_call),
)
summary: Final = record_fixtures(
(target,),
tmp_path,
examples=1,
concurrency=1,
case_type=_ParityCase,
)
assert summary.recorded == ()
assert summary.cached == ()
assert len(summary.failed) == 1
assert str(summary.failed[0].error) == (
f"fixture {path} has schema_version 0, expected {FIXTURE_SCHEMA_VERSION}; "
"delete it and regenerate the fixture bundle"
)
def test_cached_fixture_is_reported_without_provider_call(tmp_path: Path) -> None:
case_input: Final = _case("cached")
with _controlled_upstream() as upstream:
target: Final = RecordingTarget(
name="cached-target",
upstream=UpstreamEndpoint(base_url=upstream.url),
strategy=st.just(case_input),
invocation=_Invocation(_sdk_call),
)
first: Final = record_fixtures((target,), tmp_path, 1, 1, _ParityCase)
second: Final = record_fixtures((target,), tmp_path, 1, 1, _ParityCase)
assert len(first.recorded) == 1
assert len(second.cached) == 1
assert upstream.request_count == 1
def test_streaming_response_records_and_replays_chunks() -> None:
with _controlled_upstream() as upstream:
responses: Final = record_upstream_responses(
UpstreamEndpoint(base_url=upstream.url),
_case("stream"),
_stream_sdk_call,
)
response: Final = responses[0]
assert isinstance(response, RecordedHttpStreamResponse)
assert tuple(chunk.data_bytes() for chunk in response.chunks) == _SSE_CHUNKS
assert isinstance(response.model_dump(mode="json")["chunks"], list)
with replay_server() as provider:
provider.enqueue_response(response)
with httpx.stream("POST", f"{provider.url}/v1/chat/completions", json={}) as replayed:
replayed_chunks: Final = tuple(replayed.iter_raw())
provider.take_requests(1)
assert replayed_chunks == _SSE_CHUNKS
def test_non_successful_provider_response_is_recorded() -> None:
with _controlled_upstream() as upstream:
responses: Final = record_upstream_responses(
UpstreamEndpoint(base_url=upstream.url),
_case("provider-error"),
_error_sdk_call,
)
response: Final = responses[0]
assert response.status_code == 429
def test_sensitive_response_headers_are_not_recorded() -> None:
with _controlled_upstream() as upstream:
responses: Final = record_upstream_responses(
UpstreamEndpoint(base_url=upstream.url),
_case("headers"),
_sdk_call,
)
assert all(header.name.lower() != "set-cookie" for header in responses[0].headers)
def test_recorded_requests_strip_credentials_without_changing_the_live_request() -> None:
def sdk_call(api_base: str, case_input: _FixtureInput) -> object:
return httpx.post(
f"{api_base}/credentials?api_key=query-secret&api-version=1",
headers={
"Authorization": "Bearer header-secret",
"Ocp-Apim-Subscription-Key": "azure-secret",
"Cookie": "session=cookie-secret",
"X-Test": case_input.identifier,
},
content=b"\xffdocument",
)
with _controlled_upstream() as upstream:
interactions: Final = record_upstream_interactions(
UpstreamEndpoint(upstream.url), _case("credentials"), sdk_call
)
interaction: Final = interactions[0]
assert interaction.response.status_code == 200
assert interaction.request.uri == "http://parity-provider.invalid/credentials?api-version=1"
assert interaction.request.body == b"\xffdocument"
assert interaction.request.headers["x-test"] == "credentials"
assert all(
header not in interaction.request.headers for header in ("authorization", "ocp-apim-subscription-key", "cookie")
)
@pytest.mark.parametrize("method", ("PUT", "PATCH", "DELETE"))
def test_recording_and_replay_support_mutating_http_methods(method: str) -> None:
sdk_call: Final = _method_sdk_call(method)
with _controlled_upstream() as upstream:
responses: Final = record_upstream_responses(
UpstreamEndpoint(base_url=upstream.url),
_case(method),
sdk_call,
)
with replay_server() as provider:
provider.enqueue_response(responses[0])
sdk_call(provider.url, _case(method))
requests: Final = provider.take_requests(1)
assert requests[0].method == method
def test_stream_response_model_rejects_buffered_body() -> None:
with pytest.raises(ValueError, match="Extra inputs are not permitted"):
RecordedHttpStreamResponse.model_validate(
{
"kind": "http_stream",
"status_code": 200,
"headers": [HttpHeader(name="content-type", value="text/event-stream")],
"chunks": [RecordedStreamChunk.from_bytes(b"data: [DONE]\n\n")],
"body_b64": "",
}
)
@pytest.mark.parametrize("sdk_call", (_multi_sdk_call, _polling_sdk_call))
def test_multiple_provider_calls_record_and_replay_in_order(
sdk_call: Callable[[str, _FixtureInput], object],
) -> None:
with _controlled_upstream() as upstream:
responses: Final = record_upstream_responses(
UpstreamEndpoint(base_url=upstream.url),
_case(sdk_call.__name__),
sdk_call,
)
assert len(responses) == 2
with replay_server() as provider:
for response in responses:
provider.enqueue_response(response)
sdk_call(provider.url, _case(sdk_call.__name__))
requests: Final = provider.take_requests(2)
assert len(requests) == 2
@pytest.mark.asyncio
@pytest.mark.parametrize("mode", ("sync", "async"))
@pytest.mark.parametrize("failed", (False, True), ids=("completed", "application-error"))
async def test_typed_stream_recording_cassette_replay_parity(
tmp_path: Path, mode: Literal["sync", "async"], failed: bool
) -> None:
case_input: Final = _case("typed-stream")
outcomes: Final[queue.SimpleQueue[StreamOutcome]] = queue.SimpleQueue()
def record(api_base: str, sdk_input: _FixtureInput) -> None:
outcome: Final = (
consume_sync_stream(lambda: _sync_events(api_base, sdk_input))
if mode == "sync"
else asyncio.run(_consume_async_events(api_base, sdk_input))
)
outcomes.put(outcome)
with _controlled_upstream(_event_chunks(failed)) as upstream:
target: Final = RecordingTarget(
name="stream",
upstream=UpstreamEndpoint(upstream.url),
strategy=st.just(case_input),
invocation=_Invocation(record),
)
summary: Final = record_fixtures((target,), tmp_path, 1, 1, _ParityCase)
assert summary.failed == ()
assert len(summary.recorded) == 1
recorded: Final = outcomes.get_nowait()
loaded: Final = load_fixture(tmp_path / "stream", case_input, _ParityCase)
assert loaded is not None
response: Final = loaded.provider_responses[0]
assert isinstance(response, RecordedHttpStreamResponse)
assert response.status_code == 200
wire_bytes: Final = b"".join(chunk.data_bytes() for chunk in response.chunks)
assert wire_bytes == b"".join(_event_chunks(failed))
coalesced: Final = response.model_copy(update={"chunks": (RecordedStreamChunk.from_bytes(wire_bytes),)})
with replay_server() as provider:
first: Final = await _replay_events(mode, provider, response, case_input)
second: Final = await _replay_events(mode, provider, coalesced, case_input)
assert_request_parity(first.requests, second.requests)
assert len(first.requests) == 1
assert first.requests[0].body == {"id": case_input.identifier}
assert_stream_parity(recorded, first.response)
assert_stream_parity(first.response, second.response)
expected: Final = (_StreamEvent(kind="delta", value="café"), _StreamEvent(kind="delta", value="second"))
assert first.response.chunks == (expected if failed else (*expected, _StreamEvent(kind="done", value="")))
if failed:
assert isinstance(first.response.terminal, StreamFailed)
assert first.response.terminal.phase == "iteration"
assert first.response.terminal.exception_type is _StreamApplicationError
assert first.response.terminal.error.code == "invalid_input"
assert first.response.terminal.error.message == "invalid input"
else:
assert first.response.terminal == StreamCompleted()

View file

@ -0,0 +1,54 @@
from __future__ import annotations
from collections.abc import Iterable
from typing import Final
HOP_BY_HOP_HEADERS: Final[frozenset[str]] = frozenset(
{
"connection",
"keep-alive",
"proxy-authenticate",
"proxy-authorization",
"te",
"trailer",
"trailers",
"transfer-encoding",
"upgrade",
}
)
REQUEST_DROPPED_HEADERS: Final[frozenset[str]] = HOP_BY_HOP_HEADERS | {
"host",
"content-length",
"accept-encoding",
}
RESPONSE_DROPPED_HEADERS: Final[frozenset[str]] = HOP_BY_HOP_HEADERS | {
"content-encoding",
"content-length",
"set-cookie",
}
def connection_header_names(headers: Iterable[tuple[str, str]]) -> frozenset[str]:
return frozenset(
token.strip().lower()
for name, value in headers
if name.lower() == "connection"
for token in value.split(",")
if token.strip()
)
def dropped_request_headers(headers: Iterable[tuple[str, str]]) -> frozenset[str]:
materialized: Final = tuple(headers)
return REQUEST_DROPPED_HEADERS | connection_header_names(materialized)
def dropped_response_headers(headers: Iterable[tuple[str, str]]) -> frozenset[str]:
materialized: Final = tuple(headers)
return RESPONSE_DROPPED_HEADERS | connection_header_names(materialized)
def is_streaming_response(content_type: str) -> bool:
return "text/event-stream" in content_type.lower()

View file

@ -0,0 +1,47 @@
from __future__ import annotations
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from typing import Final, Generic, TypeVar
from .models import CapturedRequest
from .recorded_http import RecordedResponse
from .replay import ReplayServer
ResponseT = TypeVar("ResponseT")
@dataclass(frozen=True, slots=True)
class InProcessExecution(Generic[ResponseT]):
requests: tuple[CapturedRequest, ...]
response: ResponseT
def run_in_process(
provider: ReplayServer,
recorded_responses: tuple[RecordedResponse, ...],
call: Callable[[str], ResponseT],
) -> InProcessExecution[ResponseT]:
for recorded_response in recorded_responses:
provider.enqueue_response(recorded_response)
try:
response: Final = call(provider.url)
return InProcessExecution(requests=provider.take_requests(len(recorded_responses)), response=response)
except Exception:
provider.reset()
raise
async def run_in_process_async(
provider: ReplayServer,
recorded_responses: tuple[RecordedResponse, ...],
call: Callable[[str], Awaitable[ResponseT]],
) -> InProcessExecution[ResponseT]:
for recorded_response in recorded_responses:
provider.enqueue_response(recorded_response)
try:
response: Final = await call(provider.url)
return InProcessExecution(requests=provider.take_requests(len(recorded_responses)), response=response)
except Exception:
provider.reset()
raise

View file

@ -0,0 +1,145 @@
from __future__ import annotations
import base64
from typing import Annotated, Final, Literal, cast
from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter
class CapturedRequest(BaseModel):
model_config = ConfigDict(frozen=True)
method: str
path: str
headers: tuple[tuple[str, str], ...]
body: JsonValue
user_agent: str | None
class SDKSuccess(BaseModel):
model_config = ConfigDict(frozen=True)
status: Literal["ok"] = "ok"
response: JsonValue
class SDKError(BaseModel):
model_config = ConfigDict(frozen=True)
status: Literal["error"] = "error"
exception_type: str
message: str
status_code: int | None
code: str | None
error_type: str | None
param: str | None
model: str | None
llm_provider: str | None
class SDKJsonChunk(BaseModel):
model_config = ConfigDict(frozen=True)
kind: Literal["json"] = "json"
value: JsonValue
class SDKBytesChunk(BaseModel):
model_config = ConfigDict(frozen=True)
kind: Literal["bytes"] = "bytes"
data_b64: str
def data_bytes(self) -> bytes:
return base64.b64decode(self.data_b64, validate=True)
SDKChunk = Annotated[SDKJsonChunk | SDKBytesChunk, Field(discriminator="kind")]
class SDKStreamCompleted(BaseModel):
model_config = ConfigDict(frozen=True)
kind: Literal["completed"] = "completed"
class SDKStreamFailed(BaseModel):
model_config = ConfigDict(frozen=True)
kind: Literal["failed"] = "failed"
error: SDKError
SDKStreamTerminal = Annotated[SDKStreamCompleted | SDKStreamFailed, Field(discriminator="kind")]
class SDKStreamReport(BaseModel):
model_config = ConfigDict(frozen=True)
status: Literal["stream"] = "stream"
chunks: tuple[SDKChunk, ...]
terminal: SDKStreamTerminal
SDKReport = Annotated[SDKSuccess | SDKError | SDKStreamReport, Field(discriminator="status")]
JSON_VALUE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
def sdk_chunk(value: object) -> SDKChunk:
if isinstance(value, bytes):
return SDKBytesChunk(data_b64=base64.b64encode(value).decode("ascii"))
if isinstance(value, BaseModel):
return SDKJsonChunk(value=JSON_VALUE_ADAPTER.validate_python(value.model_dump(mode="json")))
return SDKJsonChunk(value=JSON_VALUE_ADAPTER.validate_python(value))
def _string_attribute(error: Exception, name: str) -> str | None:
value: Final = cast(object | None, getattr(error, name, None))
return None if value is None else str(value)
def sdk_error_report(error: Exception) -> SDKError:
message, _, _ = str(error).partition("\nTraceback (most recent call last):")
raw_status_code: Final = cast(object | None, getattr(error, "status_code", None))
status_code: Final = raw_status_code if isinstance(raw_status_code, int) else None
return SDKError(
exception_type=f"{type(error).__module__}.{type(error).__qualname__}",
message=message.rstrip(),
status_code=status_code,
code=_string_attribute(error, "code"),
error_type=_string_attribute(error, "type"),
param=_string_attribute(error, "param"),
model=_string_attribute(error, "model"),
llm_provider=_string_attribute(error, "llm_provider"),
)
class Execution(BaseModel):
model_config = ConfigDict(frozen=True)
requests: tuple[CapturedRequest, ...]
report: SDKReport
class SDKCommand(BaseModel):
model_config = ConfigDict(frozen=True)
case_file: str
route: str
class WorkerSuccess(BaseModel):
model_config = ConfigDict(frozen=True)
status: Literal["ok"] = "ok"
report: SDKReport
class WorkerFailure(BaseModel):
model_config = ConfigDict(frozen=True)
status: Literal["error"] = "error"
error: str
WorkerResult = Annotated[WorkerSuccess | WorkerFailure, Field(discriminator="status")]

View file

@ -0,0 +1,63 @@
from __future__ import annotations
import base64
from typing import Annotated, Literal
from pydantic import BaseModel, ConfigDict, Field
class _RecordedHttpModel(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
class HttpHeader(_RecordedHttpModel):
name: str
value: str
class RecordedHttpResponse(_RecordedHttpModel):
kind: Literal["http"]
status_code: int
headers: tuple[HttpHeader, ...]
body_b64: str
@classmethod
def from_bytes(
cls,
status_code: int,
headers: tuple[HttpHeader, ...],
body: bytes,
) -> RecordedHttpResponse:
return cls(
kind="http",
status_code=status_code,
headers=headers,
body_b64=base64.b64encode(body).decode("ascii"),
)
def body_bytes(self) -> bytes:
return base64.b64decode(self.body_b64, validate=True)
class RecordedStreamChunk(_RecordedHttpModel):
data_b64: str
@classmethod
def from_bytes(cls, data: bytes) -> RecordedStreamChunk:
return cls(data_b64=base64.b64encode(data).decode("ascii"))
def data_bytes(self) -> bytes:
return base64.b64decode(self.data_b64, validate=True)
class RecordedHttpStreamResponse(_RecordedHttpModel):
kind: Literal["http_stream"]
status_code: int
headers: tuple[HttpHeader, ...]
chunks: tuple[RecordedStreamChunk, ...]
RecordedResponse = Annotated[
RecordedHttpResponse | RecordedHttpStreamResponse,
Field(discriminator="kind"),
]

View file

@ -0,0 +1,147 @@
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 typing import Final
from pydantic import JsonValue, TypeAdapter
from .fixtures.recording import local_response_header
from .models import CapturedRequest
from .recorded_http import RecordedHttpResponse, RecordedHttpStreamResponse, RecordedResponse
JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
EXCLUDED_REQUEST_HEADERS: Final = frozenset(
{
"host",
"content-length",
"connection",
"accept-encoding",
"user-agent",
"x-litellm-parity-route",
}
)
EXCLUDED_RESPONSE_HEADERS: Final = frozenset({"content-length", "transfer-encoding", "connection"})
class ReplayServer(ThreadingHTTPServer):
daemon_threads = True
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)
def take_requests(self, expected_count: int) -> tuple[CapturedRequest, ...]:
request_count: Final = self.requests.qsize()
if request_count != expected_count:
raise AssertionError(f"expected exactly {expected_count} provider requests, received {request_count}")
return tuple(self.requests.get_nowait() for _ in range(request_count))
def reset(self) -> None:
while not self.responses.empty():
self.responses.get_nowait()
while not self.requests.empty():
self.requests.get_nowait()
class _ReplayHandler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
def do_POST(self) -> None:
self._replay()
def do_GET(self) -> None:
self._replay()
def do_PUT(self) -> None:
self._replay()
def do_PATCH(self) -> None:
self._replay()
def do_DELETE(self) -> None:
self._replay()
def _replay(self) -> None:
provider: Final = self.server
assert isinstance(provider, ReplayServer)
length: Final = int(self.headers.get("content-length") or "0")
raw_body: Final = self.rfile.read(length) if length else b""
content_type: Final = self.headers.get("content-type", "")
body: Final = (
JSON_VALUE.validate_json(raw_body)
if raw_body and content_type.lower().startswith("application/json")
else base64.b64encode(raw_body).decode("ascii")
if raw_body
else None
)
headers: Final = tuple(
sorted(
(name.lower(), value)
for name, value in self.headers.raw_items()
if name.lower() not in EXCLUDED_REQUEST_HEADERS
)
)
provider.requests.put(
CapturedRequest(
method=self.command,
path=self.path,
headers=headers,
body=body,
user_agent=self.headers.get("user-agent"),
)
)
try:
response: Final = provider.responses.get(timeout=5)
except queue.Empty:
self.send_error(500, "no replay response queued")
return
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))
if isinstance(response, RecordedHttpResponse):
response_body: Final = response.body_bytes()
self.send_header("content-length", str(len(response_body)))
self.end_headers()
self.wfile.write(response_body)
return
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
@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)

View file

@ -0,0 +1,198 @@
from __future__ import annotations
import asyncio
import os
import subprocess
import sys
from collections import deque
from collections.abc import Callable, Generator
from concurrent.futures import ThreadPoolExecutor, TimeoutError
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Final, TextIO, cast
from pydantic import TypeAdapter, ValidationError
from .models import (
Execution,
SDKCommand,
WorkerFailure,
WorkerResult,
WorkerSuccess,
)
from .recorded_http import RecordedResponse
from .replay import ReplayServer, replay_server
WORKER_RESULT_PREFIX: Final = "LITELLM_PARITY_RESULT "
WORKER_RESULT_ADAPTER: Final[TypeAdapter[WorkerResult]] = TypeAdapter(WorkerResult)
@dataclass(frozen=True, slots=True)
class SubprocessRunner:
entrypoint: Path
baseline_user_agent: str
route_label: str
def command(self, provider_url: str) -> tuple[str, ...]:
return (
sys.executable,
"-m",
".".join(
self.entrypoint.resolve().relative_to(Path(__file__).resolve().parents[4]).with_suffix("").parts
),
"--parity-worker",
provider_url,
)
@dataclass(frozen=True, slots=True)
class ExecutionVariant:
name: str
environment: tuple[tuple[str, str], ...]
class SubprocessWorker:
def __init__(self, runner: SubprocessRunner, provider: ReplayServer, variant: ExecutionVariant) -> None:
project_root: Final = str(Path(__file__).resolve().parents[4])
existing_pythonpath: Final = os.environ.get("PYTHONPATH")
env: Final = {
**os.environ,
**dict(variant.environment),
"LITELLM_USER_AGENT": runner.baseline_user_agent,
"PYTHONPATH": os.pathsep.join(path for path in (project_root, existing_pythonpath) if path),
}
self.mode: Final = variant.name
self.route_label: Final = runner.route_label
self.provider: Final = provider
self.process: Final = subprocess.Popen(
runner.command(provider.url),
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
bufsize=1,
env=env,
)
self.output_reader: Final = ThreadPoolExecutor(max_workers=1)
self.recent_output: Final[deque[str]] = deque(maxlen=100)
def execute(
self,
case_file: Path,
route: str,
responses: tuple[RecordedResponse, ...],
) -> Execution:
stdin: Final = self.process.stdin
if stdin is None or self.process.poll() is not None:
raise AssertionError(f"{self.mode} {self.route_label} worker exited before processing {case_file}")
for response in responses:
self.provider.enqueue_response(response)
command: Final = SDKCommand(case_file=str(case_file), route=route)
try:
stdin.write(f"{command.model_dump_json()}\n")
stdin.flush()
result: Final = self.output_reader.submit(self._read_result).result(timeout=60)
except TimeoutError as error:
self.provider.reset()
self.close()
raise AssertionError(
f"{self.mode} {self.route_label} worker timed out after 60s while processing {case_file}"
) from error
except AssertionError:
self.provider.reset()
raise
except (BrokenPipeError, OSError) as error:
self.provider.reset()
raise AssertionError(self._failure_message(f"worker pipe failed while processing {case_file}")) from error
if isinstance(result, WorkerFailure):
self.provider.reset()
raise AssertionError(
f"{self.mode} {self.route_label} worker failed while processing {case_file}:\n{result.error}"
)
assert isinstance(result, WorkerSuccess)
try:
return Execution(requests=self.provider.take_requests(len(responses)), report=result.report)
except AssertionError:
self.provider.reset()
raise
def _read_result(self) -> WorkerResult:
process_stdout: Final = self.process.stdout
if process_stdout is None:
raise AssertionError(self._failure_message("worker stdout is unavailable"))
stdout: Final = cast(TextIO, process_stdout)
line: Final = stdout.readline()
if not line:
raise AssertionError(self._failure_message("worker exited without returning a result"))
stripped: Final = line.rstrip()
if not stripped.startswith(WORKER_RESULT_PREFIX):
self.recent_output.append(stripped)
return self._read_result()
payload: Final = stripped.removeprefix(WORKER_RESULT_PREFIX)
try:
return WORKER_RESULT_ADAPTER.validate_json(payload)
except ValidationError as error:
raise AssertionError(self._failure_message("worker returned an invalid result")) from error
def _failure_message(self, message: str) -> str:
output: Final = "\n".join(self.recent_output)
prefix: Final = f"{self.mode} {self.route_label}"
return f"{prefix} {message}" if not output else f"{prefix} {message}\noutput:\n{output}"
def close(self) -> None:
stdin: Final = self.process.stdin
if stdin is not None and not stdin.closed:
stdin.close()
try:
self.process.wait(timeout=10)
except subprocess.TimeoutExpired:
self.process.terminate()
self.process.wait(timeout=10)
self.output_reader.shutdown(wait=True, cancel_futures=True)
@contextmanager
def execution_worker(
runner: SubprocessRunner,
variant: ExecutionVariant,
) -> Generator[SubprocessWorker]:
with replay_server() as provider:
worker: Final = SubprocessWorker(runner, provider, variant)
try:
yield worker
finally:
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,
baseline: ExecutionVariant,
candidate: ExecutionVariant,
) -> Generator[tuple[SubprocessWorker, SubprocessWorker]]:
with execution_worker(runner, baseline) as baseline_worker:
with execution_worker(runner, candidate) as candidate_worker:
yield baseline_worker, candidate_worker
def parity_worker_main(
execute_command: Callable[[str, str, asyncio.AbstractEventLoop], WorkerResult],
mock_url: str,
) -> None:
event_loop: Final = asyncio.new_event_loop()
try:
for line in sys.stdin:
sys.stdout.write(f"{WORKER_RESULT_PREFIX}{execute_command(line, mock_url, event_loop).model_dump_json()}\n")
sys.stdout.flush()
finally:
event_loop.close()

View file

@ -0,0 +1,177 @@
from __future__ import annotations
from collections.abc import AsyncIterable, Awaitable, Callable, Iterable
from dataclasses import dataclass
from typing import Final, Literal, TypeAlias
from .compare import assert_value_parity
from .models import (
SDKError,
SDKReport,
SDKStreamCompleted,
SDKStreamFailed,
SDKStreamReport,
sdk_chunk,
sdk_error_report,
)
@dataclass(frozen=True, slots=True)
class StreamCompleted:
kind: Literal["completed"] = "completed"
@dataclass(frozen=True, slots=True)
class StreamFailed:
phase: Literal["creation", "iteration"]
exception_type: type[BaseException]
error: SDKError
kind: Literal["failed"] = "failed"
StreamTerminal: TypeAlias = StreamCompleted | StreamFailed
@dataclass(frozen=True, slots=True)
class StreamOutcome:
wrapper_type: type[object] | None
supports_sync_iteration: bool | None
supports_async_iteration: bool | None
chunks: tuple[object, ...]
chunk_types: tuple[type[object], ...]
terminal: StreamTerminal
ChunkNormalizer: TypeAlias = Callable[[object], object]
def drain_sync_stream(stream: Iterable[object]) -> None:
for _ in stream:
pass
async def drain_async_stream(stream: AsyncIterable[object]) -> None:
async for _ in stream:
pass
def capture_sync_stream(create: Callable[[], Iterable[object]]) -> SDKReport:
return _stream_report(consume_sync_stream(create))
async def capture_async_stream(create: Callable[[], Awaitable[AsyncIterable[object]]]) -> SDKReport:
return _stream_report(await consume_async_stream(create))
def _stream_report(outcome: StreamOutcome) -> SDKReport:
terminal: Final = outcome.terminal
if isinstance(terminal, StreamFailed) and terminal.phase == "creation":
return terminal.error
return SDKStreamReport(
chunks=tuple(sdk_chunk(chunk) for chunk in outcome.chunks),
terminal=SDKStreamFailed(error=terminal.error) if isinstance(terminal, StreamFailed) else SDKStreamCompleted(),
)
def _failed(phase: Literal["creation", "iteration"], error: Exception) -> StreamFailed:
return StreamFailed(
phase=phase,
exception_type=type(error),
error=sdk_error_report(error),
)
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),
)
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(),
)
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),
)
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(),
)
def normalize_chunk(chunk: object) -> object:
return chunk
def assert_stream_parity(
baseline: StreamOutcome,
candidate: StreamOutcome,
*,
normalize: ChunkNormalizer = normalize_chunk,
) -> None:
assert baseline.wrapper_type is candidate.wrapper_type
assert baseline.supports_sync_iteration is candidate.supports_sync_iteration
assert baseline.supports_async_iteration is candidate.supports_async_iteration
assert baseline.chunk_types == candidate.chunk_types
assert len(baseline.chunks) == len(candidate.chunks)
for index, (baseline_chunk, candidate_chunk) in enumerate(zip(baseline.chunks, candidate.chunks, strict=True)):
assert_value_parity(normalize(baseline_chunk), normalize(candidate_chunk), path=f"$.chunks[{index}]")
assert baseline.terminal == candidate.terminal

View file

@ -0,0 +1,190 @@
from __future__ import annotations
from types import MappingProxyType
from typing import Final
import pytest
from pydantic import BaseModel, ConfigDict, JsonValue, PrivateAttr
from .compare import assert_model_parity, assert_parity
from .models import CapturedRequest, Execution, SDKError, SDKSuccess, sdk_error_report
SENTINEL: Final = "python-parity-fallback"
class _ComparableResponse(BaseModel):
value: str
_hidden_params: dict[str, object] = PrivateAttr(default_factory=dict)
def set_hidden_param(self, key: str, value: object) -> None:
self._hidden_params[key] = value
class _DifferentResponse(BaseModel):
value: str
class _FloatResponse(BaseModel):
values: list[float]
class _PublicValue(BaseModel):
model_config = ConfigDict(extra="allow")
value: object
class _PublicError(ValueError):
status_code: Final = 400
def _execution(*, body: JsonValue = None, markdown: str = "same", user_agent: str | None = None) -> Execution:
return Execution(
requests=(
CapturedRequest(
method="POST",
path="/v1/test-route?mode=test",
headers=(("authorization", "Bearer test-key"), ("content-type", "application/json")),
body={"model": "test-model"} if body is None else body,
user_agent=user_agent,
),
),
report=SDKSuccess(response={"items": [{"text": markdown}], "model": "test-model"}),
)
def test_parity_rejects_request_difference() -> None:
python: Final = _execution(user_agent=SENTINEL)
rust: Final = _execution(body={"model": "different"}, user_agent="litellm-rust")
with pytest.raises(AssertionError):
assert_parity(python, rust, SENTINEL)
def test_parity_rejects_response_difference() -> None:
python: Final = _execution(user_agent=SENTINEL)
rust: Final = _execution(markdown="different", user_agent="litellm-rust")
with pytest.raises(AssertionError):
assert_parity(python, rust, SENTINEL)
def test_parity_rejects_error_difference() -> None:
python: Final = Execution(
requests=(),
report=SDKError(
exception_type="litellm.exceptions.BadRequestError",
message="bad request",
status_code=400,
code=None,
error_type=None,
param=None,
model="test-model",
llm_provider="test-provider",
),
)
rust: Final = python.model_copy(update={"report": python.report.model_copy(update={"status_code": 500})})
with pytest.raises(AssertionError):
assert_parity(python, rust, SENTINEL)
def test_sdk_error_report_removes_traceback_but_keeps_public_fields() -> None:
error: Final = _PublicError("invalid input\nTraceback (most recent call last):\n unstable")
report: Final = sdk_error_report(error)
assert report.exception_type.endswith("._PublicError")
assert report.message == "invalid input"
assert report.status_code == 400
def test_parity_rejects_rust_fallback() -> None:
python: Final = _execution(user_agent=SENTINEL)
rust: Final = _execution(user_agent=SENTINEL)
with pytest.raises(AssertionError, match="fell back"):
assert_parity(python, rust, SENTINEL)
def test_model_parity_compares_public_values_and_ignores_private_attrs() -> None:
python: Final = _ComparableResponse(value="same")
rust: Final = _ComparableResponse(value="same")
python.set_hidden_param("litellm_call_id", "python-id")
rust.set_hidden_param("litellm_call_id", "rust-id")
assert_model_parity(python, rust)
def test_model_parity_rejects_public_value_difference() -> None:
python: Final = _ComparableResponse(value="python")
rust: Final = _ComparableResponse(value="rust")
with pytest.raises(AssertionError):
assert_model_parity(python, rust)
def test_model_parity_rejects_type_difference() -> None:
with pytest.raises(AssertionError):
assert_model_parity(_ComparableResponse(value="same"), _DifferentResponse(value="same"))
def test_model_parity_rejects_wire_float_rounding_difference() -> None:
with pytest.raises(AssertionError, match=r"\$\.values\[0\]"):
assert_model_parity(
_FloatResponse(values=[0.22590550796036835]),
_FloatResponse(values=[0.22590550796036837]),
)
def test_model_parity_rejects_meaningful_float_difference() -> None:
with pytest.raises(AssertionError, match=r"\$\.values\[0\]"):
assert_model_parity(
_FloatResponse(values=[0.22590550796036835]),
_FloatResponse(values=[0.2259]),
)
@pytest.mark.parametrize(
("baseline", "candidate"),
(
(_ComparableResponse(value="same"), {"value": "same"}),
(_ComparableResponse(value="same"), _DifferentResponse(value="same")),
(True, 1),
(1, 1.0),
(["same"], ("same",)),
({"value": "same"}, MappingProxyType({"value": "same"})),
({True: "same"}, {1: "same"}),
),
ids=("model-dict", "model-class", "bool-int", "int-float", "list-tuple", "mapping-class", "key-type"),
)
def test_model_parity_rejects_nested_type_changes(baseline: object, candidate: object) -> None:
with pytest.raises(AssertionError, match=r"\$\.value\[0\]"):
assert_model_parity(_PublicValue(value=[baseline]), _PublicValue(value=[candidate]))
def test_model_parity_ignores_nested_private_attributes() -> None:
baseline: Final = _ComparableResponse(value="same")
candidate: Final = _ComparableResponse(value="same")
baseline.set_hidden_param("request_id", "baseline")
candidate.set_hidden_param("request_id", "candidate")
assert_model_parity(_PublicValue(value={"nested": [baseline]}), _PublicValue(value={"nested": [candidate]}))
@pytest.mark.parametrize("extras", ({"provider_value": "changed"}, {}, {"provider_value": {"value": "same"}}))
def test_model_parity_compares_public_extras(extras: dict[str, object]) -> None:
baseline: Final = _PublicValue.model_validate({"value": None, "provider_value": _ComparableResponse(value="same")})
candidate: Final = _PublicValue.model_validate({"value": None, **extras})
with pytest.raises(AssertionError):
assert_model_parity(baseline, candidate)
def test_serialized_parity_rejects_boolean_integer_substitution() -> None:
with pytest.raises(AssertionError, match="type mismatch"):
assert_parity(
_execution(body={"enabled": True}, user_agent=SENTINEL),
_execution(body={"enabled": 1}, user_agent="candidate"),
SENTINEL,
)

View file

@ -0,0 +1,344 @@
from __future__ import annotations
import queue
from collections.abc import AsyncIterator, Iterator
from typing import Final, Literal, NoReturn
import pytest
from pydantic import BaseModel, PrivateAttr, ValidationError
from .models import (
SDKBytesChunk,
SDKError,
SDKJsonChunk,
SDKReport,
SDKStreamCompleted,
SDKStreamFailed,
SDKStreamReport,
sdk_error_report,
)
from .stream import (
StreamCompleted,
StreamFailed,
StreamOutcome,
assert_stream_parity,
capture_async_stream,
capture_sync_stream,
consume_async_stream,
consume_sync_stream,
drain_async_stream,
drain_sync_stream,
)
class _Chunk(BaseModel):
value: str
_hidden_params: dict[str, object] = PrivateAttr(default_factory=dict)
def set_hidden_param(self, key: str, value: object) -> None:
self._hidden_params[key] = value
class _NestedChunk(BaseModel):
value: object
class _SyncStream:
def __init__(self, chunks: tuple[object, ...], error: BaseException | None = None) -> None:
self.chunks: Final = chunks
self.error: Final = error
def __iter__(self) -> Iterator[object]:
yield from self.chunks
if self.error is not None:
raise self.error
class _AsyncStream:
def __init__(self, chunks: tuple[object, ...], error: BaseException | None = None) -> None:
self.chunks: Final = chunks
self.error: Final = error
async def __aiter__(self) -> AsyncIterator[object]:
for chunk in self.chunks:
yield chunk
if self.error is not None:
raise self.error
class _PublicStreamError(Exception):
def __init__(
self,
message: str = "invalid input",
*,
status_code: int = 400,
llm_provider: str = "test",
model: str = "test-model",
code: str = "invalid_input",
error_type: str = "validation_error",
param: str = "input",
) -> None:
super().__init__(message)
self.status_code: Final = status_code
self.llm_provider: Final = llm_provider
self.model: Final = model
self.code: Final = code
self.type: Final = error_type
self.param: Final = param
def _creation_error() -> NoReturn:
raise _PublicStreamError(status_code=429, llm_provider="test", model="test-model")
async def _async_stream(chunks: tuple[object, ...], error: BaseException | None = None) -> _AsyncStream:
return _AsyncStream(chunks, error)
async def _consume(
mode: Literal["sync", "async"], chunks: tuple[object, ...], error: Exception | None = None
) -> StreamOutcome:
if mode == "sync":
return consume_sync_stream(lambda: _SyncStream(chunks, error))
return await consume_async_stream(lambda: _async_stream(chunks, error))
async def _capture(
mode: Literal["sync", "async"], chunks: tuple[object, ...], error: Exception | None = None
) -> SDKReport:
if mode == "sync":
return capture_sync_stream(lambda: _SyncStream(chunks, error))
return await capture_async_stream(lambda: _async_stream(chunks, error))
def test_sync_stream_parity_compares_chunks_and_ignores_private_attrs() -> None:
python_chunk: Final = _Chunk(value="same")
accelerated_chunk: Final = _Chunk(value="same")
python_chunk.set_hidden_param("request_id", "python")
accelerated_chunk.set_hidden_param("request_id", "accelerated")
python: Final = consume_sync_stream(lambda: _SyncStream((python_chunk,)))
accelerated: Final = consume_sync_stream(lambda: _SyncStream((accelerated_chunk,)))
assert python.supports_sync_iteration is True
assert python.supports_async_iteration is False
assert_stream_parity(python, accelerated)
def test_stream_parity_rejects_extra_chunk() -> None:
python: Final = consume_sync_stream(lambda: _SyncStream((_Chunk(value="one"),)))
accelerated: Final = consume_sync_stream(lambda: _SyncStream((_Chunk(value="one"), _Chunk(value="two"))))
with pytest.raises(AssertionError):
assert_stream_parity(python, accelerated)
def test_stream_parity_rejects_chunk_value_difference() -> None:
python: Final = consume_sync_stream(lambda: _SyncStream((_Chunk(value="python"),)))
accelerated: Final = consume_sync_stream(lambda: _SyncStream((_Chunk(value="accelerated"),)))
with pytest.raises(AssertionError):
assert_stream_parity(python, accelerated)
def test_stream_outcome_distinguishes_creation_and_iteration_errors() -> None:
creation: Final = consume_sync_stream(_creation_error)
iteration: Final = consume_sync_stream(
lambda: _SyncStream(
(_Chunk(value="before-error"),),
_PublicStreamError(status_code=429, llm_provider="test", model="test-model"),
)
)
assert isinstance(creation.terminal, StreamFailed)
assert creation.terminal.phase == "creation"
assert creation.chunks == ()
assert isinstance(iteration.terminal, StreamFailed)
assert iteration.terminal.phase == "iteration"
assert len(iteration.chunks) == 1
with pytest.raises(AssertionError):
assert_stream_parity(creation, iteration)
@pytest.mark.asyncio
async def test_async_stream_uses_same_trace_contract() -> None:
python_error: Final = _PublicStreamError(
"invalid input\nTraceback (most recent call last):\npython detail", status_code=500
)
accelerated_error: Final = _PublicStreamError(
"invalid input\nTraceback (most recent call last):\nrust detail", status_code=500
)
python: Final = await consume_async_stream(lambda: _async_stream((_Chunk(value="same"),), python_error))
accelerated: Final = await consume_async_stream(lambda: _async_stream((_Chunk(value="same"),), accelerated_error))
assert python.supports_sync_iteration is False
assert python.supports_async_iteration is True
assert_stream_parity(python, accelerated)
def test_stream_parity_accepts_route_specific_chunk_normalizer() -> None:
python: Final = consume_sync_stream(lambda: _SyncStream((_Chunk(value="python-generated-id"),)))
accelerated: Final = consume_sync_stream(lambda: _SyncStream((_Chunk(value="rust-generated-id"),)))
assert_stream_parity(python, accelerated, normalize=lambda chunk: type(chunk))
def test_drain_sync_stream_exhausts_lazy_iterator() -> None:
consumed: Final[queue.SimpleQueue[str]] = queue.SimpleQueue()
def chunks() -> Iterator[object]:
yield _Chunk(value="one")
consumed.put("complete")
drain_sync_stream(chunks())
assert consumed.get_nowait() == "complete"
@pytest.mark.asyncio
async def test_drain_async_stream_exhausts_lazy_iterator() -> None:
consumed: Final[queue.SimpleQueue[str]] = queue.SimpleQueue()
async def chunks() -> AsyncIterator[object]:
yield b"one"
consumed.put("complete")
await drain_async_stream(chunks())
assert consumed.get_nowait() == "complete"
def test_capture_sync_stream_serializes_model_chunks_and_partial_failure() -> None:
report: Final = capture_sync_stream(
lambda: _SyncStream(
(_Chunk(value="before-error"),),
_PublicStreamError(status_code=429, llm_provider="test", model="test-model"),
)
)
assert isinstance(report, SDKStreamReport)
assert len(report.chunks) == 1
chunk: Final = report.chunks[0]
assert isinstance(chunk, SDKJsonChunk)
assert chunk.value == {"value": "before-error"}
assert isinstance(report.terminal, SDKStreamFailed)
assert report.terminal.error.status_code == 429
@pytest.mark.asyncio
async def test_capture_async_stream_serializes_message_bytes_in_order() -> None:
report: Final = await capture_async_stream(lambda: _async_stream((b"first", b"second")))
assert isinstance(report, SDKStreamReport)
assert tuple(chunk.data_bytes() for chunk in report.chunks if isinstance(chunk, SDKBytesChunk)) == (
b"first",
b"second",
)
@pytest.mark.asyncio
@pytest.mark.parametrize("mode", ("sync", "async"))
@pytest.mark.parametrize(
"candidate_error",
(
ValueError("invalid input"),
_PublicStreamError("changed message"),
_PublicStreamError(status_code=429),
_PublicStreamError(code="changed_code"),
_PublicStreamError(error_type="changed_type"),
_PublicStreamError(param="changed_param"),
_PublicStreamError(model="changed_model"),
_PublicStreamError(llm_provider="changed_provider"),
),
ids=("exception", "message", "status", "code", "type", "param", "model", "provider"),
)
async def test_stream_parity_rejects_public_error_changes(
mode: Literal["sync", "async"], candidate_error: Exception
) -> None:
chunks: Final = (_Chunk(value="partial"),)
baseline: Final = await _consume(mode, chunks, _PublicStreamError())
candidate: Final = await _consume(mode, chunks, candidate_error)
with pytest.raises(AssertionError):
assert_stream_parity(baseline, candidate)
@pytest.mark.asyncio
@pytest.mark.parametrize("mode", ("sync", "async"))
@pytest.mark.parametrize("candidate", (("one",), ("one", "two", "three"), ("two", "one"), ("one", "changed")))
async def test_stream_parity_checks_event_sequence(mode: Literal["sync", "async"], candidate: tuple[str, ...]) -> None:
baseline: Final = await _consume(mode, (_Chunk(value="one"), _Chunk(value="two")))
changed: Final = await _consume(mode, tuple(_Chunk(value=value) for value in candidate))
with pytest.raises(AssertionError):
assert_stream_parity(baseline, changed)
@pytest.mark.asyncio
@pytest.mark.parametrize("mode", ("sync", "async"))
async def test_stream_parity_preserves_nested_types_and_ignores_private_fields(mode: Literal["sync", "async"]) -> None:
first: Final = _Chunk(value="same")
second: Final = _Chunk(value="same")
first.set_hidden_param("request_id", "first")
second.set_hidden_param("request_id", "second")
baseline: Final = await _consume(mode, (_NestedChunk(value=[first]),))
candidate: Final = await _consume(mode, (_NestedChunk(value=[second]),))
assert_stream_parity(baseline, candidate)
changed: Final = await _consume(mode, (_NestedChunk(value=[{"value": "same"}]),))
with pytest.raises(AssertionError, match=r"\$\.chunks\[0\]\.value\[0\]"):
assert_stream_parity(baseline, changed)
def test_stream_parity_rejects_wrapper_and_chunk_type_changes() -> None:
baseline: Final = consume_sync_stream(lambda: _SyncStream((_Chunk(value="same"),)))
different_wrapper: Final = consume_sync_stream(lambda: iter((_Chunk(value="same"),)))
different_chunk: Final = consume_sync_stream(lambda: _SyncStream(({"value": "same"},)))
with pytest.raises(AssertionError):
assert_stream_parity(baseline, different_wrapper)
with pytest.raises(AssertionError):
assert_stream_parity(baseline, different_chunk)
@pytest.mark.asyncio
@pytest.mark.parametrize("mode", ("sync", "async"))
@pytest.mark.parametrize("error", (None, _PublicStreamError()))
async def test_capture_keeps_serialization_failures_out_of_sdk_errors(
mode: Literal["sync", "async"], error: Exception | None
) -> None:
with pytest.raises(ValidationError):
await _capture(mode, (_Chunk(value="valid"), object()), error)
@pytest.mark.asyncio
@pytest.mark.parametrize("mode", ("sync", "async"))
async def test_empty_stream_completes(mode: Literal["sync", "async"]) -> None:
outcome: Final = await _consume(mode, ())
assert outcome.chunks == ()
assert outcome.terminal == StreamCompleted()
assert await _capture(mode, ()) == SDKStreamReport(chunks=(), terminal=SDKStreamCompleted())
@pytest.mark.asyncio
async def test_async_creation_error_matches_sync_capture() -> None:
async def create() -> _AsyncStream:
_creation_error()
sync: Final = consume_sync_stream(_creation_error)
asynchronous: Final = await consume_async_stream(create)
assert_stream_parity(sync, asynchronous)
assert isinstance(sync.terminal, StreamFailed)
assert sync.terminal.phase == "creation"
assert isinstance(capture_sync_stream(_creation_error), SDKError)
assert capture_sync_stream(_creation_error) == await capture_async_stream(create) == sync.terminal.error
@pytest.mark.asyncio
@pytest.mark.parametrize("mode", ("sync", "async"))
async def test_capture_preserves_partial_output_and_complete_error(mode: Literal["sync", "async"]) -> None:
error: Final = _PublicStreamError()
report: Final = await _capture(mode, (_Chunk(value="partial"),), error)
assert report == SDKStreamReport(
chunks=(SDKJsonChunk(value={"value": "partial"}),),
terminal=SDKStreamFailed(error=sdk_error_report(error)),
)

View file

@ -44,10 +44,12 @@ class HarnessCase:
coverage: Coverage
selectors: tuple[str, ...]
note: str = ""
surface: str = "sdk"
unit_suite: str | None = None
@property
def key(self) -> str:
return f"{self.strategy_id}:{self.sdk_function}"
return f"{self.strategy_id}:{self.sdk_function}" if self.surface == "sdk" else f"{self.strategy_id}:gateway:{self.sdk_function}"
@dataclass(frozen=True)
@ -96,7 +98,7 @@ class CaseResult:
def set_initial_status(self) -> None:
if self.case.coverage is Coverage.NOT_APPLICABLE:
self.status = RunStatus.NOT_APPLICABLE
elif not self.case.selectors:
elif not self.case.selectors and not self.case.unit_suite:
self.status = RunStatus.PLANNED
else:
self.status = RunStatus.QUEUED
@ -168,12 +170,13 @@ def section_confidence(
) -> tuple[SectionConfidence, ...]:
strategy_list = tuple(strategies)
scores: list[SectionConfidence] = []
for sdk_function in SDK_FUNCTIONS:
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
if case.sdk_function == sdk_function and case.surface == surface
and case.coverage is not Coverage.NOT_APPLICABLE
)
verified = 0
@ -195,7 +198,7 @@ def section_confidence(
level = ConfidenceLevel.LOW
scores.append(
SectionConfidence(
sdk_function=sdk_function,
sdk_function=sdk_function if surface == "sdk" else f"gateway/{sdk_function}",
verified_strategies=verified,
required_strategies=required,
level=level,

View file

@ -0,0 +1,65 @@
from __future__ import annotations
from collections.abc import Callable, Sequence
from pathlib import Path
from time import monotonic
from typing import Final, Protocol
from .models import HarnessCase, HarnessRun
from .pytest_runner import UpdateCallback
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:
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],
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)
if all(report.finished_at is not None for report in reports)
else None
),
)
def run_strategies(
cases: Sequence[HarnessCase],
repo_root: Path,
on_update: UpdateCallback,
pytest_args: Sequence[str],
resolve_runner: Callable[[str], StrategyRunner],
) -> tuple[int, HarnessRun]:
strategy_ids: Final = tuple(dict.fromkeys(case.strategy_id for case in cases))
def execute(
remaining: tuple[str, ...], 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)(
selected,
repo_root,
lambda current: on_update(combine_reports((*reports, current, pending))),
pytest_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, (), ())

View file

@ -4,6 +4,7 @@ import os
from collections.abc import Callable, Sequence
from pathlib import Path
from time import monotonic
from typing import Final
import pytest
@ -148,13 +149,22 @@ def run_pytest(
return exit_code, run
plugin = HarnessPytestPlugin(run=run, on_update=on_update)
args = [*selectors, "-p", "no:terminal", *pytest_args]
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(args, plugins=[plugin]))
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()
):

View file

@ -0,0 +1,47 @@
from __future__ import annotations
from pathlib import Path
from typing import Final
from .models import Coverage, HarnessCase, RunStatus
from .orchestration import run_strategies
from .pytest_runner import run_pytest
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")
)
code, report = run_strategies(cases, tmp_path, lambda _: None, (), lambda _: run_pytest)
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]
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"),
)
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

View file

@ -12,7 +12,6 @@ from .models import (
Coverage,
HarnessRun,
RunStatus,
SDK_FUNCTIONS,
Strategy,
section_confidence,
)
@ -52,7 +51,9 @@ def _format_duration(seconds: float) -> str:
def _rerun_command(nodeid: str) -> str:
return f"poetry run pytest {shlex.quote(nodeid)} -q"
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]:
@ -67,8 +68,9 @@ def _summary(run: HarnessRun) -> tuple[int, int, int, int]:
)
def _cell_text(run: HarnessRun, strategy_id: str, sdk_function: str) -> tuple[str, str]:
result = run.results.get(f"{strategy_id}:{sdk_function}")
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 = ""
@ -101,6 +103,7 @@ class RichDashboard(AbstractContextManager["RichDashboard"]):
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)
@ -108,23 +111,23 @@ class RichDashboard(AbstractContextManager["RichDashboard"]):
table.add_column("Results", ratio=5)
for strategy in self.strategies:
values = []
for sdk_function in SDK_FUNCTIONS:
value, style = _cell_text(run, strategy.id, sdk_function)
for surface, sdk_function in columns:
value, style = _cell_text(run, strategy.id, sdk_function, surface)
if value:
values.append(
Text.assemble((f"{sdk_function} ", "dim"), (value, style))
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 × SDK function")
table = Table(box=box.ROUNDED, expand=True, title="Strategy × API")
table.add_column("Strategy", ratio=3)
for label in SDK_FUNCTIONS:
table.add_column(label, justify="center", ratio=1)
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 sdk_function in SDK_FUNCTIONS:
value, style = _cell_text(run, strategy.id, sdk_function)
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
@ -192,7 +195,7 @@ class RichDashboard(AbstractContextManager["RichDashboard"]):
from rich.table import Table
confidence_table = Table(
title="Port confidence by SDK section", box=box.ROUNDED, expand=True
title="Port confidence by API", box=box.ROUNDED, expand=True
)
confidence_table.add_column("SDK section")
confidence_table.add_column("Score", justify="right")
@ -256,9 +259,10 @@ class PlainDashboard(AbstractContextManager["PlainDashboard"]):
f"{skipped} skipped in {_format_duration(run.duration)}",
flush=True,
)
for nodeid, _ in run.failures[:5]:
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 SDK section", flush=True)
print("Port confidence by API", flush=True)
for score in section_confidence(run, self.confidence_strategies):
print(
f" {score.sdk_function:12} "

View file

@ -0,0 +1,57 @@
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)

View file

@ -0,0 +1,33 @@
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",)

View file

@ -0,0 +1,5 @@
# 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

View file

@ -0,0 +1,26 @@
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="e2e_parity")
if __name__ == "__main__":
raise SystemExit(main())

View file

@ -0,0 +1,43 @@
from __future__ import annotations
from typing import Final
import pytest
from .....shared.parity.fixtures.pytest_support import parametrize_recorded_fixtures
from .....shared.parity.fixtures.store import fixture_id
from .fixtures.config import DEFAULT_FIXTURE_DIRECTORY, FIXTURE_DIR_ENV
from .fixtures.models import OcrParityCase
def ocr_fixture_id(fixture: OcrParityCase) -> 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 fixture_id(case_input, prefix)
def ocr_fixture_marks(fixture: OcrParityCase) -> tuple[pytest.MarkDecorator, ...]:
if fixture.litellm_input.contract not in {"reducto_v3", "reducto_legacy"}:
return ()
return (
pytest.mark.xfail(
reason="Reducto does not have a Rust OCR contract",
strict=False,
),
)
def pytest_generate_tests(metafunc: pytest.Metafunc) -> None:
parametrize_recorded_fixtures(
metafunc,
fixture_name="ocr_fixture",
case_type=OcrParityCase,
env_var=FIXTURE_DIR_ENV,
default_directory=DEFAULT_FIXTURE_DIRECTORY,
regeneration_command=(
f"uv run python -m tests.rust-python-harness.strategies.e2e_parity.sdk.ocr.fixtures.record --fixture-dir {DEFAULT_FIXTURE_DIRECTORY}"
),
id_builder=ocr_fixture_id,
marks_builder=ocr_fixture_marks,
)

View file

@ -0,0 +1,74 @@
# OCR parity fixtures
The recording command runs four stages:
1. Generate deterministic SDK inputs for every configured OCR target
2. Build target-scoped, deduplicated recording jobs
3. Record upstream responses through one globally bounded worker pool
4. Persist each fixture and report whether it was recorded, cached, or failed
Run it with:
```shell
uv run python -m tests.rust-python-harness.strategies.e2e_parity.sdk.ocr.fixtures.record --examples 1000
```
`--concurrency` defaults to 2 and caps active recording jobs across all targets. Increase it explicitly when provider
quotas permit. Concurrency limits do not guarantee a request-per-minute quota; HTTP 408, 429, and 5xx responses fail
recording without being saved. Rerunning retries missing fixtures and reuses successful recordings
New recordings are VCR YAML cassettes. The corpus retains the original 31 migrated cassettes and adds live recordings
for all four providers. Original response bytes, statuses, headers, and recording timestamps are preserved
For Vertex, authenticate and select a project once:
```shell
gcloud auth login
gcloud config set project YOUR_PROJECT_ID
```
The recording command reads the project from `VERTEXAI_PROJECT`, `VERTEX_PROJECT`, or the active gcloud configuration,
then gets an OAuth access token with `gcloud auth print-access-token`. Tokens stay in memory and are removed from
recorded headers. `VERTEX_AI_ACCESS_TOKEN` or the legacy `VERTEX_AI_API_KEY` can override token lookup. The
`VERTEXT_API_KEY` express-mode key is not used as a Bearer token. Mistral defaults to `us-central1`; DeepSeek defaults
to the global host and `global` location. `VERTEX_DEEPSEEK_LOCATION` and `VERTEX_DEEPSEEK_API_BASE` override the latter
Azure accepts `AZURE_KEY` and `AZURE_ENDPOINT` as fallbacks for both Azure OCR contracts. Provider-specific variables
take precedence. Set `AZURE_DEPLOYMENT_NAME=mistral-ocr-4-0` to record that deployment instead of the default
`mistral-document-ai-2512`. Mistral and Reducto use `MISTRAL_API_KEY` and `REDUCTO_API_KEY`
To migrate an existing JSON fixture directory locally:
```shell
uv run python -m tests.rust-python-harness.strategies.e2e_parity.sdk.ocr.fixtures.migrate --fixture-dir tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/fixtures/data
```
The migration replays each old response through the Python SDK to reconstruct missing requests, writes and validates
the YAML cassette, then removes its JSON predecessor. It calls only local recording/replay servers and needs no provider
credentials. Reconstructed requests are labeled `python_replay`; they are not historical wire captures. Filenames use
the current normalized SDK input hash, including the fixture contract
OCR strategies generate public `litellm.ocr()` and `litellm.aocr()` inputs. Every case contains the normalized model,
document, optional provider override, and LiteLLM keyword arguments. The fixture-only `contract` literal selects the
input schema and is removed before calling the SDK. Strategies never build provider wire payloads
Each contract's strategy contains baselines and cases for its supported top-level OCR parameters. The
contracts are Mistral, Azure-hosted Mistral, Vertex-hosted Mistral, Azure Document Intelligence, Vertex DeepSeek,
Reducto v3, and Reducto legacy. Credentials and endpoints only control target discovery, so a machine records the
contracts it has configured and skips the rest
`--examples 1000` exhausts the current finite strategies: 32 Mistral, 15 Azure Mistral, 17 Azure Document Intelligence,
16 Vertex Mistral, 2 Vertex DeepSeek, 55 Reducto v3, and 3 Reducto legacy cases, including fixed rejected inputs.
This covers the defined strategy choices, not every possible value accepted by the schemas. The live run recorded
139 of these 140 cases. Vertex Mistral's standalone `document_annotation_format` case repeatedly returned HTTP 500
and remains pending. Its bounding-box annotation, annotation-prompt, and confidence cases record upstream 404/422
rejections; schema acceptance does not imply support by the hosted model
Reducto fixtures record upload and parse responses. Their parity cases remain non-strict expected failures until the
Rust OCR bridge supports Reducto. Azure and Vertex generation paths are unit-tested without credentials in CI, so the
committed corpus does not need live recordings for every target
Every recording target owns a small fixed provider-rejected corpus, independent of replay implementation support.
Those inputs are recorded separately from generated valid inputs. Local validation failures use no recorded response;
the parity suite checks those unsupported providers and models, malformed documents, invalid request formats, invalid
Azure Document Intelligence parameters, and invalid headers in sync and async SDK calls

View file

@ -0,0 +1,216 @@
from __future__ import annotations
from collections.abc import Mapping
from typing import Final, Literal, cast
from hypothesis import strategies as st
from hypothesis.strategies import SearchStrategy
from pydantic import StrictInt, StrictStr, TypeAdapter, field_validator
from ......shared.parity.fixtures.recording import UpstreamEndpoint
from .base import OcrDocument, OcrSdkInputBase
from .common import (
OcrFixtureClient,
OcrRecordingTarget,
image_document,
invoke_with_api_key,
pdf_document,
)
from .mistral import (
MistralCompatibleOcrSdkInput,
mistral_input_values_strategy,
)
AzureMistralModel = Literal["azure_ai/mistral-document-ai-2512", "azure_ai/mistral-ocr-4-0"]
AzureMistralFixtureModel = AzureMistralModel | Literal["azure_ai/invalid-ocr-model-for-parity"]
AzureDocumentIntelligenceModel = Literal[
"azure_ai/doc-intelligence/prebuilt-read",
"azure_ai/doc-intelligence/prebuilt-layout",
"azure_ai/doc-intelligence/prebuilt-document",
]
AzureDocumentIntelligenceFixtureModel = (
AzureDocumentIntelligenceModel | Literal["azure_ai/doc-intelligence/invalid-ocr-model-for-parity"]
)
AZURE_MISTRAL_MODELS: Final[tuple[AzureMistralModel, ...]] = ("azure_ai/mistral-document-ai-2512",)
AZURE_DOCUMENT_INTELLIGENCE_MODELS: Final[tuple[AzureDocumentIntelligenceModel, ...]] = (
"azure_ai/doc-intelligence/prebuilt-read",
"azure_ai/doc-intelligence/prebuilt-layout",
"azure_ai/doc-intelligence/prebuilt-document",
)
# API v4 replaces prebuilt-document with prebuilt-layout plus keyValuePairs. Keep
# the broader fixture model above so existing recordings remain loadable.
AZURE_DOCUMENT_INTELLIGENCE_RECORDING_MODELS: Final[tuple[AzureDocumentIntelligenceModel, ...]] = (
"azure_ai/doc-intelligence/prebuilt-read",
"azure_ai/doc-intelligence/prebuilt-layout",
)
class AzureMistralOcrSdkInput(MistralCompatibleOcrSdkInput):
contract: Literal["azure_mistral"] = "azure_mistral"
model: AzureMistralFixtureModel
custom_llm_provider: Literal["azure_ai"] | None = None
@field_validator("model")
@classmethod
def validate_model_namespace(cls, model: str) -> str:
if not model.startswith("azure_ai/"):
raise ValueError("Azure Mistral models must use the azure_ai/ LiteLLM namespace")
return model
class AzureDocumentIntelligenceOcrSdkInput(OcrSdkInputBase):
contract: Literal["azure_document_intelligence"] = "azure_document_intelligence"
model: AzureDocumentIntelligenceFixtureModel
document: OcrDocument
custom_llm_provider: Literal["azure_ai"] | None = None
pages: str | list[StrictInt] | list[StrictStr] | None = None
features: str | list[str] | None = None
req_format: Literal["litellm"] = "litellm"
AZURE_MISTRAL_PROVIDER_REJECTED_INPUTS: Final[tuple[AzureMistralOcrSdkInput, ...]] = (
AzureMistralOcrSdkInput(
model="azure_ai/invalid-ocr-model-for-parity",
document=pdf_document(),
),
)
AZURE_DOCUMENT_INTELLIGENCE_PROVIDER_REJECTED_INPUTS: Final[tuple[AzureDocumentIntelligenceOcrSdkInput, ...]] = (
AzureDocumentIntelligenceOcrSdkInput(
model="azure_ai/doc-intelligence/invalid-ocr-model-for-parity",
document=pdf_document(),
),
)
def _azure_mistral_input(values: dict[str, object], model: AzureMistralModel) -> AzureMistralOcrSdkInput:
return AzureMistralOcrSdkInput.model_validate({**values, "model": model})
def azure_mistral_input_strategy(
inline_image_data_uri: str,
models: tuple[AzureMistralModel, ...] = AZURE_MISTRAL_MODELS,
) -> SearchStrategy[AzureMistralOcrSdkInput]:
# Foundry's active gateway schema rejects 2512-only controls and
# document_annotation_prompt, even though native Mistral accepts them.
return st.builds(
_azure_mistral_input,
values=mistral_input_values_strategy("2505", inline_image_data_uri, include_document_annotation_prompt=False),
model=st.sampled_from(models),
)
_AZURE_DOCUMENT_INTELLIGENCE_CANONICAL_MODEL: Final[AzureDocumentIntelligenceModel] = (
"azure_ai/doc-intelligence/prebuilt-layout"
)
def _document_intelligence_input(
model: AzureDocumentIntelligenceModel,
document: OcrDocument,
optional_params: Mapping[str, object] | None = None,
) -> AzureDocumentIntelligenceOcrSdkInput:
return AzureDocumentIntelligenceOcrSdkInput.model_validate(
{"model": model, "document": document, **(optional_params or {})}
)
def azure_document_intelligence_input_strategy() -> SearchStrategy[AzureDocumentIntelligenceOcrSdkInput]:
document: Final = pdf_document()
pages: Final = st.one_of(
st.sampled_from(((0,), (2, 0, 0, 1))).map(list),
st.just(["1", "2-4"]),
st.just("1-4, 5"),
).map(lambda value: {"pages": value})
features: Final = st.one_of(
st.sampled_from(
(
("languages",),
("ocrHighResolution",),
("barcodes",),
("formulas",),
("styleFont",),
("keyValuePairs",),
)
).map(list),
st.just("languages, styleFont"),
).map(lambda value: {"features": value})
combined_query: Final = st.just({"pages": (0, 1), "features": ("languages", "styleFont")})
return st.one_of(
st.sampled_from(AZURE_DOCUMENT_INTELLIGENCE_RECORDING_MODELS).map(
lambda model: _document_intelligence_input(model, document)
),
st.just(
_document_intelligence_input(
_AZURE_DOCUMENT_INTELLIGENCE_CANONICAL_MODEL,
image_document("invoice 123", 24),
)
),
pages.map(
lambda optional_params: _document_intelligence_input(
_AZURE_DOCUMENT_INTELLIGENCE_CANONICAL_MODEL, document, optional_params
)
),
features.map(
lambda optional_params: _document_intelligence_input(
_AZURE_DOCUMENT_INTELLIGENCE_CANONICAL_MODEL, document, optional_params
)
),
combined_query.map(
lambda optional_params: _document_intelligence_input(
_AZURE_DOCUMENT_INTELLIGENCE_CANONICAL_MODEL, document, optional_params
)
),
st.just(
_document_intelligence_input(
_AZURE_DOCUMENT_INTELLIGENCE_CANONICAL_MODEL,
document,
{"req_format": "litellm"},
)
),
)
def azure_mistral_recording_targets(
environ: Mapping[str, str], client: OcrFixtureClient, inline_image_data_uri: str
) -> tuple[OcrRecordingTarget, ...]:
api_key: Final = environ.get("AZURE_AI_API_KEY") or environ.get("AZURE_KEY")
base_url: Final = environ.get("AZURE_AI_API_BASE") or environ.get("AZURE_ENDPOINT")
if not api_key or not base_url:
return ()
deployment: Final = environ.get("AZURE_DEPLOYMENT_NAME")
models: Final = (
(TypeAdapter(AzureMistralModel).validate_python(f"azure_ai/{deployment.removeprefix('azure_ai/')}"),)
if deployment
else AZURE_MISTRAL_MODELS
)
return (
OcrRecordingTarget(
name="azure-mistral",
upstream=UpstreamEndpoint(base_url=base_url.rstrip("/")),
strategy=cast(
SearchStrategy[OcrSdkInputBase],
azure_mistral_input_strategy(inline_image_data_uri, models),
),
invocation=invoke_with_api_key(client, api_key),
required_inputs=AZURE_MISTRAL_PROVIDER_REJECTED_INPUTS,
),
)
def azure_document_intelligence_recording_targets(
environ: Mapping[str, str], client: OcrFixtureClient
) -> tuple[OcrRecordingTarget, ...]:
api_key: Final = environ.get("AZURE_DOCUMENT_INTELLIGENCE_API_KEY") or environ.get("AZURE_KEY")
base_url: Final = environ.get("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT") or environ.get("AZURE_ENDPOINT")
if not api_key or not base_url:
return ()
return (
OcrRecordingTarget(
name="azure-document-intelligence",
upstream=UpstreamEndpoint(base_url=base_url.rstrip("/")),
strategy=cast(SearchStrategy[OcrSdkInputBase], azure_document_intelligence_input_strategy()),
invocation=invoke_with_api_key(client, api_key),
required_inputs=AZURE_DOCUMENT_INTELLIGENCE_PROVIDER_REJECTED_INPUTS,
),
)

View file

@ -0,0 +1,48 @@
from __future__ import annotations
from typing import Annotated, Literal
from pydantic import Field
from ......shared.parity.fixture_models import (
FixtureModel,
JsonSchemaDefinition,
JsonSchemaResponseFormat,
SdkInputBase,
)
__all__ = (
"DocumentUrlDocument",
"ImageUrlDocument",
"ImageUrlValue",
"JsonSchemaDefinition",
"JsonSchemaResponseFormat",
"OcrDocument",
"OcrSdkInputBase",
)
class OcrSdkInputBase(SdkInputBase):
fixture_only_fields = ("contract",)
class ImageUrlValue(FixtureModel):
url: str
detail: Literal["low", "auto", "high"] | None = None
class ImageUrlDocument(FixtureModel):
type: Literal["image_url"]
image_url: str | ImageUrlValue
class DocumentUrlDocument(FixtureModel):
type: Literal["document_url"]
document_url: str
document_name: str | None = None
OcrDocument = Annotated[
ImageUrlDocument | DocumentUrlDocument,
Field(discriminator="type"),
]

View file

@ -0,0 +1,101 @@
from __future__ import annotations
from dataclasses import dataclass, field
from functools import cache
from typing import Final, Literal, Protocol
from hypothesis import strategies as st
from hypothesis.strategies import SearchStrategy
from ......shared.parity.fixtures.pipeline import RecordingTarget
from ......shared.parity.fixtures.media import dummy_image_url, structured_pdf_data_uri
from .base import (
DocumentUrlDocument,
ImageUrlDocument,
JsonSchemaDefinition,
JsonSchemaResponseFormat,
OcrSdkInputBase,
)
OcrRecordingTarget = RecordingTarget[OcrSdkInputBase]
class OcrFixtureClient(Protocol):
def execute(self, api_base: str, api_key: str, case_input: OcrSdkInputBase) -> None: ...
class OcrSdkCall(Protocol):
def __call__(self, **kwargs: object) -> object: ...
@dataclass(frozen=True, slots=True)
class ApiKeyOcrInvocation:
client: OcrFixtureClient
api_key: str = field(repr=False)
def execute(self, provider_url: str, case_input: OcrSdkInputBase) -> None:
self.client.execute(provider_url, self.api_key, case_input)
def image_document(text: str, font_size: int) -> ImageUrlDocument:
return ImageUrlDocument(type="image_url", image_url=dummy_image_url(text, font_size))
def image_data_document(data_uri: str) -> ImageUrlDocument:
return ImageUrlDocument(type="image_url", image_url=data_uri)
@cache
def remote_pdf_document() -> DocumentUrlDocument:
return DocumentUrlDocument(
type="document_url",
document_url="https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf",
)
@cache
def pdf_document() -> DocumentUrlDocument:
return DocumentUrlDocument(type="document_url", document_url=structured_pdf_data_uri())
def document_transport_strategy(inline_image_data_uri: str) -> SearchStrategy[ImageUrlDocument | DocumentUrlDocument]:
transports: Final[tuple[Literal["remote_image", "inline_image", "remote_pdf", "inline_pdf"], ...]] = (
"remote_image",
"inline_image",
"remote_pdf",
"inline_pdf",
)
def as_document(
transport: Literal["remote_image", "inline_image", "remote_pdf", "inline_pdf"],
) -> ImageUrlDocument | DocumentUrlDocument:
if transport == "remote_image":
return image_document("invoice 123", 24)
if transport == "inline_image":
return image_data_document(inline_image_data_uri)
if transport == "remote_pdf":
return remote_pdf_document()
return pdf_document()
return st.sampled_from(transports).map(as_document)
def annotation_format(name: str) -> JsonSchemaResponseFormat:
return JsonSchemaResponseFormat(
type="json_schema",
json_schema=JsonSchemaDefinition(
name=name,
description="Extract the visible document fields",
schema={
"type": "object",
"properties": {"title": {"type": "string"}},
"required": ["title"],
"additionalProperties": False,
},
strict=True,
),
)
def invoke_with_api_key(client: OcrFixtureClient, api_key: str) -> ApiKeyOcrInvocation:
return ApiKeyOcrInvocation(client=client, api_key=api_key)

View file

@ -0,0 +1,44 @@
from __future__ import annotations
import os
import subprocess
from collections.abc import Callable, Mapping
from pathlib import Path
from typing import Final
FIXTURE_DIR_ENV: Final = "LITELLM_OCR_FIXTURE_DIR"
DEFAULT_FIXTURE_DIRECTORY: Final = Path(__file__).with_name("data")
def read_gcloud(arguments: tuple[str, ...]) -> str:
try:
result: Final = subprocess.run(("gcloud", *arguments), capture_output=True, text=True, timeout=45, check=False)
except (OSError, subprocess.TimeoutExpired):
return ""
return result.stdout.strip() if result.returncode == 0 else ""
def recording_environment(
environ: Mapping[str, str],
command_reader: Callable[[tuple[str, ...]], str] = read_gcloud,
) -> Mapping[str, str]:
project: Final = (
environ.get("VERTEXAI_PROJECT")
or environ.get("VERTEX_PROJECT")
or command_reader(("config", "get-value", "project"))
)
if not project or project == "(unset)":
return environ
token: Final = (
environ.get("VERTEX_AI_ACCESS_TOKEN")
or environ.get("VERTEX_AI_API_KEY")
or command_reader(("auth", "print-access-token"))
)
if not token:
raise SystemExit("Vertex OCR needs OAuth credentials. Run `gcloud auth login` or set VERTEX_AI_ACCESS_TOKEN")
return {**environ, "VERTEXAI_PROJECT": project, "VERTEX_AI_API_KEY": token}
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

Some files were not shown because too many files have changed in this diff Show more