mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
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:
parent
8fc0663198
commit
ee08c36fc0
188 changed files with 10419 additions and 5452 deletions
1
.github/workflows/test-unit.yml
vendored
1
.github/workflows/test-unit.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
22
litellm-rust/Cargo.lock
generated
22
litellm-rust/Cargo.lock
generated
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"] }
|
||||
|
|
|
|||
|
|
@ -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"] }
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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>;
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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>,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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>,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
65
litellm-rust/crates/ai-gateway/src/trace_parity.rs
Normal file
65
litellm-rust/crates/ai-gateway/src/trace_parity.rs
Normal 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,
|
||||
})
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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";
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
215
litellm-rust/crates/core/src/observability/function_trace.rs
Normal file
215
litellm-rust/crates/core/src/observability/function_trace.rs
Normal 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")]
|
||||
);
|
||||
}
|
||||
}
|
||||
59
litellm-rust/crates/core/src/observability/mod.rs
Normal file
59
litellm-rust/crates/core/src/observability/mod.rs
Normal 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());
|
||||
}
|
||||
}
|
||||
|
|
@ -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>,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -134,6 +134,8 @@ impl OcrProviderConfig for MistralOcrConfig {
|
|||
document_annotation,
|
||||
usage_info,
|
||||
object: "ocr".to_string(),
|
||||
extra_fields: Map::new(),
|
||||
provider_native_response: None,
|
||||
})
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -4,4 +4,5 @@ pub mod azure_ai;
|
|||
pub mod bedrock;
|
||||
pub mod mistral;
|
||||
pub mod openai;
|
||||
pub mod reducto;
|
||||
pub mod vertex_ai;
|
||||
|
|
|
|||
1
litellm-rust/crates/core/src/providers/reducto/mod.rs
Normal file
1
litellm-rust/crates/core/src/providers/reducto/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod ocr;
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
pub mod transformation;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
202
litellm-rust/crates/core/src/providers/reducto/ocr/tests.rs
Normal file
202
litellm-rust/crates/core/src/providers/reducto/ocr/tests.rs
Normal 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())]
|
||||
);
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -1 +0,0 @@
|
|||
pub(crate) const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace";
|
||||
|
|
@ -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,
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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)?)
|
||||
}
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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`
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
113
tests/rust-python-harness/cli/__init__.py
Normal file
113
tests/rust-python-harness/cli/__init__.py
Normal 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
|
||||
116
tests/rust-python-harness/cli/catalog.py
Normal file
116
tests/rust-python-harness/cli/catalog.py
Normal 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)))
|
||||
45
tests/rust-python-harness/cli/commands.py
Normal file
45
tests/rust-python-harness/cli/commands.py
Normal 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
|
||||
458
tests/rust-python-harness/cli/test_cli.py
Normal file
458
tests/rust-python-harness/cli/test_cli.py
Normal 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"
|
||||
29
tests/rust-python-harness/conftest.py
Normal file
29
tests/rust-python-harness/conftest.py
Normal 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
|
||||
92
tests/rust-python-harness/shared/native_build.py
Normal file
92
tests/rust-python-harness/shared/native_build.py
Normal 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
|
||||
|
|
@ -1,3 +0,0 @@
|
|||
import pytest
|
||||
|
||||
pytest.register_assert_rewrite("tests.rust-python-harness.shared.parity.compare")
|
||||
|
|
@ -1 +1,7 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
JSON_OBJECT_ADAPTER: Final = TypeAdapter(dict[str, object])
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
{
|
||||
|
|
|
|||
|
|
@ -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]):
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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]}"
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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 ''}"
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
56
tests/rust-python-harness/shared/parity/local_server.py
Normal file
56
tests/rust-python-harness/shared/parity/local_server.py
Normal 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)
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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), (), ())
|
||||
|
|
|
|||
|
|
@ -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
|
||||
29
tests/rust-python-harness/shared/reporting/rendering.py
Normal file
29
tests/rust-python-harness/shared/reporting/rendering.py
Normal 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)
|
||||
92
tests/rust-python-harness/shared/reporting/strategy.py
Normal file
92
tests/rust-python-harness/shared/reporting/strategy.py
Normal 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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
116
tests/rust-python-harness/shared/test_native_build.py
Normal file
116
tests/rust-python-harness/shared/test_native_build.py
Normal 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
|
||||
|
|
@ -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)
|
||||
38
tests/rust-python-harness/shared/tracing/native.py
Normal file
38
tests/rust-python-harness/shared/tracing/native.py
Normal 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
|
||||
)
|
||||
117
tests/rust-python-harness/shared/tracing/profiler.py
Normal file
117
tests/rust-python-harness/shared/tracing/profiler.py
Normal 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)
|
||||
335
tests/rust-python-harness/shared/tracing/pytest_usage.py
Normal file
335
tests/rust-python-harness/shared/tracing/pytest_usage.py
Normal 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())
|
||||
271
tests/rust-python-harness/shared/tracing/steps.py
Normal file
271
tests/rust-python-harness/shared/tracing/steps.py
Normal 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,
|
||||
)
|
||||
|
|
@ -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",)
|
||||
108
tests/rust-python-harness/shared/tracing/test_profiler.py
Normal file
108
tests/rust-python-harness/shared/tracing/test_profiler.py
Normal 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}
|
||||
168
tests/rust-python-harness/shared/tracing/test_pytest_usage.py
Normal file
168
tests/rust-python-harness/shared/tracing/test_pytest_usage.py
Normal 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"
|
||||
157
tests/rust-python-harness/shared/tracing/test_steps.py
Normal file
157
tests/rust-python-harness/shared/tracing/test_steps.py
Normal 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
|
||||
|
|
@ -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__":
|
||||
253
tests/rust-python-harness/shared/unit_runners/rust_runner.py
Normal file
253
tests/rust-python-harness/shared/unit_runners/rust_runner.py
Normal 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))
|
||||
|
|
@ -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
|
||||
|
|
@ -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",))
|
||||
|
|
@ -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")
|
||||
)
|
||||
|
|
@ -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,)}
|
||||
|
|
@ -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.
|
||||
|
|
@ -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
|
||||
|
|
@ -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,
|
||||
)
|
||||
12
tests/rust-python-harness/strategies/e2e_parity/reporting.py
Normal file
12
tests/rust-python-harness/strategies/e2e_parity/reporting.py
Normal 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",)),)
|
||||
|
|
@ -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
Loading…
Add table
Reference in a new issue