refactor(tests): restructure rust python harness around strategy definitions (#39628)

* wip

* refactor(tests): move sdk function tracing into rust python harness

* dead code

* fix: handle harness keyboard interrupts

* refactor(tests): deduplicate rust python harness helpers

* fix(harness): expose validated strategy choices

* wip

* refactor(harness): let strategies own parity reports

* docs(harness): update strategy structure

* refactor(harness): localize strategy report views

* wip

* fix(harness): satisfy mapping runner type checks

* fix(harness): clarify trace parity output

* wip

* fix(harness): clarify unit mapping report

* fix(harness): finalize trace parity contracts

* refactor(harness): structure parity contracts

* feat: derive unit test mapping from traces

* feat(harness): map rstest test families

* feat(ocr): port Azure document intelligence tests

* feat(harness): enforce complete unit mappings

* feat(ocr): add reducto core transforms

* feat(harness): classify host-only unit tests

* fix(ocr): complete Rust provider plumbing

* fix(harness): reuse OCR parity workers
This commit is contained in:
yujonglee 2026-09-03 21:15:01 -07:00 committed by GitHub
parent 8fc0663198
commit ee08c36fc0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
188 changed files with 10419 additions and 5452 deletions

View file

@ -96,7 +96,6 @@ jobs:
- shard: misc
artifact-name: misc
test-path: >-
tests/sdk_function_trace
tests/test_litellm/batches
tests/test_litellm/secret_managers
tests/test_litellm/a2a_protocol

View file

@ -1415,6 +1415,7 @@ dependencies = [
"litellm-core",
"pyo3",
"reqwest",
"rstest",
"serde",
"serde_json",
"sha2 0.10.9",
@ -1435,14 +1436,17 @@ dependencies = [
"aws-sigv4",
"aws-smithy-runtime-api",
"aws-types",
"base64",
"rand 0.8.7",
"reqwest",
"rstest",
"serde",
"serde_json",
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
"tracing",
"tracing-subscriber",
]
[[package]]
@ -1461,7 +1465,6 @@ dependencies = [
"tokio",
"tokio-tungstenite",
"tracing",
"tracing-subscriber",
]
[[package]]
@ -1511,6 +1514,16 @@ version = "0.3.17"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a"
[[package]]
name = "mime_guess"
version = "2.0.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f7c44f8e672c00fe5308fa235f821cb4198414e1c77935c1ab6948d3fd78550e"
dependencies = [
"mime",
"unicase",
]
[[package]]
name = "mio"
version = "1.2.2"
@ -1969,6 +1982,7 @@ dependencies = [
"hyper-util",
"js-sys",
"log",
"mime_guess",
"percent-encoding",
"pin-project-lite",
"quinn",
@ -2736,6 +2750,12 @@ version = "1.20.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20"
[[package]]
name = "unicase"
version = "2.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dbc4bc3a9f746d862c45cb89d705aa10f187bb96c76001afab07a0d35ce60142"
[[package]]
name = "unicode-ident"
version = "1.0.24"

View file

@ -24,7 +24,7 @@ pyo3 = "0.29.2"
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
pythonize = "0.29.0"
rand = "0.8"
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] }
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] }
rstest = "0.26.1"
serde = { version = "1.0", features = ["derive"] }
serde_json = { version = "1.0", features = ["float_roundtrip"] }

View file

@ -32,6 +32,7 @@ subtle = { workspace = true, optional = true }
# SHA-256 hash_token) so the plaintext credential never enters a log payload.
sha2 = { workspace = true, optional = true }
pyo3 = { workspace = true, features = ["auto-initialize"], optional = true }
tower = { version = "0.5.3", features = ["util"], optional = true }
[features]
default = []
@ -39,7 +40,9 @@ server = ["dep:axum", "dep:subtle", "dep:sha2"]
# Build the gateway's config from the proxy YAML via an embedded Python
# interpreter (links libpython; requires `litellm` importable at runtime).
python-config = ["dep:pyo3"]
trace-parity = ["server", "dep:tower", "litellm-core/observability"]
[dev-dependencies]
futures-channel = "0.3"
rstest.workspace = true
tower = { version = "0.5.3", features = ["util"] }

View file

@ -28,6 +28,8 @@ pub mod auth;
pub mod routes;
#[cfg(feature = "server")]
pub mod state;
#[cfg(feature = "trace-parity")]
pub mod trace_parity;
mod constants;
pub mod integrations;

View file

@ -12,6 +12,7 @@ use litellm_core::providers::azure_ai::ocr::transformation::{
AZURE_AI_OCR_CONFIG, AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG,
};
use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
use litellm_core::providers::reducto::ocr::transformation as reducto;
use litellm_core::providers::vertex_ai::ocr::transformation as vertex_ai;
use litellm_core::providers::vertex_ai::ocr::transformation::{
VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG,
@ -39,6 +40,7 @@ pub(super) fn ocr_provider_config(
) -> Option<&'static dyn OcrProviderConfig> {
match provider {
"mistral" => Some(&MISTRAL_OCR_CONFIG),
"reducto" => reducto::config_for_model(model),
"azure_ai" if is_azure_document_intelligence_model(model) => {
Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG)
}
@ -334,6 +336,7 @@ fn operation_status(response_json: &Value) -> Result<&str, Error> {
}
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(super) async fn poll_document_intelligence(
operation_url: &str,
original_url: &str,

View file

@ -50,7 +50,11 @@ pub(crate) async fn execute_ocr_provider_call(
.await?;
return Ok(request
.config
.transform_ocr_response(&request.model, response_json)?
.transform_ocr_response_with_params(
&request.model,
response_json,
&request.optional_params,
)?
.into_json());
}
@ -71,6 +75,10 @@ pub(crate) async fn execute_ocr_provider_call(
Ok(request
.config
.transform_ocr_response(&request.model, response_json)?
.transform_ocr_response_with_params(
&request.model,
response_json,
&request.optional_params,
)?
.into_json())
}

View file

@ -1,11 +1,15 @@
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
use litellm_core::error::Error;
use litellm_core::providers::reducto::ocr::transformation::{
build_upload_request, extract_document_source, extract_upload_file_id,
};
use serde_json::{Map, Value, json};
use std::future::Future;
use std::pin::Pin;
use super::common_utils::{convert_document_url_to_data_uri, string_headers};
use super::common_utils::{convert_document_url_to_data_uri, string_headers, truncate_error_body};
use super::types::{PreparedOcrRequest, ProviderOcrRequest};
use crate::client::http_client;
use crate::integrations::custom_guardrail::{
CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest,
};
@ -89,22 +93,39 @@ impl OcrLifecycleHooks {
)?;
let model = request.model.clone();
let custom_llm_provider = request.custom_llm_provider.clone();
let document = if config.requires_data_uri_document() {
let is_reducto = custom_llm_provider == "reducto";
let document = if is_reducto {
let guarded_document = self
.run_during_call_guardrails(&model, &custom_llm_provider, &url, request.document)
.await?;
upload_reducto_document(
&guarded_document,
request.api_base.as_deref(),
request.timeout,
&upstream_headers,
)
.await?
} else if config.requires_data_uri_document() {
convert_document_url_to_data_uri(request.document).await?
} else {
request.document
};
let optional_params = request.optional_params;
let body = config
.transform_ocr_request(&request.model, document, request.optional_params)?
.transform_ocr_request(&request.model, document, optional_params.clone())?
.data;
let body = self
.run_during_call_guardrails(&model, &custom_llm_provider, &url, body)
.await?;
let body = if is_reducto {
body
} else {
self.run_during_call_guardrails(&model, &custom_llm_provider, &url, body)
.await?
};
Ok(ProviderOcrRequest {
model,
config,
url,
body,
optional_params,
upstream_headers,
timeout: request.timeout,
})
@ -165,6 +186,63 @@ impl OcrLifecycleHooks {
}
}
async fn upload_reducto_document(
document: &Value,
api_base: Option<&str>,
timeout: Option<std::time::Duration>,
upstream_headers: &[(String, String)],
) -> Result<Value, Error> {
let source = extract_document_source(document)?;
let Some(authorization) = upstream_headers
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case("authorization"))
.map(|(_, value)| value.as_str())
else {
return Err(Error::Auth(
"Reducto upload requires an Authorization header".to_string(),
));
};
let Some(upload) = build_upload_request(source, authorization, api_base) else {
return Ok(document.clone());
};
let part = reqwest::multipart::Part::bytes(upload.bytes)
.file_name(upload.file_name)
.mime_str(&upload.mime_type)
.map_err(|error| Error::InvalidRequest(error.to_string()))?;
let form = reqwest::multipart::Form::new().part("file", part);
let mut request_builder = http_client().post(upload.url).multipart(form);
for (name, value) in upstream_headers {
if !name.eq_ignore_ascii_case("content-type")
&& !name.eq_ignore_ascii_case("content-length")
{
request_builder = request_builder.header(name, value);
}
}
if let Some(timeout) = timeout {
request_builder = request_builder.timeout(timeout);
}
let response = request_builder
.send()
.await
.map_err(|error| Error::Network(error.to_string()))?;
let status = response.status();
let body = response
.text()
.await
.map_err(|error| Error::Network(error.to_string()))?;
if !status.is_success() {
return Err(Error::Http {
status: status.as_u16(),
body: truncate_error_body(&body),
});
}
let response_json: Value = serde_json::from_str(&body).map_err(|error| {
Error::InvalidResponse(format!("invalid Reducto upload response JSON: {error}"))
})?;
let file_id = extract_upload_file_id(&response_json)?;
Ok(json!({"type": "document_url", "document_url": file_id}))
}
impl CallLifecycleHooks<PreparedOcrRequest, PreparedOcrRequest, Value> for OcrLifecycleHooks {
type PreCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>;
type DuringCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>;

View file

@ -2,6 +2,7 @@ use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
use serde_json::{Map, Value};
use super::common_utils::ocr_provider_config;
use super::hooks::OcrLifecycleHooks;
@ -28,17 +29,33 @@ pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall {
let model = provider_info.model.to_string();
let custom_llm_provider = provider_info.custom_llm_provider.to_string();
let config = ocr_provider_config(&custom_llm_provider, &model)
.ok_or_else(|| litellm_core::Error::InvalidProvider(custom_llm_provider.clone()));
.ok_or_else(|| litellm_core::Error::InvalidProvider(custom_llm_provider.clone()))
.and_then(|config| {
validate_request_format(config, &request.optional_params, &custom_llm_provider)?;
Ok(config)
});
let optional_params = match &config {
Ok(config) => {
let supported = config.supported_ocr_params();
config.map_ocr_params(
let mut mapped = config.map_ocr_params(
&request
.optional_params
.into_iter()
.iter()
.filter(|(name, _)| supported.contains(&name.as_str()))
.map(|(name, value)| (name.clone(), value.clone()))
.collect(),
)
);
for name in [
"vertex_project",
"vertex_ai_project",
"vertex_location",
"vertex_ai_location",
] {
if let Some(value) = request.optional_params.get(name) {
mapped.insert(name.to_string(), value.clone());
}
}
mapped
}
Err(_) => request.optional_params,
};
@ -64,6 +81,26 @@ pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall {
}
}
fn validate_request_format(
config: &'static dyn litellm_core::ocr::transformation::OcrProviderConfig,
optional_params: &Map<String, Value>,
provider: &str,
) -> Result<(), litellm_core::Error> {
let Some(format) = optional_params.get("req_format") else {
return Ok(());
};
match format.as_str() {
Some("litellm") => Ok(()),
Some("native") if config.supported_ocr_params().contains(&"req_format") => Ok(()),
Some("native") => Err(litellm_core::Error::InvalidRequest(format!(
"`req_format=native` is not supported for provider {provider}"
))),
_ => Err(litellm_core::Error::InvalidRequest(format!(
"Invalid `req_format`: {format}. Expected `litellm` or `native`"
))),
}
}
fn new_ocr_call_id() -> String {
static COUNTER: AtomicU64 = AtomicU64::new(1);
let sequence = COUNTER.fetch_add(1, Ordering::Relaxed);

View file

@ -53,6 +53,7 @@ pub(crate) struct ProviderOcrRequest {
pub(crate) config: &'static dyn OcrProviderConfig,
pub(crate) url: String,
pub(crate) body: Value,
pub(crate) optional_params: Map<String, Value>,
pub(crate) upstream_headers: Vec<(String, String)>,
pub(crate) timeout: Option<Duration>,
}

View file

@ -21,6 +21,12 @@ pub fn router() -> Router<AppState> {
Router::new().route(MESSAGES_ROUTE_PATH, post(handle))
}
#[tracing::instrument(
name = "messages_gateway_route",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
async fn handle(
_auth: RequireMasterKey,
State(state): State<AppState>,

View file

@ -12,6 +12,12 @@ pub(crate) enum MessagesResponse {
Stream(reqwest::Response),
}
#[tracing::instrument(
name = "messages_gateway_service",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
pub async fn run(
router: &Arc<Router>,
body: Value,

View file

@ -0,0 +1,65 @@
//! Harness-only in-process adapters. Never mounted as production routes.
use std::sync::Arc;
use axum::body::{Body, to_bytes};
use axum::http::header::{AUTHORIZATION, CONTENT_TYPE};
use axum::http::{Request, StatusCode};
use litellm_core::Error;
use litellm_core::router::{Deployment, LiteLLMParams, Router as ModelRouter};
use serde::Serialize;
use serde_json::Value;
use tower::ServiceExt;
use crate::io::realtime_pool::RealtimePool;
use crate::routes;
use crate::state::AppState;
#[derive(Debug, Serialize)]
pub struct GatewayResponse {
pub status: u16,
pub body: Value,
}
pub async fn messages_request(
model_alias: String,
provider_model: String,
api_base: String,
body: Value,
) -> Result<GatewayResponse, Error> {
let state = AppState {
router: Arc::new(ModelRouter::new(vec![Deployment {
model_name: model_alias,
litellm_params: LiteLLMParams {
model: provider_model,
api_key: Some("trace-provider-key".to_string()),
api_base: Some(api_base),
},
}])),
master_key: Some(Arc::from("trace-master-key")),
loggers: Arc::new(Vec::new()),
realtime_pool: RealtimePool::disabled(),
};
let request = Request::builder()
.method("POST")
.uri("/v1/messages")
.header(AUTHORIZATION, "Bearer trace-master-key")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(body.to_string()))
.map_err(|error| Error::InvalidRequest(error.to_string()))?;
let response = routes::app(state)
.oneshot(request)
.await
.map_err(|error| match error {})?;
let status: StatusCode = response.status();
let bytes = to_bytes(response.into_body(), usize::MAX)
.await
.map_err(|error| Error::InvalidResponse(error.to_string()))?;
let body = serde_json::from_slice(&bytes).map_err(|error| {
Error::InvalidResponse(format!("gateway returned invalid JSON: {error}"))
})?;
Ok(GatewayResponse {
status: status.as_u16(),
body,
})
}

View file

@ -6,12 +6,14 @@ license.workspace = true
repository.workspace = true
[dependencies]
base64.workspace = true
rand.workspace = true
reqwest.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
tracing.workspace = true
tracing-subscriber = { workspace = true, optional = true }
sha2.workspace = true
aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }
aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"], optional = true }
@ -30,6 +32,9 @@ bedrock-auth = [
"dep:aws-types",
"dep:aws-smithy-runtime-api",
]
observability = ["dep:tracing-subscriber"]
[dev-dependencies]
rstest.workspace = true
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
tracing-subscriber.workspace = true

View file

@ -41,3 +41,5 @@ pub const CHAT_COMPLETION_OBJECT: &str = "chat.completion";
/// `litellm/litellm_core_utils/prompt_templates/factory.py`.
pub const EMPTY_TEXT_PLACEHOLDER: &str =
"[System: Empty message content sanitised to satisfy protocol]";
pub const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace";

View file

@ -6,6 +6,8 @@ pub mod constants;
pub mod error;
pub mod http_utils;
pub mod messages;
#[cfg(any(feature = "observability", test))]
pub mod observability;
pub mod ocr;
pub mod providers;
pub mod realtime;

View file

@ -0,0 +1,215 @@
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use serde::Serialize;
use tracing::span::{Attributes, Id};
use tracing::{Dispatch, Subscriber};
use tracing_subscriber::layer::Context;
use tracing_subscriber::prelude::*;
use tracing_subscriber::registry::LookupSpan;
use tracing_subscriber::{Layer, Registry};
use super::function_trace_filter;
#[derive(Clone, Debug, PartialEq, Serialize)]
pub struct FunctionTraceEvent {
pub id: usize,
pub parent_id: Option<usize>,
pub function: &'static str,
pub module_path: Option<&'static str>,
pub file: Option<&'static str>,
pub line: Option<u32>,
}
#[derive(Clone, Default)]
pub struct FunctionTrace {
events: Arc<Mutex<Vec<FunctionTraceEvent>>>,
span_events: Arc<Mutex<HashMap<Id, usize>>>,
}
impl FunctionTrace {
pub fn dispatcher(&self) -> Dispatch {
Dispatch::new(
Registry::default().with(
FunctionTraceLayer {
trace: self.clone(),
}
.with_filter(function_trace_filter()),
),
)
}
pub fn events(&self) -> Vec<FunctionTraceEvent> {
self.events
.lock()
.unwrap_or_else(|error| error.into_inner())
.clone()
}
}
struct FunctionTraceLayer {
trace: FunctionTrace,
}
impl<S> Layer<S> for FunctionTraceLayer
where
S: Subscriber + for<'lookup> LookupSpan<'lookup>,
{
fn on_new_span(&self, attributes: &Attributes<'_>, id: &Id, context: Context<'_, S>) {
let parent_id = context.span(id).and_then(|span| {
let span_events = self
.trace
.span_events
.lock()
.unwrap_or_else(|error| error.into_inner());
span.scope()
.skip(1)
.find_map(|ancestor| span_events.get(&ancestor.id()).copied())
});
let mut events = self
.trace
.events
.lock()
.unwrap_or_else(|error| error.into_inner());
let event_id = events.len();
events.push(FunctionTraceEvent {
id: event_id,
parent_id,
function: attributes.metadata().name(),
module_path: attributes.metadata().module_path(),
file: attributes.metadata().file(),
line: attributes.metadata().line(),
});
self.trace
.span_events
.lock()
.unwrap_or_else(|error| error.into_inner())
.insert(id.clone(), event_id);
}
}
#[cfg(test)]
mod tests {
use crate::constants::FUNCTION_TRACE_TARGET;
use super::*;
fn event(
id: usize,
parent_id: Option<usize>,
function: &'static str,
) -> (usize, Option<usize>, &'static str) {
(id, parent_id, function)
}
fn structural_events(
events: &[FunctionTraceEvent],
) -> Vec<(usize, Option<usize>, &'static str)> {
events
.iter()
.map(|event| (event.id, event.parent_id, event.function))
.collect()
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
async fn outer() {
tokio::task::yield_now().await;
inner().await;
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
async fn inner() {
tokio::task::yield_now().await;
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
async fn concurrent_parent() {
tokio::join!(inner(), inner());
}
#[tokio::test]
async fn concurrent_futures_keep_separate_traces_across_yields() {
use tracing::instrument::WithSubscriber;
let first = FunctionTrace::default();
let second = FunctionTrace::default();
let outside = FunctionTrace::default();
async {
tokio::join!(
outer().with_subscriber(first.dispatcher()),
inner().with_subscriber(second.dispatcher()),
);
inner().await;
}
.with_subscriber(outside.dispatcher())
.await;
assert_eq!(
structural_events(&first.events()),
vec![event(0, None, "outer"), event(1, Some(0), "inner")],
);
assert_eq!(
structural_events(&second.events()),
vec![event(0, None, "inner")],
);
assert_eq!(
structural_events(&outside.events()),
vec![event(0, None, "inner")],
);
}
#[tokio::test]
async fn concurrent_siblings_keep_the_same_parent() {
use tracing::instrument::WithSubscriber;
let trace = FunctionTrace::default();
concurrent_parent()
.with_subscriber(trace.dispatcher())
.await;
assert_eq!(
structural_events(&trace.events()),
vec![
event(0, None, "concurrent_parent"),
event(1, Some(0), "inner"),
event(2, Some(0), "inner"),
]
);
}
#[test]
fn records_matching_spans_in_creation_order() {
let trace = FunctionTrace::default();
let dispatch = trace.dispatcher();
tracing::dispatcher::with_default(&dispatch, || {
let _ignored = tracing::trace_span!(target: "other", "ignored");
let _first = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name");
let _wrong_level = tracing::debug_span!(target: FUNCTION_TRACE_TARGET, "wrong_level");
let _second = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name");
});
assert_eq!(
structural_events(&trace.events()),
vec![event(0, None, "same_name"), event(1, None, "same_name")]
);
}
#[test]
fn records_matching_span_nesting_depth() {
let trace = FunctionTrace::default();
let dispatch = trace.dispatcher();
tracing::dispatcher::with_default(&dispatch, || {
let outer = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "outer");
let _outer_guard = outer.enter();
let _inner = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "inner");
});
assert_eq!(
structural_events(&trace.events()),
vec![event(0, None, "outer"), event(1, Some(0), "inner")]
);
}
}

View file

@ -0,0 +1,59 @@
use tracing::span::Id;
use tracing::{Level, Metadata, Subscriber};
use tracing_subscriber::filter::{FilterFn, LevelFilter, filter_fn};
use tracing_subscriber::layer::Context;
use tracing_subscriber::registry::LookupSpan;
use crate::constants::FUNCTION_TRACE_TARGET;
pub mod function_trace;
pub use function_trace::{FunctionTrace, FunctionTraceEvent};
pub fn function_trace_filter() -> FilterFn<impl Fn(&Metadata<'_>) -> bool> {
filter_fn(|metadata| {
metadata.is_span()
&& metadata.target() == FUNCTION_TRACE_TARGET
&& *metadata.level() == Level::TRACE
})
.with_max_level_hint(LevelFilter::TRACE)
}
pub fn span_depth<S>(context: &Context<'_, S>, id: &Id) -> usize
where
S: Subscriber + for<'lookup> LookupSpan<'lookup>,
{
context
.span(id)
.map(|span| span.scope().skip(1).count())
.unwrap_or_default()
}
#[cfg(test)]
mod tests {
use tracing::instrument::WithSubscriber;
use super::*;
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
async fn instrumented_with_literal_target() {}
#[tokio::test]
async fn literal_instrument_target_matches_filter_constant() {
assert_eq!(FUNCTION_TRACE_TARGET, "litellm::function_trace");
let trace = FunctionTrace::default();
instrumented_with_literal_target()
.with_subscriber(trace.dispatcher())
.await;
let events = trace.events();
assert_eq!(events.len(), 1);
assert_eq!(events[0].id, 0);
assert_eq!(events[0].parent_id, None);
assert_eq!(events[0].function, "instrumented_with_literal_target");
assert_eq!(events[0].module_path, Some(module_path!()));
assert_eq!(events[0].file, Some(file!()));
assert!(events[0].line.is_some());
}
}

View file

@ -51,6 +51,15 @@ pub trait OcrProviderConfig: Sync {
response_json: Value,
) -> Result<OcrResponseData, Error>;
fn transform_ocr_response_with_params(
&self,
model: &str,
response_json: Value,
_optional_params: &Map<String, Value>,
) -> Result<OcrResponseData, Error> {
self.transform_ocr_response(model, response_json)
}
fn complete_url(
&self,
api_base: Option<&str>,

View file

@ -1,5 +1,5 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
use serde_json::{Map, Value};
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct OcrRequestData {
@ -14,16 +14,25 @@ pub struct OcrResponseData {
pub document_annotation: Option<Value>,
pub usage_info: Option<Value>,
pub object: String,
pub extra_fields: Map<String, Value>,
pub provider_native_response: Option<Value>,
}
impl OcrResponseData {
pub fn into_json(self) -> Value {
serde_json::json!({
let mut response = serde_json::json!({
"pages": self.pages,
"model": self.model,
"document_annotation": self.document_annotation,
"usage_info": self.usage_info,
"object": self.object,
})
});
if let Value::Object(object) = &mut response {
object.extend(self.extra_fields);
if let Some(native_response) = self.provider_native_response {
object.insert("provider_native_response".to_string(), native_response);
}
}
response
}
}

View file

@ -134,6 +134,8 @@ impl OcrProviderConfig for MistralOcrConfig {
document_annotation,
usage_info,
object: "ocr".to_string(),
extra_fields: Map::new(),
provider_native_response: None,
})
}

View file

@ -4,4 +4,5 @@ pub mod azure_ai;
pub mod bedrock;
pub mod mistral;
pub mod openai;
pub mod reducto;
pub mod vertex_ai;

View file

@ -0,0 +1 @@
pub mod ocr;

View file

@ -0,0 +1,4 @@
pub mod transformation;
#[cfg(test)]
mod tests;

View file

@ -0,0 +1,202 @@
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use super::transformation::*;
use crate::ocr::transformation::OcrProviderConfig;
#[fixture]
fn parse_response() -> Value {
json!({
"job_id": "job_123",
"usage": {"num_pages": 3, "credits": 3},
"result": {
"chunks": [
{
"content": "Page 1 block A",
"blocks": [{
"content": "Page 1 block A",
"bbox": {"page": 1},
"kind": "text",
}],
},
{
"content": "Page 2 block A",
"blocks": [{
"content": "Page 2 block A",
"bbox": {"page": 2},
"kind": "table",
}],
},
{
"content": "Page 1 block B",
"blocks": [{
"content": "Page 1 block B",
"bbox": {"page": 1},
"kind": "text",
}],
},
{
"content": "Page 3 block A",
"blocks": [{
"content": "Page 3 block A",
"bbox": {"page": 3},
"kind": "figure",
}],
},
],
},
})
}
#[rstest]
fn test_parse_v3_file_upload_and_response_mapping(parse_response: Value) {
let source = classify_document_source("data:application/pdf;base64,JVBERi0xLjQ=")
.expect("PDF data URI should be valid");
let upload = build_upload_request(
source,
"Bearer test-key",
Some("https://platform.reducto.ai"),
)
.expect("data URI should require upload");
assert_eq!(upload.url, "https://platform.reducto.ai/upload");
assert_eq!(upload.authorization, "Bearer test-key");
assert_eq!(upload.file_name, "document");
assert_eq!(upload.mime_type, "application/pdf");
assert_eq!(upload.bytes, b"%PDF-1.4");
let optional_params = json!({
"formatting": {"table_output_format": "html"},
"retrieval": {"chunk_mode": "section"},
"settings": {"ocr_system": "standard"},
})
.as_object()
.expect("params should be an object")
.clone();
let request = build_parse_v3_request("reducto://uploaded.pdf", optional_params);
assert_eq!(
request.data,
json!({
"input": "reducto://uploaded.pdf",
"formatting": {"table_output_format": "html"},
"retrieval": {"chunk_mode": "section"},
"settings": {"ocr_system": "standard"},
})
);
let transformed = transform_reducto_response("parse-v3", parse_response.clone())
.expect("response should transform");
assert_eq!(
transformed.usage_info,
Some(json!({"pages_processed": 3, "credits": 3}))
);
assert_eq!(transformed.pages.len(), 3);
assert_eq!(
transformed.pages[0],
json!({
"index": 0,
"markdown": "Page 1 block A\n\nPage 1 block B",
"blocks": [
{"content": "Page 1 block A", "bbox": {"page": 1}, "kind": "text"},
{"content": "Page 1 block B", "bbox": {"page": 1}, "kind": "text"},
],
})
);
assert_eq!(transformed.pages[1]["markdown"], "Page 2 block A");
assert_eq!(transformed.pages[2]["markdown"], "Page 3 block A");
assert_eq!(transformed.provider_native_response, Some(parse_response));
}
#[rstest]
fn test_parse_v3_reducto_id_passthrough_skips_upload(parse_response: Value) {
let document = json!({
"type": "document_url",
"document_url": "reducto://already-uploaded.pdf",
});
let source = extract_document_source(&document).expect("Reducto ID should be valid");
assert!(build_upload_request(source.clone(), "Bearer test-key", None).is_none());
assert_eq!(
source,
ReductoDocumentSource::FileId("reducto://already-uploaded.pdf".to_string())
);
let request = REDUCTO_PARSE_V3_CONFIG
.transform_ocr_request(
"parse-v3",
document,
json!({"retrieval": {"chunk_mode": "section"}})
.as_object()
.expect("params should be object")
.clone(),
)
.expect("direct ID should transform");
assert_eq!(request.data["input"], "reducto://already-uploaded.pdf");
assert_eq!(request.data["retrieval"]["chunk_mode"], "section");
let response = REDUCTO_PARSE_V3_CONFIG
.transform_ocr_response("parse-v3", parse_response)
.expect("response should transform");
assert!(
response.pages[0]["markdown"]
.as_str()
.expect("markdown should be string")
.starts_with("Page 1 block A")
);
}
#[rstest]
fn test_parse_legacy_wraps_enhance_under_options() {
let request = build_parse_legacy_request(
"reducto://legacy.pdf",
json!({"enhance": {"agentic": [{"type": "table"}]}})
.as_object()
.expect("params should be object"),
);
assert_eq!(
request.data,
json!({
"document_url": "reducto://legacy.pdf",
"options": {"enhance": {"agentic": [{"type": "table"}]}},
})
);
}
#[rstest]
fn test_parse_v3_image_data_uri_upload_uses_image_mime() {
let source = classify_document_source("data:image/png;base64,iVBORw0KGgo=")
.expect("PNG data URI should be valid");
let upload = build_upload_request(
source,
"Bearer programmatic-key",
Some("https://custom.reducto.test/"),
)
.expect("data URI should require upload");
assert_eq!(upload.url, "https://custom.reducto.test/upload");
assert_eq!(upload.authorization, "Bearer programmatic-key");
assert_eq!(upload.mime_type, "image/png");
assert_eq!(upload.bytes, b"\x89PNG\r\n\x1a\n");
}
#[rstest]
#[case::http("http://example.com/document.pdf")]
#[case::https("https://example.com/document.pdf")]
fn test_parse_v3_rejects_plain_http_urls(#[case] source: &str) {
let error = classify_document_source(source).expect_err("plain URL should be rejected");
assert!(error.to_string().contains("upload the file first"));
}
#[rstest]
fn test_parse_v3_uses_programmatic_api_key_over_env() {
let key = resolve_api_key(Some("passed-key"), &|_| Some("env-reducto-key".to_string()))
.expect("explicit key should resolve");
assert_eq!(key, "passed-key");
let headers = REDUCTO_PARSE_V3_CONFIG
.validate_environment(Vec::new(), Some("passed-key"), &|_| {
Some("env-reducto-key".to_string())
})
.expect("headers should validate");
assert_eq!(
headers,
vec![("Authorization".to_string(), "Bearer passed-key".to_string())]
);
}

View file

@ -0,0 +1,407 @@
use std::collections::BTreeMap;
use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use serde_json::{Map, Value, json};
use crate::error::{Error, json_type_name};
use crate::ocr::transformation::OcrProviderConfig;
use crate::ocr::types::{OcrRequestData, OcrResponseData};
pub const REDUCTO_API_BASE: &str = "https://platform.reducto.ai";
pub const REDUCTO_API_KEY_ENV: &str = "REDUCTO_API_KEY";
pub const REDUCTO_ID_PREFIX: &str = "reducto://";
const PARSE_V3_SUPPORTED_OCR_PARAMS: &[&str] = &["formatting", "retrieval", "settings"];
const PARSE_LEGACY_SUPPORTED_OCR_PARAMS: &[&str] = &["enhance"];
const MISSING_KEY_MESSAGE: &str = "Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()";
const DATA_URI_UPLOAD_REQUIRED: &str =
"Reducto data URI upload must complete before OCR request transformation";
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ReductoDocumentSource {
FileId(String),
Upload { bytes: Vec<u8>, mime_type: String },
}
#[derive(Clone, PartialEq, Eq)]
pub struct ReductoUploadRequest {
pub url: String,
pub authorization: String,
pub file_name: &'static str,
pub bytes: Vec<u8>,
pub mime_type: String,
}
pub struct ReductoParseV3Config;
pub struct ReductoParseLegacyConfig;
pub const REDUCTO_PARSE_V3_CONFIG: ReductoParseV3Config = ReductoParseV3Config;
pub const REDUCTO_PARSE_LEGACY_CONFIG: ReductoParseLegacyConfig = ReductoParseLegacyConfig;
pub fn config_for_model(model: &str) -> Option<&'static dyn OcrProviderConfig> {
match model {
"parse-v3" => Some(&REDUCTO_PARSE_V3_CONFIG),
"parse-legacy" => Some(&REDUCTO_PARSE_LEGACY_CONFIG),
_ => None,
}
}
pub fn normalize_api_base(api_base: Option<&str>) -> String {
api_base
.map(str::trim)
.filter(|base| !base.is_empty())
.unwrap_or(REDUCTO_API_BASE)
.trim_end_matches('/')
.to_string()
}
pub fn parse_url(api_base: Option<&str>) -> String {
format!("{}/parse", normalize_api_base(api_base))
}
pub fn upload_url(api_base: Option<&str>) -> String {
format!("{}/upload", normalize_api_base(api_base))
}
pub fn resolve_api_key(
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
api_key
.map(str::trim)
.filter(|key| !key.is_empty())
.map(str::to_string)
.or_else(|| {
env_lookup(REDUCTO_API_KEY_ENV)
.map(|key| key.trim().to_string())
.filter(|key| !key.is_empty())
})
.ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string()))
}
pub fn extract_document_source(document: &Value) -> Result<ReductoDocumentSource, Error> {
let document = document.as_object().ok_or_else(|| Error::InvalidType {
expected: "object",
actual: json_type_name(document),
})?;
let source = document
.get("document_url")
.and_then(Value::as_str)
.filter(|source| !source.is_empty())
.or_else(|| document.get("image_url").and_then(Value::as_str))
.ok_or_else(|| {
Error::InvalidRequest(
"Reducto expected OCR preprocessing to produce document_url or image_url"
.to_string(),
)
})?;
classify_document_source(source)
}
pub fn classify_document_source(source: &str) -> Result<ReductoDocumentSource, Error> {
if source.starts_with(REDUCTO_ID_PREFIX) {
return Ok(ReductoDocumentSource::FileId(source.to_string()));
}
if source.starts_with("http://") || source.starts_with("https://") {
return Err(Error::InvalidRequest(
"Reducto requires type='file' (auto-uploaded) or a reducto:// id. Plain http(s) URLs are not supported; upload the file first."
.to_string(),
));
}
if !source.starts_with("data:") {
return Err(Error::InvalidRequest(
"Reducto requires a reducto:// id or a base64 data URI after OCR preprocessing."
.to_string(),
));
}
let (header, encoded) = source
.split_once(',')
.ok_or_else(|| Error::InvalidRequest("Invalid Reducto data URI provided.".to_string()))?;
if !header.split(';').any(|part| part == "base64") {
return Err(Error::InvalidRequest(
"Reducto only supports base64-encoded data URIs.".to_string(),
));
}
let mime_type = header
.strip_prefix("data:")
.and_then(|header| header.split(';').next())
.filter(|mime| !mime.is_empty())
.unwrap_or("application/octet-stream")
.to_string();
let bytes = BASE64_STANDARD.decode(encoded).map_err(|_| {
Error::InvalidRequest("Invalid Reducto base64 payload provided.".to_string())
})?;
Ok(ReductoDocumentSource::Upload { bytes, mime_type })
}
pub fn build_upload_request(
source: ReductoDocumentSource,
authorization: &str,
api_base: Option<&str>,
) -> Option<ReductoUploadRequest> {
let ReductoDocumentSource::Upload { bytes, mime_type } = source else {
return None;
};
Some(ReductoUploadRequest {
url: upload_url(api_base),
authorization: authorization.to_string(),
file_name: "document",
bytes,
mime_type,
})
}
pub fn extract_upload_file_id(response_json: &Value) -> Result<&str, Error> {
response_json
.as_object()
.and_then(|response| response.get("file_id"))
.and_then(Value::as_str)
.filter(|file_id| !file_id.is_empty())
.ok_or_else(|| {
Error::InvalidResponse(format!(
"Reducto /upload returned 200 without a file_id; got payload={response_json}"
))
})
}
pub fn build_parse_v3_request(
file_id: &str,
optional_params: Map<String, Value>,
) -> OcrRequestData {
let data = std::iter::once(("input".to_string(), Value::String(file_id.to_string())))
.chain(optional_params)
.collect();
OcrRequestData {
data: Value::Object(data),
files: None,
}
}
pub fn build_parse_legacy_request(
file_id: &str,
optional_params: &Map<String, Value>,
) -> OcrRequestData {
let options = optional_params
.get("enhance")
.filter(|enhance| !enhance.is_null())
.map(|enhance| json!({"options": {"enhance": enhance}}));
let data = match options {
Some(Value::Object(options)) => std::iter::once((
"document_url".to_string(),
Value::String(file_id.to_string()),
))
.chain(options)
.collect(),
_ => Map::from_iter([(
"document_url".to_string(),
Value::String(file_id.to_string()),
)]),
};
OcrRequestData {
data: Value::Object(data),
files: None,
}
}
fn source_file_id(document: &Value) -> Result<String, Error> {
match extract_document_source(document)? {
ReductoDocumentSource::FileId(file_id) => Ok(file_id),
ReductoDocumentSource::Upload { .. } => Err(Error::Unsupported(DATA_URI_UPLOAD_REQUIRED)),
}
}
fn page_number(block: &Map<String, Value>) -> Option<i64> {
let page = block.get("bbox")?.as_object()?.get("page")?;
page.as_i64()
.or_else(|| page.as_u64().and_then(|page| i64::try_from(page).ok()))
.or_else(|| page.as_str().and_then(|page| page.parse().ok()))
}
fn chunks(result: &Map<String, Value>) -> &[Value] {
result
.get("chunks")
.and_then(Value::as_array)
.map(Vec::as_slice)
.unwrap_or_default()
}
fn build_pages(result: &Map<String, Value>) -> Vec<Value> {
let blocks_by_page = chunks(result)
.iter()
.filter_map(Value::as_object)
.filter_map(|chunk| chunk.get("blocks").and_then(Value::as_array))
.flatten()
.filter_map(|block| block.as_object().map(|object| (block, object)))
.filter_map(|(block, object)| page_number(object).map(|page| (page, block.clone())))
.fold(
BTreeMap::<i64, Vec<Value>>::new(),
|mut pages, (page, block)| {
pages.entry(page).or_default().push(block);
pages
},
);
if blocks_by_page.is_empty() {
let markdown = chunks(result)
.iter()
.filter_map(Value::as_object)
.filter_map(|chunk| chunk.get("content").and_then(Value::as_str))
.filter(|content| !content.is_empty())
.collect::<Vec<_>>()
.join("\n\n");
return if markdown.is_empty() {
Vec::new()
} else {
vec![json!({"index": 0, "markdown": markdown})]
};
}
blocks_by_page
.into_iter()
.map(|(page, blocks)| {
let markdown = blocks
.iter()
.filter_map(Value::as_object)
.filter_map(|block| block.get("content").and_then(Value::as_str))
.filter(|content| !content.is_empty())
.collect::<Vec<_>>()
.join("\n\n");
json!({
"index": page.saturating_sub(1).max(0),
"markdown": markdown,
"blocks": blocks,
})
})
.collect()
}
pub fn transform_reducto_response(
model: &str,
response_json: Value,
) -> Result<OcrResponseData, Error> {
let response = response_json
.as_object()
.ok_or_else(|| Error::InvalidType {
expected: "object",
actual: json_type_name(&response_json),
})?;
let empty_result = Map::new();
let result = match response.get("result") {
Some(Value::Object(result)) => result,
Some(Value::Null) => &empty_result,
Some(_) => {
return Err(Error::InvalidResponse(
"Reducto result must be an object".to_string(),
));
}
None => response,
};
let usage = response
.get("usage")
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
let usage_info = Some(json!({
"pages_processed": usage.get("num_pages").cloned().unwrap_or(Value::Null),
"credits": usage.get("credits").cloned().unwrap_or(Value::Null),
}));
Ok(OcrResponseData {
pages: build_pages(result),
model: model.to_string(),
document_annotation: None,
usage_info,
object: "ocr".to_string(),
extra_fields: Map::new(),
provider_native_response: Some(response_json),
})
}
impl OcrProviderConfig for ReductoParseV3Config {
fn supported_ocr_params(&self) -> &'static [&'static str] {
PARSE_V3_SUPPORTED_OCR_PARAMS
}
fn transform_ocr_request(
&self,
_model: &str,
document: Value,
optional_params: Map<String, Value>,
) -> Result<OcrRequestData, Error> {
let file_id = source_file_id(&document)?;
Ok(build_parse_v3_request(&file_id, optional_params))
}
fn transform_ocr_response(
&self,
model: &str,
response_json: Value,
) -> Result<OcrResponseData, Error> {
transform_reducto_response(model, response_json)
}
fn complete_url(
&self,
api_base: Option<&str>,
_model: &str,
_optional_params: &Map<String, Value>,
_env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
Ok(parse_url(api_base))
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
resolve_api_key(api_key, env_lookup)
}
}
impl OcrProviderConfig for ReductoParseLegacyConfig {
fn supported_ocr_params(&self) -> &'static [&'static str] {
PARSE_LEGACY_SUPPORTED_OCR_PARAMS
}
fn transform_ocr_request(
&self,
_model: &str,
document: Value,
optional_params: Map<String, Value>,
) -> Result<OcrRequestData, Error> {
let file_id = source_file_id(&document)?;
Ok(build_parse_legacy_request(&file_id, &optional_params))
}
fn transform_ocr_response(
&self,
model: &str,
response_json: Value,
) -> Result<OcrResponseData, Error> {
transform_reducto_response(model, response_json)
}
fn complete_url(
&self,
api_base: Option<&str>,
_model: &str,
_optional_params: &Map<String, Value>,
_env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
Ok(parse_url(api_base))
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
resolve_api_key(api_key, env_lookup)
}
}

