scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-16 04:29:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 138
litellm/litellm-rust/crates/host-python/tests/driver.rs
yujonglee 8302f006dd
refactor(rust-bridge): return an unstarted Execution from native entrypoints and drive it in Python (#46299) 
* refactor(rust-bridge): return an unstarted Execution from native entrypoints and drive it in Python

Rust's run_call built Python streams itself and imported
litellm.rust_bridge.streams by name, so sync calls were driven in Rust
while async calls came back as a Rust-made coroutine. Every native
inference entrypoint now returns one unstarted Execution and
public_execution.py drives it with lifecycle.run or lifecycle.drive,
wrapping streams and _hidden_params on the Python side. Rust only
imports the lifecycle step classes through StepTypes, and
python-bridge/src/lifecycle.rs, PythonLifecycle and streams.py are gone.

The adapter trusts the _native.pyi contract with a single cast instead
of a per-request runtime_checkable isinstance. Binding tests assert the
adapter starts the execution and returns its completed value rather
than checking callable identity.

* refactor(rust-bridge): fold lifecycle and public_execution into one execution module

The public stream wrappers only forwarded to the lifecycle streams and added the
header dict, so the streams now carry _hidden_params themselves. lifecycle.py no
longer held any call lifecycle, just the Execution driver, so it is renamed to
execution.py and absorbs the synchronous and asynchronous adapters

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(rust-bridge): keep execution.py free of third-party imports

The host-python tests load execution.py into a bare interpreter, so the
pydantic header validation broke them. The stream head comes from the
binding's encode_stream_head, which always builds a str-keyed dict

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test(host-python): preserve pre-async hook coverage after execution handoff

* style(rust-bridge): sort litellm_host_python imports

Co-Authored-By: Yujong Lee <yujong@berri.ai>

* test(rust-bridge): drive unstarted native calls through execution in the split lifecycle tests

Co-Authored-By: Yujong Lee <yujong@berri.ai>

* fix(rust-bridge): mark execution TypeVar as Final

Co-Authored-By: Yujong Lee <yujong@berri.ai>

* test(rust-bridge): cover stream close, read failure, and sync suspension in execution

Co-Authored-By: Yujong Lee <yujong@berri.ai>

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-10-14 17:50:44 -07:00

770 lines
24 KiB
Rust
Raw Blame History

mod support;
use std::num::NonZeroUsize;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use litellm_host::{
hooks::CallHooks,
lifecycle::{CallEvent, ExecutionEvent, FailureOrigin},
machine::CallMachine,
observation::observation_channel,
};
use litellm_host_python::{CallOptions, HookChain, PythonOwned, PythonRuntime, run_call};
use pyo3::{
exceptions::{
PyBaseException, PyException, PyLookupError, PyRuntimeError, PyStopIteration, PyTypeError,
PyValueError, asyncio::CancelledError,
},
gc::{PyTraverseError, PyVisit},
prelude::*,
types::PyDict,
};
use rstest::rstest;
use support::{
HookScript, Interpreter, Log, MachineFactory, OpScript, Synthetic, binding, call_options,
failing_machine, handed, hooks, interpreter, lifecycle_module, open_stream, read_all,
run_composed, run_scripted, settled, success_machine,
};
fn observed(capacity: usize) -> (CallOptions, impl FnMut() -> Vec<CallEvent>) {
let (sender, mut receiver) = observation_channel(NonZeroUsize::new(capacity).unwrap());
let options = CallOptions {
asynchronous: false,
steps: lifecycle_module,
observers: Some(sender),
};
(options, move || {
std::iter::from_fn(|| receiver.try_recv().ok()).collect()
})
}
#[rstest]
#[case::success(OpScript::AwaitAnswer)]
#[case::failure(OpScript::AwaitFailure)]
fn asynchronous_host_operations_resume_the_same_machine(
#[from(interpreter)] _python: Interpreter,
#[case] op: OpScript,
) {
Python::attach(|py| {
let machine: MachineFactory = Box::new(|_| {
CallMachine::<Synthetic>::new(None, |host| {
Box::pin(async move { host.services.call(|reply| ("read", reply)).await })
})
});
let (result, log) = run_scripted(py, machine, op, HookScript::Plain, true);
match op {
OpScript::AwaitAnswer => {
assert_eq!(result.unwrap().extract::<String>(py).unwrap(), "awaited")
}
OpScript::AwaitFailure => {
assert!(result.unwrap_err().is_instance_of::<PyLookupError>(py));
assert_eq!(
log.iter()
.filter(|entry| entry.starts_with("failed:"))
.count(),
1
);
assert!(!log.iter().any(|entry| entry.starts_with("succeeded:")));
}
_ => unreachable!(),
}
});
}
#[rstest]
fn composed_hooks_share_prepared_arguments_and_one_execution(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let log = Log::default();
let chain = HookChain::new()
.with(hooks(&log, HookScript::ObserveArguments))
.with(hooks(&log, HookScript::RewriteArguments))
.with(hooks(&log, HookScript::ReplaceResponse));
let result = handed(
py,
run_call(
py,
move |_, _, request| Ok(success_machine()(request)),
binding(&log, OpScript::Answer),
chain,
PyDict::new(py).unbind(),
call_options(asynchronous),
),
asynchronous,
);
if asynchronous {
assert!(log.entries().is_empty());
}
let response = settled(py, result, asynchronous).unwrap();
assert_eq!(response.extract::<String>(py).unwrap(), "replaced");
let entries = log.entries();
assert!(entries.iter().any(|entry| entry == "adopted:hook"));
assert_eq!(
entries.iter().filter(|entry| *entry == "project").count(),
1
);
assert_eq!(
entries.iter().filter(|entry| *entry == "op:sign").count(),
1
);
assert_eq!(
entries
.iter()
.filter(|entry| *entry == "succeeded:replaced")
.count(),
3
);
assert!(!entries.iter().any(|entry| entry.starts_with("failed:")));
});
}
#[rstest]
#[case::draining(4, false)]
#[case::full(1, false)]
#[case::closed(4, true)]
fn observation_delivery_preserves_callback_order_and_response(
#[from(interpreter)] _python: Interpreter,
#[case] capacity: usize,
#[case] closed: bool,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let (expected, expected_log) = run_scripted(
py,
success_machine(),
OpScript::Answer,
HookScript::ReplaceResponse,
asynchronous,
);
let (sender, mut receiver) = observation_channel(NonZeroUsize::new(capacity).unwrap());
if closed {
receiver.close();
}
let log = Log::default();
let (actual, actual_log) = run_composed(
py,
success_machine(),
binding(&log, OpScript::Answer),
hooks(&log, HookScript::ReplaceResponse),
std::convert::identity,
CallOptions {
asynchronous,
steps: lifecycle_module,
observers: Some(sender.clone()),
},
);
assert_eq!(
actual.unwrap().extract::<String>(py).unwrap(),
expected.unwrap().extract::<String>(py).unwrap()
);
assert_eq!(actual_log, expected_log);
let events: Vec<_> = std::iter::from_fn(|| receiver.try_recv().ok()).collect();
assert_eq!(events.len() as u64 + sender.dropped_events(), 3);
if !closed && capacity >= 3 {
let [
CallEvent::Started { start_time },
CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }),
CallEvent::Succeeded { timing, .. },
] = events.as_slice()
else {
panic!("expected start, provider response and success: {events:?}");
};
assert_eq!(raw.body, "raw");
assert_eq!(*start_time, timing.start_time);
assert!(timing.end_time >= timing.start_time);
}
});
}
#[rstest]
#[case::preparation(HookScript::FailBegin, false, 1)]
#[case::transformation(HookScript::FailAfterSuccess, false, 1)]
#[case::terminal_cancellation(HookScript::CancelTerminal, true, 0)]
#[case::terminal_failure(HookScript::FailTerminal, false, 0)]
fn a_terminal_hook_failure_is_reported_once_and_never_replays_callbacks(
#[from(interpreter)] _python: Interpreter,
#[case] script: HookScript,
#[case] cancelled: bool,
#[case] failure_callbacks: usize,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let (options, mut events) = observed(4);
let log = Log::default();
let (result, log) = run_composed(
py,
success_machine(),
binding(&log, OpScript::Answer),
hooks(&log, script),
std::convert::identity,
CallOptions {
asynchronous,
..options
},
);
let error = result.unwrap_err();
assert_eq!(error.is_instance_of::<CancelledError>(py), cancelled);
let events = events();
assert!(matches!(events.first(), Some(CallEvent::Started { .. })));
assert_eq!(
matches!(events.last(), Some(CallEvent::Cancelled { .. })),
cancelled
);
assert_eq!(
matches!(
events.last(),
Some(CallEvent::Failed {
origin: FailureOrigin::Host,
..
})
),
!cancelled
);
let terminal = events.iter().filter(|event| {
matches!(
event,
CallEvent::Failed { .. }
| CallEvent::Succeeded { .. }
| CallEvent::Cancelled { .. }
)
});
assert_eq!(terminal.count(), 1);
assert_eq!(
log.iter()
.filter(|entry| entry.starts_with("failed:"))
.count(),
failure_callbacks
);
assert!(
!log.iter()
.any(|entry| entry.starts_with("succeeded:") || entry == "cancelled")
);
});
}
#[rstest]
#[case::preparation(HookScript::FailBegin, OpScript::Answer)]
#[case::native_decode(HookScript::Plain, OpScript::RejectRequestNatively)]
#[case::python_decode(HookScript::Plain, OpScript::RaiseRequestPython)]
fn startup_failures_release_the_factory_without_constructing_a_machine(
#[from(interpreter)] _python: Interpreter,
#[case] script: HookScript,
#[case] op: OpScript,
#[values(false, true)] asynchronous: bool,
) {
struct Release(Arc<AtomicBool>);
impl Drop for Release {
fn drop(&mut self) {
self.0.store(true, Ordering::SeqCst);
}
}
Python::attach(|py| {
let constructed = Arc::new(AtomicBool::new(false));
let did_construct = constructed.clone();
let released = Arc::new(AtomicBool::new(false));
let release = Release(released.clone());
let machine: MachineFactory = Box::new(move |request| {
let _release = release;
did_construct.store(true, Ordering::SeqCst);
success_machine()(request)
});
let (result, log) = run_scripted(py, machine, op, script, asynchronous);
assert!(result.is_err());
assert!(!constructed.load(Ordering::SeqCst));
assert!(released.load(Ordering::SeqCst));
assert_eq!(
log.iter()
.filter(|event| event.starts_with("failed:"))
.count(),
1
);
assert_eq!(
log.iter().filter(|event| *event == "adapter.close").count(),
1
);
assert_eq!(log.iter().filter(|event| *event == "host.close").count(), 1);
});
}
#[rstest]
fn success_runs_every_step_in_order_and_returns_the_public_response(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
success_machine(),
OpScript::Answer,
HookScript::Plain,
asynchronous,
);
assert_eq!(
result.unwrap().extract::<String>(py).unwrap(),
"project:1|sign|rewritten"
);
assert_eq!(
log,
[
"started",
"begin",
"project",
"op:sign",
"before_provider_request",
"response:raw",
"complete",
"after_success",
"succeeded:project:1|sign|rewritten",
"adapter.close",
"host.close",
]
);
});
}
#[rstest]
fn explicitly_closing_a_stream_dispatches_success_for_delivered_chunks(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let log = Log::default();
let (options, mut events) = observed(3);
let stream = open_stream(
py,
hooks(&log, HookScript::Plain),
CallOptions {
asynchronous,
..options
},
);
if asynchronous {
for method in ["__anext__", "aclose", "aclose"] {
let stop = stream
.call_method0(method)
.unwrap()
.call_method1("send", (py.None(),))
.unwrap_err();
assert!(stop.is_instance_of::<PyStopIteration>(py));
}
} else {
assert_eq!(
stream
.call_method0("__next__")
.unwrap()
.extract::<String>()
.unwrap(),
"first"
);
stream.call_method0("close").unwrap();
stream.call_method0("close").unwrap();
}
assert!(matches!(
events().as_slice(),
[CallEvent::Started { .. }, CallEvent::Succeeded { .. }]
));
assert_eq!(
log.entries(),
[
"started",
"begin",
"opened",
"delivered",
"succeeded:None",
"adapter.close"
]
);
});
}
#[rstest]
fn a_stream_carries_its_head_before_the_first_chunk(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let log = Log::default();
let stream = open_stream(
py,
hooks(&log, HookScript::Plain),
call_options(asynchronous),
);
let hidden: std::collections::HashMap<String, std::collections::HashMap<String, String>> =
stream.getattr("_hidden_params").unwrap().extract().unwrap();
assert_eq!(
hidden["additional_headers"],
std::collections::HashMap::from([("request-id".to_string(), "req_1".to_string())])
);
assert_eq!(log.entries(), ["started", "begin", "opened"]);
assert_eq!(read_all(py, &stream, asynchronous), ["first", "second"]);
});
}
struct Classification {
provider_fails: bool,
op: OpScript,
classifier_fails: bool,
message: &'static str,
native_context: Option<&'static str>,
steps: &'static [&'static str],
}
#[rstest]
#[case::provider_failure(Classification {
provider_fails: true,
op: OpScript::Answer,
classifier_fails: false,
message: "classified: provider exploded",
native_context: None,
steps: &["project", "classify:provider exploded"],
})]
#[case::host_op_rejection(Classification {
provider_fails: false,
op: OpScript::RejectNatively,
classifier_fails: false,
message: "classified: op rejected",
native_context: None,
steps: &["project", "op:sign", "classify:op rejected"],
})]
#[case::host_op_python_exception(Classification {
provider_fails: false,
op: OpScript::RaisePython,
classifier_fails: false,
message: "op failed",
native_context: None,
steps: &["project", "op:sign"],
})]
#[case::failing_classifier(Classification {
provider_fails: true,
op: OpScript::Answer,
classifier_fails: true,
message: "classifier failed",
native_context: Some("provider exploded"),
steps: &["project", "classify:provider exploded"],
})]
fn a_native_failure_is_classified_once_and_a_python_exception_is_raised_as_is(
#[from(interpreter)] _python: Interpreter,
#[case] case: Classification,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let Classification {
provider_fails,
op,
classifier_fails,
message,
native_context,
steps,
} = case;
let log = Log::default();
let binding = if classifier_fails {
binding(&log, op).with_failing_classifier()
} else {
binding(&log, op)
};
let machine = if provider_fails {
failing_machine()
} else {
success_machine()
};
let (result, log) = run_composed(
py,
machine,
binding,
hooks(&log, HookScript::Plain),
std::convert::identity,
call_options(asynchronous),
);
let error = result.unwrap_err();
assert_eq!(error.value(py).to_string(), message);
match native_context {
Some(context) => {
assert!(error.is_instance_of::<PyTypeError>(py));
let chained = error.context(py).unwrap();
assert!(chained.is_instance_of::<PyRuntimeError>(py));
assert_eq!(chained.value(py).to_string(), context);
}
None => assert!(error.is_instance_of::<PyValueError>(py)),
}
let expected: Vec<String> = ["started", "begin"]
.iter()
.chain(steps)
.map(|step| step.to_string())
.chain([
format!("failed:Call:{message}"),
"adapter.close".into(),
"host.close".into(),
])
.collect();
assert_eq!(log, expected);
});
}
#[rstest]
fn begin_failures_are_host_failures_without_provider_mapping(
#[from(interpreter)] _python: Interpreter,
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
success_machine(),
OpScript::Answer,
HookScript::FailBegin,
false,
);
assert_eq!(result.unwrap_err().value(py).to_string(), "begin failed");
assert_eq!(
log,
[
"started",
"begin",
"failed:Host:begin failed",
"adapter.close",
"host.close"
]
);
});
}
enum ArgumentPolicy {
Inherit,
Reject(Py<PyBaseException>),
}
impl CallHooks<PythonRuntime> for ArgumentPolicy {
fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> PyResult<()> {
match self {
Self::Inherit => arguments.bind(py).set_item("api_key", "inherited"),
Self::Reject(error) => Err(PyErr::from_value(error.bind(py).clone().into_any())),
}
}
}
impl PythonOwned for ArgumentPolicy {
fn close(&mut self, _: Python<'_>) {}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
match self {
Self::Inherit => Ok(()),
Self::Reject(error) => visit.call(error),
}
}
}
#[rstest]
fn resource_setup_uses_prepared_arguments_and_failures_are_terminal(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
#[values(false, true)] fail_setup: bool,
) {
Python::attach(|py| {
let log = Log::default();
let setup_log = log.clone();
let error = PyValueError::new_err("setup failed").into_value(py);
let setup_error = error.clone_ref(py);
let arguments = PyDict::new(py);
arguments.set_item("api_key", "original").unwrap();
let result = handed(
py,
run_call(
py,
move |py, prepared, request| {
let source: String = prepared.get_item("prepared")?.unwrap().extract()?;
let key: String = prepared.get_item("api_key")?.unwrap().extract()?;
setup_log.push(format!("setup:{source}:{key}"));
if fail_setup {
return Err(PyErr::from_value(setup_error.into_bound(py).into_any()));
}
Ok(success_machine()(request))
},
binding(&log, OpScript::Answer),
HookChain::new()
.with(hooks(&log, HookScript::RewriteArguments))
.with(ArgumentPolicy::Inherit),
arguments.clone().unbind(),
call_options(asynchronous),
),
asynchronous,
);
if asynchronous {
assert!(log.entries().is_empty());
}
let settled = settled(py, result, asynchronous);
assert_eq!(
arguments
.get_item("api_key")
.unwrap()
.unwrap()
.extract::<String>()
.unwrap(),
"original"
);
assert_eq!(
&log.entries()[..4],
["started", "begin", "project", "setup:hook:inherited"]
);
if fail_setup {
assert!(settled.unwrap_err().value(py).is(error.bind(py)));
assert_eq!(
log.entries(),
[
"started",
"begin",
"project",
"setup:hook:inherited",
"failed:Host:setup failed",
"adapter.close",
"host.close",
]
);
} else {
assert!(settled.is_ok());
assert!(
log.entries()
.iter()
.any(|entry| entry.starts_with("succeeded:"))
);
}
});
}
#[rstest]
fn closing_an_unstarted_call_never_initializes_resources(
#[from(interpreter)] _python: Interpreter,
) {
Python::attach(|py| {
let log = Log::default();
let setup_log = log.clone();
let pending = handed(
py,
run_call(
py,
move |_, _, request| {
setup_log.push("setup");
Ok(success_machine()(request))
},
binding(&log, OpScript::Answer),
hooks(&log, HookScript::Plain),
PyDict::new(py).unbind(),
call_options(true),
),
true,
)
.unwrap();
assert!(log.entries().is_empty());
pending.call_method0(py, "close").unwrap();
drop(pending);
assert_eq!(log.entries(), ["adapter.close", "host.close"]);
});
}
#[rstest]
fn an_argument_policy_rejection_is_the_callers_error_and_the_machine_never_starts(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let raised = PyValueError::new_err("over budget").into_value(py);
let log = Log::default();
let (result, log) = run_composed(
py,
success_machine(),
binding(&log, OpScript::Answer),
hooks(&log, HookScript::Plain),
|hooks| {
HookChain::new()
.with(hooks)
.with(ArgumentPolicy::Reject(raised.clone_ref(py)))
},
call_options(asynchronous),
);
assert!(result.unwrap_err().value(py).is(raised.bind(py)));
assert_eq!(
log,
[
"started",
"begin",
"failed:Host:over budget",
"adapter.close",
"host.close"
]
);
});
}
#[rstest]
fn the_host_projects_from_the_keyword_view_the_argument_policy_rewrote(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let log = Log::default();
let (result, _) = run_composed(
py,
success_machine(),
binding(&log, OpScript::Answer),
hooks(&log, HookScript::Plain),
|hooks| HookChain::new().with(hooks).with(ArgumentPolicy::Inherit),
call_options(asynchronous),
);
assert_eq!(
result.unwrap().extract::<String>(py).unwrap(),
"project:2|sign|rewritten"
);
});
}
#[rstest]
fn a_failure_while_finalizing_fails_the_call_instead_of_succeeding(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
success_machine(),
OpScript::Answer,
HookScript::FailAfterSuccess,
asynchronous,
);
assert_eq!(
result.unwrap_err().value(py).to_string(),
"after_success failed"
);
assert_eq!(
&log[log.len() - 4..],
[
"after_success",
"failed:Host:after_success failed",
"adapter.close",
"host.close"
]
);
assert!(!log.iter().any(|entry| entry.starts_with("succeeded")));
});
}
#[rstest]
fn cancellation_ends_the_call_without_terminal_dispatch(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
success_machine(),
OpScript::CancelRequest,
HookScript::Plain,
asynchronous,
);
let error = result.unwrap_err();
assert!(!error.is_instance_of::<PyException>(py));
assert_eq!(
log,
["started", "begin", "project", "adapter.close", "host.close"]
);
});
}
Reference in a new issue View git blame Copy permalink