View file

@ -212,6 +212,7 @@ impl OcrProviderConfig for VertexAiOcrConfig {
MISTRAL_OCR_CONFIG.supported_ocr_params()
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn transform_ocr_request(
&self,
model: &str,
@ -229,6 +230,7 @@ impl OcrProviderConfig for VertexAiOcrConfig {
MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn complete_url(
&self,
api_base: Option<&str>,
@ -253,10 +255,21 @@ impl OcrProviderConfig for VertexAiOcrConfig {
}
impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn supported_ocr_params(&self) -> &'static [&'static str] {
DEEPSEEK_SUPPORTED_OCR_PARAMS
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn map_ocr_params(&self, non_default_params: &Map<String, Value>) -> Map<String, Value> {
non_default_params
.iter()
.filter(|(name, _)| DEEPSEEK_SUPPORTED_OCR_PARAMS.contains(&name.as_str()))
.map(|(name, value)| (name.clone(), value.clone()))
.collect()
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn transform_ocr_request(
&self,
model: &str,
@ -283,6 +296,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
})
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn transform_ocr_response(
&self,
model: &str,
@ -335,9 +349,12 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
document_annotation: object.get("document_annotation").cloned(),
usage_info,
object: "ocr".to_string(),
extra_fields: Map::new(),
provider_native_response: None,
})
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn complete_url(
&self,
api_base: Option<&str>,
@ -360,6 +377,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
#[cfg(test)]
mod tests {
use super::*;
use rstest::rstest;
#[test]
fn vertex_mistral_url_uses_project_location_and_model() {
@ -411,6 +429,22 @@ mod tests {
);
}
#[rstest]
#[case::bare_model("deepseek-ocr-maas")]
#[case::namespaced_model("deepseek-ai/deepseek-ocr-maas")]
fn vertex_deepseek_request_uses_single_provider_namespace(#[case] model: &str) {
let body = VERTEX_AI_DEEPSEEK_OCR_CONFIG
.transform_ocr_request(
model,
json!({"type": "image_url", "image_url": "data:image/png;base64,AA=="}),
Map::new(),
)
.expect("request transforms")
.data;
assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas");
}
#[test]
fn vertex_deepseek_response_wraps_markdown_content() {
let response = VERTEX_AI_DEEPSEEK_OCR_CONFIG

View file

@ -14,11 +14,15 @@ default = ["abi3"]
abi3 = ["pyo3/abi3-py310"]
extension-module = ["pyo3/extension-module"]
panic-test = []
trace-parity = [
"dep:tracing",
"litellm-core/observability",
"litellm-ai-gateway/trace-parity",
]
[dependencies]
futures-util.workspace = true
tracing.workspace = true
tracing-subscriber.workspace = true
tracing = { workspace = true, optional = true }
litellm-core = { workspace = true, features = ["bedrock-auth"] }
litellm-ai-gateway = { workspace = true, default-features = false }
litellm-python-interop.workspace = true
@ -31,6 +35,7 @@ tokio.workspace = true
[dev-dependencies]
criterion = "0.8.2"
tokio-tungstenite.workspace = true
tracing.workspace = true
[[bench]]
name = "serialization"

View file

@ -1 +0,0 @@
pub(crate) const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace";

View file

@ -1,216 +1,22 @@
use std::future::Future;
use std::sync::{Arc, Mutex};
use litellm_core::observability::{FunctionTrace, FunctionTraceEvent};
use serde::Serialize;
use tracing::instrument::WithSubscriber;
use tracing::span::{Attributes, Id};
use tracing::{Dispatch, Level, Subscriber};
use tracing_subscriber::filter::{LevelFilter, filter_fn};
use tracing_subscriber::layer::Context;
use tracing_subscriber::prelude::*;
use tracing_subscriber::registry::LookupSpan;
use tracing_subscriber::{Layer, Registry};
use crate::constants::FUNCTION_TRACE_TARGET;
#[derive(Serialize)]
#[serde(untagged)]
pub(crate) enum TraceResponse<T> {
Plain(T),
Traced {
response: T,
trace: Vec<FunctionTraceEvent>,
},
pub(crate) struct TracedResponse<T> {
response: T,
trace: Vec<FunctionTraceEvent>,
}
pub(crate) async fn trace_call<T, E>(
pub(crate) async fn capture<T, E>(
future: impl Future<Output = Result<T, E>>,
enabled: bool,
) -> Result<TraceResponse<T>, E> {
if !enabled {
return future.await.map(TraceResponse::Plain);
}
) -> Result<TracedResponse<T>, E> {
let trace = FunctionTrace::default();
let response = future.with_subscriber(trace.dispatcher()).await?;
Ok(TraceResponse::Traced {
Ok(TracedResponse {
response,
trace: trace.events(),
})
}
#[derive(Clone, Debug, PartialEq, Serialize)]
pub struct FunctionTraceEvent {
pub function: &'static str,
pub depth: usize,
}
#[derive(Clone, Default)]
pub struct FunctionTrace {
events: Arc<Mutex<Vec<FunctionTraceEvent>>>,
}
impl FunctionTrace {
pub fn dispatcher(&self) -> Dispatch {
let filter = filter_fn(|metadata| {
metadata.is_span()
&& metadata.target() == FUNCTION_TRACE_TARGET
&& *metadata.level() == Level::TRACE
})
.with_max_level_hint(LevelFilter::TRACE);
Dispatch::new(
Registry::default().with(
FunctionTraceLayer {
trace: self.clone(),
}
.with_filter(filter),
),
)
}
pub fn events(&self) -> Vec<FunctionTraceEvent> {
self.events
.lock()
.unwrap_or_else(|error| error.into_inner())
.clone()
}
}
struct FunctionTraceLayer {
trace: FunctionTrace,
}
impl<S> Layer<S> for FunctionTraceLayer
where
S: Subscriber + for<'lookup> LookupSpan<'lookup>,
{
fn on_new_span(&self, attributes: &Attributes<'_>, id: &Id, context: Context<'_, S>) {
let depth = context
.span(id)
.map(|span| span.scope().skip(1).count())
.unwrap_or_default();
self.trace
.events
.lock()
.unwrap_or_else(|error| error.into_inner())
.push(FunctionTraceEvent {
function: attributes.metadata().name(),
depth,
});
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
async fn outer() {
tokio::task::yield_now().await;
inner().await;
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
async fn inner() {
tokio::task::yield_now().await;
}
#[tokio::test]
async fn concurrent_futures_keep_separate_traces_across_yields() {
use tracing::instrument::WithSubscriber;
let first = FunctionTrace::default();
let second = FunctionTrace::default();
let outside = FunctionTrace::default();
async {
tokio::join!(
outer().with_subscriber(first.dispatcher()),
inner().with_subscriber(second.dispatcher()),
);
inner().await;
}
.with_subscriber(outside.dispatcher())
.await;
assert_eq!(
first.events(),
vec![
FunctionTraceEvent {
function: "outer",
depth: 0
},
FunctionTraceEvent {
function: "inner",
depth: 1
},
],
);
assert_eq!(
second.events(),
vec![FunctionTraceEvent {
function: "inner",
depth: 0
}],
);
assert_eq!(
outside.events(),
vec![FunctionTraceEvent {
function: "inner",
depth: 0
}],
);
}
#[test]
fn records_matching_spans_in_creation_order() {
let trace = FunctionTrace::default();
let dispatch = trace.dispatcher();
tracing::dispatcher::with_default(&dispatch, || {
let _ignored = tracing::trace_span!(target: "other", "ignored");
let _first = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name");
let _wrong_level = tracing::debug_span!(target: FUNCTION_TRACE_TARGET, "wrong_level");
let _second = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name");
});
assert_eq!(
trace.events(),
vec![
FunctionTraceEvent {
function: "same_name",
depth: 0,
},
FunctionTraceEvent {
function: "same_name",
depth: 0,
},
]
);
}
#[test]
fn records_matching_span_nesting_depth() {
let trace = FunctionTrace::default();
let dispatch = trace.dispatcher();
tracing::dispatcher::with_default(&dispatch, || {
let outer = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "outer");
let _outer_guard = outer.enter();
let _inner = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "inner");
});
assert_eq!(
trace.events(),
vec![
FunctionTraceEvent {
function: "outer",
depth: 0,
},
FunctionTraceEvent {
function: "inner",
depth: 1,
},
]
);
}
}

View file

@ -1,8 +1,8 @@
mod constants;
mod diagnostics;
mod errors;
mod execution;
pub mod function_trace;
#[cfg(feature = "trace-parity")]
mod function_trace;
mod marshal;
mod routes;
@ -115,9 +115,43 @@ mod tests {
.extract::<Vec<String>>()
.expect("module names should be strings")
.into_iter()
.filter(|name| !name.starts_with("__"))
.filter(|name| !name.starts_with('_'))
.collect();
assert_eq!(public_names, expected);
#[cfg(not(feature = "trace-parity"))]
assert!(!module.hasattr("_trace").expect("module lookup should work"));
#[cfg(feature = "trace-parity")]
{
let trace = module
.getattr("_trace")
.expect("trace build should expose its diagnostic namespace");
let trace_names: Vec<String> = trace
.cast::<PyModule>()
.expect("trace namespace should be a module")
.dict()
.keys()
.extract::<Vec<String>>()
.expect("trace names should be strings")
.into_iter()
.filter(|name| !name.starts_with("__"))
.collect();
assert_eq!(
trace_names,
[
"ocr",
"aocr",
"transcription",
"atranscription",
"messages",
"amessages",
"chat_completions",
"achat_completions",
"gateway_messages",
]
);
}
});
}

View file

@ -54,16 +54,16 @@ bridge_route! {
required = {
model: String,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
audio: Value,
audio: serde_json::Value,
},
optional = {
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
extra_headers: Option<Value>,
extra_headers: Option<serde_json::Value>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
optional_params: Option<Value>,
optional_params: Option<serde_json::Value>,
timeout_seconds: Option<f64>,
},
prepare = prepare_transcription,

View file

@ -73,16 +73,16 @@ bridge_route! {
required = {
model: String,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
messages: Value,
messages: serde_json::Value,
},
optional = {
#[pyo3(from_py_with = litellm_python_interop::from_py)]
optional_params: Option<Value>,
optional_params: Option<serde_json::Value>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
extra_headers: Option<Value>,
extra_headers: Option<serde_json::Value>,
timeout_seconds: Option<f64>,
},
prepare = prepare_chat_completions,

View file

@ -20,43 +20,33 @@ macro_rules! bridge_route {
}
#[pyfunction]
#[pyo3(signature = ($($required_name),*, $($optional_name=None,)* trace=false))]
#[pyo3(signature = ($($required_name),*, $($optional_name=None),*))]
#[allow(clippy::too_many_arguments)]
fn $sync_name(
py: pyo3::Python<'_>,
$($(#[$required_attr])* $required_name: $required_type,)*
$($(#[$optional_attr])* $optional_name: $optional_type,)*
trace: bool,
) -> pyo3::PyResult<pyo3::Py<pyo3::PyAny>> {
let future = $prepare($inputs {
$($required_name,)*
$($optional_name),*
})?;
$crate::execution::run_sync(
py,
$crate::function_trace::trace_call(future, trace),
$map_error,
)
$crate::execution::run_sync(py, future, $map_error)
}
#[pyfunction]
#[pyo3(signature = ($($required_name),*, $($optional_name=None,)* trace=false))]
#[pyo3(signature = ($($required_name),*, $($optional_name=None),*))]
#[allow(clippy::too_many_arguments)]
fn $async_name(
py: pyo3::Python<'_>,
$($(#[$required_attr])* $required_name: $required_type,)*
$($(#[$optional_attr])* $optional_name: $optional_type,)*
trace: bool,
) -> pyo3::PyResult<pyo3::Bound<'_, pyo3::PyAny>> {
let future = $prepare($inputs {
$($required_name,)*
$($optional_name),*
})?;
$crate::execution::run_async(
py,
$crate::function_trace::trace_call(future, trace),
$map_error,
)
$crate::execution::run_async(py, future, $map_error)
}
pub(super) fn register(
@ -67,6 +57,71 @@ macro_rules! bridge_route {
$crate::routes::definition::add_function(module, pyo3::wrap_pyfunction!($async_name, module)?)?;
Ok(())
}
#[cfg(feature = "trace-parity")]
mod trace {
use pyo3::prelude::*;
use super::{$inputs, $map_error, $prepare};
#[pyfunction]
#[pyo3(signature = ($($required_name),*, $($optional_name=None),*))]
#[allow(clippy::too_many_arguments)]
fn $sync_name(
py: pyo3::Python<'_>,
$($(#[$required_attr])* $required_name: $required_type,)*
$($(#[$optional_attr])* $optional_name: $optional_type,)*
) -> pyo3::PyResult<pyo3::Py<pyo3::PyAny>> {
let future = $prepare($inputs {
$($required_name,)*
$($optional_name),*
})?;
$crate::execution::run_sync(
py,
$crate::function_trace::capture(future),
$map_error,
)
}
#[pyfunction]
#[pyo3(signature = ($($required_name),*, $($optional_name=None),*))]
#[allow(clippy::too_many_arguments)]
fn $async_name(
py: pyo3::Python<'_>,
$($(#[$required_attr])* $required_name: $required_type,)*
$($(#[$optional_attr])* $optional_name: $optional_type,)*
) -> pyo3::PyResult<pyo3::Bound<'_, pyo3::PyAny>> {
let future = $prepare($inputs {
$($required_name,)*
$($optional_name),*
})?;
$crate::execution::run_async(
py,
$crate::function_trace::capture(future),
$map_error,
)
}
pub(super) fn register(
module: &pyo3::Bound<'_, pyo3::types::PyModule>,
) -> pyo3::PyResult<()> {
$crate::routes::definition::add_function(
module,
pyo3::wrap_pyfunction!($sync_name, module)?,
)?;
$crate::routes::definition::add_function(
module,
pyo3::wrap_pyfunction!($async_name, module)?,
)?;
Ok(())
}
}
#[cfg(feature = "trace-parity")]
pub(super) fn register_trace(
module: &pyo3::Bound<'_, pyo3::types::PyModule>,
) -> pyo3::PyResult<()> {
trace::register(module)
}
};
}
@ -130,20 +185,26 @@ mod tests {
) -> PyResult<impl Future<Output = Result<String, Error>> + Send + 'static> {
FUTURE_DROPPED.store(false, Ordering::SeqCst);
let drop_guard = (inputs.value == "pending").then_some(DropGuard);
Ok(async move {
let _drop_guard = drop_guard;
tokio::task::yield_now().await;
match inputs.value.as_str() {
"error" => Err(Error::InvalidRequest("synthetic error".to_string())),
"map_panic" => Err(Error::InvalidRequest("panic in mapper".to_string())),
"panic" => panic!("synthetic panic"),
"pending" => {
pending::<()>().await;
unreachable!()
}
_ => Ok(inputs.value),
Ok(execute_echo(inputs, drop_guard))
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
async fn execute_echo(
inputs: EchoInputs,
drop_guard: Option<DropGuard>,
) -> Result<String, Error> {
let _drop_guard = drop_guard;
tokio::task::yield_now().await;
match inputs.value.as_str() {
"error" => Err(Error::InvalidRequest("synthetic error".to_string())),
"map_panic" => Err(Error::InvalidRequest("panic in mapper".to_string())),
"panic" => panic!("synthetic panic"),
"pending" => {
pending::<()>().await;
unreachable!()
}
})
_ => Ok(inputs.value),
}
}
fn map_error(error: Error) -> PyErr {
@ -164,22 +225,22 @@ mod tests {
(
"ocr",
"aocr",
"(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=False)",
"(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None)",
),
(
"transcription",
"atranscription",
"(model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=False)",
"(model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None)",
),
(
"messages",
"amessages",
"(model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=False)",
"(model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None)",
),
(
"chat_completions",
"achat_completions",
"(model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=False)",
"(model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None)",
),
];
@ -411,6 +472,32 @@ asyncio.run(exercise())
});
}
#[cfg(feature = "trace-parity")]
#[test]
fn diagnostic_route_returns_the_response_and_filtered_trace() {
Python::initialize();
Python::attach(|py| {
let module = PyModule::new(py, "synthetic").expect("module should be created");
synthetic::register_trace(&module).expect("trace routes should register");
let locals = PyDict::new(py);
locals
.set_item("routes", &module)
.expect("module should enter Python locals");
let code = CString::new(
r#"
result = routes.echo("traced")
assert result == {
"response": "traced",
"trace": [{"function": "execute_echo", "depth": 0}],
}
"#,
)
.expect("Python source should not contain null bytes");
py.run(&code, Some(&locals), Some(&locals))
.expect("diagnostic route should return its response and trace");
});
}
#[test]
fn route_registration_rejects_duplicate_python_names() {
Python::initialize();

View file

@ -0,0 +1,29 @@
use pyo3::prelude::*;
use serde_json::Value;
use crate::errors::core_error_to_pyerr;
#[pyfunction]
fn gateway_messages<'py>(
py: Python<'py>,
model_alias: String,
provider_model: String,
api_base: String,
#[pyo3(from_py_with = litellm_python_interop::from_py)] body: Value,
) -> PyResult<Bound<'py, PyAny>> {
let future = litellm_ai_gateway::trace_parity::messages_request(
model_alias,
provider_model,
api_base,
body,
);
crate::execution::run_async(
py,
crate::function_trace::capture(future),
core_error_to_pyerr,
)
}
pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> {
super::definition::add_function(module, wrap_pyfunction!(gateway_messages, module)?)
}

View file

@ -50,14 +50,14 @@ bridge_route! {
required = {
model: String,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
body: Value,
body: serde_json::Value,
},
optional = {
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
extra_headers: Option<Value>,
extra_headers: Option<serde_json::Value>,
timeout_seconds: Option<f64>,
},
prepare = prepare_messages,

View file

@ -3,6 +3,9 @@ use pyo3::prelude::*;
#[macro_use]
mod definition;
#[cfg(feature = "trace-parity")]
mod gateway_messages;
mod audio_transcription;
mod chat_completions;
mod messages;
@ -12,5 +15,16 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
ocr::register(module)?;
audio_transcription::register(module)?;
messages::register(module)?;
chat_completions::register(module)
chat_completions::register(module)?;
#[cfg(feature = "trace-parity")]
{
let trace = PyModule::new(module.py(), "_trace")?;
ocr::register_trace(&trace)?;
audio_transcription::register_trace(&trace)?;
messages::register_trace(&trace)?;
chat_completions::register_trace(&trace)?;
gateway_messages::register_trace(&trace)?;
module.add_submodule(&trace)?;
}
Ok(())
}

View file

@ -56,16 +56,16 @@ bridge_route! {
required = {
model: String,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
document: Value,
document: serde_json::Value,
},
optional = {
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
extra_headers: Option<Value>,
extra_headers: Option<serde_json::Value>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
optional_params: Option<Value>,
optional_params: Option<serde_json::Value>,
timeout_seconds: Option<f64>,
},
prepare = prepare_ocr,

View file

@ -1,423 +0,0 @@
use std::future::Future;
use std::panic::AssertUnwindSafe;
use std::time::Duration;
use futures_util::FutureExt;
use litellm_core::error::Error;
use litellm_python_interop::{Pythonized, panic_to_pyerr, release_gil};
use pyo3::exceptions::PyRuntimeError;
use pyo3::prelude::*;
use serde::Serialize;
use tokio::runtime::{Handle, Runtime};
use tokio::time::{self, MissedTickBehavior};
pub(super) fn run_sync<T, F>(
py: Python<'_>,
future: F,
map_error: fn(Error) -> PyErr,
) -> PyResult<Py<PyAny>>
where
T: Serialize + Send + 'static,
F: Future<Output = Result<T, Error>> + Send + 'static,
{
run_sync_on(
py,
pyo3_async_runtimes::tokio::get_runtime(),
future,
map_error,
)
}
fn run_sync_on<T, F>(
py: Python<'_>,
runtime: &Runtime,
future: F,
map_error: fn(Error) -> PyErr,
) -> PyResult<Py<PyAny>>
where
T: Serialize + Send + 'static,
F: Future<Output = Result<T, Error>> + Send + 'static,
{
if Handle::try_current().is_ok() {
return Err(PyRuntimeError::new_err(
"synchronous native routes cannot run from a Tokio context; use the async route",
));
}
let result = release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?;
let result = map_core_result(result, map_error)?;
Pythonized(result).into_pyobject(py).map(Bound::unbind)
}
pub(super) fn run_async<T, F>(
py: Python<'_>,
future: F,
map_error: fn(Error) -> PyErr,
) -> PyResult<Bound<'_, PyAny>>
where
T: Serialize + Send + 'static,
F: Future<Output = Result<T, Error>> + Send + 'static,
{
pyo3_async_runtimes::tokio::future_into_py(py, async move {
let result = catch_route_panic(future).await?;
let result = map_core_result(result, map_error)?;
Ok(Pythonized(result))
})
}
fn map_core_result<T>(result: Result<T, Error>, map_error: fn(Error) -> PyErr) -> PyResult<T> {
match result {
Ok(value) => Ok(value),
Err(error) => Err(
std::panic::catch_unwind(AssertUnwindSafe(|| map_error(error)))
.map_err(panic_to_pyerr)?,
),
}
}
async fn catch_route_panic<T, F>(future: F) -> PyResult<Result<T, Error>>
where
F: Future<Output = Result<T, Error>>,
{
AssertUnwindSafe(future)
.catch_unwind()
.await
.map_err(panic_to_pyerr)
}
async fn wait_for_sync_result<T, F>(future: F) -> PyResult<Result<T, Error>>
where
F: Future<Output = Result<T, Error>>,
{
let future = catch_route_panic(future);
tokio::pin!(future);
let signal_interval = Duration::from_millis(50);
let mut signal_checks =
time::interval_at(time::Instant::now() + signal_interval, signal_interval);
signal_checks.set_missed_tick_behavior(MissedTickBehavior::Delay);
loop {
tokio::select! {
result = &mut future => return result,
_ = signal_checks.tick() => Python::attach(|py| py.check_signals())?,
}
}
}
#[cfg(test)]
mod tests {
use std::ffi::CString;
use std::future::poll_fn;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, mpsc};
use std::task::Poll;
use std::thread;
use std::time::Instant;
use pyo3::panic::PanicException;
use pyo3::types::{PyDict, PyModule};
use serde::Serializer;
use tokio::runtime::Builder;
use super::*;
fn runtime_error(error: Error) -> PyErr {
PyRuntimeError::new_err(error.to_string())
}
fn panicking_error_mapper(_error: Error) -> PyErr {
panic!("error mapper panicked")
}
struct PanickingOutput;
static ASYNC_PROBE_COMPLETED: AtomicUsize = AtomicUsize::new(0);
impl Serialize for PanickingOutput {
fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
panic!("serializer panicked")
}
}
#[pyfunction]
fn async_serialization_panic(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
run_async(py, async { Ok(PanickingOutput) }, runtime_error)
}
#[pyfunction]
fn async_runtime_probe(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
run_async(
py,
async {
ASYNC_PROBE_COMPLETED.fetch_add(1, Ordering::SeqCst);
Ok(true)
},
runtime_error,
)
}
#[pyfunction]
fn runtime_worker_count() -> usize {
pyo3_async_runtimes::tokio::get_runtime()
.metrics()
.num_workers()
}
#[pyfunction]
fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> bool {
let completion_deadline = Instant::now() + Duration::from_secs(2);
while ASYNC_PROBE_COMPLETED.load(Ordering::SeqCst) < expected_completions {
if Instant::now() >= completion_deadline {
return false;
}
thread::sleep(Duration::from_millis(1));
}
let (heartbeat_tx, heartbeat_rx) = mpsc::sync_channel(1);
pyo3_async_runtimes::tokio::get_runtime().spawn(async move {
let _ = heartbeat_tx.send(());
});
heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok()
}
fn extract_bool(py: Python<'_>, result: PyResult<Py<PyAny>>) -> bool {
result
.expect("route should complete")
.bind(py)
.extract()
.expect("result should convert")
}
#[test]
fn sync_runner_polls_future_on_the_caller_thread() {
Python::initialize();
Python::attach(|py| {
let caller_thread = std::thread::current().id();
let result = run_sync(
py,
async move { Ok(std::thread::current().id() == caller_thread) },
runtime_error,
);
assert!(extract_bool(py, result));
});
}
#[test]
fn sync_runner_releases_gil_while_waiting() {
Python::initialize();
Python::attach(|py| {
let result = run_sync(
py,
async {
let gil_acquired = tokio::time::timeout(
Duration::from_secs(2),
tokio::task::spawn_blocking(|| Python::attach(|_| true)),
)
.await;
Ok(matches!(gil_acquired, Ok(Ok(true))))
},
runtime_error,
);
assert!(extract_bool(py, result));
});
}
#[test]
fn sync_runner_rejects_calls_from_a_tokio_context() {
Python::initialize();
let runtime = Builder::new_current_thread()
.enable_all()
.build()
.expect("runtime should build");
let error = runtime.block_on(async {
Python::attach(|py| {
run_sync::<bool, _>(py, async { Ok(true) }, runtime_error)
.expect_err("sync route should reject a nested Tokio runtime")
})
});
assert_eq!(
error.to_string(),
"RuntimeError: synchronous native routes cannot run from a Tokio context; use the async route"
);
}
#[test]
fn sync_runner_can_drive_a_current_thread_runtime() {
Python::initialize();
let runtime = Builder::new_current_thread()
.enable_all()
.build()
.expect("runtime should build");
Python::attach(|py| {
let result = run_sync_on(
py,
&runtime,
async {
tokio::task::yield_now().await;
Ok(true)
},
runtime_error,
);
assert!(extract_bool(py, result));
});
}
#[test]
fn sync_runner_maps_a_panicked_future() {
Python::initialize();
Python::attach(|py| {
let error = run_sync::<bool, _>(
py,
poll_fn(|_| -> Poll<Result<bool, Error>> { panic!("route future panicked") }),
runtime_error,
)
.expect_err("panicked route should become a Python exception");
assert!(error.is_instance_of::<PanicException>(py));
assert_eq!(error.to_string(), "PanicException: route future panicked");
});
}
#[test]
fn sync_runner_maps_a_panicked_error_mapper() {
Python::initialize();
Python::attach(|py| {
let error = run_sync::<bool, _>(
py,
async { Err(Error::InvalidRequest("invalid".to_string())) },
panicking_error_mapper,
)
.expect_err("panicked mapper should become a Python exception");
assert!(error.is_instance_of::<PanicException>(py));
assert_eq!(error.to_string(), "PanicException: error mapper panicked");
});
}
#[test]
fn sync_runner_surfaces_serializer_panics() {
Python::initialize();
Python::attach(|py| {
let error = run_sync(py, async { Ok(PanickingOutput) }, runtime_error)
.expect_err("serializer panic should become a Python exception");
assert!(error.is_instance_of::<PanicException>(py));
assert_eq!(error.to_string(), "PanicException: serializer panicked");
});
}
#[test]
fn sync_runner_supports_concurrent_callers_on_the_shared_runtime() {
Python::initialize();
let barrier = Arc::new(tokio::sync::Barrier::new(2));
let callers: Vec<_> = (0..2)
.map(|_| {
let barrier = Arc::clone(&barrier);
thread::spawn(move || {
Python::attach(|py| {
extract_bool(
py,
run_sync(
py,
async move {
Ok(tokio::time::timeout(Duration::from_secs(2), barrier.wait())
.await
.is_ok())
},
runtime_error,
),
)
})
})
})
.collect();
let results: Vec<_> = callers
.into_iter()
.map(|caller| caller.join().expect("caller should not panic"))
.collect();
assert_eq!(results, vec![true, true]);
}
#[test]
fn async_runner_surfaces_serializer_panics() {
Python::initialize();
Python::attach(|py| {
let module = PyModule::new(py, "runtime").expect("module should be created");
module
.add_function(
wrap_pyfunction!(async_serialization_panic, &module)
.expect("function should wrap"),
)
.expect("function should register");
let locals = PyDict::new(py);
locals
.set_item("runtime", &module)
.expect("module should enter Python locals");
let code = CString::new(
r#"
import asyncio
async def exercise():
try:
await runtime.async_serialization_panic()
except BaseException as error:
assert type(error).__name__ == "PanicException"
assert str(error) == "serializer panicked"
else:
raise AssertionError("serializer panic was not raised")
asyncio.run(exercise())
"#,
)
.expect("Python source should not contain null bytes");
py.run(&code, Some(&locals), Some(&locals))
.expect("serializer panic should reach the Python awaiter");
});
}
#[test]
fn async_result_delivery_does_not_stall_tokio_workers() {
Python::initialize();
ASYNC_PROBE_COMPLETED.store(0, Ordering::SeqCst);
Python::attach(|py| {
let module = PyModule::new(py, "runtime").expect("module should be created");
for function in [
wrap_pyfunction!(async_runtime_probe, &module).expect("function should wrap"),
wrap_pyfunction!(runtime_worker_count, &module).expect("function should wrap"),
wrap_pyfunction!(runtime_is_responsive, &module).expect("function should wrap"),
] {
module
.add_function(function)
.expect("function should register");
}
let locals = PyDict::new(py);
locals
.set_item("runtime", &module)
.expect("module should enter Python locals");
let code = CString::new(
r#"
import asyncio
async def exercise():
worker_count = runtime.runtime_worker_count()
awaitables = [runtime.async_runtime_probe() for _ in range(worker_count)]
assert runtime.runtime_is_responsive(worker_count)
assert await asyncio.gather(*awaitables) == [True] * worker_count
asyncio.run(exercise())
"#,
)
.expect("Python source should not contain null bytes");
py.run(&code, Some(&locals), Some(&locals))
.expect("result delivery should leave Tokio workers responsive");
});
}
}

View file

@ -4,6 +4,7 @@ from litellm.rust_bridge.configuration import use_litellm_rust
from litellm.rust_bridge.loader import (
get_native_bridge,
native_bridge_available,
reset_native_bridge_cache,
)
__all__ = ["get_native_bridge", "native_bridge_available", "use_litellm_rust"]
__all__ = ["get_native_bridge", "native_bridge_available", "reset_native_bridge_cache", "use_litellm_rust"]

View file

@ -24,6 +24,12 @@ def get_native_bridge() -> ModuleType | None:
return _native
def reset_native_bridge_cache() -> None:
"""Forget the cached extension so the next lookup reimports it from disk."""
global _cached_bridge
_cached_bridge = _BRIDGE_SENTINEL
def native_bridge_available() -> bool:
"""Whether the packaged Rust extension is importable."""
return get_native_bridge() is not None

View file

@ -3,42 +3,75 @@
```text
tests/rust-python-harness/
├── __main__.py
├── cli/
│ ├── __init__.py
│ ├── catalog.py
│ └── commands.py
├── strategies/
│ ├── e2e_parity/
│ │ ├── runner.py
│ │ ├── __init__.py
│ │ ├── reporting.py
│ │ ├── sdk/
│ │ │ ├── ocr/
│ │ │ ├── messages/
│ │ │ ├── chat_completions/
│ │ │ └── responses/
│ │ └── gateway/
│ │ │ └── ocr/
│ │
│ ├── trace_parity/
│ │ ├── runner.py
│ │ ├── sdk/
│ │ └── gateway/
│ │ ├── __init__.py
│ │ ├── models.py
│ │ ├── reporting.py
│ │ └── sdk/
│ │ ├── chat_completions/
│ │ ├── messages/
│ │ ├── ocr/
│ │ └── transcription/
│ │
│ └── unit_tests/
│ ├── runner.py
│ ├── mapping_validator.py
│ ├── python_runner.py
│ └── rust_runner.py
│ ├── unit_tests_mapping/
│ │ ├── __init__.py
│ │ ├── contracts.py
│ │ ├── cases/
│ │ │ └── ocr.py
│ │ ├── mapping_report.py
│ │ ├── mappings.py
│ │ ├── mapping_validator.py
│ │ ├── reporting.py
│ │ └── runner.py
│ │
│ ├── unit_tests_parity/
│ │ ├── __init__.py
│ │ ├── reporting.py
│ │ └── runner.py
│ │
│ └── unit_tests_rust/
│ ├── __init__.py
│ ├── reporting.py
│ └── runner.py
└── shared/
├── parity/
├── tracing/
└── reporting/
├── reporting/
│ └── strategy.py
└── unit_runners/
└── suite_runner.py
```
- A strategy is a folder under `strategies/` with a one-line `AGENTS.md` and an `__init__.py` exporting exactly one `STRATEGY: StrategyDefinition`; its id must equal the folder name
- `shared/reporting/strategy.py` is the contract: runnable module/suite specs, not-implemented/skipped specs, the runner protocol, and `StrategyDefinition`
- Every `STRATEGY` explicitly classifies every SDK function; surface-aware strategies declare their surfaces and classify the complete surface-by-function matrix
- Run locally only; no CI integration
- `__main__.py` selects strategies and combines their reports; each strategy also runs independently
- `python -m tests.rust-python-harness run <strategy>|all` runs the selected strategy; `--function` is common, while each strategy exposes only its supported options
- Examples: `run e2e_parity --surface sdk --function ocr`, `run unit_tests_parity --function ocr --pytest-arg=-x`, or `run all --function ocr`
- `cli/catalog.py` discovers strategies, validates their Python definitions, and orders them; `cli/__init__.py` builds the Click command tree; `cli/commands.py` runs selected cases
- `e2e_parity/` compares SDK objects, exceptions, callbacks, and streams, or gateway HTTP responses
- `trace_parity/` compares mapped operations, call counts, and required execution ordering
- E2E and trace runners share orchestration across `sdk/` and `gateway/`; surface-specific execution lives in those folders
- `unit_tests/runner.py` combines mapping validation, Python test runs, and native Rust test runs
- `mapping_validator.py` matches Python/Rust tests by agreed names or annotations and reports missing or ambiguous counterparts
- `python_runner.py` runs existing Python tests with Rust disabled and enabled in separate processes, verifies backend selection, and compares results
- `rust_runner.py` runs Cargo tests; native Rust unit tests stay beside their implementation
- `shared/` contains reusable parity, tracing, and reporting machinery
- `trace_parity/` compares mapped operations, call counts, and required execution ordering; before running it rebuilds the native bridge with the `trace-parity` feature whenever `litellm-rust` sources are newer than the installed extension (`shared/native_build.py`)
- E2E and trace strategies load their registered module cases and run surface-specific execution from their folders
- `unit_tests_mapping/contracts.py` owns typed harness-side mapping contracts, per-function contracts live below `cases/`, and `mappings.py` exports the registry; live test discovery derives unmapped Python and Rust-only tests without an exhaustive manifest
- `unit_tests_mapping/runner.py` validates confirmed mappings against the live Python and Rust inventories and attaches the derived status report
- `unit_tests_parity/runner.py` runs each contract's `unit_parity_scope` with `LITELLM_RUST=0` and `LITELLM_RUST=1` in separate processes and requires matching outcomes, including failures; exclusions require a reason in the contract
- `unit_tests_rust/runner.py` runs each contract's focused Cargo test suite; native Rust unit tests stay beside their implementation
- `shared/unit_runners/suite_runner.py` runs typed suites registered in code with nodeids of the form `suite:<strategy_id>:<function>:<suite>`
- Every strategy declares its report sections and presentation in its own `reporting.py`; shared reporting code only provides reusable models and cell-formatting primitives
- `shared/` contains reusable parity, tracing, reporting primitives, and unit-runner machinery
- Keep fixtures with their owning API and existing Python tests in their current locations
- Each strategy folder carries an `AGENTS.md` one-liner stating what it should be doing
- Run the harness's own checks with `uv run pytest -o consider_namespace_packages=true tests/rust-python-harness/shared tests/rust-python-harness/cli tests/rust-python-harness/strategies/unit_tests_mapping tests/rust-python-harness/strategies/unit_tests_parity tests/rust-python-harness/strategies/unit_tests_rust tests/test_rust_python_harness.py -q`

View file

@ -1,105 +0,0 @@
# Rust/Python migration harness
This local harness follows [the agreed structure](AGENTS.md). The root command selects strategies and combines their reports. Each strategy has an independent entry point
```text
strategies/
e2e_parity/runner.py
sdk/ocr/fixtures/
sdk/messages/
sdk/chat_completions/
sdk/responses/
gateway/
existing_e2e_test_sdk/runner.py
trace_parity/runner.py
sdk/
gateway/
unit_tests/
runner.py
mapping_validator.py
python_runner.py
rust_runner.py
shared/
parity/
tracing/
reporting/
```
## Run locally
```bash
uv run python -m tests.rust-python-harness --list
uv run python -m tests.rust-python-harness --function ocr --plain
uv run python -m tests.rust-python-harness --strategy e2e_parity --surface sdk --function ocr --plain
uv run python -m tests.rust-python-harness.strategies.e2e_parity.runner --function ocr --plain
uv run python -m tests.rust-python-harness.strategies.trace_parity.runner --plain
uv run python -m tests.rust-python-harness.strategies.unit_tests.runner --plain
uv run python -m tests.rust-python-harness.strategies.existing_e2e_test_sdk.runner --function transcription --plain
```
Use `--interactive` for strategy and function selection, `--pytest-arg=-x` to stop pytest on its first failure, and `--coverage` to write Python coverage under `target/rust-python-harness/`. The harness enables pytest namespace-package discovery only for its own invocations
This harness has no CI execution. A configured test that fails or disappears makes the command fail. An unconfigured strategy cell remains planned and contributes no passing evidence. Interruptions and collection errors stop execution; ordinary test failures remain in the combined report while later strategies run
## Strategy responsibilities
E2E parity compares SDK objects, exceptions, callbacks, streams, and provider requests. Gateway tests compare HTTP responses. Both surfaces use the same strategy runner and keep execution details and fixtures in their own folders. OCR has recorded sync/async SDK coverage; the existing Messages and Responses bridge checks remain partial
Trace parity compares operation names through an explicit Python/Rust mapping, call counts, and required completion-before-start ordering with `shared/tracing/compare.py`. Surface tests supply captured operation intervals. No production trace instrumentation or trace case is configured yet
Unit testing combines test mapping validation, separate Python processes with Rust disabled and enabled, backend verification, result comparison, and native Cargo tests. Native tests stay beside their Rust implementation. Existing Python tests stay at their original paths. No complete Python/native unit mapping is configured yet, so these cells remain planned
The existing E2E SDK strategy retains the live provider tests configured upstream. It runs OCR, Chat Completions, and Transcription checks from their existing paths and reports them separately from parity tests. These tests require provider credentials
## Configure cases
Each strategy has a `strategy.json`. Its `functions` object defines SDK cases for OCR, Messages, Responses, Count Tokens, Chat Completions, and Transcription. E2E and trace manifests also accept a `gateway` object keyed by API name. A case has `coverage`, `selectors`, and an optional `note`
```json
{
"coverage": "partial",
"selectors": ["tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py"]
}
```
Selectors use pytest file or node syntax. A selector ending in `/` includes tests recursively from that directory
Use `planned` with no selectors until an executable contract exists, `partial` for incomplete coverage, `complete` for the full contract, and `not_applicable` when a strategy does not apply. The dashboard shows passing evidence separately from coverage completeness and LOC coverage
Unit cases use `unit_suite` instead of `selectors`, pointing to a repository-relative JSON file with this shape:
```json
{
"python_selectors": ["tests/test_api.py::test_decode"],
"cargo_manifest": "litellm-rust/Cargo.toml",
"cargo_package": "litellm-core",
"cargo_filter": "ocr::",
"backend": {
"environment_variable": "LITELLM_USE_RUST_OCR",
"probe": "tests.rust-python-harness.strategies.unit_tests.python_runner:ocr_backend"
},
"mappings": [{"python": "tests/test_api.py::test_decode", "rust": "ocr::test_decode"}]
}
```
Names match automatically when the collected Python and Rust test names agree. Explicit `mappings` handle different names, class names, and parametrized cases. Missing or ambiguous counterparts fail validation in either direction. The Cargo filter must select the same behavior as the Python selectors
The backend probe returns `python` or `rust` and runs at startup and before every test call, after fixtures have run. The OCR probe verifies the dispatch flag and native extension availability. Surface tests must also assert that calls reach their intended implementation to catch per-call fallback. Python outcomes must agree, and failed runs remain failures even if both backends fail identically
## OCR fixtures
Fixtures, provider configuration, input strategies, and recording commands live in [the OCR package](strategies/e2e_parity/sdk/ocr/fixtures/README.md). Record with provider credentials:
```bash
uv run python -m tests.rust-python-harness.strategies.e2e_parity.sdk.ocr.fixtures.record --examples 1000
```
`LITELLM_OCR_FIXTURE_DIR` and `--fixture-dir` override the default directory. Shared recording, replay, comparison, streaming, and cassette persistence live in `shared/parity/`
Run the harness's own checks locally:
```bash
uv run pytest -o consider_namespace_packages=true tests/rust-python-harness/shared tests/rust-python-harness/strategies/unit_tests tests/test_rust_python_harness.py -q
```
Existing OCR parity gaps remain visible: invalid-model provider errors differ, Reducto lacks a native contract, and the expanded Azure corpus exposes duplicate Content-Type headers. Moving the harness does not change provider responses or weaken assertions

View file

@ -1,5 +1,4 @@
"""Interactive Rust/Python SDK parity test harness."""
from .cli import main
from .cli.catalog import load_catalog
from .catalog import load_catalog
__all__ = ["load_catalog"]
__all__ = ["load_catalog", "main"]

View file

@ -1,75 +0,0 @@
from __future__ import annotations
import json
from pathlib import Path
from typing import Final
from pydantic import BaseModel, ConfigDict, ValidationError
from .shared.reporting.models import Coverage, HarnessCase, SDK_FUNCTIONS, Strategy
STRATEGIES_ROOT: Final = Path(__file__).parent / "strategies"
class CaseSpec(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
coverage: Coverage
selectors: tuple[str, ...] = ()
note: str = ""
unit_suite: str | None = None
class StrategySpec(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
order: int
id: str
label: str
description: str
functions: dict[str, CaseSpec]
gateway: dict[str, CaseSpec] = {}
def _load_strategy(source: Path) -> Strategy:
data: Final = StrategySpec.model_validate_json(source.read_text(encoding="utf-8"))
if set(data.functions) != set(SDK_FUNCTIONS):
raise ValueError(f"{source}: functions must exactly match {SDK_FUNCTIONS}")
cases: Final = tuple(
HarnessCase(
strategy_id=data.id,
strategy_label=data.label,
sdk_function=name,
coverage=case.coverage,
selectors=case.selectors,
note=case.note,
surface=surface,
unit_suite=case.unit_suite,
)
for surface, functions in (("sdk", data.functions), ("gateway", data.gateway))
for name in (SDK_FUNCTIONS if surface == "sdk" else functions)
for case in (functions[name],)
)
for case in cases:
if case.coverage in {Coverage.PLANNED, Coverage.NOT_APPLICABLE} and (case.selectors or case.unit_suite):
raise ValueError(f"{source}: {case.coverage.value} case {case.key} cannot configure tests")
if any(not selector.strip() for selector in case.selectors):
raise ValueError(f"{source}: empty selector in {case.key}")
if data.id == "unit_tests" and case.selectors:
raise ValueError(f"{source}: unit_tests must configure unit_suite instead of pytest selectors")
if data.id != "unit_tests" and case.unit_suite:
raise ValueError(f"{source}: unit_suite is only valid for unit_tests")
return Strategy(data.order, data.id, data.label, data.description, source.parent, cases)
def load_catalog(root: Path = STRATEGIES_ROOT) -> tuple[Strategy, ...]:
sources: Final = tuple(sorted(root.glob("*/strategy.json")))
if not sources:
raise ValueError(f"No strategy manifests found below {root}")
try:
strategies: Final = tuple(sorted((_load_strategy(source) for source in sources), key=lambda item: item.order))
except (ValidationError, json.JSONDecodeError) as error:
raise ValueError(str(error)) from error
if len({strategy.id for strategy in strategies}) != len(strategies):
raise ValueError(f"Duplicate strategy id in {root}")
return strategies

View file

@ -1,245 +0,0 @@
from __future__ import annotations
import argparse
import importlib.util
from collections.abc import Sequence
from pathlib import Path
from .catalog import load_catalog
from .shared.reporting.models import SDK_FUNCTIONS, HarnessCase, Strategy
from .shared.reporting.orchestration import StrategyRunner, run_strategies
from .shared.reporting.ui import make_dashboard
from .strategies.e2e_parity.runner import run as run_e2e
from .strategies.existing_e2e_test_sdk.runner import run as run_existing
from .strategies.trace_parity.runner import run as run_trace
from .strategies.unit_tests.mapping_validator import FunctionReport, build_function_report
from .strategies.unit_tests.runner import run as run_units
REPO_ROOT = Path(__file__).resolve().parents[2]
COVERAGE_ROOT = REPO_ROOT / "target" / "rust-python-harness"
def _parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
prog="rust-python-harness",
description="Run Rust/Python parity tests with a live strategy-by-SDK-function dashboard.",
)
parser.add_argument(
"-i",
"--interactive",
action="store_true",
help="pick strategies and SDK functions in a guided terminal menu",
)
parser.add_argument(
"--list", action="store_true", help="show the catalog without running tests"
)
parser.add_argument(
"--strategy",
action="append",
default=[],
metavar="ID",
help="run only this strategy",
)
parser.add_argument(
"--function",
action="append",
default=[],
dest="sdk_functions",
choices=SDK_FUNCTIONS,
help="run only this SDK function",
)
parser.add_argument("--surface", choices=("sdk", "gateway"), help="run only this API surface")
parser.add_argument(
"--validate-ledger",
action="store_true",
help=(
"report Python<->Rust test-parity ledger gaps and drift instead of "
"running the dashboard; narrow with --function"
),
)
parser.add_argument(
"--plain",
action="store_true",
help="disable the interactive terminal dashboard",
)
parser.add_argument(
"--coverage",
action="store_true",
help="write Python reference LOC reports (HTML, JSON, and XML)",
)
parser.add_argument(
"--pytest-arg",
action="append",
default=[],
metavar="ARG",
help="append an argument to pytest (repeatable, for example --pytest-arg=-x)",
)
return parser
def _coverage_pytest_args(output_root: Path = COVERAGE_ROOT) -> tuple[str, ...]:
output_root.mkdir(parents=True, exist_ok=True)
return (
"--cov=litellm",
"--cov-context=test",
f"--cov-report=json:{output_root / 'python.json'}",
f"--cov-report=xml:{output_root / 'python.xml'}",
f"--cov-report=html:{output_root / 'python-html'}",
)
def _pick_values(
title: str, options: Sequence[tuple[str, str]], input_fn=input
) -> set[str]:
print(f"\n{title} (Enter = all)")
for index, (value, label) in enumerate(options, start=1):
print(f" {index:>2}. {label} [{value}]")
while True:
answer = input_fn("Choose numbers, comma-separated: ").strip()
if not answer:
return set()
try:
indexes = {int(part.strip()) for part in answer.split(",")}
except ValueError:
print("Please enter numbers separated by commas.")
continue
if indexes and all(1 <= index <= len(options) for index in indexes):
return {options[index - 1][0] for index in indexes}
print(f"Choose values from 1 to {len(options)}.")
def _interactive_filters(strategies: Sequence[Strategy]) -> tuple[set[str], set[str]]:
strategy_ids = _pick_values(
"Testing strategies", [(strategy.id, strategy.label) for strategy in strategies]
)
sdk_functions = _pick_values(
"SDK functions",
[(name, name) for name in SDK_FUNCTIONS],
)
return strategy_ids, sdk_functions
def _select(
strategies: Sequence[Strategy], strategy_ids: set[str], sdk_functions: set[str]
) -> tuple[HarnessCase, ...]:
known_ids = {strategy.id for strategy in strategies}
unknown = strategy_ids - known_ids
if unknown:
raise ValueError(f"Unknown strategy: {', '.join(sorted(unknown))}")
return tuple(
case
for strategy in strategies
if not strategy_ids or strategy.id in strategy_ids
for case in strategy.cases
if not sdk_functions or case.sdk_function in sdk_functions
)
def _print_catalog(strategies: Sequence[Strategy]) -> None:
for strategy in strategies:
print(f"{strategy.id:20} {strategy.label}")
for case in strategy.cases:
selectors = (
", ".join(case.selectors) if case.selectors else case.unit_suite or "no test configured"
)
print(f" {case.surface}/{case.sdk_function:12} {case.coverage.value:14} {selectors}")
def _print_function_report(report: FunctionReport) -> None:
print(f"\n{report.sdk_function}")
if report.ledger is None or report.audit is None:
print(" no ledger yet")
return
ledger, audit = report.ledger, report.audit
print(
f" {ledger.mapped_count}/{ledger.total_count} python tests mapped to rust "
f"({ledger.percentage}%)"
)
print(f" {len(ledger.rust_only_tests)} rust-only tests with no python counterpart")
if audit.is_clean:
print(" ledger is in sync with the live test files")
return
for label, items in (
("ledger references a python test that no longer exists", audit.missing_python_tests),
("python test exists but is not tracked in the ledger", audit.stale_python_tests),
("ledger references a rust test that no longer exists", audit.missing_rust_tests),
("rust test exists but is not tracked in the ledger", audit.stale_rust_tests),
):
for item in items:
print(f" {label}: {item}")
def _validate_ledger(sdk_functions: set[str]) -> int:
functions = sdk_functions or set(SDK_FUNCTIONS)
reports = tuple(build_function_report(function) for function in sorted(functions))
for report in reports:
_print_function_report(report)
return 0 if all(report.is_clean for report in reports) else 1
def _resolve_runner(strategy_id: str) -> StrategyRunner:
match strategy_id:
case "e2e_parity":
return run_e2e
case "trace_parity":
return run_trace
case "unit_tests":
return run_units
case "existing_e2e_test_sdk":
return run_existing
case _:
raise ValueError(f"Unknown strategy: {strategy_id}")
def main(argv: Sequence[str] | None = None, *, strategy_id: str | None = None) -> int:
args = _parser().parse_args(argv)
if args.coverage and importlib.util.find_spec("pytest_cov") is None:
_parser().error(
"--coverage requires the project's pytest-cov dependency; run with "
"`poetry run python -m tests.rust-python-harness --coverage`"
)
if args.validate_ledger:
return _validate_ledger(set(args.sdk_functions))
catalog = load_catalog()
strategies = tuple(strategy for strategy in catalog if strategy_id is None or strategy.id == strategy_id)
if args.list:
_print_catalog(strategies)
return 0
strategy_ids = set(args.strategy)
sdk_functions = set(args.sdk_functions)
if args.interactive:
picked_strategies, picked_functions = _interactive_filters(strategies)
strategy_ids = strategy_ids or picked_strategies
sdk_functions = sdk_functions or picked_functions
try:
selected = _select(strategies, strategy_ids, sdk_functions)
cases = tuple(case for case in selected if args.surface is None or case.surface == args.surface)
except ValueError as exc:
_parser().error(str(exc))
selected_strategy_ids = {case.strategy_id for case in cases}
visible_strategies = tuple(
strategy for strategy in strategies if strategy.id in selected_strategy_ids
)
dashboard = make_dashboard(
visible_strategies,
plain=args.plain,
confidence_strategies=strategies,
)
pytest_args = [*args.pytest_arg]
if args.coverage:
pytest_args.extend(_coverage_pytest_args())
with dashboard:
exit_code, run = run_strategies(
cases=cases,
repo_root=REPO_ROOT,
on_update=dashboard.update,
pytest_args=pytest_args,
resolve_runner=_resolve_runner,
)
dashboard.finish(run, exit_code)
if args.coverage and (COVERAGE_ROOT / "python.json").exists():
print(f"Python LOC heatmap: {COVERAGE_ROOT / 'python-html' / 'index.html'}")
print(f"Machine-readable coverage: {COVERAGE_ROOT / 'python.json'}")
return exit_code

View file

@ -0,0 +1,113 @@
from __future__ import annotations
import sys
from collections.abc import Sequence
from typing import Final, cast
import click
from ..shared.reporting.models import SDK_FUNCTIONS, SdkFunction, Strategy, Surface
from .catalog import load_catalog
from .commands import run_command, select_cases
__all__ = ["load_catalog", "main"]
_INTERRUPTED_EXIT_CODE: Final = 130
def _function_option() -> click.Option:
return click.Option(
("--function", "sdk_functions"),
type=click.Choice(SDK_FUNCTIONS),
multiple=True,
help="run only this SDK function; repeat to select more than one",
)
def _run_all_command(strategies: Sequence[Strategy]) -> click.Command:
def run_all(sdk_functions: tuple[str, ...]) -> int:
selected_functions: Final = cast(frozenset[SdkFunction], frozenset(sdk_functions))
cases: Final = select_cases(strategies, selected_functions)
return run_command(strategies, cases)
return click.Command(
"all",
params=[_function_option()],
callback=run_all,
help="run every strategy",
)
def _strategy_command(strategy: Strategy) -> click.Command:
params: list[click.Parameter] = [_function_option()]
if strategy.definition.surfaces:
params.append(
click.Option(
("--surface",),
type=click.Choice(strategy.definition.surfaces),
help="run only this API surface; omit to run every surface",
)
)
runner_argument: Final = strategy.definition.runner_argument
if runner_argument is not None:
params.append(
click.Option(
(runner_argument.option, "runner_args"),
multiple=True,
metavar=runner_argument.metavar,
help=runner_argument.help,
)
)
def run_strategy(
sdk_functions: tuple[str, ...],
surface: str | None = None,
runner_args: tuple[str, ...] = (),
) -> int:
selected_functions: Final = cast(frozenset[SdkFunction], frozenset(sdk_functions))
selected_surface: Final = cast(Surface | None, surface)
cases: Final = select_cases((strategy,), selected_functions, selected_surface)
return run_command((strategy,), cases, runner_args)
return click.Command(
strategy.id,
params=params,
callback=run_strategy,
help=strategy.description,
)
def _build_cli(strategies: Sequence[Strategy]) -> click.Group:
root: Final = click.Group(
"rust-python-harness",
help="Run Rust/Python parity tests with raw progress and strategy reports.",
)
run: Final = click.Group("run", help="run one strategy or the complete harness")
run.add_command(_run_all_command(strategies))
for strategy in strategies:
run.add_command(_strategy_command(strategy))
root.add_command(run)
return root
def main(argv: Sequence[str] | None = None) -> int:
try:
strategies: Final = load_catalog()
result: Final = _build_cli(strategies).main(
args=None if argv is None else list(argv),
prog_name="rust-python-harness",
standalone_mode=False,
)
exit_code: Final = result if isinstance(result, int) else 0
except click.ClickException as error:
error.show()
return error.exit_code
except click.Abort:
click.echo("Aborted!", err=True)
return 1
except KeyboardInterrupt:
sys.stderr.write("\nInterrupted\n")
return _INTERRUPTED_EXIT_CODE
if exit_code == _INTERRUPTED_EXIT_CODE:
sys.stderr.write("Interrupted\n")
return exit_code

View file

@ -0,0 +1,116 @@
from __future__ import annotations
import hashlib
import importlib
import importlib.util
import pkgutil
import sys
from pathlib import Path
from types import ModuleType
from typing import Final
from .. import strategies as _strategies_package
from ..shared.reporting.models import SDK_FUNCTIONS, SURFACES, CaseDisposition, HarnessCase, Strategy
from ..shared.reporting.strategy import StrategyDefinition
_STRATEGIES_PACKAGE: Final = _strategies_package
STRATEGIES_ROOT: Final = Path(_STRATEGIES_PACKAGE.__path__[0])
def _load_strategy_module(name: str, folder: Path, prefix: str | None) -> ModuleType:
if prefix is not None:
return importlib.import_module(f"{prefix}.{name}")
module_name: Final = _synthetic_module_name(folder)
spec: Final = importlib.util.spec_from_file_location(
module_name, folder / "__init__.py"
)
if spec is None or spec.loader is None:
raise ValueError(f"{folder}: cannot load strategy package")
module: Final = importlib.util.module_from_spec(spec)
sys.modules[module_name] = module
try:
spec.loader.exec_module(module)
except Exception as error:
del sys.modules[module_name]
raise ValueError(f"{folder}: cannot import strategy package: {error}") from error
return module
def _synthetic_module_name(folder: Path) -> str:
digest: Final = hashlib.sha1(str(folder.resolve()).encode()).hexdigest()[:8]
return f"_harness_strategy_{folder.name}_{digest}"
def _load_strategy(name: str, folder: Path, prefix: str | None) -> Strategy:
module: Final = _load_strategy_module(name, folder, prefix)
definition: Final = getattr(module, "STRATEGY", None)
if not isinstance(definition, StrategyDefinition):
raise ValueError(f"{folder}: __init__.py must export STRATEGY: StrategyDefinition")
if definition.id != name:
raise ValueError(f"{folder}: strategy id {definition.id!r} must match folder name {name!r}")
if definition.directory.resolve() != folder.resolve():
raise ValueError(f"{folder}: strategy directory must be {folder}")
if len(set(definition.surfaces)) != len(definition.surfaces) or any(
surface not in SURFACES for surface in definition.surfaces
):
raise ValueError(f"{folder}: invalid strategy surfaces: {definition.surfaces}")
keys: Final = tuple((case.surface, case.sdk_function) for case in definition.cases)
duplicates: Final = tuple(sorted(key for key in set(keys) if keys.count(key) > 1))
if duplicates:
raise ValueError(f"{folder}: duplicate strategy cases: {duplicates}")
expected: Final = frozenset(
(surface, function)
for surface in (definition.surfaces or (None,))
for function in SDK_FUNCTIONS
)
actual: Final = frozenset(keys)
if actual != expected:
missing: Final = tuple(sorted(expected - actual))
extra: Final = tuple(sorted(actual - expected))
raise ValueError(
f"{folder}: strategy cases must exactly match its declared matrix; missing={missing}, extra={extra}"
)
incompatible: Final = tuple(
(case.surface, case.sdk_function)
for case in definition.cases
if case.spec.disposition is CaseDisposition.RUNNABLE
and not isinstance(case.spec, definition.runnable_spec)
)
if incompatible:
raise ValueError(f"{folder}: runnable cases do not match {definition.runnable_spec.__name__}: {incompatible}")
cases: Final = tuple(
HarnessCase(
strategy_id=definition.id,
strategy_label=definition.label,
sdk_function=case.sdk_function,
spec=case.spec,
surface=case.surface,
)
for case in definition.cases
)
return Strategy(
definition.order,
definition.id,
definition.label,
definition.description,
definition.directory,
cases,
definition,
)
def load_catalog(root: Path | None = None) -> tuple[Strategy, ...]:
resolved: Final = STRATEGIES_ROOT if root is None else root
prefix: Final = _STRATEGIES_PACKAGE.__name__ if resolved == STRATEGIES_ROOT else None
folders: Final = tuple(
info.name for info in pkgutil.iter_modules([str(resolved)]) if info.ispkg
)
if not folders:
raise ValueError(f"No strategy packages found below {resolved}")
strategies: Final = tuple(
_load_strategy(name, resolved / name, prefix) for name in sorted(folders)
)
ids: Final = [strategy.id for strategy in strategies]
if len(set(ids)) != len(ids):
raise ValueError(f"Duplicate strategy id in {resolved}")
return tuple(sorted(strategies, key=lambda strategy: (strategy.order, strategy.id)))

View file

@ -0,0 +1,45 @@
from __future__ import annotations
from collections.abc import Sequence, Set
from dataclasses import replace
from pathlib import Path
from typing import Final
from ..shared.reporting.models import HarnessCase, SdkFunction, Strategy, Surface
from ..shared.reporting.orchestration import run_strategies
from ..shared.reporting.ui import make_dashboard
REPO_ROOT: Final = Path(__file__).resolve().parents[3]
def select_cases(
strategies: Sequence[Strategy],
sdk_functions: Set[SdkFunction],
surface: Surface | None = None,
) -> tuple[HarnessCase, ...]:
return tuple(
case
for strategy in strategies
for case in strategy.cases
if (not sdk_functions or case.sdk_function in sdk_functions)
and (surface is None or case.surface == surface)
)
def run_command(
strategies: Sequence[Strategy],
cases: Sequence[HarnessCase],
runner_args: Sequence[str] = (),
) -> int:
grouped: Final = {
strategy.id: tuple(case for case in cases if case.strategy_id == strategy.id)
for strategy in strategies
}
visible: Final = tuple(strategy for strategy in strategies if grouped[strategy.id])
runners: Final = tuple(replace(strategy, cases=grouped[strategy.id]) for strategy in visible)
dashboard: Final = make_dashboard(visible)
with dashboard:
exit_code, run = run_strategies(runners, REPO_ROOT, dashboard.update, runner_args)
if exit_code != 130:
dashboard.finish(run, exit_code)
return exit_code

View file

@ -0,0 +1,458 @@
from __future__ import annotations
import importlib
from collections.abc import Callable, Sequence
from dataclasses import replace
from pathlib import Path
from typing import Final
import pytest
from ..shared.reporting.models import (
SDK_FUNCTIONS,
SURFACES,
CaseDisposition,
HarnessCase,
HarnessRun,
RunStatus,
Strategy,
)
from ..shared.reporting.strategy import NotImplementedCaseSpec, SkippedCaseSpec, StrategyDefinition
from ..shared.reporting.ui import PlainDashboard, final_report, make_dashboard
from ..strategies.unit_tests_mapping.mappings import UNIT_TEST_CONTRACTS
from ..strategies.unit_tests_parity import UNIT_PARITY_SUITES
from ..strategies.unit_tests_rust import RUST_SUITES
from . import main
from .catalog import STRATEGIES_ROOT, load_catalog
from .commands import REPO_ROOT, select_cases
def _strategy_source(
*,
strategy_id: str = "example",
surfaces: tuple[str, ...] = (),
drop: tuple[str | None, str] | None = None,
duplicate: tuple[str | None, str] | None = None,
incompatible: tuple[str | None, str] | None = None,
) -> str:
cells: Final = tuple(
(surface, function)
for surface in (surfaces or (None,))
for function in SDK_FUNCTIONS
if (surface, function) != drop
)
definitions: Final = tuple(
(
f"strategy.CaseDefinition({function!r}, "
"strategy.ModuleCaseSpec(coverage=models.Coverage.COMPLETE, module='tests.example'), "
f"surface={surface!r})"
if (surface, function) == incompatible
else (
f"strategy.CaseDefinition({function!r}, "
"strategy.NotImplementedCaseSpec(reason='Not implemented yet'), "
f"surface={surface!r})"
)
)
for surface, function in (*cells, *((duplicate,) if duplicate is not None else ()))
)
return (
"import importlib\n"
"from pathlib import Path\n"
"strategy = importlib.import_module('tests.rust-python-harness.shared.reporting.strategy')\n"
"models = importlib.import_module('tests.rust-python-harness.shared.reporting.models')\n"
"runner = importlib.import_module('tests.rust-python-harness.strategies.trace_parity.runner')\n"
"rendering = importlib.import_module('tests.rust-python-harness.shared.reporting.rendering')\n"
"def render(results):\n"
" return (rendering.ReportSection('Example outcomes', "
"tuple(rendering.render_case_outcome(r) for r in results)),)\n"
f"CASES = ({','.join(definitions)},)\n"
"STRATEGY = strategy.StrategyDefinition("
f"id={strategy_id!r}, order=1, label='Example strategy', description='Example description', "
"directory=Path(__file__).parent, runnable_spec=strategy.SuiteCaseSpec, cases=CASES, "
f"run=runner.run_trace_cases, render=render, surfaces={surfaces!r})\n"
)
def _write_strategy_folder(
root: Path,
name: str = "example",
*,
init_source: str | None = None,
) -> Path:
folder: Final = root / name
folder.mkdir(parents=True)
(folder / "__init__.py").write_text(init_source or _strategy_source(), encoding="utf-8")
return folder
def test_should_load_surface_aware_and_function_only_strategies() -> None:
strategies: Final = load_catalog()
assert [strategy.id for strategy in strategies] == [
"e2e_parity",
"trace_parity",
"unit_tests_mapping",
"unit_tests_parity",
"unit_tests_rust",
]
for strategy in strategies:
expected: Final = tuple(
(surface, function) for surface in (strategy.definition.surfaces or (None,)) for function in SDK_FUNCTIONS
)
assert tuple((case.surface, case.sdk_function) for case in strategy.cases) == expected
def test_unit_strategies_use_function_only_cases() -> None:
strategies: Final = {
strategy.id: strategy
for strategy in load_catalog()
if strategy.id in {"unit_tests_mapping", "unit_tests_parity", "unit_tests_rust"}
}
for sdk_function in SDK_FUNCTIONS:
cases: Final = tuple(
case for strategy in strategies.values() for case in strategy.cases if case.sdk_function == sdk_function
)
assert len(cases) == 3
assert all(case.surface is None for case in cases)
expected_mapping: Final = (
CaseDisposition.RUNNABLE if sdk_function in UNIT_TEST_CONTRACTS else CaseDisposition.NOT_IMPLEMENTED
)
assert cases[0].spec.disposition is expected_mapping
expected_parity: Final = (
CaseDisposition.RUNNABLE if sdk_function in UNIT_PARITY_SUITES else CaseDisposition.NOT_IMPLEMENTED
)
expected_rust: Final = (
CaseDisposition.RUNNABLE if sdk_function in RUST_SUITES else CaseDisposition.NOT_IMPLEMENTED
)
assert cases[1].spec.disposition is expected_parity
assert cases[2].spec.disposition is expected_rust
def test_raw_dashboard_is_always_the_default() -> None:
assert isinstance(make_dashboard(load_catalog()), PlainDashboard)
def test_every_strategy_folder_complies() -> None:
strategies: Final = load_catalog()
folders: Final = {
path.name for path in STRATEGIES_ROOT.iterdir() if path.is_dir() and (path / "__init__.py").exists()
}
assert folders == {strategy.id for strategy in strategies}
for strategy in strategies:
definition: Final = strategy.definition
assert isinstance(definition, StrategyDefinition)
assert definition.directory == strategy.directory
assert not (strategy.directory / "strategy.json").exists()
assert (strategy.directory / "AGENTS.md").exists()
for case in strategy.cases:
if case.spec.disposition is CaseDisposition.RUNNABLE:
assert isinstance(case.spec, definition.runnable_spec)
@pytest.mark.parametrize("surfaces", ((), SURFACES))
def test_should_reject_a_registry_missing_a_declared_matrix_cell(tmp_path: Path, surfaces: tuple[str, ...]) -> None:
surface: Final = surfaces[0] if surfaces else None
_write_strategy_folder(
tmp_path,
init_source=_strategy_source(surfaces=surfaces, drop=(surface, "count_tokens")),
)
with pytest.raises(ValueError, match="must exactly match its declared matrix"):
load_catalog(tmp_path)
def test_should_reject_a_duplicate_matrix_cell(tmp_path: Path) -> None:
_write_strategy_folder(tmp_path, init_source=_strategy_source(duplicate=(None, "ocr")))
with pytest.raises(ValueError, match="duplicate strategy cases"):
load_catalog(tmp_path)
def test_should_reject_invalid_declared_surfaces(tmp_path: Path) -> None:
_write_strategy_folder(tmp_path, init_source=_strategy_source(surfaces=("sdk", "sdk")))
with pytest.raises(ValueError, match="invalid strategy surfaces"):
load_catalog(tmp_path)
def test_should_reject_a_folder_without_a_strategy_definition(tmp_path: Path) -> None:
folder: Final = tmp_path / "example"
folder.mkdir()
(folder / "__init__.py").write_text("VALUE = 1\n", encoding="utf-8")
with pytest.raises(ValueError, match="STRATEGY"):
load_catalog(tmp_path)
def test_should_reject_a_strategy_id_that_differs_from_its_folder(tmp_path: Path) -> None:
_write_strategy_folder(tmp_path, init_source=_strategy_source(strategy_id="other"))
with pytest.raises(ValueError, match="must match folder name"):
load_catalog(tmp_path)
def test_should_reject_a_runnable_case_incompatible_with_the_strategy(tmp_path: Path) -> None:
_write_strategy_folder(tmp_path, init_source=_strategy_source(incompatible=(None, "ocr")))
with pytest.raises(ValueError, match="runnable cases do not match SuiteCaseSpec"):
load_catalog(tmp_path)
@pytest.mark.parametrize("case_type", (NotImplementedCaseSpec, SkippedCaseSpec))
def test_should_reject_an_unavailable_case_with_a_blank_reason(
case_type: type[NotImplementedCaseSpec] | type[SkippedCaseSpec],
) -> None:
with pytest.raises(ValueError, match="at least 1 character"):
case_type(reason=" ")
def test_should_select_functions_and_surfaces() -> None:
strategy: Final = next(strategy for strategy in load_catalog() if strategy.id == "e2e_parity")
assert tuple(case.key for case in select_cases((strategy,), {"messages"})) == (
"e2e_parity:messages",
"e2e_parity:gateway:messages",
)
assert tuple(case.display_name for case in select_cases((strategy,), {"ocr"}, "gateway")) == ("gateway/ocr",)
def _assert_unavailable_cell(strategy: Strategy, case: HarnessCase, section_title: str) -> None:
spec: Final = case.spec
assert isinstance(spec, (NotImplementedCaseSpec, SkippedCaseSpec))
scoped: Final = replace(strategy, cases=(case,))
exit_code, run = strategy.definition.run((case,), REPO_ROOT, lambda _: None)
result: Final = run.results[case.key]
expected: Final = (
RunStatus.NOT_IMPLEMENTED if spec.disposition is CaseDisposition.NOT_IMPLEMENTED else RunStatus.SKIPPED
)
report: Final = final_report(run, exit_code, (scoped,))
assert exit_code == 0
assert result.status is expected
assert spec.reason in report
assert section_title in report
expected_result: Final = "NOT RUN" if expected is RunStatus.NOT_IMPLEMENTED else "SKIPPED"
expected_implemented: Final = 0 if expected is RunStatus.NOT_IMPLEMENTED else 1
assert f"Result: {expected_result}" in report
assert f"Harness support: {expected_implemented}/1 cases implemented" in report
def test_every_unavailable_case_finishes_and_explains_itself() -> None:
section_titles: Final = {
"e2e_parity": "End-to-end parity outcomes",
"trace_parity": "trace comparisons",
"unit_tests_mapping": "Python/Rust unit-test mappings",
"unit_tests_parity": "Python backend parity outcomes",
"unit_tests_rust": "Native Rust unit-test outcomes",
}
unavailable: Final = tuple(
(strategy, case)
for strategy in load_catalog()
for case in strategy.cases
if case.spec.disposition is not CaseDisposition.RUNNABLE
)
for strategy, case in unavailable:
_assert_unavailable_cell(strategy, case, section_titles[strategy.id])
@pytest.mark.parametrize(
("strategy_id", "present", "absent"),
(
("e2e_parity", "--surface", "--pytest-arg"),
("trace_parity", "--surface", "--pytest-arg"),
("unit_tests_parity", "--pytest-arg", "--surface"),
("unit_tests_mapping", "--detail", "--surface"),
("unit_tests_rust", "--function", "--surface"),
),
)
def test_strategy_help_only_lists_supported_options(
strategy_id: str,
present: str,
absent: str,
capsys: pytest.CaptureFixture[str],
) -> None:
exit_code: Final = main(["run", strategy_id, "--help"])
captured: Final = capsys.readouterr()
assert exit_code == 0
assert present in captured.out
assert absent not in captured.out
def test_run_help_lists_all_and_every_strategy(capsys: pytest.CaptureFixture[str]) -> None:
exit_code: Final = main(["run", "--help"])
captured: Final = capsys.readouterr()
assert exit_code == 0
for command in (
"all",
"e2e_parity",
"trace_parity",
"unit_tests_mapping",
"unit_tests_parity",
"unit_tests_rust",
):
assert command in captured.out
@pytest.mark.parametrize(
"argv",
(
("list",),
("check",),
("run", "--strategy", "unit_tests_parity"),
("run", "unit_tests_parity", "--surface", "sdk"),
("run", "unit_tests_parity", "--plain"),
("run", "unit_tests_parity", "--runner-arg=-x"),
("run", "all", "--pytest-arg=-x"),
),
)
def test_removed_commands_and_options_are_rejected(argv: tuple[str, ...], capsys: pytest.CaptureFixture[str]) -> None:
exit_code: Final = main(argv)
captured: Final = capsys.readouterr()
assert exit_code == 2
assert captured.err
def test_strategy_command_forwards_repeated_filters_and_runner_arguments(
monkeypatch: pytest.MonkeyPatch,
) -> None:
cli: Final = importlib.import_module("tests.rust-python-harness.cli")
captured: list[tuple[tuple[str, ...], tuple[str, ...], tuple[str, ...]]] = []
def capture_run(
strategies: Sequence[Strategy],
cases: Sequence[HarnessCase],
runner_args: Sequence[str] = (),
) -> int:
captured.append(
(
tuple(strategy.id for strategy in strategies),
tuple(case.display_name for case in cases),
tuple(runner_args),
)
)
return 0
monkeypatch.setattr(cli, "run_command", capture_run)
assert (
main(
[
"run",
"unit_tests_parity",
"--function",
"ocr",
"--function",
"messages",
"--pytest-arg=-x",
]
)
== 0
)
assert captured == [
(("unit_tests_parity",), ("ocr", "messages"), ("-x",)),
]
def test_omitted_surface_selects_every_strategy_surface(monkeypatch: pytest.MonkeyPatch) -> None:
cli: Final = importlib.import_module("tests.rust-python-harness.cli")
selected: list[str] = []
def capture_run(
strategies: Sequence[Strategy],
cases: Sequence[HarnessCase],
runner_args: Sequence[str] = (),
) -> int:
del strategies, runner_args
selected.extend(case.display_name for case in cases)
return 0
monkeypatch.setattr(cli, "run_command", capture_run)
assert main(["run", "e2e_parity", "--function", "ocr"]) == 0
assert selected == ["sdk/ocr", "gateway/ocr"]
def test_run_all_selects_every_declared_case_once(monkeypatch: pytest.MonkeyPatch) -> None:
cli: Final = importlib.import_module("tests.rust-python-harness.cli")
selected: list[HarnessCase] = []
def capture_run(
strategies: Sequence[Strategy],
cases: Sequence[HarnessCase],
runner_args: Sequence[str] = (),
) -> int:
del strategies, runner_args
selected.extend(cases)
return 0
monkeypatch.setattr(cli, "run_command", capture_run)
assert main(["run", "all", "--function", "ocr"]) == 0
assert len(selected) == 7
assert sum(case.surface is None for case in selected) == 3
assert sum(case.surface is not None for case in selected) == 4
def test_run_reports_not_implemented_surface_as_not_run(
capsys: pytest.CaptureFixture[str],
) -> None:
exit_code: Final = main(["run", "trace_parity", "--surface", "gateway", "--function", "ocr"])
captured: Final = capsys.readouterr()
assert exit_code == 0
assert "Result: NOT RUN" in captured.out
assert "Harness support: 0/1 cases implemented" in captured.out
assert "Cases: 1 selected, 1 not implemented, 0 skipped" in captured.out
assert "Not implemented" in captured.out
assert "No gateway OCR trace-parity case is registered." in captured.out
def test_keyboard_interrupt_exits_cleanly(
monkeypatch: pytest.MonkeyPatch,
capsys: pytest.CaptureFixture[str],
) -> None:
cli: Final = importlib.import_module("tests.rust-python-harness.cli")
def interrupt() -> tuple[object, ...]:
raise KeyboardInterrupt
monkeypatch.setattr(cli, "load_catalog", interrupt)
exit_code: Final = main(["run", "all"])
captured: Final = capsys.readouterr()
assert exit_code == 130
assert captured.out == ""
assert captured.err == "\nInterrupted\n"
def test_runner_interrupt_skips_the_completion_report(
monkeypatch: pytest.MonkeyPatch,
capsys: pytest.CaptureFixture[str],
) -> None:
commands: Final = importlib.import_module("tests.rust-python-harness.cli.commands")
def interrupt_run(
strategies: Sequence[Strategy],
repo_root: Path,
on_update: Callable[[HarnessRun], None],
runner_args: Sequence[str] = (),
) -> tuple[int, HarnessRun]:
del repo_root, on_update, runner_args
run: Final = HarnessRun.from_cases(case for strategy in strategies for case in strategy.cases)
return 130, run
monkeypatch.setattr(commands, "run_strategies", interrupt_run)
exit_code: Final = main(["run", "trace_parity", "--surface", "gateway"])
captured: Final = capsys.readouterr()
assert exit_code == 130
assert "Rust <-> Python parity report" not in captured.out
assert captured.err == "Interrupted\n"

View file

@ -0,0 +1,29 @@
from __future__ import annotations
from collections.abc import Callable
from pathlib import Path
from typing import Final
import pytest
HARNESS_ROOT: Final = Path(__file__).resolve().parents[2]
@pytest.fixture(autouse=True)
def subprocess_test_environment(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("PYTHONPATH", str(HARNESS_ROOT))
monkeypatch.setenv("PYTEST_DISABLE_PLUGIN_AUTOLOAD", "1")
@pytest.fixture
def cargo_project(tmp_path: Path) -> Callable[[str, str], Path]:
def create(package: str, source: str) -> Path:
manifest: Final = tmp_path / "Cargo.toml"
manifest.write_text(
f'[package]\nname = "{package}"\nversion = "0.1.0"\nedition = "2021"\n[workspace]\n'
)
(tmp_path / "src").mkdir()
(tmp_path / "src/lib.rs").write_text(source)
return manifest
return create

View file

@ -0,0 +1,92 @@
from __future__ import annotations
import importlib.util
import os
import subprocess
import sys
from collections.abc import Iterator
from pathlib import Path
from typing import Final
from litellm.rust_bridge import get_native_bridge, reset_native_bridge_cache
MATURIN_SPEC: Final = "maturin==1.15.0"
BRIDGE_FEATURE: Final = "trace-parity"
_RUST_ROOT: Final = "litellm-rust"
_LOCKFILE: Final = "Cargo.lock"
_SOURCE_SUFFIXES: Final = frozenset({".rs", ".toml"})
_FAILURE_OUTPUT_LINES: Final = 15
def needs_rebuild(native_mtime: float | None, newest_source_mtime: float | None) -> bool:
if native_mtime is None:
return True
if newest_source_mtime is None:
return False
return newest_source_mtime > native_mtime
def _source_files(rust_root: Path) -> Iterator[Path]:
for path in rust_root.rglob("*"):
relative: Final = path.relative_to(rust_root)
if "target" in relative.parts or not path.is_file():
continue
if path.name == _LOCKFILE or path.suffix in _SOURCE_SUFFIXES:
yield path
def _newest_source_mtime(repo_root: Path) -> float | None:
rust_root: Final = repo_root / _RUST_ROOT
if not rust_root.is_dir():
return None
return max((path.stat().st_mtime for path in _source_files(rust_root)), default=None)
def _native_module_path() -> Path | None:
try:
spec: Final = importlib.util.find_spec("litellm.rust_bridge._native")
except (ImportError, ValueError):
return None
origin: Final = getattr(spec, "origin", None)
return Path(origin) if origin else None
def _drop_imported_bridge() -> None:
reset_native_bridge_cache()
for name in tuple(sys.modules):
if name.startswith("litellm.rust_bridge._native"):
del sys.modules[name]
def _rebuild(repo_root: Path) -> tuple[bool, str]:
command: Final = ("uvx", "--from", MATURIN_SPEC, "maturin", "develop", "--features", BRIDGE_FEATURE)
completed: Final = subprocess.run(
command,
cwd=repo_root,
env={**os.environ, "VIRTUAL_ENV": sys.prefix},
capture_output=True,
text=True,
check=False,
)
output: Final = f"{completed.stdout}\n{completed.stderr}".strip()
lines: Final = tuple(output.splitlines())
return completed.returncode == 0, "\n".join(lines[-_FAILURE_OUTPUT_LINES:])
def ensure_trace_bridge(repo_root: Path) -> str | None:
native_path: Final = _native_module_path()
native_mtime: Final = native_path.stat().st_mtime if native_path is not None and native_path.exists() else None
if needs_rebuild(native_mtime, _newest_source_mtime(repo_root)):
print(f"Rebuilding native Rust bridge ({BRIDGE_FEATURE} feature)...", flush=True)
succeeded: Final
output: Final
succeeded, output = _rebuild(repo_root)
if not succeeded:
return f"native Rust bridge rebuild failed:\n{output}"
_drop_imported_bridge()
bridge: Final = get_native_bridge()
if bridge is None:
return "native Rust bridge is not importable"
if getattr(bridge, "_trace", None) is None:
return f"native Rust bridge does not expose _trace; it must be built with the {BRIDGE_FEATURE} feature"
return None

View file

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

View file

@ -1 +1,7 @@
from __future__ import annotations
from typing import Final
from pydantic import TypeAdapter
JSON_OBJECT_ADAPTER: Final = TypeAdapter(dict[str, object])

View file

@ -5,11 +5,10 @@ from datetime import datetime
from itertools import accumulate
from typing import Final, Literal
from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, TypeAdapter
from pydantic import AwareDatetime, BaseModel, ConfigDict, Field
from vcr.serialize import serialize
from vcr.serializers import yamlserializer
from .recording import RecordedInteraction
from ..recorded_http import (
HttpHeader,
RecordedHttpResponse,
@ -17,8 +16,8 @@ from ..recorded_http import (
RecordedResponse,
RecordedStreamChunk,
)
_OBJECT: Final = TypeAdapter(dict[str, object])
from . import JSON_OBJECT_ADAPTER
from .recording import RecordedInteraction
class _CassetteModel(BaseModel):
@ -118,7 +117,7 @@ def serialize_cassette(
recorded_at: datetime,
request_source: Literal["recorded", "python_replay"],
) -> str:
normalized: Final = _OBJECT.validate_python(
normalized: Final = JSON_OBJECT_ADAPTER.validate_python(
yamlserializer.deserialize(
serialize(
{

View file

@ -8,11 +8,11 @@ from types import MappingProxyType
from typing import Final, Generic, Literal, Protocol, TypeVar
from hypothesis.strategies import SearchStrategy
from pydantic import BaseModel
from .inputs import generate_case_inputs
from .recording import UpstreamEndpoint, record_upstream_interactions
from .store import (
CaseT,
FixtureInput,
canonical_json,
fixture_cache_key,
@ -25,7 +25,6 @@ from .store import (
LOGGER: Final = logging.getLogger(__name__)
InputT = TypeVar("InputT", bound=FixtureInput)
InputT_contra = TypeVar("InputT_contra", bound=FixtureInput, contravariant=True)
CaseT = TypeVar("CaseT", bound=BaseModel)
class RecordingInvocation(Protocol[InputT_contra]):

View file

@ -3,14 +3,12 @@ from __future__ import annotations
import os
from collections.abc import Callable
from pathlib import Path
from typing import Final, TypeVar
from typing import Final
import pytest
from pydantic import BaseModel, ValidationError
from pydantic import ValidationError
from .store import recorded_fixtures
CaseT = TypeVar("CaseT", bound=BaseModel)
from .store import CaseT, fixture_directory, recorded_fixtures
def parametrize_recorded_fixtures(
@ -27,9 +25,10 @@ def parametrize_recorded_fixtures(
if fixture_name not in metafunc.fixturenames:
return
configured: Final = os.environ.get(env_var)
if configured == "":
raise pytest.UsageError(f"{env_var} is set but empty")
directory: Final = Path(configured).expanduser() if configured is not None else default_directory
try:
directory: Final = fixture_directory(None, configured, default_directory)
except ValueError as error:
raise pytest.UsageError(f"{env_var} is set but empty") from error
try:
fixtures: Final = recorded_fixtures(directory, case_type)
except (ValidationError, ValueError) as error:

View file

@ -1,23 +1,24 @@
from __future__ import annotations
import queue
import threading
from collections.abc import Callable, Generator, Iterable
from contextlib import contextmanager
from contextlib import AbstractContextManager
from dataclasses import dataclass
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Final, TypeVar, cast
from urllib.parse import urlsplit, urlunsplit
import httpx
from vcr.filters import remove_query_parameters
from vcr.request import Request
from ..http import (
PARITY_PROVIDER_HOST,
dropped_request_headers,
dropped_response_headers,
is_streaming_response,
local_response_header,
normalized_response_header,
)
from ..local_server import LocalHttpHandler, LocalHttpServer, serve_in_thread
from ..recorded_http import (
HttpHeader,
RecordedHttpResponse,
@ -26,7 +27,6 @@ from ..recorded_http import (
RecordedStreamChunk,
)
_PARITY_PROVIDER_HOST: Final = "parity-provider.invalid"
_SECRET_HEADERS: Final = frozenset(
{
"authorization",
@ -61,42 +61,18 @@ def _end_to_end_headers(headers: httpx.Headers) -> tuple[HttpHeader, ...]:
decoded: Final = tuple((name.decode("ascii"), value.decode("latin-1")) for name, value in headers.raw)
excluded: Final = dropped_response_headers(decoded)
return tuple(
HttpHeader(name=name, value=_normalized_response_header(name, value))
HttpHeader(name=name, value=normalized_response_header(name, value))
for name, value in decoded
if name.lower() not in excluded
)
def _normalized_response_header(name: str, value: str) -> str:
if name.lower() not in {"location", "operation-location"}:
return value
parsed: Final = urlsplit(value)
if not parsed.netloc:
return value
return urlunsplit(("http", _PARITY_PROVIDER_HOST, parsed.path, parsed.query, parsed.fragment))
def local_response_header(name: str, value: str, provider_url: str) -> str:
if name.lower() not in {"location", "operation-location"}:
return value
parsed: Final = urlsplit(value)
if parsed.hostname != _PARITY_PROVIDER_HOST:
return value
return f"{provider_url}{parsed.path}{'?' + parsed.query if parsed.query else ''}"
class _RecordingProvider(ThreadingHTTPServer):
daemon_threads = True
class _RecordingProvider(LocalHttpServer):
def __init__(self, spec: UpstreamEndpoint) -> None:
super().__init__(("127.0.0.1", 0), _RecordingHandler)
self.spec: Final = spec
self.interactions: queue.Queue[RecordedInteraction] = queue.Queue()
@property
def url(self) -> str:
return f"http://127.0.0.1:{self.server_address[1]}"
def take_interactions(self) -> tuple[RecordedInteraction, ...]:
try:
first: Final = self.interactions.get(timeout=5)
@ -106,9 +82,7 @@ class _RecordingProvider(ThreadingHTTPServer):
return (first, *remaining)
class _RecordingHandler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
class _RecordingHandler(LocalHttpHandler):
def do_POST(self) -> None:
self._forward()
@ -151,7 +125,7 @@ class _RecordingHandler(BaseHTTPRequestHandler):
recorded_request: Final = remove_query_parameters(
Request(
self.command,
f"http://{_PARITY_PROVIDER_HOST}{self.path}",
f"http://{PARITY_PROVIDER_HOST}{self.path}",
request_body,
{name: value for name, value in forwarded_headers if name.lower() not in _SECRET_HEADERS},
),
@ -191,8 +165,10 @@ class _RecordingHandler(BaseHTTPRequestHandler):
self.send_header("transfer-encoding", "chunked")
self.end_headers()
chunks: Final = tuple(self._relay_chunks(upstream.iter_bytes()))
self.wfile.write(b"0\r\n\r\n")
self.wfile.flush()
try:
self.finish_chunked()
except (BrokenPipeError, ConnectionResetError):
pass
return RecordedHttpStreamResponse(
kind="http_stream",
status_code=upstream.status_code,
@ -202,10 +178,7 @@ class _RecordingHandler(BaseHTTPRequestHandler):
def _relay_chunks(self, chunks: Iterable[bytes]) -> Generator[RecordedStreamChunk, None, None]:
for chunk in chunks:
self.wfile.write(f"{len(chunk):X}\r\n".encode("ascii"))
self.wfile.write(chunk)
self.wfile.write(b"\r\n")
self.wfile.flush()
self.write_chunk(chunk)
yield RecordedStreamChunk.from_bytes(chunk)
def _send_response(self, status_code: int, headers: tuple[HttpHeader, ...], body: bytes) -> None:
@ -218,21 +191,8 @@ class _RecordingHandler(BaseHTTPRequestHandler):
self.end_headers()
self.wfile.write(body)
def log_message(self, format: str, *args: object) -> None:
return
@contextmanager
def _recording_provider(spec: UpstreamEndpoint) -> Generator[_RecordingProvider]:
server: Final = _RecordingProvider(spec)
thread: Final = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
try:
yield server
finally:
server.shutdown()
server.server_close()
thread.join(timeout=5)
def _recording_provider(spec: UpstreamEndpoint) -> AbstractContextManager[_RecordingProvider]:
return serve_in_thread(_RecordingProvider(spec))
def _invoke_and_take_interactions(

View file

@ -8,15 +8,13 @@ from datetime import datetime, timezone
from pathlib import Path
from typing import Final, Literal, Protocol, TypeVar, cast
from pydantic import AwareDatetime, BaseModel, ConfigDict, TypeAdapter, ValidationError
from pydantic import AwareDatetime, BaseModel, ConfigDict, ValidationError
from . import JSON_OBJECT_ADAPTER
from .cassette import deserialize_cassette, serialize_cassette
from .recording import RecordedInteraction
FIXTURE_SCHEMA_VERSION: Final = 1
JSON_OBJECT: Final = TypeAdapter(dict[str, object])
class FixtureInput(Protocol):
def canonical_input(self) -> dict[str, object]: ...
@ -40,10 +38,13 @@ def fixture_cache_key(case_input: FixtureInput) -> dict[str, object]:
return case_input.canonical_input()
def fixture_path(directory: Path, case_input: FixtureInput) -> Path:
def _fixture_digest(case_input: FixtureInput) -> str:
input_json: Final = canonical_json(fixture_cache_key(case_input))
digest: Final = hashlib.sha256(input_json.encode("utf-8")).hexdigest()
return directory / f"{digest}.yaml"
return hashlib.sha256(input_json.encode("utf-8")).hexdigest()
def fixture_path(directory: Path, case_input: FixtureInput) -> Path:
return directory / f"{_fixture_digest(case_input)}.yaml"
def load_fixture(directory: Path, case_input: FixtureInput, case_type: type[CaseT]) -> CaseT | None:
@ -87,7 +88,7 @@ def save_fixture(
def read_fixture(path: Path, case_type: type[CaseT]) -> CaseT:
contents: Final = path.read_text(encoding="utf-8")
if path.suffix == ".json":
return _load_fixture(JSON_OBJECT.validate_json(contents), path, case_type)
return _load_fixture(JSON_OBJECT_ADAPTER.validate_json(contents), path, case_type)
try:
cassette: Final = deserialize_cassette(contents)
return case_type.model_validate(cassette.case_data())
@ -117,10 +118,12 @@ def recorded_fixtures(directory: Path, case_type: type[CaseT]) -> tuple[CaseT, .
def fixture_directory(configured: Path | None, env_value: str | None, default: Path) -> Path:
return (configured or Path(env_value or default)).expanduser()
if configured is not None:
return configured.expanduser()
if env_value == "":
raise ValueError("fixture directory environment variable is set but empty")
return Path(env_value).expanduser() if env_value is not None else default.expanduser()
def fixture_id(case_input: FixtureInput, prefix: str) -> str:
input_json: Final = canonical_json(case_input.canonical_input())
digest: Final = hashlib.sha256(input_json.encode("utf-8")).hexdigest()[:8]
return f"{prefix}-{digest}"
return f"{prefix}-{_fixture_digest(case_input)[:8]}"

View file

@ -10,9 +10,6 @@ from vcr import VCR
from vcr.request import Request
from ..fixture_models import ParityCase, SdkInputBase
from .cassette import deserialize_cassette
from .recording import RecordedInteraction
from .store import load_fixture, save_fixture
from ..recorded_http import (
HttpHeader,
RecordedHttpResponse,
@ -21,6 +18,9 @@ from ..recorded_http import (
RecordedStreamChunk,
)
from ..replay import replay_server
from .cassette import deserialize_cassette
from .recording import RecordedInteraction
from .store import load_fixture, save_fixture
_URI: Final = "http://parity-provider.invalid/operation?api-version=1"
@ -93,3 +93,18 @@ def test_cassette_preserves_duplicate_response_headers(tmp_path: Path) -> None:
save_fixture(tmp_path, sdk_input, case, (RecordedInteraction(Request("POST", _URI, b"", {}), response),))
assert load_fixture(tmp_path, sdk_input, ParityCase[_Input]) == case
def test_local_replay_skips_recorded_retry_delay() -> None:
response: Final = RecordedHttpResponse.from_bytes(
200,
(HttpHeader(name="retry-after", value="5"),),
b"{}",
)
with replay_server() as server:
server.enqueue_response(response)
replayed: Final = httpx.post(f"{server.url}/operation", content=b"{}")
server.take_requests(1)
assert replayed.headers["retry-after"] == "0"

View file

@ -2,10 +2,8 @@ from __future__ import annotations
import logging
import threading
from collections.abc import Generator
from contextlib import contextmanager
from contextlib import AbstractContextManager
from dataclasses import dataclass
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Final, Literal
@ -14,6 +12,8 @@ import pytest
from hypothesis import strategies as st
from pydantic import BaseModel, ConfigDict
from ..local_server import LocalHttpHandler, LocalHttpServer, serve_in_thread
from ..recorded_http import RecordedResponse
from .pipeline import (
RecordingInvocation,
RecordingTarget,
@ -22,7 +22,6 @@ from .pipeline import (
)
from .recording import UpstreamEndpoint
from .store import fixture_path
from ..recorded_http import RecordedResponse
class _FixtureInput(BaseModel):
@ -41,20 +40,12 @@ class _ParityCase(BaseModel):
provider_responses: tuple[RecordedResponse, ...]
class _Upstream(ThreadingHTTPServer):
daemon_threads = True
class _Upstream(LocalHttpServer):
def __init__(self, status: int = 200) -> None:
super().__init__(("127.0.0.1", 0), _UpstreamHandler)
self.response_status: Final = status
@property
def url(self) -> str:
return f"http://127.0.0.1:{self.server_address[1]}"
class _UpstreamHandler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
class _UpstreamHandler(LocalHttpHandler):
def do_POST(self) -> None:
length: Final = int(self.headers.get("content-length") or "0")
@ -68,21 +59,8 @@ class _UpstreamHandler(BaseHTTPRequestHandler):
self.end_headers()
self.wfile.write(body)
def log_message(self, format: str, *args: object) -> None:
return
@contextmanager
def _upstream(status: int = 200) -> Generator[_Upstream]:
server: Final = _Upstream(status)
thread: Final = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
try:
yield server
finally:
server.shutdown()
server.server_close()
thread.join(timeout=5)
def _upstream(status: int = 200) -> AbstractContextManager[_Upstream]:
return serve_in_thread(_Upstream(status))
@dataclass(frozen=True, slots=True)

View file

@ -3,10 +3,9 @@ from __future__ import annotations
import asyncio
import queue
import threading
from collections.abc import AsyncIterator, Callable, Generator, Iterator
from contextlib import contextmanager
from collections.abc import AsyncIterator, Callable, Iterator
from contextlib import AbstractContextManager
from dataclasses import dataclass
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Final, Literal
@ -17,19 +16,8 @@ from openai._streaming import SSEDecoder
from pydantic import BaseModel, ConfigDict
from ..compare import assert_request_parity
from .pipeline import RecordingTarget, record_fixtures
from .recording import (
UpstreamEndpoint,
record_upstream_interactions,
record_upstream_responses,
)
from .store import (
FIXTURE_SCHEMA_VERSION,
fixture_path,
load_fixture,
recorded_fixtures,
)
from ..inprocess import InProcessExecution, run_in_process, run_in_process_async
from ..local_server import LocalHttpHandler, LocalHttpServer, serve_in_thread
from ..recorded_http import (
HttpHeader,
RecordedHttpStreamResponse,
@ -45,6 +33,18 @@ from ..stream import (
consume_async_stream,
consume_sync_stream,
)
from .pipeline import RecordingTarget, record_fixtures
from .recording import (
UpstreamEndpoint,
record_upstream_interactions,
record_upstream_responses,
)
from .store import (
FIXTURE_SCHEMA_VERSION,
fixture_path,
load_fixture,
recorded_fixtures,
)
_SSE_CHUNKS: Final = (
b'data: {"choices":[{"delta":{"content":"hello"}}]}\n\n',
@ -146,9 +146,7 @@ class _Invocation:
self.sdk_call(provider_url, case_input)
class _ControlledUpstream(ThreadingHTTPServer):
daemon_threads = True
class _ControlledUpstream(LocalHttpServer):
def __init__(self, stream_chunks: tuple[bytes, ...]) -> None:
super().__init__(("127.0.0.1", 0), _ControlledUpstreamHandler)
self.stream_chunks: Final = stream_chunks
@ -158,10 +156,6 @@ class _ControlledUpstream(ThreadingHTTPServer):
self.max_active_requests: int = 0
self.request_count: int = 0
@property
def url(self) -> str:
return f"http://127.0.0.1:{self.server_address[1]}"
def start_request(self) -> None:
with self.lock:
self.active_requests += 1
@ -176,9 +170,7 @@ class _ControlledUpstream(ThreadingHTTPServer):
self.active_requests -= 1
class _ControlledUpstreamHandler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
class _ControlledUpstreamHandler(LocalHttpHandler):
def do_POST(self) -> None:
upstream: Final = self.server
assert isinstance(upstream, _ControlledUpstream)
@ -207,13 +199,7 @@ class _ControlledUpstreamHandler(BaseHTTPRequestHandler):
self.send_header("content-type", "text/event-stream")
self.send_header("transfer-encoding", "chunked")
self.end_headers()
for chunk in upstream.stream_chunks:
self.wfile.write(f"{len(chunk):X}\r\n".encode("ascii"))
self.wfile.write(chunk)
self.wfile.write(b"\r\n")
self.wfile.flush()
self.wfile.write(b"0\r\n\r\n")
self.wfile.flush()
self.write_chunked(upstream.stream_chunks)
return
if self.path == "/error":
self._send_json(429, b'{"error":{"message":"rate limited"}}')
@ -252,21 +238,10 @@ class _ControlledUpstreamHandler(BaseHTTPRequestHandler):
self.end_headers()
self.wfile.write(body)
def log_message(self, format: str, *args: object) -> None:
return
@contextmanager
def _controlled_upstream(stream_chunks: tuple[bytes, ...] = _SSE_CHUNKS) -> Generator[_ControlledUpstream]:
server: Final = _ControlledUpstream(stream_chunks)
thread: Final = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
try:
yield server
finally:
server.shutdown()
server.server_close()
thread.join(timeout=5)
def _controlled_upstream(
stream_chunks: tuple[bytes, ...] = _SSE_CHUNKS,
) -> AbstractContextManager[_ControlledUpstream]:
return serve_in_thread(_ControlledUpstream(stream_chunks))
def _case(identifier: str) -> _FixtureInput:

View file

@ -2,6 +2,9 @@ from __future__ import annotations
from collections.abc import Iterable
from typing import Final
from urllib.parse import urlsplit, urlunsplit
PARITY_PROVIDER_HOST: Final = "parity-provider.invalid"
HOP_BY_HOP_HEADERS: Final[frozenset[str]] = frozenset(
{
@ -52,3 +55,21 @@ def dropped_response_headers(headers: Iterable[tuple[str, str]]) -> frozenset[st
def is_streaming_response(content_type: str) -> bool:
return "text/event-stream" in content_type.lower()
def normalized_response_header(name: str, value: str) -> str:
if name.lower() not in {"location", "operation-location"}:
return value
parsed: Final = urlsplit(value)
if not parsed.netloc:
return value
return urlunsplit(("http", PARITY_PROVIDER_HOST, parsed.path, parsed.query, parsed.fragment))
def local_response_header(name: str, value: str, provider_url: str) -> str:
if name.lower() not in {"location", "operation-location"}:
return value
parsed: Final = urlsplit(value)
if parsed.hostname != PARITY_PROVIDER_HOST:
return value
return f"{provider_url}{parsed.path}{'?' + parsed.query if parsed.query else ''}"

View file

@ -1,136 +0,0 @@
from __future__ import annotations
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any
@dataclass(frozen=True, slots=True)
class LedgerEntry:
python_file: str
python_test: str
status: str
rust_file: str
rust_test: str
justification: str
reason: str
@dataclass(frozen=True, slots=True)
class RustOnlyEntry:
rust_file: str
rust_test: str
reason: str
@dataclass(frozen=True, slots=True)
class TestLedger:
sdk_function: str
python_scope: tuple[str, ...]
rust_scope: tuple[str, ...]
entries: tuple[LedgerEntry, ...]
rust_only_tests: tuple[RustOnlyEntry, ...]
@property
def mapped_count(self) -> int:
return sum(1 for entry in self.entries if entry.status == "mapped")
@property
def total_count(self) -> int:
return len(self.entries)
@property
def percentage(self) -> float:
if self.total_count == 0:
return 0.0
return round(100.0 * self.mapped_count / self.total_count, 1)
def _require_string(value: Any, field: str, source: Path) -> str:
if not isinstance(value, str) or not value.strip():
raise ValueError(f"{source}: {field} must be a non-empty string")
return value
def _require_string_list(value: Any, field: str, source: Path) -> tuple[str, ...]:
if not isinstance(value, list) or not all(isinstance(item, str) and item for item in value):
raise ValueError(f"{source}: {field} must be a list of non-empty strings")
return tuple(value)
def _load_entry(data: Any, index: int, source: Path) -> LedgerEntry:
if not isinstance(data, dict):
raise ValueError(f"{source}: entries[{index}] must be an object")
python_file = _require_string(data.get("python_file"), f"entries[{index}].python_file", source)
python_test = _require_string(data.get("python_test"), f"entries[{index}].python_test", source)
status = data.get("status")
if status not in ("mapped", "unmapped"):
raise ValueError(f"{source}: entries[{index}].status must be 'mapped' or 'unmapped'")
if status == "mapped":
rust_file = _require_string(data.get("rust_file"), f"entries[{index}].rust_file", source)
rust_test = _require_string(data.get("rust_test"), f"entries[{index}].rust_test", source)
justification = _require_string(
data.get("justification"), f"entries[{index}].justification", source
)
return LedgerEntry(
python_file=python_file,
python_test=python_test,
status=status,
rust_file=rust_file,
rust_test=rust_test,
justification=justification,
reason="",
)
reason = _require_string(data.get("reason"), f"entries[{index}].reason", source)
return LedgerEntry(
python_file=python_file,
python_test=python_test,
status=status,
rust_file="",
rust_test="",
justification="",
reason=reason,
)
def _load_rust_only_entry(data: Any, index: int, source: Path) -> RustOnlyEntry:
if not isinstance(data, dict):
raise ValueError(f"{source}: rust_only_tests[{index}] must be an object")
return RustOnlyEntry(
rust_file=_require_string(data.get("rust_file"), f"rust_only_tests[{index}].rust_file", source),
rust_test=_require_string(data.get("rust_test"), f"rust_only_tests[{index}].rust_test", source),
reason=_require_string(data.get("reason"), f"rust_only_tests[{index}].reason", source),
)
def load_ledger(path: Path) -> TestLedger:
with path.open(encoding="utf-8") as stream:
data = json.load(stream)
sdk_function = _require_string(data.get("sdk_function"), "sdk_function", path)
python_scope = _require_string_list(data.get("python_scope"), "python_scope", path)
rust_scope = _require_string_list(data.get("rust_scope"), "rust_scope", path)
entries_data = data.get("entries")
if not isinstance(entries_data, list):
raise ValueError(f"{path}: entries must be a list")
entries = tuple(
_load_entry(entry, index, path) for index, entry in enumerate(entries_data)
)
rust_only_data = data.get("rust_only_tests")
if not isinstance(rust_only_data, list):
raise ValueError(f"{path}: rust_only_tests must be a list")
rust_only_tests = tuple(
_load_rust_only_entry(entry, index, path) for index, entry in enumerate(rust_only_data)
)
return TestLedger(
sdk_function=sdk_function,
python_scope=python_scope,
rust_scope=rust_scope,
entries=entries,
rust_only_tests=rust_only_tests,
)

View file

@ -0,0 +1,56 @@
from __future__ import annotations
import threading
from collections.abc import Generator, Iterable
from contextlib import contextmanager
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Final, TypeVar
class LocalHttpServer(ThreadingHTTPServer):
daemon_threads = True
@property
def url(self) -> str:
return f"http://127.0.0.1:{self.server_address[1]}"
class LocalHttpHandler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
def write_chunk(self, chunk: bytes) -> None:
self.wfile.write(f"{len(chunk):X}\r\n".encode("ascii"))
self.wfile.write(chunk)
self.wfile.write(b"\r\n")
self.wfile.flush()
def finish_chunked(self) -> None:
self.wfile.write(b"0\r\n\r\n")
self.wfile.flush()
def write_chunked(self, chunks: Iterable[bytes]) -> None:
for chunk in chunks:
self.write_chunk(chunk)
self.finish_chunked()
def log_message(self, format: str, *args: object) -> None:
return
ServerT = TypeVar("ServerT", bound=LocalHttpServer)
@contextmanager
def serve_in_thread(server: ServerT, poll_interval: float = 0.5) -> Generator[ServerT]:
thread: Final = threading.Thread(
target=server.serve_forever,
kwargs={"poll_interval": poll_interval},
daemon=True,
)
thread.start()
try:
yield server
finally:
server.shutdown()
server.server_close()
thread.join(timeout=5)

View file

@ -2,15 +2,13 @@ from __future__ import annotations
import base64
import queue
import threading
from collections.abc import Generator
from contextlib import contextmanager
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from contextlib import AbstractContextManager
from typing import Final
from pydantic import JsonValue, TypeAdapter
from .fixtures.recording import local_response_header
from .http import local_response_header
from .local_server import LocalHttpHandler, LocalHttpServer, serve_in_thread
from .models import CapturedRequest
from .recorded_http import RecordedHttpResponse, RecordedHttpStreamResponse, RecordedResponse
@ -28,18 +26,18 @@ EXCLUDED_REQUEST_HEADERS: Final = frozenset(
EXCLUDED_RESPONSE_HEADERS: Final = frozenset({"content-length", "transfer-encoding", "connection"})
class ReplayServer(ThreadingHTTPServer):
daemon_threads = True
def _replay_response_header(name: str, value: str, provider_url: str) -> str:
if name.lower() == "retry-after":
return "0"
return local_response_header(name, value, provider_url)
class ReplayServer(LocalHttpServer):
def __init__(self) -> None:
super().__init__(("127.0.0.1", 0), _ReplayHandler)
self.responses: queue.Queue[RecordedResponse] = queue.Queue()
self.requests: queue.Queue[CapturedRequest] = queue.Queue()
@property
def url(self) -> str:
return f"http://127.0.0.1:{self.server_address[1]}"
def enqueue_response(self, response: RecordedResponse) -> None:
self.responses.put(response)
@ -56,9 +54,7 @@ class ReplayServer(ThreadingHTTPServer):
self.requests.get_nowait()
class _ReplayHandler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
class _ReplayHandler(LocalHttpHandler):
def do_POST(self) -> None:
self._replay()
@ -111,7 +107,7 @@ class _ReplayHandler(BaseHTTPRequestHandler):
self.send_response_only(response.status_code)
for header in response.headers:
if header.name.lower() not in EXCLUDED_RESPONSE_HEADERS:
self.send_header(header.name, local_response_header(header.name, header.value, provider.url))
self.send_header(header.name, _replay_response_header(header.name, header.value, provider.url))
if isinstance(response, RecordedHttpResponse):
response_body: Final = response.body_bytes()
self.send_header("content-length", str(len(response_body)))
@ -121,27 +117,8 @@ class _ReplayHandler(BaseHTTPRequestHandler):
assert isinstance(response, RecordedHttpStreamResponse)
self.send_header("transfer-encoding", "chunked")
self.end_headers()
for chunk in response.chunks:
data = chunk.data_bytes()
self.wfile.write(f"{len(data):X}\r\n".encode("ascii"))
self.wfile.write(data)
self.wfile.write(b"\r\n")
self.wfile.flush()
self.wfile.write(b"0\r\n\r\n")
self.wfile.flush()
def log_message(self, format: str, *args: object) -> None:
return
self.write_chunked(chunk.data_bytes() for chunk in response.chunks)
@contextmanager
def replay_server() -> Generator[ReplayServer]:
server: Final = ReplayServer()
thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.01}, daemon=True)
thread.start()
try:
yield server
finally:
server.shutdown()
server.server_close()
thread.join(timeout=5)
def replay_server() -> AbstractContextManager[ReplayServer]:
return serve_in_thread(ReplayServer(), poll_interval=0.01)

View file

@ -26,6 +26,7 @@ from .replay import ReplayServer, replay_server
WORKER_RESULT_PREFIX: Final = "LITELLM_PARITY_RESULT "
WORKER_RESULT_ADAPTER: Final[TypeAdapter[WorkerResult]] = TypeAdapter(WorkerResult)
PROJECT_ROOT: Final = Path(__file__).resolve().parents[4]
@dataclass(frozen=True, slots=True)
@ -39,7 +40,7 @@ class SubprocessRunner:
sys.executable,
"-m",
".".join(
self.entrypoint.resolve().relative_to(Path(__file__).resolve().parents[4]).with_suffix("").parts
self.entrypoint.resolve().relative_to(PROJECT_ROOT).with_suffix("").parts
),
"--parity-worker",
provider_url,
@ -54,7 +55,7 @@ class ExecutionVariant:
class SubprocessWorker:
def __init__(self, runner: SubprocessRunner, provider: ReplayServer, variant: ExecutionVariant) -> None:
project_root: Final = str(Path(__file__).resolve().parents[4])
project_root: Final = str(PROJECT_ROOT)
existing_pythonpath: Final = os.environ.get("PYTHONPATH")
env: Final = {
**os.environ,
@ -165,15 +166,6 @@ def execution_worker(
worker.close()
def run_execution(
worker: SubprocessWorker,
case_file: Path,
route: str,
responses: tuple[RecordedResponse, ...],
) -> Execution:
return worker.execute(case_file, route, responses)
@contextmanager
def execution_worker_pair(
runner: SubprocessRunner,

View file

@ -81,80 +81,61 @@ def _failed(phase: Literal["creation", "iteration"], error: Exception) -> Stream
)
def _creation_failure(error: Exception) -> StreamOutcome:
return StreamOutcome(
wrapper_type=None,
supports_sync_iteration=None,
supports_async_iteration=None,
chunks=(),
chunk_types=(),
terminal=_failed("creation", error),
)
def _stream_outcome(
stream: object,
chunks: Iterable[object],
terminal: StreamTerminal,
) -> StreamOutcome:
recorded: Final = tuple(chunks)
return StreamOutcome(
wrapper_type=type(stream),
supports_sync_iteration=hasattr(stream, "__iter__"),
supports_async_iteration=hasattr(stream, "__aiter__"),
chunks=recorded,
chunk_types=tuple(type(chunk) for chunk in recorded),
terminal=terminal,
)
def consume_sync_stream(create: Callable[[], Iterable[object]]) -> StreamOutcome:
try:
stream: Final = create()
except Exception as error:
return StreamOutcome(
wrapper_type=None,
supports_sync_iteration=None,
supports_async_iteration=None,
chunks=(),
chunk_types=(),
terminal=_failed("creation", error),
)
return _creation_failure(error)
chunks: list[object] = [] # mutable-ok: iterator consumption builds an ordered trace
try:
for chunk in stream:
chunks.append(chunk) # noqa: PERF402 # partial trace is required if iteration raises
except Exception as error:
recorded: Final = tuple(chunks)
return StreamOutcome(
wrapper_type=type(stream),
supports_sync_iteration=hasattr(stream, "__iter__"),
supports_async_iteration=hasattr(stream, "__aiter__"),
chunks=recorded,
chunk_types=tuple(type(chunk) for chunk in recorded),
terminal=_failed("iteration", error),
)
completed_chunks: Final = tuple(chunks)
return StreamOutcome(
wrapper_type=type(stream),
supports_sync_iteration=hasattr(stream, "__iter__"),
supports_async_iteration=hasattr(stream, "__aiter__"),
chunks=completed_chunks,
chunk_types=tuple(type(chunk) for chunk in completed_chunks),
terminal=StreamCompleted(),
)
return _stream_outcome(stream, chunks, _failed("iteration", error))
return _stream_outcome(stream, chunks, StreamCompleted())
async def consume_async_stream(create: Callable[[], Awaitable[AsyncIterable[object]]]) -> StreamOutcome:
try:
stream: Final = await create()
except Exception as error:
return StreamOutcome(
wrapper_type=None,
supports_sync_iteration=None,
supports_async_iteration=None,
chunks=(),
chunk_types=(),
terminal=_failed("creation", error),
)
return _creation_failure(error)
chunks: list[object] = [] # mutable-ok: iterator consumption builds an ordered trace
try:
async for chunk in stream:
chunks.append(chunk)
except Exception as error:
recorded: Final = tuple(chunks)
return StreamOutcome(
wrapper_type=type(stream),
supports_sync_iteration=hasattr(stream, "__iter__"),
supports_async_iteration=hasattr(stream, "__aiter__"),
chunks=recorded,
chunk_types=tuple(type(chunk) for chunk in recorded),
terminal=_failed("iteration", error),
)
completed_chunks: Final = tuple(chunks)
return StreamOutcome(
wrapper_type=type(stream),
supports_sync_iteration=hasattr(stream, "__iter__"),
supports_async_iteration=hasattr(stream, "__aiter__"),
chunks=completed_chunks,
chunk_types=tuple(type(chunk) for chunk in completed_chunks),
terminal=StreamCompleted(),
)
return _stream_outcome(stream, chunks, _failed("iteration", error))
return _stream_outcome(stream, chunks, StreamCompleted())
def normalize_chunk(chunk: object) -> object:

View file

@ -1,17 +1,25 @@
from __future__ import annotations
from collections.abc import Iterable
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from time import monotonic
from typing import Iterable
from typing import TYPE_CHECKING, Final, Literal, TypeAlias, assert_never
if TYPE_CHECKING:
from .strategy import CaseSpec, StrategyDefinition
class Coverage(str, Enum):
COMPLETE = "complete"
PARTIAL = "partial"
PLANNED = "planned"
NOT_APPLICABLE = "not_applicable"
class CaseDisposition(str, Enum):
RUNNABLE = "runnable"
NOT_IMPLEMENTED = "not_implemented"
SKIPPED = "skipped"
class RunStatus(str, Enum):
@ -23,33 +31,45 @@ class RunStatus(str, Enum):
SKIPPED = "skipped"
ERROR = "error"
MISSING = "missing"
PLANNED = "planned"
NOT_APPLICABLE = "not_applicable"
NOT_IMPLEMENTED = "not_implemented"
class ConfidenceLevel(str, Enum):
HIGH = "HIGH"
MEDIUM = "MEDIUM"
LOW = "LOW"
SDK_FUNCTIONS = ("ocr", "messages", "responses", "count_tokens", "chat_completions", "transcription")
SdkFunction: TypeAlias = Literal["ocr", "messages", "responses", "count_tokens", "chat_completions", "transcription"]
Surface: TypeAlias = Literal["sdk", "gateway"]
SURFACES: Final[tuple[Surface, ...]] = ("sdk", "gateway")
SDK_FUNCTIONS: Final[tuple[SdkFunction, ...]] = (
"ocr",
"messages",
"responses",
"count_tokens",
"chat_completions",
"transcription",
)
@dataclass(frozen=True)
class HarnessCase:
strategy_id: str
strategy_label: str
sdk_function: str
coverage: Coverage
selectors: tuple[str, ...]
note: str = ""
surface: str = "sdk"
unit_suite: str | None = None
sdk_function: SdkFunction
spec: CaseSpec
surface: Surface | None = None
@property
def key(self) -> str:
return f"{self.strategy_id}:{self.sdk_function}" if self.surface == "sdk" else f"{self.strategy_id}:gateway:{self.sdk_function}"
return (
f"{self.strategy_id}:{self.sdk_function}"
if self.surface in {None, "sdk"}
else f"{self.strategy_id}:gateway:{self.sdk_function}"
)
@property
def display_name(self) -> str:
return self.sdk_function if self.surface is None else f"{self.surface}/{self.sdk_function}"
@property
def coverage(self) -> Coverage | None:
return self.spec.coverage
@dataclass(frozen=True)
@ -60,6 +80,7 @@ class Strategy:
description: str
directory: Path
cases: tuple[HarnessCase, ...]
definition: StrategyDefinition
@dataclass
@ -74,6 +95,7 @@ class CaseResult:
errors: int = 0
outcomes: dict[str, RunStatus] = field(default_factory=dict)
durations: dict[str, float] = field(default_factory=dict)
artifacts: dict[str, tuple[ResultArtifact, ...]] = field(default_factory=dict)
@property
def total(self) -> int:
@ -83,10 +105,18 @@ class CaseResult:
def duration(self) -> float:
return sum(self.durations.values())
def record(self, nodeid: str, status: RunStatus, duration: float = 0.0) -> None:
def record(
self,
nodeid: str,
status: RunStatus,
duration: float = 0.0,
artifacts: tuple[ResultArtifact, ...] = (),
) -> None:
"""Record a terminal outcome, allowing teardown errors to replace a pass."""
self.outcomes[nodeid] = status
self.durations[nodeid] = self.durations.get(nodeid, 0.0) + duration
self.add_duration(nodeid, duration)
if artifacts:
self.artifacts[nodeid] = artifacts
self.completed = set(self.outcomes)
values = tuple(self.outcomes.values())
self.passed = values.count(RunStatus.PASSED)
@ -95,16 +125,25 @@ class CaseResult:
self.errors = values.count(RunStatus.ERROR)
self.finalize()
def add_duration(self, nodeid: str, duration: float) -> None:
self.durations[nodeid] = self.durations.get(nodeid, 0.0) + duration
def set_initial_status(self) -> None:
if self.case.coverage is Coverage.NOT_APPLICABLE:
self.status = RunStatus.NOT_APPLICABLE
elif not self.case.selectors and not self.case.unit_suite:
self.status = RunStatus.PLANNED
else:
self.status = RunStatus.QUEUED
disposition: Final = self.case.spec.disposition
match disposition:
case CaseDisposition.RUNNABLE:
self.status = RunStatus.QUEUED
return
case CaseDisposition.NOT_IMPLEMENTED:
self.status = RunStatus.NOT_IMPLEMENTED
return
case CaseDisposition.SKIPPED:
self.status = RunStatus.SKIPPED
return
assert_never(disposition)
def finalize(self) -> None:
if self.status in {RunStatus.NOT_APPLICABLE, RunStatus.PLANNED}:
if self.status in {RunStatus.NOT_IMPLEMENTED, RunStatus.SKIPPED} and not self.collected:
return
if not self.collected:
self.status = RunStatus.MISSING
@ -118,11 +157,18 @@ class CaseResult:
self.status = RunStatus.SKIPPED
@dataclass(frozen=True, slots=True)
class ResultArtifact:
kind: str
body: str
@dataclass
class HarnessRun:
results: dict[str, CaseResult]
current_nodeid: str | None = None
failures: list[tuple[str, str]] = field(default_factory=list)
strategy_durations: dict[str, float] = field(default_factory=dict)
started_at: float = field(default_factory=monotonic)
finished_at: float | None = None
@ -131,92 +177,20 @@ class HarnessRun:
return (self.finished_at or monotonic()) - self.started_at
@property
def unique_tests(self) -> int:
def unique_checks(self) -> int:
return len(
{nodeid for result in self.results.values() for nodeid in result.collected}
)
@property
def completed_tests(self) -> int:
def completed_checks(self) -> int:
return len(
{nodeid for result in self.results.values() for nodeid in result.completed}
)
@classmethod
def from_cases(cls, cases: Iterable[HarnessCase]) -> "HarnessRun":
def from_cases(cls, cases: Iterable[HarnessCase]) -> HarnessRun:
results = {case.key: CaseResult(case=case) for case in cases}
for result in results.values():
result.set_initial_status()
return cls(results=results)
@dataclass(frozen=True)
class SectionConfidence:
sdk_function: str
verified_strategies: int
required_strategies: int
level: ConfidenceLevel
details: tuple[str, ...]
@property
def percentage(self) -> int:
if not self.required_strategies:
return 0
return round(100 * self.verified_strategies / self.required_strategies)
def section_confidence(
run: HarnessRun, strategies: Iterable[Strategy]
) -> tuple[SectionConfidence, ...]:
strategy_list = tuple(strategies)
scores: list[SectionConfidence] = []
sections = tuple(dict.fromkeys((case.surface, case.sdk_function) for strategy in strategy_list for case in strategy.cases))
for surface, sdk_function in sections:
cases = tuple(
case
for strategy in strategy_list
for case in strategy.cases
if case.sdk_function == sdk_function and case.surface == surface
and case.coverage is not Coverage.NOT_APPLICABLE
)
verified = 0
details: list[str] = []
for case in cases:
result = run.results.get(case.key)
status = result.status if result is not None else RunStatus.NOT_RUN
if status is RunStatus.PASSED:
verified += 1
details.append(
f"{STATUS_LABELS[status]} {case.strategy_id} ({case.coverage.value})"
)
required = len(cases)
if required and verified == required:
level = ConfidenceLevel.HIGH
elif verified:
level = ConfidenceLevel.MEDIUM
else:
level = ConfidenceLevel.LOW
scores.append(
SectionConfidence(
sdk_function=sdk_function if surface == "sdk" else f"gateway/{sdk_function}",
verified_strategies=verified,
required_strategies=required,
level=level,
details=tuple(details),
)
)
return tuple(scores)
STATUS_LABELS = {
RunStatus.NOT_RUN: "·",
RunStatus.QUEUED: "",
RunStatus.RUNNING: "",
RunStatus.PASSED: "",
RunStatus.FAILED: "",
RunStatus.SKIPPED: "",
RunStatus.ERROR: "!",
RunStatus.MISSING: "?",
RunStatus.PLANNED: "",
RunStatus.NOT_APPLICABLE: "n/a",
}

View file

@ -1,29 +1,31 @@
from __future__ import annotations
from collections.abc import Callable, Sequence
from collections.abc import Sequence
from pathlib import Path
from time import monotonic
from typing import Final, Protocol
from typing import Final
from .models import HarnessCase, HarnessRun
from .pytest_runner import UpdateCallback
from .models import HarnessRun, Strategy
from .strategy import StrategyRunner, UpdateCallback
__all__ = ["StrategyRunner", "run_strategies"]
class StrategyRunner(Protocol):
def __call__(
self,
cases: Sequence[HarnessCase],
repo_root: Path,
on_update: UpdateCallback,
pytest_args: Sequence[str] = (),
) -> tuple[int, HarnessRun]: ...
def combine_reports(reports: Sequence[HarnessRun]) -> HarnessRun:
def combine_reports(
reports: Sequence[HarnessRun],
*,
timed_reports: Sequence[HarnessRun] | None = None,
) -> HarnessRun:
duration_sources: Final = reports if timed_reports is None else timed_reports
return HarnessRun(
results={key: result for report in reports for key, result in report.results.items()},
current_nodeid=next((report.current_nodeid for report in reversed(reports) if report.current_nodeid), None),
failures=[failure for report in reports for failure in report.failures],
strategy_durations={
strategy_id: report.duration
for report in duration_sources
for strategy_id in {result.case.strategy_id for result in report.results.values()}
},
started_at=min((report.started_at for report in reports), default=monotonic()),
finished_at=(
max((report.finished_at for report in reports if report.finished_at is not None), default=None)
@ -34,32 +36,38 @@ def combine_reports(reports: Sequence[HarnessRun]) -> HarnessRun:
def run_strategies(
cases: Sequence[HarnessCase],
strategies: Sequence[Strategy],
repo_root: Path,
on_update: UpdateCallback,
pytest_args: Sequence[str],
resolve_runner: Callable[[str], StrategyRunner],
runner_args: Sequence[str] = (),
) -> tuple[int, HarnessRun]:
strategy_ids: Final = tuple(dict.fromkeys(case.strategy_id for case in cases))
cases: Final = tuple(case for strategy in strategies for case in strategy.cases)
def execute(
remaining: tuple[str, ...], reports: tuple[HarnessRun, ...], codes: tuple[int, ...]
remaining: tuple[Strategy, ...], reports: tuple[HarnessRun, ...], codes: tuple[int, ...]
) -> tuple[int, HarnessRun]:
if not remaining:
combined: Final = combine_reports(reports)
on_update(combined)
return next((code for code in codes if code), 0), combined
strategy_id, *tail = remaining
selected: Final = tuple(case for case in cases if case.strategy_id == strategy_id)
pending: Final = HarnessRun.from_cases(case for case in cases if case.strategy_id in tail)
code, report = resolve_runner(strategy_id)(
strategy, *tail = remaining
selected: Final = tuple(case for case in cases if case.strategy_id == strategy.id)
pending: Final = HarnessRun.from_cases(
case for case in cases if case.strategy_id in {later.id for later in tail}
)
code, report = strategy.definition.run(
selected,
repo_root,
lambda current: on_update(combine_reports((*reports, current, pending))),
pytest_args,
lambda current: on_update(
combine_reports(
(*reports, current, pending),
timed_reports=(*reports, current),
)
),
runner_args,
)
if code in {2, 3, 4}:
return code, combine_reports((*reports, report, pending))
return execute(tuple(tail), (*reports, report), (*codes, code))
return execute(strategy_ids, (), ())
return execute(tuple(strategies), (), ())

View file

@ -1,172 +0,0 @@
from __future__ import annotations
import os
from collections.abc import Callable, Sequence
from pathlib import Path
from time import monotonic
from typing import Final
import pytest
from .models import CaseResult, HarnessCase, HarnessRun, RunStatus
UpdateCallback = Callable[[HarnessRun], None]
def selector_matches_node(selector: str, nodeid: str) -> bool:
normalized_selector = selector.replace("\\", "/")
normalized_nodeid = nodeid.replace("\\", "/")
if normalized_selector.endswith("/"):
return normalized_nodeid.startswith(normalized_selector)
if "::" in normalized_selector:
return normalized_nodeid == normalized_selector or normalized_nodeid.startswith(
f"{normalized_selector}["
)
return normalized_nodeid == normalized_selector or normalized_nodeid.startswith(
f"{normalized_selector}::"
)
def selector_path(selector: str) -> Path:
return Path(selector.split("::", 1)[0])
def runnable_selectors(
cases: Sequence[HarnessCase], repo_root: Path
) -> tuple[str, ...]:
selectors = {
selector
for case in cases
for selector in case.selectors
if (repo_root / selector_path(selector)).exists()
}
return tuple(sorted(selectors))
class HarnessPytestPlugin:
def __init__(self, run: HarnessRun, on_update: UpdateCallback) -> None:
self.run = run
self.on_update = on_update
self.node_to_results: dict[str, list[CaseResult]] = {}
def _notify(self) -> None:
self.on_update(self.run)
def pytest_collection_modifyitems(self, items: list[pytest.Item]) -> None:
for item in items:
matched_results: list[CaseResult] = []
for result in self.run.results.values():
if any(
selector_matches_node(selector, item.nodeid)
for selector in result.case.selectors
):
result.collected.add(item.nodeid)
matched_results.append(result)
if matched_results:
self.node_to_results[item.nodeid] = matched_results
for result in self.run.results.values():
if result.status is RunStatus.QUEUED and not result.collected:
result.status = RunStatus.MISSING
self._notify()
def pytest_runtest_logstart(
self, nodeid: str, location: tuple[str, int | None, str]
) -> None:
del location
self.run.current_nodeid = nodeid
for result in self.node_to_results.get(nodeid, []):
if result.status not in {RunStatus.FAILED, RunStatus.ERROR}:
result.status = RunStatus.RUNNING
self._notify()
def pytest_runtest_logreport(self, report: pytest.TestReport) -> None:
if report.when not in {"setup", "call", "teardown"}:
return
results = self.node_to_results.get(report.nodeid, [])
if not results:
return
terminal = report.when == "call" or report.failed or report.skipped
if not terminal:
for result in results:
result.durations[report.nodeid] = (
result.durations.get(report.nodeid, 0.0) + report.duration
)
return
for result in results:
if report.when == "teardown" and not report.failed:
result.durations[report.nodeid] = (
result.durations.get(report.nodeid, 0.0) + report.duration
)
continue
if report.skipped:
status = RunStatus.SKIPPED
elif report.failed and report.when in {"setup", "teardown"}:
status = RunStatus.ERROR
elif report.failed:
status = RunStatus.FAILED
else:
status = RunStatus.PASSED
result.record(report.nodeid, status, report.duration)
if report.failed:
failure = (report.nodeid, str(report.longrepr))
if failure not in self.run.failures:
self.run.failures.append(failure)
self._notify()
def pytest_sessionfinish(
self, session: pytest.Session, exitstatus: int | pytest.ExitCode
) -> None:
del session, exitstatus
self.run.current_nodeid = None
self.run.finished_at = monotonic()
for result in self.run.results.values():
result.finalize()
self._notify()
def run_pytest(
cases: Sequence[HarnessCase],
repo_root: Path,
on_update: UpdateCallback,
pytest_args: Sequence[str] = (),
) -> tuple[int, HarnessRun]:
run = HarnessRun.from_cases(cases)
selectors = runnable_selectors(cases, repo_root)
if not selectors:
for result in run.results.values():
result.finalize()
run.finished_at = monotonic()
on_update(run)
has_missing_test = any(
result.status is RunStatus.MISSING for result in run.results.values()
)
exit_code = (
int(pytest.ExitCode.TESTS_FAILED)
if has_missing_test
else int(pytest.ExitCode.OK)
)
return exit_code, run
plugin = HarnessPytestPlugin(run=run, on_update=on_update)
args: Final = (*selectors, "-q", "--tb=no", "--no-summary", "-o", "consider_namespace_packages=true", *pytest_args)
previous_directory = Path.cwd()
try:
os.chdir(repo_root)
exit_code = int(pytest.main(list(args), plugins=[plugin]))
finally:
os.chdir(previous_directory)
for result in run.results.values():
missing = tuple(
selector for selector in result.case.selectors
if not any(selector_matches_node(selector, node) for node in result.collected)
)
if missing:
result.status = RunStatus.MISSING
run.failures.extend((selector, "Configured selector collected no tests") for selector in missing)
on_update(run)
if exit_code == 0 and any(
result.status is RunStatus.MISSING for result in run.results.values()
):
exit_code = int(pytest.ExitCode.TESTS_FAILED)
return exit_code, run

View file

@ -0,0 +1,29 @@
from __future__ import annotations
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Final, Protocol, assert_never
from .models import CaseDisposition, CaseResult
@dataclass(frozen=True, slots=True)
class ReportSection:
title: str
blocks: tuple[str, ...]
class StrategyRenderer(Protocol):
def __call__(self, results: Sequence[CaseResult]) -> tuple[ReportSection, ...]: ...
def render_case_outcome(result: CaseResult) -> str:
prefix: Final = f"- {result.case.display_name}: {result.status.value}"
spec: Final = result.case.spec
match spec.disposition:
case CaseDisposition.RUNNABLE:
progress: Final = f", {len(result.completed)}/{result.total} checks" if result.total else ""
return f"{prefix}{progress}, {spec.coverage.value} coverage"
case CaseDisposition.NOT_IMPLEMENTED | CaseDisposition.SKIPPED:
return f"{prefix}, {spec.reason}"
assert_never(spec.disposition)

View file

@ -0,0 +1,92 @@
from __future__ import annotations
from collections.abc import Callable, Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Annotated, Literal, Protocol, TypeAlias
from pydantic import BaseModel, ConfigDict, StringConstraints
from .models import CaseDisposition, Coverage, HarnessCase, HarnessRun, SdkFunction, Surface
from .rendering import StrategyRenderer
UpdateCallback: TypeAlias = Callable[[HarnessRun], None]
NonBlankString: TypeAlias = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)]
class SuiteCaseSpec(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
disposition: Literal[CaseDisposition.RUNNABLE] = CaseDisposition.RUNNABLE
coverage: Coverage
suite: NonBlankString
note: str = ""
class ModuleCaseSpec(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
disposition: Literal[CaseDisposition.RUNNABLE] = CaseDisposition.RUNNABLE
coverage: Coverage
module: NonBlankString
note: str = ""
class NotImplementedCaseSpec(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
disposition: Literal[CaseDisposition.NOT_IMPLEMENTED] = CaseDisposition.NOT_IMPLEMENTED
coverage: None = None
reason: NonBlankString
class SkippedCaseSpec(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
disposition: Literal[CaseDisposition.SKIPPED] = CaseDisposition.SKIPPED
coverage: None = None
reason: NonBlankString
RunnableCaseSpec: TypeAlias = SuiteCaseSpec | ModuleCaseSpec
UnavailableCaseSpec: TypeAlias = NotImplementedCaseSpec | SkippedCaseSpec
CaseSpec: TypeAlias = RunnableCaseSpec | UnavailableCaseSpec
@dataclass(frozen=True, slots=True)
class CaseDefinition:
sdk_function: SdkFunction
spec: CaseSpec
surface: Surface | None = None
@dataclass(frozen=True, slots=True)
class RunnerArgumentDefinition:
option: str
help: str
metavar: str = "ARG"
class StrategyRunner(Protocol):
def __call__(
self,
cases: Sequence[HarnessCase],
repo_root: Path,
on_update: UpdateCallback,
runner_args: Sequence[str] = (),
) -> tuple[int, HarnessRun]: ...
@dataclass(frozen=True, slots=True)
class StrategyDefinition:
id: str
order: int
label: str
description: str
directory: Path
runnable_spec: type[SuiteCaseSpec] | type[ModuleCaseSpec]
cases: tuple[CaseDefinition, ...]
run: StrategyRunner
render: StrategyRenderer
surfaces: tuple[Surface, ...] = ()
runner_argument: RunnerArgumentDefinition | None = None

View file

@ -1,47 +1,161 @@
from __future__ import annotations
import logging
from collections.abc import Sequence
from pathlib import Path
from time import monotonic
from typing import Final
from .models import Coverage, HarnessCase, RunStatus
from .models import CaseResult, Coverage, HarnessCase, HarnessRun, RunStatus, Strategy
from .orchestration import run_strategies
from .pytest_runner import run_pytest
from .rendering import ReportSection, StrategyRenderer, render_case_outcome
from .strategy import (
CaseDefinition,
ModuleCaseSpec,
NotImplementedCaseSpec,
StrategyDefinition,
UpdateCallback,
)
from .ui import HarnessOutputFilter, final_report
def test_combines_independent_strategy_reports_and_keeps_failures(tmp_path: Path) -> None:
(tmp_path / "test_first.py").write_text("def test_first():\n assert 1 == 2\n")
(tmp_path / "test_second.py").write_text("def test_second():\n assert True\n")
cases: Final = tuple(
HarnessCase(
strategy_id=name,
strategy_label=name,
sdk_function="ocr",
coverage=Coverage.COMPLETE,
selectors=(f"test_{name}.py",),
)
for name in ("first", "second")
def _run_cases(
cases: Sequence[HarnessCase],
repo_root: Path,
on_update: UpdateCallback,
runner_args: Sequence[str] = (),
) -> tuple[int, HarnessRun]:
del repo_root, runner_args
run: Final = HarnessRun.from_cases(cases)
for case in cases:
_record_case(run, case, on_update)
run.finished_at = monotonic()
return int(bool(run.failures)), run
def _record_case(run: HarnessRun, case: HarnessCase, on_update: UpdateCallback) -> None:
result: Final = run.results[case.key]
nodeid: Final = f"check:{case.strategy_id}:{case.sdk_function}"
result.collected.add(nodeid)
failed: Final = isinstance(case.spec, ModuleCaseSpec) and case.spec.module == "fail"
result.record(nodeid, RunStatus.FAILED if failed else RunStatus.PASSED)
if failed:
run.failures.append((nodeid, "comparison failed"))
on_update(run)
def _render_test_results(results: Sequence[CaseResult]) -> tuple[ReportSection, ...]:
return (ReportSection("Test outcomes", tuple(render_case_outcome(result) for result in results)),)
def _strategy(name: str, module: str, *, render: StrategyRenderer = _render_test_results) -> Strategy:
case_definition: Final = CaseDefinition("ocr", ModuleCaseSpec(coverage=Coverage.COMPLETE, module=module))
definition: Final = StrategyDefinition(
id=name,
order=1,
label=name,
description="Example strategy",
directory=Path.cwd(),
runnable_spec=ModuleCaseSpec,
cases=(case_definition,),
run=_run_cases,
render=render,
)
code, report = run_strategies(cases, tmp_path, lambda _: None, (), lambda _: run_pytest)
case: Final = HarnessCase(
strategy_id=name,
strategy_label=name,
sdk_function="ocr",
spec=case_definition.spec,
)
return Strategy(1, name, name, "", Path.cwd(), (case,), definition)
def test_combines_strategy_reports_and_delegates_rendering() -> None:
strategies: Final = (_strategy("first", "fail"), _strategy("second", "pass"))
code, report = run_strategies(strategies, Path.cwd(), lambda _: None)
assert code == 1
assert report.results["first:ocr"].status is RunStatus.FAILED
assert report.results["second:ocr"].status is RunStatus.PASSED
assert report.completed_tests == 2
assert len(report.failures) == 1
assert "assert 1 == 2" in report.failures[0][1]
assert "terminalreporter" not in report.failures[0][1]
assert report.completed_checks == 2
rendered: Final = final_report(report, code, strategies)
assert "Result: FAILED" in rendered
assert rendered.count("Test outcomes") == 2
assert "- ocr: failed, 1/1 checks, complete coverage" in rendered
assert "- ocr: passed, 1/1 checks, complete coverage" in rendered
assert "Failures (showing 1 of 1)" in rendered
assert "Port confidence" not in rendered
assert "Slowest tests" not in rendered
def test_missing_selector_cannot_hide_behind_a_passing_surface(tmp_path: Path) -> None:
(tmp_path / "test_present.py").write_text("def test_present():\n assert True\n")
case: Final = HarnessCase(
strategy_id="e2e_parity",
strategy_label="End-to-end parity",
sdk_function="ocr",
surface="gateway",
coverage=Coverage.PARTIAL,
selectors=("test_present.py", "test_missing.py"),
def test_strategy_can_replace_the_generic_result_view() -> None:
def render_custom(results: Sequence[CaseResult]) -> tuple[ReportSection, ...]:
del results
return (ReportSection("Custom comparison", ("domain-owned diff",)),)
strategy: Final = _strategy("custom", "pass", render=render_custom)
code, report = run_strategies((strategy,), Path.cwd(), lambda _: None)
rendered: Final = final_report(report, code, (strategy,))
assert "Custom comparison\ndomain-owned diff" in rendered
assert "sdk/ocr" not in rendered
def test_report_separates_successful_execution_from_incomplete_coverage() -> None:
runnable: Final = _strategy("mixed", "pass")
unavailable: Final = HarnessCase(
strategy_id="mixed",
strategy_label="mixed",
sdk_function="messages",
spec=NotImplementedCaseSpec(reason="No Messages case is registered."),
)
code, report = run_pytest((case,), tmp_path, lambda _: None)
assert code == 1
assert report.results["e2e_parity:gateway:ocr"].status is RunStatus.MISSING
assert ("test_missing.py", "Configured selector collected no tests") in report.failures
code, executed = run_strategies((runnable,), Path.cwd(), lambda _: None)
unavailable_run: Final = HarnessRun.from_cases((unavailable,))
combined: Final = HarnessRun(
results={**executed.results, **unavailable_run.results},
started_at=executed.started_at,
finished_at=executed.finished_at,
)
rendered: Final = final_report(combined, code, (runnable,))
assert code == 0
assert "Result: PASSED" in rendered
assert "Harness support: 1/2 cases implemented" in rendered
assert "Cases: 2 selected, 1 not implemented, 0 skipped" in rendered
def test_harness_output_filter_suppresses_expected_harness_warnings() -> None:
output_filter: Final = HarnessOutputFilter()
ocr_cost_warning: Final = logging.LogRecord(
"LiteLLM",
logging.WARNING,
"/repo/litellm/cost_calculator.py",
1953,
"OCR cost: model=%s has no pricing",
("example",),
None,
)
other_warning: Final = logging.LogRecord(
"LiteLLM",
logging.WARNING,
"/repo/litellm/main.py",
1,
"Provider warning",
(),
None,
)
loop_warning: Final = logging.LogRecord(
"LiteLLM",
logging.WARNING,
"/repo/litellm/litellm_core_utils/logging_worker.py",
129,
"LoggingWorker: event loop changed; carried %d pending and revived %d dequeued logging task(s) onto the new loop",
(1, 0),
None,
)
assert output_filter.filter(ocr_cost_warning) is False
assert output_filter.filter(loop_warning) is False
assert output_filter.filter(other_warning) is True

View file

@ -1,45 +1,39 @@
from __future__ import annotations
import os
import shlex
import sys
import logging
from collections.abc import Sequence
from contextlib import AbstractContextManager
from pathlib import Path
from typing import Any
from textwrap import indent
from types import TracebackType
from typing import TYPE_CHECKING, Final
from .models import (
Coverage,
HarnessRun,
RunStatus,
Strategy,
section_confidence,
)
from litellm._logging import handler as litellm_log_handler
STATUS_GLYPHS = {
RunStatus.NOT_RUN: "·",
RunStatus.QUEUED: "",
RunStatus.RUNNING: "",
RunStatus.PASSED: "",
RunStatus.FAILED: "",
RunStatus.SKIPPED: "",
RunStatus.ERROR: "!",
RunStatus.MISSING: "?",
RunStatus.PLANNED: "",
RunStatus.NOT_APPLICABLE: "n/a",
}
from .models import HarnessRun, RunStatus, Strategy
from .rendering import ReportSection
STATUS_STYLES = {
RunStatus.QUEUED: "dim",
RunStatus.RUNNING: "bold cyan",
RunStatus.PASSED: "bold green",
RunStatus.FAILED: "bold red",
RunStatus.SKIPPED: "yellow",
RunStatus.ERROR: "bold red",
RunStatus.MISSING: "magenta",
RunStatus.PLANNED: "dim",
RunStatus.NOT_APPLICABLE: "dim",
}
if TYPE_CHECKING:
from rich.live import Live
class HarnessOutputFilter(logging.Filter):
def filter(self, record: logging.LogRecord) -> bool:
noisy_prefixes: Final = (
"OCR cost:",
"LoggingWorker: event loop changed;",
)
return not (record.name == "LiteLLM" and record.getMessage().startswith(noisy_prefixes))
_HARNESS_OUTPUT_FILTER: Final = HarnessOutputFilter()
def _start_output_filtering() -> None:
litellm_log_handler.addFilter(_HARNESS_OUTPUT_FILTER)
def _stop_output_filtering() -> None:
litellm_log_handler.removeFilter(_HARNESS_OUTPUT_FILTER)
def _format_duration(seconds: float) -> str:
@ -50,250 +44,221 @@ def _format_duration(seconds: float) -> str:
return f"{int(seconds // 60)}m {seconds % 60:.0f}s"
def _rerun_command(nodeid: str) -> str:
if nodeid.startswith("unit-suite:"):
return "uv run python -m tests.rust-python-harness.strategies.unit_tests.runner --plain"
return f"poetry run pytest {shlex.quote(nodeid)} -q -o consider_namespace_packages=true"
def _summary(run: HarnessRun) -> tuple[int, int, int, int]:
outcomes: dict[str, RunStatus] = {}
for result in run.results.values():
outcomes.update(result.outcomes)
values: Final = tuple(outcomes.values())
return (
list(outcomes.values()).count(RunStatus.PASSED),
list(outcomes.values()).count(RunStatus.FAILED),
list(outcomes.values()).count(RunStatus.ERROR),
list(outcomes.values()).count(RunStatus.SKIPPED),
values.count(RunStatus.PASSED),
values.count(RunStatus.FAILED),
values.count(RunStatus.ERROR),
values.count(RunStatus.SKIPPED),
)
def _cell_text(run: HarnessRun, strategy_id: str, sdk_function: str, surface: str = "sdk") -> tuple[str, str]:
key = f"{strategy_id}:{sdk_function}" if surface == "sdk" else f"{strategy_id}:gateway:{sdk_function}"
result = run.results.get(key)
if result is None:
return "", ""
counts = ""
if result.total:
counts = f" {len(result.completed)}/{result.total}"
coverage = "" if result.case.coverage is Coverage.PARTIAL else ""
return f"{STATUS_GLYPHS[result.status]}{counts}{coverage}", STATUS_STYLES.get(
result.status, ""
def _strategy_state(statuses: tuple[RunStatus, ...], outcomes: tuple[RunStatus, ...]) -> str:
for status in (
RunStatus.ERROR,
RunStatus.FAILED,
RunStatus.MISSING,
RunStatus.RUNNING,
RunStatus.QUEUED,
):
if status in statuses:
return status.value
if RunStatus.NOT_IMPLEMENTED in statuses:
return RunStatus.NOT_IMPLEMENTED.value
if outcomes and all(outcome is RunStatus.SKIPPED for outcome in outcomes):
return RunStatus.SKIPPED.value
if statuses and all(status is RunStatus.SKIPPED for status in statuses):
return RunStatus.SKIPPED.value
for status in (RunStatus.PASSED, RunStatus.SKIPPED):
if status in statuses:
return status.value
return RunStatus.NOT_RUN.value
def _strategy_line(strategy: Strategy, run: HarnessRun) -> str:
results: Final = tuple(run.results[case.key] for case in strategy.cases if case.key in run.results)
outcomes: dict[str, RunStatus] = {}
collected: set[str] = set()
for result in results:
outcomes.update(result.outcomes)
collected.update(result.collected)
values: Final = tuple(outcomes.values())
statuses: Final = tuple(result.status for result in results)
state: Final = _strategy_state(statuses, values)
completed: Final = len(outcomes)
total: Final = len(collected)
progress: Final = f", {completed}/{total} checks" if total else ""
counts: Final = (
f", {values.count(RunStatus.PASSED)} passed, "
f"{values.count(RunStatus.FAILED) + values.count(RunStatus.ERROR)} failed, "
f"{values.count(RunStatus.SKIPPED)} skipped"
if total
else ""
)
duration: Final = run.strategy_durations.get(strategy.id, 0.0)
return f"- {strategy.label}: {state}{progress}{counts}, {_format_duration(duration)}"
def _rendered_sections(run: HarnessRun, strategies: Sequence[Strategy]) -> tuple[ReportSection, ...]:
return tuple(
section
for strategy in strategies
if any(case.key in run.results for case in strategy.cases)
for section in strategy.definition.render(
tuple(run.results[case.key] for case in strategy.cases if case.key in run.results)
)
)
def _format_section(section: ReportSection) -> str:
return f"{section.title}\n" + ("\n\n".join(section.blocks) or "- No results")
def _run_result(run: HarnessRun, exit_code: int) -> str:
statuses: Final = tuple(result.status for result in run.results.values())
if exit_code:
return "FAILED"
if not statuses or all(status is RunStatus.NOT_IMPLEMENTED for status in statuses):
return "NOT RUN"
if all(status is RunStatus.SKIPPED for status in statuses):
return "SKIPPED"
return "PASSED"
def final_report(run: HarnessRun, exit_code: int, strategies: Sequence[Strategy]) -> str:
passed, failed, errors, skipped = _summary(run)
run_result: Final = _run_result(run, exit_code)
statuses: Final = tuple(case_result.status for case_result in run.results.values())
not_implemented: Final = statuses.count(RunStatus.NOT_IMPLEMENTED)
implemented: Final = len(statuses) - not_implemented
skipped_cells: Final = statuses.count(RunStatus.SKIPPED)
failure_lines: Final = tuple(
f"{index}. {nodeid}\n{indent(detail.strip(), ' ')}"
for index, (nodeid, detail) in enumerate(run.failures[:5], start=1)
)
rendered: Final = tuple(_format_section(section) for section in _rendered_sections(run, strategies))
summary: Final = (
"Rust <-> Python parity report\n\n"
f"Result: {run_result}\n"
f"Harness support: {implemented}/{len(statuses)} cases implemented\n"
f"Cases: {len(statuses)} selected, {not_implemented} not implemented, {skipped_cells} skipped\n"
f"Checks: {run.completed_checks}/{run.unique_checks} completed, {passed} passed, "
f"{failed} failed, {errors} errors, {skipped} skipped\n"
f"Duration: {_format_duration(run.duration)}\n"
f"Exit code: {exit_code}"
)
failures: Final = (
(f"Failures (showing {len(failure_lines)} of {len(run.failures)})\n" + "\n\n".join(failure_lines))
if failure_lines
else ""
)
return "\n\n".join((summary, *rendered, *((failures,) if failures else ())))
class RichDashboard(AbstractContextManager["RichDashboard"]):
def __init__(
self,
strategies: Sequence[Strategy],
confidence_strategies: Sequence[Strategy],
) -> None:
def __init__(self, strategies: Sequence[Strategy]) -> None:
from rich.console import Console
from rich.live import Live
self.strategies = strategies
self.confidence_strategies = confidence_strategies
self.console = Console()
self.live: Any = Live(
console=self.console, refresh_per_second=12, transient=False
)
self.live: Live = Live(console=self.console, refresh_per_second=12, transient=True)
self._live_active = False
def _table(self, run: HarnessRun) -> Any:
from rich import box
from rich.table import Table
from rich.text import Text
columns = tuple(dict.fromkeys((case.surface, case.sdk_function) for strategy in self.strategies for case in strategy.cases))
narrow = self.console.width < 96
if narrow:
table = Table(box=box.SIMPLE_HEAVY, expand=True, show_header=False)
table.add_column("Strategy", ratio=3)
table.add_column("Results", ratio=5)
for strategy in self.strategies:
values = []
for surface, sdk_function in columns:
value, style = _cell_text(run, strategy.id, sdk_function, surface)
if value:
values.append(
Text.assemble((f"{surface}/{sdk_function} ", "dim"), (value, style))
)
table.add_row(strategy.label, Text(" ").join(values))
return table
table = Table(box=box.ROUNDED, expand=True, title="Strategy × API")
table.add_column("Strategy", ratio=3)
for surface, label in columns:
table.add_column(label if surface == "sdk" else f"gateway/{label}", justify="center", ratio=1)
for strategy in self.strategies:
cells = []
for surface, sdk_function in columns:
value, style = _cell_text(run, strategy.id, sdk_function, surface)
cells.append(Text(value, style=style))
table.add_row(strategy.label, *cells)
return table
def __enter__(self) -> "RichDashboard":
def __enter__(self) -> RichDashboard:
_start_output_filtering()
self.live.__enter__()
self._live_active = True
return self
def __exit__(self, *args: object) -> None:
self.live.__exit__(*args)
def __exit__(
self,
exc_type: type[BaseException] | None,
exc_value: BaseException | None,
traceback: TracebackType | None,
) -> None:
_stop_output_filtering()
if self._live_active:
self.live.__exit__(exc_type, exc_value, traceback)
self._live_active = False
def update(self, run: HarnessRun) -> None:
from rich.markup import escape
from rich.panel import Panel
active = run.current_nodeid or "Waiting for test events…"
if len(active) > max(40, self.console.width - 16):
active = f"{active[-(self.console.width - 17):]}"
active: Final = run.current_nodeid or "Waiting for test events..."
available_width: Final = max(40, self.console.width - 10)
visible_active: Final = active if len(active) <= available_width else f"...{active[-(available_width - 3) :]}"
passed, failed, errors, skipped = _summary(run)
progress = (
f"[bold]{run.completed_tests}/{run.unique_tests}[/bold] tests "
f"[green]{passed} passed[/green] [red]{failed + errors} failed[/red] "
f"[yellow]{skipped} skipped[/yellow] [dim]{_format_duration(run.duration)}[/dim]"
)
legend = "✓ pass ✗ fail ! error ↷ skip\n? configured test missing — planned ◐ partial coverage"
total: Final = run.unique_checks
percentage: Final = round(100 * run.completed_checks / total) if total else 0
strategy_lines: Final = "\n".join(escape(_strategy_line(strategy, run)) for strategy in self.strategies)
self.live.update(
Panel(
self._table(run),
title="⚡ Rust ↔ Python parity lab",
subtitle=f"{progress}\n[dim]{escape(active)}[/dim]\n{legend}",
border_style="cyan",
)
"[bold]Running Rust <-> Python parity[/bold]\n"
f"Progress: [bold]{run.completed_checks}/{total} ({percentage}%)[/bold] | "
f"[green]{passed} passed[/green] | [red]{failed + errors} failed[/red] | "
f"[yellow]{skipped} skipped[/yellow] | [dim]{_format_duration(run.duration)}[/dim]\n"
f"Strategies:\n{strategy_lines}\n"
f"Current: [dim]{escape(visible_active)}[/dim]"
)
def finish(self, run: HarnessRun, exit_code: int) -> None:
self.update(run)
if run.failures:
from rich.markup import escape
from rich.panel import Panel
for nodeid, detail in run.failures[:5]:
rerun = _rerun_command(nodeid)
self.console.print(
Panel(
f"{escape(detail)}\n\n[bold]Rerun just this test[/bold]\n"
f"[cyan]{escape(rerun)}[/cyan]",
title=f"{escape(nodeid)}",
border_style="red",
)
)
durations: dict[str, float] = {}
for result in run.results.values():
for nodeid, duration in result.durations.items():
durations[nodeid] = max(duration, durations.get(nodeid, 0.0))
if durations:
slow = sorted(durations.items(), key=lambda item: item[1], reverse=True)[:3]
self.console.print(
"[bold]Slowest tests[/bold] "
+ "".join(
f"{Path(nodeid).name} [dim]{_format_duration(duration)}[/dim]"
for nodeid, duration in slow
)
)
from rich import box
from rich.table import Table
confidence_table = Table(
title="Port confidence by API", box=box.ROUNDED, expand=True
)
confidence_table.add_column("SDK section")
confidence_table.add_column("Score", justify="right")
confidence_table.add_column("Confidence")
confidence_table.add_column("Strategy evidence", ratio=4)
confidence_styles = {"HIGH": "green", "MEDIUM": "yellow", "LOW": "red"}
for score in section_confidence(run, self.confidence_strategies):
confidence_table.add_row(
score.sdk_function,
f"{score.verified_strategies}/{score.required_strategies} {score.percentage}%",
f"[{confidence_styles[score.level.value]}]{score.level.value}[/]",
" ".join(score.details),
)
self.console.print(confidence_table)
self.console.print(
"[dim]Score = required strategies with passing evidence. "
"LOC coverage remains a separate report.[/dim]"
)
style = "green" if exit_code == 0 else "red"
self.console.print(
f"[{style}]Harness finished in {_format_duration(run.duration)} "
f"(exit {exit_code})[/{style}]"
)
if self._live_active:
self.live.stop()
self._live_active = False
print(final_report(run, exit_code, self.strategies), flush=True) # noqa: T201 # CLI output
class PlainDashboard(AbstractContextManager["PlainDashboard"]):
def __init__(
self,
strategies: Sequence[Strategy],
confidence_strategies: Sequence[Strategy],
) -> None:
def __init__(self, strategies: Sequence[Strategy]) -> None:
self.strategies = strategies
self.confidence_strategies = confidence_strategies
self._seen: dict[str, tuple[RunStatus, int]] = {}
def __enter__(self) -> "PlainDashboard":
print("Rust <-> Python SDK parity harness", flush=True)
def __enter__(self) -> PlainDashboard:
_start_output_filtering()
print("Running Rust <-> Python parity", flush=True) # noqa: T201 # CLI output
return self
def __exit__(self, *args: object) -> None:
return None
def __exit__(
self,
exc_type: type[BaseException] | None,
exc_value: BaseException | None,
traceback: TracebackType | None,
) -> None:
del exc_type, exc_value, traceback
_stop_output_filtering()
def update(self, run: HarnessRun) -> None:
for key, result in run.results.items():
state = (result.status, len(result.completed))
if self._seen.get(key) != state:
self._seen[key] = state
progress = (
f" {len(result.completed)}/{result.total}" if result.total else ""
)
print(
f"{STATUS_GLYPHS[result.status]} {key}: {result.status.value}{progress}",
flush=True,
)
self._update_result(key, result.case.display_name, result.status, len(result.completed), result.total)
def _update_result(self, key: str, label: str, status: RunStatus, completed: int, total: int) -> None:
state: Final = (status, completed)
previous: Final = self._seen.get(key)
self._seen[key] = state
visible: Final = status not in {
RunStatus.NOT_RUN,
RunStatus.QUEUED,
RunStatus.NOT_IMPLEMENTED,
RunStatus.SKIPPED,
}
should_print: Final = visible and (
previous is None or previous[0] is not status or (completed > 0 and completed % 25 == 0)
)
if should_print:
progress: Final = f" {completed}/{total}" if total else ""
print( # noqa: T201 # CLI output
f"{label}: {status.value}{progress}", flush=True
)
def finish(self, run: HarnessRun, exit_code: int) -> None:
self.update(run)
passed, failed, errors, skipped = _summary(run)
print(
f"Summary: {passed} passed, {failed} failed, {errors} errors, "
f"{skipped} skipped in {_format_duration(run.duration)}",
flush=True,
print( # noqa: T201 # CLI output
f"\n{final_report(run, exit_code, self.strategies)}", flush=True
)
for nodeid, detail in run.failures[:5]:
print(f"{nodeid}: {detail}", flush=True)
print(f"Rerun: {_rerun_command(nodeid)}", flush=True)
print("Port confidence by API", flush=True)
for score in section_confidence(run, self.confidence_strategies):
print(
f" {score.sdk_function:12} "
f"{score.verified_strategies}/{score.required_strategies} "
f"{score.percentage:3}% {score.level.value:6} "
f"{' | '.join(score.details)}",
flush=True,
)
print(
" Score = required strategies with passing evidence; LOC is reported separately.",
flush=True,
)
print(f"Harness finished with exit code {exit_code}", flush=True)
def make_dashboard(
strategies: Sequence[Strategy],
plain: bool = False,
confidence_strategies: Sequence[Strategy] | None = None,
) -> RichDashboard | PlainDashboard:
confidence_strategies = confidence_strategies or strategies
interactive_terminal = (
sys.stdout.isatty()
and not os.environ.get("CI")
and os.environ.get("TERM") != "dumb"
)
if not plain and interactive_terminal:
try:
import rich # noqa: F401
return RichDashboard(strategies, confidence_strategies)
except ImportError:
pass
return PlainDashboard(strategies, confidence_strategies)
def make_dashboard(strategies: Sequence[Strategy]) -> PlainDashboard:
return PlainDashboard(strategies)

View file

@ -0,0 +1,116 @@
from __future__ import annotations
import os
from types import SimpleNamespace
from typing import Final
import pytest
from . import native_build
def test_needs_rebuild_when_bridge_is_missing() -> None:
assert native_build.needs_rebuild(None, 1.0)
def test_needs_rebuild_when_sources_are_newer_than_bridge() -> None:
assert native_build.needs_rebuild(1.0, 2.0)
def test_fresh_bridge_with_older_sources_needs_no_rebuild() -> None:
assert not native_build.needs_rebuild(2.0, 1.0)
def test_bridge_without_rust_sources_needs_no_rebuild() -> None:
assert not native_build.needs_rebuild(2.0, None)
def test_newest_source_mtime_tracks_rust_sources_and_skips_target(tmp_path: Final) -> None:
source: Final = tmp_path / "litellm-rust" / "crates" / "bridge" / "src"
source.mkdir(parents=True)
(source / "lib.rs").write_text("fn main() {}\n")
os.utime(source / "lib.rs", (1_000, 1_000))
manifest: Final = tmp_path / "litellm-rust" / "crates" / "bridge" / "Cargo.toml"
manifest.write_text("[package]\n")
os.utime(manifest, (2_000, 2_000))
lockfile: Final = tmp_path / "litellm-rust" / "Cargo.lock"
lockfile.write_text("")
os.utime(lockfile, (1_500, 1_500))
target: Final = tmp_path / "litellm-rust" / "target" / "debug" / "junk.rs"
target.parent.mkdir(parents=True)
target.write_text("fn main() {}\n")
os.utime(target, (9_999, 9_999))
assert native_build._newest_source_mtime(tmp_path) == 2_000.0
def test_newest_source_mtime_is_none_without_rust_workspace(tmp_path: Final) -> None:
assert native_build._newest_source_mtime(tmp_path) is None
def test_ensure_trace_bridge_rebuilds_when_stale(
tmp_path: Final, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
native: Final = tmp_path / "_native.abi3.so"
native.write_bytes(b"")
os.utime(native, (1_000, 1_000))
source: Final = tmp_path / "litellm-rust" / "crates" / "bridge" / "src" / "lib.rs"
source.parent.mkdir(parents=True)
source.write_text("fn main() {}\n")
os.utime(source, (2_000, 2_000))
state: Final = SimpleNamespace(rebuilt=False)
def fake_rebuild(repo_root: object) -> tuple[bool, str]:
state.rebuilt = True
return True, ""
monkeypatch.setattr(native_build, "_native_module_path", lambda: native)
monkeypatch.setattr(native_build, "_rebuild", fake_rebuild)
monkeypatch.setattr(native_build, "_drop_imported_bridge", lambda: None)
monkeypatch.setattr(native_build, "get_native_bridge", lambda: SimpleNamespace(_trace=object()))
assert native_build.ensure_trace_bridge(tmp_path) is None
assert state.rebuilt is True
assert "Rebuilding native Rust bridge" in capsys.readouterr().out
def test_ensure_trace_bridge_reports_failed_rebuild(tmp_path: Final, monkeypatch: pytest.MonkeyPatch) -> None:
source: Final = tmp_path / "litellm-rust" / "crates" / "bridge" / "src" / "lib.rs"
source.parent.mkdir(parents=True)
source.write_text("fn main() {}\n")
monkeypatch.setattr(native_build, "_native_module_path", lambda: None)
monkeypatch.setattr(native_build, "_rebuild", lambda repo_root: (False, "boom"))
message: Final = native_build.ensure_trace_bridge(tmp_path)
assert message is not None
assert "rebuild failed" in message
assert "boom" in message
def test_ensure_trace_bridge_flags_missing_trace_feature_without_rebuild(
tmp_path: Final, monkeypatch: pytest.MonkeyPatch
) -> None:
native: Final = tmp_path / "_native.abi3.so"
native.write_bytes(b"")
os.utime(native, (9_999, 9_999))
source: Final = tmp_path / "litellm-rust" / "crates" / "bridge" / "src" / "lib.rs"
source.parent.mkdir(parents=True)
source.write_text("fn main() {}\n")
os.utime(source, (1_000, 1_000))
state: Final = SimpleNamespace(rebuilt=False)
def fake_rebuild(repo_root: object) -> tuple[bool, str]:
state.rebuilt = True
return True, ""
monkeypatch.setattr(native_build, "_native_module_path", lambda: native)
monkeypatch.setattr(native_build, "_rebuild", fake_rebuild)
monkeypatch.setattr(native_build, "get_native_bridge", lambda: SimpleNamespace(_trace=None))
message: Final = native_build.ensure_trace_bridge(tmp_path)
assert message is not None
assert "_trace" in message
assert state.rebuilt is False

View file

@ -1,57 +0,0 @@
from __future__ import annotations
from collections import Counter
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Final
@dataclass(frozen=True, slots=True)
class Operation:
name: str
started: int
finished: int
def compare_traces(
python: Sequence[Operation],
rust: Sequence[Operation],
mapping: Mapping[str, str],
required_order: Sequence[tuple[str, str]] = (),
) -> tuple[str, ...]:
python_names: Final = {operation.name for operation in python}
rust_names: Final = {operation.name for operation in rust}
problems: Final = (
*(f"unmapped Python operation: {name}" for name in sorted(python_names - mapping.keys())),
*(f"unmapped Rust operation: {name}" for name in sorted(rust_names - set(mapping.values()))),
*(f"ambiguous Rust operation: {name}" for name, count in Counter(mapping.values()).items() if count > 1),
*(
f"invalid interval: {operation.name}"
for operation in (*python, *rust)
if operation.started > operation.finished
),
)
if problems:
return problems
python_counts: Final = Counter(operation.name for operation in python)
rust_counts: Final = Counter(operation.name for operation in rust)
counts: Final = tuple(
f"call count differs for {name}: Python={python_counts[name]}, Rust={rust_counts[target]}"
for name, target in mapping.items()
if python_counts[name] != rust_counts[target]
)
ordering: Final = tuple(
f"{label}: required order {before} before {after} was not observed"
for before, after in required_order
for label, operations, first, second in (
("Python", python, before, after),
("Rust", rust, mapping.get(before), mapping.get(after)),
)
if not first
or not second
or not any(operation.name == first for operation in operations)
or not any(operation.name == second for operation in operations)
or max(operation.finished for operation in operations if operation.name == first)
> min(operation.started for operation in operations if operation.name == second)
)
return (*counts, *ordering)

View file

@ -0,0 +1,38 @@
from __future__ import annotations
from typing import Final
from pydantic import BaseModel, ConfigDict
from .profiler import FunctionTraceEvent
class _TraceEventPayload(BaseModel):
model_config = ConfigDict(strict=True, extra="forbid")
id: int
parent_id: int | None
function: str
module_path: str | None = None
file: str | None = None
line: int | None = None
class TraceResponsePayload(BaseModel):
model_config = ConfigDict(strict=True, extra="forbid")
response: object
trace: tuple[_TraceEventPayload, ...] | list[_TraceEventPayload]
def native_trace_events(payload: object) -> tuple[FunctionTraceEvent, ...]:
response: Final = TraceResponsePayload.model_validate(payload)
return tuple(
FunctionTraceEvent(
event.id,
event.parent_id,
event.function,
event.module_path,
event.file,
event.line,
)
for event in response.trace
)

View file

@ -0,0 +1,117 @@
from __future__ import annotations
import sys
import threading
from collections.abc import Generator
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from types import CodeType, FrameType
from typing import Final
@dataclass(frozen=True, slots=True)
class FunctionTraceEvent:
id: int
parent_id: int | None
function: str
module_path: str | None = None
file: str | None = None
line: int | None = None
@property
def raw(self) -> str:
location: Final = f"{self.file}:{self.line}" if self.file is not None and self.line is not None else ""
qualified: Final = f"{self.module_path}::{self.function}" if self.module_path is not None else self.function
return f"{location} {qualified}" if location else qualified
class PythonProfiler:
def __init__(self, source_root: Path) -> None:
self._source_root: Final = str(source_root.resolve()) + "/"
self._seen_frames: Final[set[FrameType]] = set()
self._event_ids: Final[dict[FrameType, int]] = {}
self.events: Final[list[FunctionTraceEvent]] = []
def __call__(self, frame: FrameType, event: str, _arg: object) -> None:
if event != "call" or frame in self._seen_frames:
return
function_name: Final = self.function_name(frame.f_code)
if function_name is None:
return
event_id: Final = len(self.events)
parent_id: Final = next(
(self._event_ids[ancestor] for ancestor in _frame_ancestors(frame) if ancestor in self._event_ids),
None,
)
self._seen_frames.add(frame)
self._event_ids[frame] = event_id
self.events.append(FunctionTraceEvent(id=event_id, parent_id=parent_id, function=function_name))
def function_name(self, code: CodeType) -> str | None:
if not code.co_filename.startswith(self._source_root):
return None
relative: Final = code.co_filename.removeprefix(self._source_root)
return f"{relative}:{code.co_firstlineno} {getattr(code, 'co_qualname', code.co_name)}"
class PythonFunctionUsageProfiler:
def __init__(self, source_root: Path, functions: frozenset[str]) -> None:
self._source_root: Final = str(source_root.resolve()) + "/"
self._functions: Final = functions
self.called: Final[set[str]] = set()
def __call__(self, frame: FrameType, event: str, _arg: object) -> None:
if event != "call":
return
code: Final = frame.f_code
if not code.co_filename.startswith(self._source_root):
return
relative: Final = code.co_filename.removeprefix(self._source_root)
function: Final = f"{relative}:{code.co_firstlineno} {getattr(code, 'co_qualname', code.co_name)}"
if function in self._functions:
self.called.add(function)
def _frame_ancestors(frame: FrameType) -> Generator[FrameType]:
ancestor: Final = frame.f_back
if ancestor is not None:
yield ancestor
yield from _frame_ancestors(ancestor)
@contextmanager
def profile_python(source_root: Path, *, threads: bool = False) -> Generator[PythonProfiler]:
profiler: Final = PythonProfiler(source_root)
previous_thread: Final = threading.getprofile()
if threads:
threading.setprofile(profiler)
previous: Final = sys.getprofile()
sys.setprofile(profiler)
try:
yield profiler
finally:
sys.setprofile(previous)
if threads:
threading.setprofile(previous_thread)
@contextmanager
def profile_python_function_usage(
source_root: Path,
functions: frozenset[str],
*,
threads: bool = False,
) -> Generator[PythonFunctionUsageProfiler]:
profiler: Final = PythonFunctionUsageProfiler(source_root, functions)
previous_thread: Final = threading.getprofile()
if threads:
threading.setprofile(profiler)
previous: Final = sys.getprofile()
sys.setprofile(profiler)
try:
yield profiler
finally:
sys.setprofile(previous)
if threads:
threading.setprofile(previous_thread)

View file

@ -0,0 +1,335 @@
from __future__ import annotations
import argparse
import ast
import importlib
import inspect
import os
import subprocess
import sys
import tempfile
import warnings
from collections.abc import Generator, Sequence
from pathlib import Path
from typing import TYPE_CHECKING, Final
from pluggy import HookimplMarker
from pydantic import BaseModel, ConfigDict
from .profiler import profile_python_function_usage
if TYPE_CHECKING:
import pytest
hookimpl: Final = HookimplMarker("pytest")
class PythonFunctionIdentity(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
file: str
line: int
qualname: str
@property
def raw(self) -> str:
return f"{self.file}:{self.line} {self.qualname}"
@property
def key(self) -> str:
return f"{self.file}::{self.qualname}"
@classmethod
def from_trace(cls, raw: str) -> PythonFunctionIdentity:
location, separator, qualname = raw.partition(" ")
file, line_separator, line = location.rpartition(":")
if not separator or not line_separator or not file or not line.isdigit() or not qualname:
raise ValueError(f"Unrecognized Python trace function: {raw}")
return cls(file=file, line=int(line), qualname=qualname)
class PythonFunctionReference(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
module: str
qualname: str
@property
def owner(self) -> str:
return self.qualname.partition(".")[0]
def resolve(self, source_root: Path) -> PythonFunctionIdentity:
value: object = importlib.import_module(self.module)
for component in self.qualname.split("."):
value = getattr(value, component)
function: Final = inspect.unwrap(value)
code: Final = getattr(function, "__code__", None)
if code is None:
raise ValueError(f"Python function has no code object: {self.module}:{self.qualname}")
source: Final = Path(code.co_filename).resolve()
try:
relative: Final = source.relative_to(source_root.resolve())
except ValueError as error:
raise ValueError(f"Python function is outside {source_root}: {source}") from error
return PythonFunctionIdentity(
file=relative.as_posix(),
line=code.co_firstlineno,
qualname=code.co_qualname,
)
class RustFunctionIdentity(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
file: str
line: int
module_path: str
function: str
@property
def test_module(self) -> str:
_, separator, module = self.module_path.partition("::")
if not separator:
raise ValueError(f"Rust function has no crate-qualified module: {self.module_path}")
return f"{module}::tests"
@classmethod
def from_trace(cls, raw: str) -> RustFunctionIdentity:
location, separator, qualified = raw.partition(" ")
file, line_separator, line = location.rpartition(":")
module_path, function_separator, function = qualified.rpartition("::")
if (
not separator
or not line_separator
or not function_separator
or not file
or not line.isdigit()
or not module_path
or not function
):
raise ValueError(f"Unrecognized Rust trace function: {raw}")
return cls(file=file, line=int(line), module_path=module_path, function=function)
class PythonFunctionUsage(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
function: PythonFunctionIdentity
tests: tuple[str, ...]
class PythonUsageReport(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
usages: tuple[PythonFunctionUsage, ...]
collected_tests: tuple[str, ...]
exit_code: int
problems: tuple[str, ...] = ()
def candidate_test_files(
functions: Sequence[PythonFunctionReference | PythonFunctionIdentity],
search_roots: Sequence[str],
repo_root: Path,
*,
exclude_roots: Sequence[str] = (),
) -> tuple[str, ...]:
owners: Final = frozenset(
function.owner if isinstance(function, PythonFunctionReference) else function.qualname.partition(".")[0]
for function in functions
if "." in function.qualname
and (
isinstance(function, PythonFunctionReference)
or function.file.startswith("ocr/")
or "/ocr/" in function.file
)
)
top_level_functions: Final = frozenset(
function.qualname for function in functions if "." not in function.qualname and function.qualname.isidentifier()
)
candidates: Final = tuple(
path.relative_to(repo_root).as_posix()
for root in search_roots
for path in sorted((repo_root / root).rglob("test*.py"))
if not any(
path == repo_root / excluded or path.is_relative_to(repo_root / excluded) for excluded in exclude_roots
)
if _references_function(path, owners, top_level_functions)
)
return tuple(dict.fromkeys(candidates))
def _references_function(path: Path, owners: frozenset[str], top_level_functions: frozenset[str]) -> bool:
contents: Final = path.read_text(errors="ignore")
if any(owner in contents for owner in owners):
return True
try:
with warnings.catch_warnings():
warnings.simplefilter("ignore", SyntaxWarning)
tree: Final = ast.parse(contents)
except SyntaxError:
return False
aliases: Final = frozenset(
alias.asname or alias.name
for node in ast.walk(tree)
if isinstance(node, ast.ImportFrom)
for alias in node.names
if alias.name in top_level_functions
)
names: Final = top_level_functions | aliases
return any(
isinstance(node, ast.Call)
and (
(isinstance(node.func, ast.Name) and node.func.id in names)
or (isinstance(node.func, ast.Attribute) and node.func.attr in top_level_functions)
)
for node in ast.walk(tree)
)
class _WorkerConfig(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
functions: tuple[PythonFunctionIdentity, ...]
source_root: Path
output: Path
pytest_args: tuple[str, ...]
class _FunctionUsagePlugin:
def __init__(self, functions: tuple[PythonFunctionIdentity, ...], source_root: Path) -> None:
self._functions: Final = functions
self._function_names: Final = frozenset(function.raw for function in functions)
self._source_root: Final = source_root
self._tests_by_function: Final[dict[str, set[str]]] = {function.raw: set() for function in functions}
self.collected_tests: tuple[str, ...] = ()
self.problems: tuple[str, ...] = ()
def pytest_collection_finish(self, session: pytest.Session) -> None:
self.collected_tests = tuple(item.nodeid for item in session.items)
def pytest_collectreport(self, report: pytest.CollectReport) -> None:
if report.failed:
self.problems = (*self.problems, str(report.longrepr))
@hookimpl(hookwrapper=True)
def pytest_runtest_protocol(self, item: pytest.Item, nextitem: pytest.Item | None) -> Generator[None, object, None]:
del nextitem
with profile_python_function_usage(self._source_root, self._function_names, threads=True) as profiler:
yield
for function in self._functions:
if function.raw in profiler.called:
self._tests_by_function[function.raw].add(item.nodeid)
def usages(self) -> tuple[PythonFunctionUsage, ...]:
return tuple(
PythonFunctionUsage(
function=function,
tests=tuple(sorted(self._tests_by_function[function.raw])),
)
for function in self._functions
)
def collect_python_function_tests(
functions: Sequence[PythonFunctionIdentity],
selectors: Sequence[str],
repo_root: Path,
*,
source_root: Path | None = None,
exclusions: Sequence[str] = (),
) -> PythonUsageReport:
selected_functions: Final = tuple(dict.fromkeys(functions))
if not selected_functions:
raise ValueError("Python function discovery needs at least one function")
if not selectors:
raise ValueError("Python function discovery needs at least one test selector")
with tempfile.TemporaryDirectory(prefix="litellm-function-tests-") as directory:
temporary: Final = Path(directory)
config_path: Final = temporary / "config.json"
output_path: Final = temporary / "report.json"
config: Final = _WorkerConfig(
functions=selected_functions,
source_root=source_root or repo_root / "litellm",
output=output_path,
pytest_args=tuple(
(
"-o",
"consider_namespace_packages=true",
"-p",
"no:cacheprovider",
*selectors,
*(f"--deselect={nodeid}" for nodeid in exclusions),
)
),
)
config_path.write_text(config.model_dump_json())
import_roots: Final = tuple(
dict.fromkeys(
(
str(repo_root),
str(source_root or repo_root / "litellm"),
*(
str(path.parent if path.suffix == ".py" else path)
for selector in selectors
if (path := repo_root / selector.partition("::")[0]).exists()
),
os.environ.get("PYTHONPATH", ""),
)
)
)
env: Final = {
**os.environ,
"PYTHONPATH": os.pathsep.join(import_roots),
}
try:
result: Final = subprocess.run(
(sys.executable, "-m", __name__, str(config_path)),
cwd=repo_root,
env=env,
capture_output=True,
text=True,
timeout=600,
check=False,
)
except (OSError, subprocess.TimeoutExpired) as error:
return PythonUsageReport(usages=(), collected_tests=(), exit_code=1, problems=(str(error),))
if not output_path.exists():
return PythonUsageReport(
usages=(),
collected_tests=(),
exit_code=result.returncode or 1,
problems=((result.stdout + result.stderr).strip(),),
)
report: Final = PythonUsageReport.model_validate_json(output_path.read_text())
process_output: Final = (result.stdout + result.stderr).strip()
if result.returncode and not report.problems and process_output:
return report.model_copy(update={"problems": (process_output,)})
return report
def _run_worker(config: _WorkerConfig) -> int:
import pytest
plugin: Final = _FunctionUsagePlugin(config.functions, config.source_root)
exit_code: Final = int(pytest.main(list(config.pytest_args), plugins=[plugin]))
report: Final = PythonUsageReport(
usages=plugin.usages(),
collected_tests=plugin.collected_tests,
exit_code=exit_code,
problems=plugin.problems,
)
config.output.write_text(report.model_dump_json())
return exit_code
def main(argv: Sequence[str] | None = None) -> int:
parser: Final = argparse.ArgumentParser()
parser.add_argument("config", type=Path)
namespace: Final = parser.parse_args(argv)
config: Final = _WorkerConfig.model_validate_json(namespace.config.read_text())
return _run_worker(config)
if __name__ == "__main__":
raise SystemExit(main())

View file

@ -0,0 +1,271 @@
from __future__ import annotations
import re
from collections import Counter
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Final, Literal
from .profiler import FunctionTraceEvent
Engine = Literal["python", "rust"]
@dataclass(frozen=True, slots=True)
class TraceMapping:
span: str
python: re.Pattern[str] | None
rust: str | None
def mapping(
*,
python_frame: str | None = None,
rust_span: str | None = None,
span: str | None = None,
) -> TraceMapping:
if rust_span is None:
if python_frame is None:
raise ValueError("mapping needs a python_frame pattern, a rust_span name, or both")
if span is None:
raise ValueError("a python-only mapping needs an explicit span to compare under")
return TraceMapping(span, re.compile(python_frame), None)
if python_frame is None:
return TraceMapping(rust_span, None, rust_span)
if span is not None and span != rust_span:
raise ValueError(f"span {span!r} disagrees with rust_span {rust_span!r}")
return TraceMapping(rust_span, re.compile(python_frame), rust_span)
@dataclass(frozen=True, slots=True)
class TraceContract:
unordered_children_of: frozenset[str] = frozenset()
@dataclass(frozen=True, slots=True)
class PipelineStep:
id: int
parent_id: int | None
span: str
raw: str
@dataclass(frozen=True, slots=True)
class PipelineProjection:
steps: tuple[PipelineStep, ...] = ()
unmatched: int = 0
def _span_for(engine: Engine, function: str, mappings: Sequence[TraceMapping]) -> str | None:
matches: Final = tuple(
item.span
for item in mappings
if (
engine == "python"
and item.python is not None
and item.python.search(function)
or engine == "rust"
and item.rust == function
)
)
if len(matches) > 1:
raise ValueError(f"{engine} event {function!r} matches multiple trace mappings: {matches}")
if matches:
return matches[0]
return function if engine == "rust" else None
def pipeline_projection(
engine: Engine, events: Sequence[FunctionTraceEvent], mappings: Sequence[TraceMapping]
) -> PipelineProjection:
raw_parents: dict[int, int | None] = {}
projected_ids: set[int] = set()
shown: list[PipelineStep] = []
unmatched: int = 0
for event in events:
if event.id in raw_parents:
raise ValueError(f"duplicate trace event id {event.id}")
if event.parent_id is not None and event.parent_id not in raw_parents:
raise ValueError(f"trace event {event.id} references unknown or later parent {event.parent_id}")
raw_parents[event.id] = event.parent_id
span = _span_for(engine, event.function, mappings)
if span is None:
unmatched += 1
continue
parent_id: int | None = event.parent_id
while parent_id is not None and parent_id not in projected_ids:
parent_id = raw_parents[parent_id]
shown.append(PipelineStep(event.id, parent_id, span, event.raw))
projected_ids.add(event.id)
return PipelineProjection(tuple(shown), unmatched)
@dataclass(frozen=True, slots=True)
class TraceNode:
id: int
span: str
children: tuple[TraceNode, ...]
def trace_depths(steps: Sequence[PipelineStep]) -> dict[int, int]:
depths: dict[int, int] = {}
for step in steps:
depths[step.id] = 0 if step.parent_id is None else depths[step.parent_id] + 1
return depths
def _forest(steps: Sequence[PipelineStep]) -> tuple[TraceNode, ...]:
children: dict[int | None, list[PipelineStep]] = {}
known: set[int] = set()
for step in steps:
if step.id in known:
raise ValueError(f"duplicate projected event id {step.id}")
if step.parent_id is not None and step.parent_id not in known:
raise ValueError(f"projected event {step.id} references unknown or later parent {step.parent_id}")
known.add(step.id)
children.setdefault(step.parent_id, []).append(step)
def node(step: PipelineStep) -> TraceNode:
return TraceNode(step.id, step.span, tuple(node(child) for child in children.get(step.id, ())))
return tuple(node(step) for step in children.get(None, ()))
def _exclusive_spans(engine: Engine, mappings: Sequence[TraceMapping]) -> frozenset[str]:
return frozenset(
item.span
for item in mappings
if (engine == "python" and item.rust is None) or (engine == "rust" and item.python is None)
)
def _comparable_steps(
engine: Engine, steps: Sequence[PipelineStep], mappings: Sequence[TraceMapping]
) -> tuple[PipelineStep, ...]:
exclusive: Final = _exclusive_spans(engine, mappings)
raw_parents: Final = {step.id: step.parent_id for step in steps}
included: Final = {step.id for step in steps if step.span not in exclusive}
comparable: list[PipelineStep] = []
for step in steps:
if step.id not in included:
continue
parent_id: int | None = step.parent_id
while parent_id is not None and parent_id not in included:
parent_id = raw_parents[parent_id]
comparable.append(PipelineStep(step.id, parent_id, step.span, step.raw))
return tuple(comparable)
def _signature(node: TraceNode, contract: TraceContract) -> tuple[object, ...]:
children: tuple[tuple[object, ...], ...] = tuple(_signature(child, contract) for child in node.children)
normalized: Final = tuple(sorted(children, key=repr)) if node.span in contract.unordered_children_of else children
return (node.span, normalized)
def trace_signature(
engine: Engine,
steps: Sequence[PipelineStep],
mappings: Sequence[TraceMapping],
contract: TraceContract,
) -> tuple[tuple[object, ...], ...]:
return tuple(_signature(root, contract) for root in _forest(_comparable_steps(engine, steps, mappings)))
@dataclass(frozen=True, slots=True)
class TraceDiff:
python_only: tuple[str, ...]
rust_only: tuple[str, ...]
shared_order_matches: bool
missing_mappings: tuple[str, ...] = ()
first_difference: str | None = None
@property
def matches(self) -> bool:
return (
not self.python_only
and not self.rust_only
and not self.missing_mappings
and self.shared_order_matches
)
def _missing_mappings(
python: Sequence[PipelineStep], rust: Sequence[PipelineStep], mappings: Sequence[TraceMapping]
) -> tuple[str, ...]:
python_seen: Final = frozenset(step.span for step in python)
rust_seen: Final = frozenset(step.span for step in rust)
return tuple(
item.span
for item in mappings
if (item.python is not None and item.span not in python_seen)
or (item.rust is not None and item.span not in rust_seen)
)
def _first_difference(
python: Sequence[PipelineStep],
rust: Sequence[PipelineStep],
mappings: Sequence[TraceMapping],
contract: TraceContract,
) -> str | None:
python_forest: Final = _forest(_comparable_steps("python", python, mappings))
rust_forest: Final = _forest(_comparable_steps("rust", rust, mappings))
def compare_children(
python_nodes: Sequence[TraceNode], rust_nodes: Sequence[TraceNode], path: str, *, unordered: bool
) -> str | None:
if unordered:
python_signatures: Final = Counter(_signature(node, contract) for node in python_nodes)
rust_signatures: Final = Counter(_signature(node, contract) for node in rust_nodes)
if python_signatures != rust_signatures:
return f"{path}: unordered child subtree multiset differs"
return None
for index in range(max(len(python_nodes), len(rust_nodes))):
child_path = f"{path}/child[{index + 1}]"
if index >= len(python_nodes):
return f"{child_path}: Rust has extra {rust_nodes[index].span!r}"
if index >= len(rust_nodes):
return f"{child_path}: Python has extra {python_nodes[index].span!r}"
python_node = python_nodes[index]
rust_node = rust_nodes[index]
if python_node.span != rust_node.span:
return f"{child_path}: Python={python_node.span!r}, Rust={rust_node.span!r}"
difference = compare_children(
python_node.children,
rust_node.children,
f"{child_path}/{python_node.span}",
unordered=python_node.span in contract.unordered_children_of,
)
if difference is not None:
return difference
return None
return compare_children(python_forest, rust_forest, "root", unordered=False)
def trace_diff(
python: Sequence[PipelineStep],
rust: Sequence[PipelineStep],
mappings: Sequence[TraceMapping] = (),
contract: TraceContract = TraceContract(),
) -> TraceDiff:
python_comparable: Final = _comparable_steps("python", python, mappings)
rust_comparable: Final = _comparable_steps("rust", rust, mappings)
python_spans: Final = tuple(step.span for step in python_comparable)
rust_spans: Final = tuple(step.span for step in rust_comparable)
python_counts: Final = Counter(python_spans)
rust_counts: Final = Counter(rust_spans)
python_only_counts: Final = python_counts - rust_counts
rust_only_counts: Final = rust_counts - python_counts
python_only: Final = tuple(
span for span, count in python_only_counts.items() for _ in range(count)
)
rust_only: Final = tuple(span for span, count in rust_only_counts.items() for _ in range(count))
first_difference: Final = _first_difference(python, rust, mappings, contract)
return TraceDiff(
python_only=python_only,
rust_only=rust_only,
shared_order_matches=bool(python_comparable or rust_comparable) and first_difference is None,
missing_mappings=_missing_mappings(python, rust, mappings),
first_difference=first_difference,
)

View file

@ -1,33 +0,0 @@
from __future__ import annotations
import pytest
from .compare import Operation, compare_traces
@pytest.mark.parametrize(
("rust", "message"),
(
((Operation("decode", 0, 1), Operation("send", 2, 3)), None),
((Operation("decode", 0, 1), Operation("send", 2, 3), Operation("send", 4, 5)), "call count differs"),
((Operation("send", 0, 1), Operation("decode", 2, 3)), "required order"),
((Operation("decode", 0, 4), Operation("send", 2, 3)), "required order"),
((Operation("decode", 0, 1), Operation("unknown", 2, 3)), "unmapped Rust"),
),
)
def test_compare_mapped_calls_and_required_completion_order(rust: tuple[Operation, ...], message: str | None) -> None:
problems = compare_traces(
(Operation("parse", 0, 1), Operation("request", 2, 3)),
rust,
{"parse": "decode", "request": "send"},
(("parse", "request"),),
)
if message is None:
assert problems == ()
else:
assert any(message in problem for problem in problems)
def test_missing_required_operations_and_ambiguous_mappings_fail() -> None:
assert compare_traces((), (), {"parse": "decode"}, (("parse", "request"),))
assert compare_traces((), (), {"parse": "decode", "request": "decode"}) == ("ambiguous Rust operation: decode",)

View file

@ -0,0 +1,108 @@
from __future__ import annotations
import asyncio
import sys
import threading
from pathlib import Path
from typing import Final
import pytest
from .profiler import FunctionTraceEvent, PythonProfiler, profile_python, profile_python_function_usage
def _events_named(profiler: PythonProfiler, name: str) -> tuple[FunctionTraceEvent, ...]:
return tuple(event for event in profiler.events if event.function.endswith(name))
def test_profiler_keeps_repeated_calls() -> None:
def called() -> None:
return None
with profile_python(Path(__file__).parent) as profiler:
called()
called()
assert len(_events_named(profiler, "called")) == 2
def test_profiler_records_real_frame_ancestry() -> None:
def called() -> None:
return None
def outer() -> None:
called()
with profile_python(Path(__file__).parent) as profiler:
outer()
outer_event, called_event = (event for event in profiler.events if event.function.endswith(("outer", "called")))
assert called_event.parent_id == outer_event.id
def test_profiler_restores_previous_profiler_after_failure() -> None:
previous: Final = sys.getprofile()
with pytest.raises(RuntimeError, match="stop"):
with profile_python(Path(__file__).parent):
raise RuntimeError("stop")
assert sys.getprofile() is previous
def test_profiler_does_not_count_coroutine_resumption_as_another_call() -> None:
async def suspended() -> None:
await asyncio.sleep(0)
await asyncio.sleep(0)
with profile_python(Path(__file__).parent) as profiler:
asyncio.run(suspended())
assert len(_events_named(profiler, "suspended")) == 1
def test_profiler_preserves_parent_across_coroutine_suspension() -> None:
def called() -> None:
return None
async def suspended() -> None:
await asyncio.sleep(0)
called()
with profile_python(Path(__file__).parent) as profiler:
asyncio.run(suspended())
suspended_event: Final = _events_named(profiler, "suspended")[0]
called_event: Final = _events_named(profiler, "called")[0]
assert called_event.parent_id == suspended_event.id
def test_profiler_captures_worker_threads_when_enabled() -> None:
def called() -> None:
return None
with profile_python(Path(__file__).parent, threads=True) as profiler:
thread: Final = threading.Thread(target=called)
thread.start()
thread.join()
called_event: Final = _events_named(profiler, "called")[0]
assert called_event.parent_id is None
def test_function_usage_profiler_records_only_selected_functions() -> None:
def selected() -> None:
return None
def ignored() -> None:
return None
source_root: Final = Path(__file__).parent
function: Final = PythonProfiler(source_root).function_name(selected.__code__)
assert function is not None
with profile_python_function_usage(source_root, frozenset((function,))) as profiler:
selected()
ignored()
assert profiler.called == {function}

View file

@ -0,0 +1,168 @@
from __future__ import annotations
from pathlib import Path
from typing import Final
import pytest
from .pytest_usage import (
PythonFunctionIdentity,
PythonFunctionReference,
RustFunctionIdentity,
candidate_test_files,
collect_python_function_tests,
)
def test_collects_parameterized_tests_that_execute_function(tmp_path: Path) -> None:
(tmp_path / "pytest.ini").write_text("[pytest]\n")
(tmp_path / "source.py").write_text("def target():\n return 1\n\ndef other():\n return 2\n")
(tmp_path / "test_source.py").write_text(
"import pytest\n"
"from source import other, target\n"
"@pytest.mark.parametrize('value', [1, 2])\n"
"def test_target(value): assert target() + value > 0\n"
"def test_other(): assert other() == 2\n"
)
target: Final = PythonFunctionIdentity(file="source.py", line=1, qualname="target")
report: Final = collect_python_function_tests(
(target,),
("test_source.py",),
tmp_path,
source_root=tmp_path,
)
assert report.exit_code == 0, report.problems
assert report.usages[0].tests == (
"test_source.py::test_target[1]",
"test_source.py::test_target[2]",
)
def test_collects_async_and_threaded_function_calls(tmp_path: Path) -> None:
(tmp_path / "pytest.ini").write_text("[pytest]\n")
(tmp_path / "source.py").write_text(
"async def async_target():\n return 1\n\ndef threaded_target():\n return 2\n"
)
(tmp_path / "test_source.py").write_text(
"import asyncio\n"
"from threading import Thread\n"
"from source import async_target, threaded_target\n"
"def test_async(): assert asyncio.run(async_target()) == 1\n"
"def test_thread():\n"
" thread = Thread(target=threaded_target)\n"
" thread.start()\n"
" thread.join()\n"
)
functions: Final = (
PythonFunctionIdentity(file="source.py", line=1, qualname="async_target"),
PythonFunctionIdentity(file="source.py", line=4, qualname="threaded_target"),
)
report: Final = collect_python_function_tests(
functions,
("test_source.py",),
tmp_path,
source_root=tmp_path,
)
assert report.exit_code == 0, report.problems
assert report.usages[0].tests == ("test_source.py::test_async",)
assert report.usages[1].tests == ("test_source.py::test_thread",)
def test_adds_candidate_directory_to_worker_import_path(tmp_path: Path) -> None:
(tmp_path / "pytest.ini").write_text("[pytest]\n")
source: Final = tmp_path / "source"
tests: Final = tmp_path / "tests"
source.mkdir()
tests.mkdir()
(source / "implementation.py").write_text("def target():\n return 1\n")
(tests / "helper.py").write_text("VALUE = 1\n")
(tests / "test_source.py").write_text(
"from helper import VALUE\nfrom implementation import target\ndef test_target(): assert target() == VALUE\n"
)
target: Final = PythonFunctionIdentity(file="implementation.py", line=1, qualname="target")
report: Final = collect_python_function_tests(
(target,),
("tests/test_source.py",),
tmp_path,
source_root=source,
)
assert report.exit_code == 0, report.problems
assert report.usages[0].tests == ("tests/test_source.py::test_target",)
def test_parses_function_identity_from_trace() -> None:
function: Final = PythonFunctionIdentity.from_trace("llms/mistral/ocr/transformation.py:72 Config.map")
assert function.file == "llms/mistral/ocr/transformation.py"
assert function.line == 72
assert function.qualname == "Config.map"
def test_resolves_function_and_finds_candidate_test_files(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
package: Final = tmp_path / "package"
tests: Final = tmp_path / "tests"
package.mkdir()
tests.mkdir()
(package / "__init__.py").write_text("")
(package / "implementation.py").write_text("class Config:\n def transform(self):\n return 1\n")
(tests / "test_implementation.py").write_text("from package.implementation import Config\n")
(tests / "test_unrelated.py").write_text("def test_other(): pass\n")
monkeypatch.syspath_prepend(tmp_path)
reference: Final = PythonFunctionReference(module="package.implementation", qualname="Config.transform")
function: Final = reference.resolve(tmp_path)
candidates: Final = candidate_test_files((reference,), ("tests",), tmp_path)
assert function.file == "package/implementation.py"
assert function.qualname == "Config.transform"
assert candidates == ("tests/test_implementation.py",)
def test_candidate_test_files_excludes_harness_roots(tmp_path: Path) -> None:
tests: Final = tmp_path / "tests"
harness: Final = tests / "harness"
harness.mkdir(parents=True)
(tests / "test_implementation.py").write_text("from package.implementation import Config\n")
(harness / "test_fixture.py").write_text("from package.implementation import Config\n")
function: Final = PythonFunctionReference(module="package.implementation", qualname="Config.transform")
candidates: Final = candidate_test_files(
(function,),
("tests",),
tmp_path,
exclude_roots=("tests/harness",),
)
assert candidates == ("tests/test_implementation.py",)
def test_candidate_test_files_finds_top_level_calls_and_import_aliases(tmp_path: Path) -> None:
tests: Final = tmp_path / "tests"
tests.mkdir()
(tests / "test_attribute.py").write_text("import package\ndef test_call(): package.ocr()\n")
(tests / "test_alias.py").write_text("from package import ocr as run_ocr\ndef test_call(): run_ocr()\n")
(tests / "test_unrelated.py").write_text("def test_call(): return 'ocr'\n")
function: Final = PythonFunctionIdentity(file="ocr/main.py", line=1, qualname="ocr")
candidates: Final = candidate_test_files((function,), ("tests",), tmp_path)
assert candidates == (
"tests/test_alias.py",
"tests/test_attribute.py",
)
def test_parses_rust_function_identity_and_derives_test_module() -> None:
function: Final = RustFunctionIdentity.from_trace(
"crates/core/src/providers/mistral/ocr/transformation.rs:73 "
"litellm_core::providers::mistral::ocr::transformation::supported_ocr_params"
)
assert function.file == "crates/core/src/providers/mistral/ocr/transformation.rs"
assert function.test_module == "providers::mistral::ocr::transformation::tests"

View file

@ -0,0 +1,157 @@
from __future__ import annotations
from typing import Final
import pytest
from .profiler import FunctionTraceEvent
from .steps import Engine, TraceContract, mapping, pipeline_projection, trace_depths, trace_diff
MAPPINGS: Final = (
mapping(rust_span="route", python_frame=r"entry$"),
mapping(rust_span="provider", python_frame=r"provider$"),
mapping(rust_span="request", python_frame=r"request$"),
mapping(rust_span="http", python_frame=r"post$"),
mapping(rust_span="response", python_frame=r"response$"),
)
def event(event_id: int, function: str, parent_id: int | None = None) -> FunctionTraceEvent:
return FunctionTraceEvent(event_id, parent_id, function)
def test_python_projection_collapses_unmapped_parents_and_counts_noise() -> None:
events: Final = (
event(0, "module.py:1 entry"),
event(1, "noise", 0),
event(2, "module.py:2 provider", 1),
event(3, "module.py:3 request", 0),
event(4, "client.py:4 post", 3),
event(5, "module.py:5 response", 0),
)
projection: Final = pipeline_projection("python", events, MAPPINGS)
assert projection.unmatched == 1
assert [(step.id, step.parent_id, step.span, step.raw) for step in projection.steps] == [
(0, None, "route", "module.py:1 entry"),
(2, 0, "provider", "module.py:2 provider"),
(3, 0, "request", "module.py:3 request"),
(4, 3, "http", "client.py:4 post"),
(5, 0, "response", "module.py:5 response"),
]
def test_rust_projection_keeps_unknown_spans() -> None:
projection: Final = pipeline_projection("rust", (event(0, "route"), event(1, "new_span", 0)), MAPPINGS)
assert [(step.span, step.parent_id) for step in projection.steps] == [("route", None), ("new_span", 0)]
def test_projection_preserves_repeated_occurrences() -> None:
projection: Final = pipeline_projection(
"rust",
(event(0, "route"), event(1, "http", 0), event(2, "http", 0)),
MAPPINGS,
)
assert [step.span for step in projection.steps] == ["route", "http", "http"]
def test_projection_preserves_multiple_roots() -> None:
projection: Final = pipeline_projection("rust", (event(0, "route"), event(1, "request")), MAPPINGS)
assert trace_depths(projection.steps) == {0: 0, 1: 0}
def test_projection_rejects_duplicate_and_unknown_parent_ids() -> None:
with pytest.raises(ValueError, match="duplicate trace event id"):
pipeline_projection("rust", (event(0, "route"), event(0, "request")), MAPPINGS)
with pytest.raises(ValueError, match="unknown or later parent"):
pipeline_projection("rust", (event(1, "request", 0),), MAPPINGS)
@pytest.mark.parametrize("engine", ("python", "rust"))
def test_rust_only_mappings_do_not_swallow_python_frames(engine: Engine) -> None:
projection: Final = pipeline_projection(
engine,
(event(0, "anything"),),
(mapping(rust_span="rust_only_span"),),
)
if engine == "python":
assert projection.unmatched == 1
assert projection.steps == ()
else:
assert projection.unmatched == 0
assert projection.steps[0].span == "anything"
def test_mapping_builder_rejects_empty_and_ambiguous_declarations() -> None:
with pytest.raises(ValueError, match="mapping needs"):
mapping()
with pytest.raises(ValueError, match="python-only mapping needs"):
mapping(python_frame=r"frame$")
with pytest.raises(ValueError, match="disagrees with"):
mapping(rust_span="span_a", python_frame=r"frame$", span="span_b")
def test_projection_rejects_ambiguous_python_mapping() -> None:
mappings: Final = (
mapping(rust_span="first", python_frame=r"same$"),
mapping(rust_span="second", python_frame=r"same$"),
)
with pytest.raises(ValueError, match="multiple trace mappings"):
pipeline_projection("python", (event(0, "module.py:1 same"),), mappings)
def test_trace_diff_matches_identical_occurrence_trees() -> None:
mappings: Final = (MAPPINGS[0], MAPPINGS[2])
steps: Final = pipeline_projection(
"rust", (event(0, "route"), event(1, "request", 0), event(2, "request", 0)), mappings
).steps
assert trace_diff(steps, steps, mappings).matches
def test_trace_diff_rejects_missing_occurrence_and_parent_drift() -> None:
python: Final = pipeline_projection(
"rust", (event(0, "route"), event(1, "request", 0), event(2, "request", 0)), MAPPINGS
).steps
missing: Final = pipeline_projection("rust", (event(0, "route"), event(1, "request", 0)), MAPPINGS).steps
reparented: Final = pipeline_projection(
"rust", (event(0, "route"), event(1, "request", 0), event(2, "request", 1)), MAPPINGS
).steps
assert trace_diff(python, missing, MAPPINGS).python_only == ("request",)
assert not trace_diff(python, reparented, MAPPINGS).matches
def test_trace_diff_rejects_sequential_reorder() -> None:
first: Final = pipeline_projection(
"rust", (event(0, "route"), event(1, "request", 0), event(2, "response", 0)), MAPPINGS
).steps
second: Final = pipeline_projection(
"rust", (event(0, "route"), event(1, "response", 0), event(2, "request", 0)), MAPPINGS
).steps
diff: Final = trace_diff(first, second, MAPPINGS)
assert not diff.matches
assert diff.first_difference == "root/child[1]/route/child[1]: Python='request', Rust='response'"
def test_trace_diff_allows_reordered_concurrent_children() -> None:
mappings: Final = (MAPPINGS[0], MAPPINGS[2], MAPPINGS[4])
first: Final = pipeline_projection(
"rust", (event(0, "route"), event(1, "request", 0), event(2, "response", 0)), mappings
).steps
second: Final = pipeline_projection(
"rust", (event(0, "route"), event(1, "response", 0), event(2, "request", 0)), mappings
).steps
contract: Final = TraceContract(frozenset({"route"}))
assert trace_diff(first, second, mappings, contract).matches
def test_trace_diff_prunes_declared_engine_only_nodes_but_requires_them() -> None:
mappings: Final = (MAPPINGS[0], mapping(rust_span="rust_prepare"))
python: Final = pipeline_projection("python", (event(0, "module.py:1 entry"),), mappings).steps
rust: Final = pipeline_projection(
"rust", (event(0, "route"), event(1, "rust_prepare", 0)), mappings
).steps
assert trace_diff(python, rust, mappings).matches
assert trace_diff(python, rust[:1], mappings).missing_mappings == ("rust_prepare",)
def test_trace_diff_does_not_claim_empty_traces_match() -> None:
assert not trace_diff((), ()).matches

View file

@ -1,7 +1,6 @@
from __future__ import annotations
import argparse
import ast
import importlib
import os
import subprocess
@ -9,11 +8,16 @@ import sys
import tempfile
from collections.abc import Callable, Sequence
from pathlib import Path
from typing import Final, Literal, cast
from typing import TYPE_CHECKING, Final, Literal, cast
import pytest
from pluggy import HookimplMarker
from pydantic import BaseModel, ConfigDict
if TYPE_CHECKING:
import pytest
hookimpl: Final = HookimplMarker("pytest")
Backend = Literal["python", "rust"]
@ -32,22 +36,20 @@ class BackendSpec(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
environment_variable: str
probe: str = ""
class WorkerArgs(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
backend: Backend
probe: str
def ocr_backend() -> Backend:
from litellm.rust_bridge import native_bridge_available
from litellm.rust_bridge.configuration import rust_ocr_enabled
if not rust_ocr_enabled():
return "python"
if not native_bridge_available():
raise RuntimeError("Rust OCR was enabled but the native extension is unavailable")
return "rust"
output: Path
pytest_args: tuple[str, ...]
class ResultPlugin:
def __init__(self, backend: Backend, probe: Callable[[], object]) -> None:
def __init__(self, backend: Backend, probe: Callable[[], object] | None) -> None:
self.backend: Final = backend
self.probe: Final = probe
self.tests: tuple[str, ...] = ()
@ -55,13 +57,13 @@ class ResultPlugin:
self.problems: tuple[str, ...] = ()
def verify(self) -> None:
if self.probe() != self.backend:
if self.probe is not None and self.probe() != self.backend:
raise RuntimeError(f"backend probe did not select {self.backend}")
def pytest_collection_finish(self, session: pytest.Session) -> None:
self.tests = tuple(item.nodeid for item in session.items)
@pytest.hookimpl(tryfirst=True)
@hookimpl(tryfirst=True)
def pytest_runtest_call(self, item: pytest.Item) -> None:
del item
self.verify()
@ -91,8 +93,7 @@ def run_python_tests(
__name__,
"--backend",
backend,
"--probe",
spec.probe,
*(("--probe", spec.probe) if spec.probe else ()),
"--output",
str(output),
"--",
@ -118,6 +119,9 @@ def run_python_tests(
problems=(result.stdout + result.stderr,),
)
report: Final = PythonReport.model_validate_json(output.read_text())
process_output: Final = (result.stdout + result.stderr).strip()
if result.returncode and not report.problems and process_output:
return report.model_copy(update={"problems": (process_output,)})
if report.exit_code != result.returncode:
return report.model_copy(
update={
@ -129,43 +133,42 @@ def run_python_tests(
def compare_python_runs(python: PythonReport, rust: PythonReport) -> tuple[str, ...]:
python_only: Final = tuple(sorted(set(python.outcomes) - set(rust.outcomes)))
rust_only: Final = tuple(sorted(set(rust.outcomes) - set(python.outcomes)))
return (
*(("backend selection was not verified",) if not python.verified or not rust.verified else ()),
*(("Python run used the wrong backend",) if python.backend != "python" else ()),
*(("Rust run used the wrong backend",) if rust.backend != "rust" else ()),
*(("Python/Rust test inventories differ",) if python.tests != rust.tests else ()),
*(("Python/Rust test outcomes differ",) if sorted(python.outcomes) != sorted(rust.outcomes) else ()),
*(("Python/Rust test outcomes differ",) if python_only or rust_only else ()),
*(f"Python only: {nodeid} [{stage}] {outcome}" for nodeid, stage, outcome in python_only),
*(f"Rust only: {nodeid} [{stage}] {outcome}" for nodeid, stage, outcome in rust_only),
*(f"Python run: {problem}" for problem in python.problems if not python.verified or not python.tests),
*(f"Rust run: {problem}" for problem in rust.problems if not rust.verified or not rust.tests),
*(("no Python tests collected",) if not python.tests else ()),
*(
("Python tests did not all pass",)
if set(python.tests)
!= {node for node, phase, status in python.outcomes if phase == "call" and status == "passed"}
else ()
),
*(
("Rust-enabled Python tests did not all pass",)
if set(rust.tests)
!= {node for node, phase, status in rust.outcomes if phase == "call" and status == "passed"}
else ()
),
*(("Python test run failed",) if python.exit_code else ()),
*(("Rust-enabled Python test run failed",) if rust.exit_code else ()),
*python.problems,
*rust.problems,
*(("Python/Rust exit codes differ",) if python.exit_code != rust.exit_code else ()),
)
def _load_probe(reference: str) -> Callable[[], object] | None:
if not reference:
return None
module, name = reference.rsplit(":", 1)
return cast(Callable[[], object], getattr(importlib.import_module(module), name))
def main(argv: Sequence[str] | None = None) -> int:
import pytest
parser: Final = argparse.ArgumentParser()
parser.add_argument("--backend", required=True, choices=("python", "rust"))
parser.add_argument("--probe", required=True)
parser.add_argument("--probe", default="")
parser.add_argument("--output", required=True, type=Path)
parser.add_argument("pytest_args", nargs=argparse.REMAINDER)
args: Final = parser.parse_args(argv)
namespace: Final = parser.parse_args(argv)
args: Final = WorkerArgs.model_validate(vars(namespace))
try:
module, name = args.probe.rsplit(":", 1)
probe: Final = cast(Callable[[], object], getattr(importlib.import_module(module), name))
plugin: Final = ResultPlugin(args.backend, probe)
plugin: Final = ResultPlugin(args.backend, _load_probe(args.probe))
plugin.verify()
code: Final = int(
pytest.main(["-o", "consider_namespace_packages=true", *args.pytest_args[1:]], plugins=[plugin])
@ -186,22 +189,29 @@ def main(argv: Sequence[str] | None = None) -> int:
return report.exit_code
def enumerate_python_tests(repo_root: Path, relative_path: str) -> frozenset[str]:
source = (repo_root / relative_path).read_text(encoding="utf-8")
tree = ast.parse(source, filename=relative_path)
def contract_nodeid(nodeid: str) -> str:
owner, separator, test = nodeid.rpartition("::")
function: Final = test.partition("[")[0]
if not separator or not function.startswith("test_"):
raise ValueError(f"Unrecognized pytest node id: {nodeid}")
return f"{owner}::{function}"
module_level: list[str] = []
for node in ast.iter_child_nodes(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name.startswith("test_"):
module_level.append(node.name)
elif isinstance(node, ast.ClassDef):
for child in ast.iter_child_nodes(node):
if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)) and child.name.startswith(
"test_"
):
module_level.append(f"{node.name}::{child.name}")
return frozenset(module_level)
def collect_python_tests(selectors: Sequence[str], repo_root: Path) -> frozenset[str]:
report: Final = run_python_tests(
selectors,
repo_root,
"python",
BackendSpec(environment_variable="LITELLM_RUST"),
("--collect-only", "-p", "no:cacheprovider"),
)
if report.exit_code or report.problems:
details: Final = "\n".join(report.problems) or f"pytest exited with code {report.exit_code}"
raise ValueError(f"Python test collection failed:\n{details}")
tests: Final = frozenset(contract_nodeid(nodeid) for nodeid in report.tests)
if not tests:
raise ValueError(f"pytest collected no tests for: {', '.join(selectors)}")
return tests
if __name__ == "__main__":

View file

@ -0,0 +1,253 @@
from __future__ import annotations
import subprocess
from collections.abc import Callable
from dataclasses import dataclass
from itertools import groupby
from pathlib import Path
from typing import Annotated, Final, Literal, TypeAlias
from pydantic import BaseModel, ConfigDict, Field, model_validator
from typing_extensions import Self
CommandRunner: TypeAlias = Callable[[tuple[str, ...], Path], str]
_MODEL_CONFIG: Final = ConfigDict(extra="forbid", frozen=True, strict=True)
class RustTarget(BaseModel):
model_config = _MODEL_CONFIG
package: str
name: str
kind: Literal["lib", "bin", "test"]
@property
def key(self) -> str:
return f"{self.package}/{self.kind}/{self.name}"
class RustTestIdentity(BaseModel):
model_config = _MODEL_CONFIG
target: RustTarget
name: str
@property
def key(self) -> str:
return f"{self.target.key}::{self.name}"
class RustTestScope(BaseModel):
model_config = _MODEL_CONFIG
target: RustTarget
modules: Annotated[tuple[str, ...], Field(min_length=1)]
features: tuple[str, ...] = ()
default_features: bool = True
@model_validator(mode="after")
def validate_scope(self) -> Self:
duplicate_features: Final = tuple(
feature for feature, values in groupby(sorted(self.features)) if sum(1 for _ in values) > 1
)
duplicate_modules: Final = tuple(
module for module, values in groupby(sorted(self.modules)) if sum(1 for _ in values) > 1
)
if duplicate_features:
raise ValueError(f"Rust features contain duplicates: {', '.join(duplicate_features)}")
if duplicate_modules:
raise ValueError(f"Rust modules contain duplicates: {', '.join(duplicate_modules)}")
if any(not module or module.endswith("::") for module in self.modules):
raise ValueError("Rust modules must be non-empty and omit the trailing :: separator")
overlaps: Final = tuple(
f"{outer} includes {inner}"
for outer in self.modules
for inner in self.modules
if inner.startswith(f"{outer}::")
)
if overlaps:
raise ValueError(f"Rust modules overlap: {', '.join(overlaps)}")
return self
def contains(self, identity: RustTestIdentity) -> bool:
return identity.target == self.target and any(
identity.name.startswith(f"{module}::") for module in self.modules
)
class _CargoPackage(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True, strict=True)
id: str
name: str
class _CargoMetadata(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True, strict=True)
packages: tuple[_CargoPackage, ...]
class _CargoMessage(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True, strict=True)
reason: str
class _CargoTarget(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True, strict=True)
name: str
kind: tuple[str, ...]
class _CargoProfile(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True, strict=True)
test: bool
class _CargoArtifact(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True, strict=True)
reason: Literal["compiler-artifact"]
package_id: str
target: _CargoTarget
profile: _CargoProfile
executable: str | None
@dataclass(frozen=True, slots=True)
class RustReport:
tests: tuple[str, ...]
exit_code: int
output: str
def run_command(command: tuple[str, ...], cwd: Path) -> str:
try:
result: Final = subprocess.run(command, cwd=cwd, capture_output=True, text=True, check=False, timeout=600)
except (OSError, subprocess.TimeoutExpired) as error:
raise ValueError(f"Rust inventory command failed: {error}") from error
if result.returncode != 0:
raise ValueError(
f"Rust inventory command failed ({result.returncode}): {' '.join(command)}\n"
f"{result.stderr}\n{result.stdout}"
)
return result.stdout
def run_rust_tests(manifest: Path, package: str | None, test_filter: str, *, collect_only: bool = False) -> RustReport:
command: Final = (
"cargo",
"test",
"--manifest-path",
str(manifest),
*(("--package", package) if package else ()),
"--lib",
test_filter,
"--",
*(("--list",) if collect_only else ("--format=pretty",)),
)
try:
result: Final = subprocess.run(command, capture_output=True, text=True, check=False, timeout=600)
except (OSError, subprocess.TimeoutExpired) as error:
return RustReport((), 1, str(error))
tests: Final = (
tuple(line.removesuffix(": test") for line in result.stdout.splitlines() if line.endswith(": test"))
if collect_only
else tuple(
line.removeprefix("test ").removesuffix(" ... ok")
for line in result.stdout.splitlines()
if line.startswith("test ") and line.endswith(" ... ok")
)
)
return RustReport(tests, result.returncode, result.stdout + result.stderr)
def _build_command(scope: RustTestScope) -> tuple[str, ...]:
selector: Final = ("--lib",) if scope.target.kind == "lib" else (f"--{scope.target.kind}", scope.target.name)
features: Final = ("--features", ",".join(scope.features)) if scope.features else ()
defaults: Final = () if scope.default_features else ("--no-default-features",)
return (
"cargo",
"test",
"--package",
scope.target.package,
*selector,
*features,
*defaults,
"--locked",
"--no-run",
"--message-format=json",
"--color",
"never",
)
def _test_names(output: str) -> frozenset[str]:
lines: Final = tuple(line for line in output.splitlines() if line)
invalid: Final = tuple(line for line in lines if not line.endswith((": test", ": benchmark")))
if invalid:
raise ValueError(f"Unrecognized libtest inventory output: {invalid!r}")
names: Final = tuple(line.removesuffix(": test") for line in lines if line.endswith(": test"))
if len(names) != len(frozenset(names)):
raise ValueError("Duplicate test names in libtest inventory")
return frozenset(names)
def _scope_tests(
scope: RustTestScope,
metadata: _CargoMetadata,
cwd: Path,
command_runner: CommandRunner,
) -> frozenset[RustTestIdentity]:
package_ids: Final = tuple(package.id for package in metadata.packages if package.name == scope.target.package)
if len(package_ids) != 1:
raise ValueError(f"Expected one Cargo package for {scope.target.package}, found {len(package_ids)}")
output: Final = command_runner(_build_command(scope), cwd)
artifacts: Final = tuple(
_CargoArtifact.model_validate_json(line)
for line in output.splitlines()
if _CargoMessage.model_validate_json(line).reason == "compiler-artifact"
)
executables: Final = frozenset(
artifact.executable
for artifact in artifacts
if artifact.package_id == package_ids[0]
and artifact.target.name == scope.target.name
and scope.target.kind in artifact.target.kind
and artifact.profile.test
and artifact.executable is not None
)
if len(executables) != 1:
raise ValueError(f"Expected one test executable for {scope.target.key}, found {len(executables)}")
executable: Final = next(iter(executables))
names: Final = _test_names(command_runner((executable, "--list", "--format", "terse"), cwd))
ignored: Final = _test_names(command_runner((executable, "--list", "--ignored", "--format", "terse"), cwd))
identities: Final = frozenset(RustTestIdentity(target=scope.target, name=name) for name in names)
scoped: Final = frozenset(identity for identity in identities if scope.contains(identity))
ignored_scoped: Final = tuple(sorted(identity.key for identity in scoped if identity.name in ignored))
if ignored_scoped:
raise ValueError(f"Ignored Rust tests cannot satisfy the mapping: {', '.join(ignored_scoped)}")
empty_modules: Final = tuple(
module for module in scope.modules if not any(identity.name.startswith(f"{module}::") for identity in scoped)
)
if empty_modules:
raise ValueError(f"No compiled tests in {scope.target.key} modules: {', '.join(empty_modules)}")
return scoped
def enumerate_rust_tests(
repo_root: Path,
scopes: tuple[RustTestScope, ...],
*,
command_runner: CommandRunner = run_command,
) -> frozenset[RustTestIdentity]:
if not scopes:
return frozenset()
cwd: Final = repo_root / "litellm-rust"
metadata: Final = _CargoMetadata.model_validate_json(
command_runner(("cargo", "metadata", "--format-version", "1", "--no-deps", "--locked"), cwd)
)
return frozenset(identity for scope in scopes for identity in _scope_tests(scope, metadata, cwd, command_runner))

View file

@ -0,0 +1,76 @@
from __future__ import annotations
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
from time import monotonic
from typing import Final, TypeVar
from pydantic import BaseModel
from ..reporting.models import HarnessCase, HarnessRun, ResultArtifact, RunStatus, SdkFunction
from ..reporting.strategy import SuiteCaseSpec, UpdateCallback
S = TypeVar("S", bound=BaseModel)
@dataclass(frozen=True, slots=True)
class SuiteExecution:
problems: tuple[str, ...] = ()
artifacts: tuple[ResultArtifact, ...] = ()
SuiteExecutor = Callable[[S, Path, Sequence[str]], SuiteExecution]
def suite_nodeid(case: HarnessCase) -> str:
spec = case.spec
suite = spec.suite if isinstance(spec, SuiteCaseSpec) else "invalid"
return f"suite:{case.strategy_id}:{case.sdk_function}:{suite}"
def run_suites(
cases: Sequence[HarnessCase],
repo_root: Path,
on_update: UpdateCallback,
runner_args: Sequence[str] = (),
*,
suites: Mapping[SdkFunction, S],
execute: SuiteExecutor[S],
) -> tuple[int, HarnessRun]:
report = HarnessRun.from_cases(cases)
for case in cases:
result = report.results[case.key]
spec = case.spec
if not isinstance(spec, SuiteCaseSpec):
continue
nodeid = suite_nodeid(case)
result.collected.add(nodeid)
result.status = RunStatus.RUNNING
on_update(report)
suite = suites.get(case.sdk_function)
if suite is None:
result.record(nodeid, RunStatus.ERROR)
report.failures.append((nodeid, f"no suite registered for {case.sdk_function}"))
continue
try:
execution: Final = execute(suite, repo_root, runner_args)
except (OSError, ValueError) as error:
result.record(nodeid, RunStatus.ERROR)
report.failures.append((nodeid, str(error)))
continue
result.record(
nodeid,
RunStatus.FAILED if execution.problems else RunStatus.PASSED,
artifacts=execution.artifacts,
)
report.failures.extend((nodeid, problem) for problem in execution.problems)
on_update(report)
report.finished_at = monotonic()
on_update(report)
return int(
any(
result.status in {RunStatus.ERROR, RunStatus.FAILED, RunStatus.MISSING}
for result in report.results.values()
)
), report

View file

@ -4,9 +4,7 @@ import os
from pathlib import Path
from typing import Final
from .python_runner import BackendSpec, compare_python_runs, run_python_tests
HARNESS_ROOT: Final = Path(__file__).resolve().parents[4]
from .python_runner import BackendSpec, collect_python_tests, compare_python_runs, run_python_tests
def _suite(root: Path, *, mismatch: bool = False) -> BackendSpec:
@ -23,9 +21,7 @@ def _suite(root: Path, *, mismatch: bool = False) -> BackendSpec:
return BackendSpec(environment_variable="TEST_USE_RUST", probe="backend_probe:selected")
def test_runs_existing_python_tests_in_separate_verified_backends(tmp_path: Path, monkeypatch) -> None:
monkeypatch.setenv("PYTHONPATH", str(HARNESS_ROOT))
monkeypatch.setenv("PYTEST_DISABLE_PLUGIN_AUTOLOAD", "1")
def test_runs_existing_python_tests_in_separate_verified_backends(tmp_path: Path) -> None:
spec: Final = _suite(tmp_path)
python: Final = run_python_tests(("test_backend.py",), tmp_path, "python", spec)
rust: Final = run_python_tests(("test_backend.py",), tmp_path, "rust", spec)
@ -35,9 +31,7 @@ def test_runs_existing_python_tests_in_separate_verified_backends(tmp_path: Path
assert (tmp_path / "python.pid").read_text() != str(os.getpid())
def test_rejects_wrong_backend_and_different_test_results(tmp_path: Path, monkeypatch) -> None:
monkeypatch.setenv("PYTHONPATH", str(HARNESS_ROOT))
monkeypatch.setenv("PYTEST_DISABLE_PLUGIN_AUTOLOAD", "1")
def test_rejects_wrong_backend_and_different_test_results(tmp_path: Path) -> None:
spec: Final = _suite(tmp_path, mismatch=True)
python: Final = run_python_tests(("test_backend.py",), tmp_path, "python", spec)
rust: Final = run_python_tests(("test_backend.py",), tmp_path, "rust", spec)
@ -51,3 +45,44 @@ def test_rejects_wrong_backend_and_different_test_results(tmp_path: Path, monkey
assert wrong.exit_code == 1
assert not wrong.verified
assert "backend probe did not select rust" in wrong.problems[0]
def test_matches_outcomes_without_a_probe_when_both_backends_fail_identically(tmp_path: Path) -> None:
(tmp_path / "pytest.ini").write_text("[pytest]\n")
(tmp_path / "test_backend.py").write_text("def test_fails():\n assert False\n")
spec: Final = BackendSpec(environment_variable="TEST_USE_RUST")
python: Final = run_python_tests(("test_backend.py",), tmp_path, "python", spec)
rust: Final = run_python_tests(("test_backend.py",), tmp_path, "rust", spec)
assert compare_python_runs(python, rust) == ()
def test_reports_worker_output_when_pytest_exits_before_collection(tmp_path: Path) -> None:
report: Final = run_python_tests(
("missing.py",),
tmp_path,
"python",
BackendSpec(environment_variable="TEST_USE_RUST"),
)
assert report.exit_code != 0
assert report.problems
assert "missing.py" in report.problems[0]
def test_collects_tests_with_pytest_semantics_and_collapses_parameters(tmp_path: Path) -> None:
(tmp_path / "pytest.ini").write_text("[pytest]\n")
(tmp_path / "test_inventory.py").write_text(
"import pytest\n"
"class Helper:\n"
" def test_not_collected(self): pass\n"
"class TestCollected:\n"
" @pytest.mark.parametrize('value', [1, 2])\n"
" def test_parameterized(self, value): pass\n",
encoding="utf-8",
)
tests: Final = collect_python_tests(("test_inventory.py",), tmp_path)
assert tests == frozenset(("test_inventory.py::TestCollected::test_parameterized",))

View file

@ -0,0 +1,61 @@
from __future__ import annotations
import shutil
from collections.abc import Callable
from pathlib import Path
from typing import Final
import pytest
from .rust_runner import RustTarget, RustTestScope, enumerate_rust_tests, run_command, run_rust_tests
@pytest.mark.skipif(shutil.which("cargo") is None, reason="Cargo is required for native runner integration")
def test_collects_and_runs_native_tests_and_propagates_failure(
tmp_path: Path,
cargo_project: Callable[[str, str], Path],
) -> None:
manifest: Final = cargo_project("harness-runner-check", "#[test] fn test_parity() { assert_eq!(2 + 2, 4); }\n")
source: Final = tmp_path / "src/lib.rs"
inventory: Final = run_rust_tests(manifest, "harness-runner-check", "test_parity", collect_only=True)
assert inventory.exit_code == 0, inventory.output
assert inventory.tests == ("test_parity",)
passing: Final = run_rust_tests(manifest, "harness-runner-check", "test_parity")
assert passing.exit_code == 0, passing.output
source.write_text("#[test] fn test_parity() { assert_eq!(2 + 2, 5); }\n")
failed: Final = run_rust_tests(manifest, "harness-runner-check", "test_parity")
assert failed.exit_code != 0
assert "test_parity" in failed.output
@pytest.mark.skipif(shutil.which("cargo") is None, reason="Cargo is required for compiled inventory tests")
def test_discovers_compiled_fully_qualified_tests(tmp_path: Path) -> None:
workspace: Final = tmp_path / "litellm-rust"
source: Final = workspace / "src"
external: Final = source / "ocr" / "external.rs"
external.parent.mkdir(parents=True)
(workspace / "Cargo.toml").write_text(
'[package]\nname = "inventory-fixture"\nversion = "0.1.0"\nedition = "2021"\n[workspace]\n',
encoding="utf-8",
)
(source / "lib.rs").write_text(
"#[cfg(test)]\n"
"mod ocr {\n"
" mod external;\n"
" #[test] fn same_name() {}\n"
" #[cfg(any())] #[test] fn compiled_out() {}\n"
" macro_rules! generate_test { ($name:ident) => { #[test] fn $name() {} }; }\n"
" generate_test!(generated_case);\n"
"}\n",
encoding="utf-8",
)
external.write_text("#[test] fn same_name() {}\n", encoding="utf-8")
run_command(("cargo", "generate-lockfile", "--offline"), workspace)
target: Final = RustTarget(package="inventory-fixture", name="inventory_fixture", kind="lib")
scope: Final = RustTestScope(target=target, modules=("ocr",))
inventory: Final = enumerate_rust_tests(tmp_path, (scope,))
assert frozenset(identity.name for identity in inventory) == frozenset(
("ocr::same_name", "ocr::external::same_name", "ocr::generated_case")
)

View file

@ -0,0 +1,88 @@
from __future__ import annotations
from collections.abc import Sequence
from pathlib import Path
from typing import Final
from pydantic import BaseModel
from ..reporting.models import Coverage, HarnessCase, ResultArtifact, RunStatus
from ..reporting.strategy import CaseSpec, NotImplementedCaseSpec, SuiteCaseSpec
from .suite_runner import SuiteExecution, run_suites
class _Suite(BaseModel):
problems: tuple[str, ...] = ()
def _execute(suite: _Suite, repo_root: Path, pytest_args: Sequence[str]) -> SuiteExecution:
del repo_root, pytest_args
return SuiteExecution(problems=suite.problems)
def _case(spec: CaseSpec) -> HarnessCase:
return HarnessCase(
strategy_id="example",
strategy_label="Example",
sdk_function="ocr",
spec=spec,
)
def test_not_implemented_cell_finalizes_without_running(tmp_path: Path) -> None:
case = _case(NotImplementedCaseSpec(reason="No suite is registered."))
code, report = run_suites((case,), tmp_path, lambda _: None, (), suites={}, execute=_execute)
assert code == 0
assert report.results[case.key].status is RunStatus.NOT_IMPLEMENTED
assert not report.failures
def test_missing_registered_suite_marks_the_cell_as_error(tmp_path: Path) -> None:
case = _case(SuiteCaseSpec(coverage=Coverage.COMPLETE, suite="ocr"))
code, report = run_suites((case,), tmp_path, lambda _: None, (), suites={}, execute=_execute)
assert code == 1
assert report.results[case.key].status is RunStatus.ERROR
assert report.failures
def test_suite_problems_mark_the_cell_as_failed(tmp_path: Path) -> None:
case = _case(SuiteCaseSpec(coverage=Coverage.COMPLETE, suite="ocr"))
code, report = run_suites(
(case,), tmp_path, lambda _: None, (), suites={"ocr": _Suite(problems=("boom",))}, execute=_execute
)
assert code == 1
assert report.results[case.key].status is RunStatus.FAILED
assert ("suite:example:ocr:ocr", "boom") in report.failures
def test_suite_without_problems_passes(tmp_path: Path) -> None:
case = _case(SuiteCaseSpec(coverage=Coverage.COMPLETE, suite="ocr"))
code, report = run_suites((case,), tmp_path, lambda _: None, (), suites={"ocr": _Suite()}, execute=_execute)
assert code == 0
assert report.results[case.key].status is RunStatus.PASSED
assert not report.failures
def test_suite_attaches_artifacts_to_passing_and_failing_results(tmp_path: Path) -> None:
case: Final = _case(SuiteCaseSpec(coverage=Coverage.COMPLETE, suite="ocr"))
artifact: Final = ResultArtifact("example", "body")
def execute(suite: _Suite, repo_root: Path, pytest_args: Sequence[str]) -> SuiteExecution:
del repo_root, pytest_args
return SuiteExecution(problems=suite.problems, artifacts=(artifact,))
_, passing = run_suites((case,), tmp_path, lambda _: None, suites={"ocr": _Suite()}, execute=execute)
_, failing = run_suites(
(case,), tmp_path, lambda _: None, suites={"ocr": _Suite(problems=("boom",))}, execute=execute
)
assert passing.results[case.key].artifacts == {"suite:example:ocr:ocr": (artifact,)}
assert failing.results[case.key].artifacts == {"suite:example:ocr:ocr": (artifact,)}

View file

@ -0,0 +1 @@
Switches between the Rust implementation and the existing Python core, then compares their observable behavior for parity across SDK objects, exceptions, callbacks, streams, and gateway HTTP responses using generated and recorded inputs.

View file

@ -1,5 +0,0 @@
# E2E Parity
Run independently with `uv run python -m tests.rust-python-harness.strategies.e2e_parity.runner --plain`. Configure SDK and gateway selectors in `strategy.json`; keep API-specific execution and fixtures in their owning surface folder
See [the harness guide](../../README.md) for coverage status and shared comparison tools

View file

@ -0,0 +1,103 @@
from pathlib import Path
from typing import Final
from ...shared.reporting.models import SURFACES, Coverage
from ...shared.reporting.strategy import (
CaseDefinition,
ModuleCaseSpec,
NotImplementedCaseSpec,
StrategyDefinition,
)
from .reporting import render_e2e_results
from .runner import run_e2e_cases
CASES: Final[tuple[CaseDefinition, ...]] = (
CaseDefinition(
"ocr",
ModuleCaseSpec(
coverage=Coverage.PARTIAL,
module="tests.rust-python-harness.strategies.e2e_parity.sdk.ocr.test_sdk_parity",
note=(
"Recorded sync/async SDK parity; invalid-model provider errors differ, "
"and Reducto lacks a Rust contract."
),
),
surface="sdk",
),
CaseDefinition(
"messages",
NotImplementedCaseSpec(
reason="Bridge unit tests exist, but no standalone end-to-end parity case is registered."
),
surface="sdk",
),
CaseDefinition(
"responses",
NotImplementedCaseSpec(
reason="Bridge unit tests exist, but no standalone end-to-end parity case is registered."
),
surface="sdk",
),
CaseDefinition(
"count_tokens",
NotImplementedCaseSpec(reason="No Rust count_tokens parity test is present yet."),
surface="sdk",
),
CaseDefinition(
"chat_completions",
NotImplementedCaseSpec(
reason="Bridge unit tests exist, but no standalone end-to-end parity case is registered."
),
surface="sdk",
),
CaseDefinition(
"transcription",
NotImplementedCaseSpec(
reason="Bridge unit tests exist, but no standalone end-to-end parity case is registered."
),
surface="sdk",
),
CaseDefinition(
"ocr",
NotImplementedCaseSpec(reason="No gateway end-to-end OCR parity case is registered."),
surface="gateway",
),
CaseDefinition(
"messages",
NotImplementedCaseSpec(reason="No gateway end-to-end Messages parity case is registered."),
surface="gateway",
),
CaseDefinition(
"responses",
NotImplementedCaseSpec(reason="No gateway end-to-end Responses parity case is registered."),
surface="gateway",
),
CaseDefinition(
"count_tokens",
NotImplementedCaseSpec(reason="No gateway end-to-end token-count parity case is registered."),
surface="gateway",
),
CaseDefinition(
"chat_completions",
NotImplementedCaseSpec(reason="No gateway end-to-end chat parity case is registered."),
surface="gateway",
),
CaseDefinition(
"transcription",
NotImplementedCaseSpec(reason="No gateway end-to-end transcription parity case is registered."),
surface="gateway",
),
)
STRATEGY: Final = StrategyDefinition(
id="e2e_parity",
order=10,
label="End-to-end parity",
description="Compare observable Python and Rust behavior over generated and recorded inputs.",
directory=Path(__file__).parent,
runnable_spec=ModuleCaseSpec,
cases=CASES,
run=run_e2e_cases,
render=render_e2e_results,
surfaces=SURFACES,
)

View file

@ -0,0 +1,12 @@
from __future__ import annotations
from collections.abc import Sequence
from typing import Final
from ...shared.reporting.models import CaseResult
from ...shared.reporting.rendering import ReportSection, render_case_outcome
def render_e2e_results(results: Sequence[CaseResult]) -> tuple[ReportSection, ...]:
blocks: Final = tuple(render_case_outcome(result) for result in results)
return (ReportSection("End-to-end parity outcomes", blocks or ("No end-to-end cases selected",)),)

View file

@ -1,26 +1,125 @@
from __future__ import annotations
from collections.abc import Sequence
import importlib
from collections.abc import Callable, Sequence
from contextlib import AbstractContextManager, nullcontext
from dataclasses import dataclass
from pathlib import Path
from time import monotonic
from typing import Final, cast
from ...shared.reporting.models import HarnessCase, HarnessRun
from ...shared.reporting.pytest_runner import UpdateCallback, run_pytest
from ...shared.reporting.models import CaseResult, HarnessCase, HarnessRun, RunStatus
from ...shared.reporting.strategy import ModuleCaseSpec, UpdateCallback
def run(
@dataclass(frozen=True, slots=True)
class E2ECheck:
name: str
execute: Callable[[], None]
@dataclass(frozen=True, slots=True)
class E2ELoadFailure:
message: str
def _load_checks(reference: str) -> AbstractContextManager[object] | E2ELoadFailure:
try:
module: Final = importlib.import_module(reference)
factory_value: Final[object] = getattr(module, "parity_checks", None)
if not callable(factory_value):
return E2ELoadFailure(f"{reference} must export parity_checks()")
factory: Final = cast(Callable[[], object], factory_value)
checks_value: Final = factory()
except Exception as error:
return E2ELoadFailure(f"cannot load {reference}: {type(error).__name__}: {error}")
if isinstance(checks_value, tuple):
return nullcontext(cast(object, checks_value))
if isinstance(checks_value, AbstractContextManager):
return cast(AbstractContextManager[object], checks_value)
return E2ELoadFailure(
f"{reference}.parity_checks() must return tuple[E2ECheck, ...] or a context manager yielding one"
)
def _validate_checks(reference: str, checks_value: object) -> tuple[E2ECheck, ...] | E2ELoadFailure:
if not isinstance(checks_value, tuple):
return E2ELoadFailure(f"{reference}.parity_checks() context manager must yield tuple[E2ECheck, ...]")
untyped_checks: Final = cast(tuple[object, ...], checks_value)
if not all(isinstance(check, E2ECheck) for check in untyped_checks):
return E2ELoadFailure(f"{reference}.parity_checks() must return tuple[E2ECheck, ...]")
return cast(tuple[E2ECheck, ...], untyped_checks)
def _run_check(
run: HarnessRun,
result: CaseResult,
check: E2ECheck,
nodeid: str,
on_update: UpdateCallback,
) -> None:
started_at: Final = monotonic()
try:
check.execute()
except Exception as error:
result.record(nodeid, RunStatus.FAILED, monotonic() - started_at)
run.failures.append((nodeid, f"{type(error).__name__}: {error}"))
else:
result.record(nodeid, RunStatus.PASSED, monotonic() - started_at)
on_update(run)
def _run_case(run: HarnessRun, harness_case: HarnessCase, on_update: UpdateCallback) -> None:
result: Final = run.results[harness_case.key]
spec: Final = harness_case.spec
if not isinstance(spec, ModuleCaseSpec):
return
loaded: Final = _load_checks(spec.module)
if isinstance(loaded, E2ELoadFailure):
load_nodeid: Final = f"e2e:{harness_case.surface}:{harness_case.sdk_function}:load"
result.collected.add(load_nodeid)
result.record(load_nodeid, RunStatus.ERROR)
run.failures.append((load_nodeid, loaded.message))
on_update(run)
return
try:
with loaded as checks_value:
checks: Final = _validate_checks(spec.module, checks_value)
if isinstance(checks, E2ELoadFailure):
raise TypeError(checks.message)
nodeids: Final = tuple(
(check, f"e2e:{harness_case.surface}:{harness_case.sdk_function}:{check.name}") for check in checks
)
result.collected.update(nodeid for _, nodeid in nodeids)
if not nodeids:
result.status = RunStatus.SKIPPED
on_update(run)
return
result.status = RunStatus.RUNNING
on_update(run)
for check, nodeid in nodeids:
_run_check(run, result, check, nodeid, on_update)
except Exception as error:
session_nodeid: Final = f"e2e:{harness_case.surface}:{harness_case.sdk_function}:session"
result.collected.add(session_nodeid)
result.record(session_nodeid, RunStatus.ERROR)
run.failures.append((session_nodeid, f"{type(error).__name__}: {error}"))
on_update(run)
def run_e2e_cases(
cases: Sequence[HarnessCase],
repo_root: Path,
on_update: UpdateCallback,
pytest_args: Sequence[str] = (),
runner_args: Sequence[str] = (),
) -> tuple[int, HarnessRun]:
return run_pytest(cases, repo_root, on_update, pytest_args)
def main(argv: Sequence[str] | None = None) -> int:
from ...cli import main as harness_main
return harness_main(argv, strategy_id="e2e_parity")
if __name__ == "__main__":
raise SystemExit(main())
del repo_root, runner_args
run: Final = HarnessRun.from_cases(cases)
for harness_case in cases:
_run_case(run, harness_case, on_update)
run.finished_at = monotonic()
on_update(run)
failed: Final = any(
result.status in {RunStatus.ERROR, RunStatus.FAILED, RunStatus.MISSING} for result in run.results.values()
)
return int(failed), run

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