litellm/litellm-rust/crates/host-python/tests/handle.rs

mod support;

use litellm_host::{hooks::CallHooks, lifecycle::Timing};
use litellm_host_python::{
    Execution, ExecutionBody, ExecutionStep, HookStep, PythonCallEvent, PythonOwned, PythonRuntime,
    Suspension, run_call,
};
use pyo3::{
    exceptions::PyBaseException,
    gc::{PyTraverseError, PyVisit},
    prelude::*,
    types::{PyCFunction, PyDict},
};
use rstest::rstest;
use support::{
    Interpreter, Log, OpScript, binding, call_options, handed, install_lifecycle_module,
    interpreter, lifecycle_module, success_machine,
};

struct RetainingHost {
    retained: Option<Py<PyAny>>,
}

impl ExecutionBody for RetainingHost {
    fn resume(&mut self, _: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
        Python::attach(|py| Ok(ExecutionStep::Return(py.None())))
    }

    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
        visit.call(&self.retained)
    }
}

#[pyfunction]
fn retaining_coroutine(py: Python<'_>, retained: Py<PyAny>) -> PyResult<Py<Execution>> {
    Py::new(
        py,
        Execution::new(
            RetainingHost {
                retained: Some(retained),
            },
            lifecycle_module,
        ),
    )
}

struct AwaitBody(Option<Py<PyAny>>);

impl ExecutionBody for AwaitBody {
    fn resume(&mut self, result: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
        match self.0.take() {
            Some(awaitable) => Ok(ExecutionStep::Suspended(Suspension::Await(awaitable))),
            None => result
                .expect("selected await completed")
                .map(ExecutionStep::Return),
        }
    }

    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
        visit.call(&self.0)
    }
}

#[pyfunction]
fn await_execution(awaitable: Py<PyAny>) -> Execution {
    Execution::new(AwaitBody(Some(awaitable)), lifecycle_module)
}

struct CallingBody(Py<PyAny>);

impl ExecutionBody for CallingBody {
    fn resume(&mut self, _: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
        Python::attach(|py| self.0.call0(py).map(ExecutionStep::Return))
    }

    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
        visit.call(&self.0)
    }
}

#[pyfunction]
fn calling_execution(callback: Py<PyAny>) -> Execution {
    Execution::new(CallingBody(callback), lifecycle_module)
}

struct ErrorBody(Option<Py<PyBaseException>>);

impl ExecutionBody for ErrorBody {
    fn resume(&mut self, _: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
        Python::attach(|py| {
            Err(PyErr::from_value(
                self.0.take().unwrap().into_bound(py).into_any(),
            ))
        })
    }

    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
        visit.call(&self.0)
    }
}

#[pyfunction]
fn error_execution(error: Bound<'_, PyBaseException>) -> Execution {
    Execution::new(ErrorBody(Some(error.unbind())), lifecycle_module)
}

/// Where a paused call keeps the test's object: in its hooks, or selected as the call's
/// response or exception, with the call paused in the hook that is told about it.
#[derive(Clone, Copy, Debug)]
enum Holder {
    Hooks,
    Response,
    Exception,
}

struct PausingHooks {
    retained: Option<Py<PyAny>>,
    pause: Option<Py<PyAny>>,
    holder: Holder,
}

impl CallHooks<PythonRuntime> for PausingHooks {
    fn transform_response(
        &mut self,
        py: Python<'_>,
        response: Py<PyAny>,
        _: Timing,
    ) -> PyResult<HookStep<Self, Py<PyAny>>> {
        match self.holder {
            Holder::Hooks => Ok(HookStep::Ready(response)),
            Holder::Response => Ok(HookStep::Ready(self.retained.take().unwrap())),
            Holder::Exception => Err(PyErr::from_value(
                self.retained.take().unwrap().into_bound(py),
            )),
        }
    }

    fn on_event(
        &mut self,
        _: Python<'_>,
        event: PythonCallEvent<'_>,
    ) -> PyResult<HookStep<Self, ()>> {
        let pauses_here = match self.holder {
            Holder::Hooks => true,
            Holder::Response => matches!(event, PythonCallEvent::Succeeded { .. }),
            Holder::Exception => matches!(event, PythonCallEvent::Failed { .. }),
        };
        Ok(match self.pause.take_if(|_| pauses_here) {
            Some(pause) => {
                HookStep::Await(pause, |_, _, answer| answer.map(|_| HookStep::Ready(())))
            }
            None => HookStep::Ready(()),
        })
    }
}

impl PythonOwned for PausingHooks {
    fn close(&mut self, _: Python<'_>) {
        self.retained = None;
        self.pause = None;
    }

    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
        visit.call(&self.retained)?;
        visit.call(&self.pause)
    }
}

fn paused_call(
    py: Python<'_>,
    retained: Py<PyAny>,
    pause: Py<PyAny>,
    holder: Holder,
) -> PyResult<Py<PyAny>> {
    let machine = success_machine();
    handed(
        py,
        run_call(
            py,
            move |_, _, request| Ok(machine(request)),
            binding(&Log::default(), OpScript::Answer),
            PausingHooks {
                retained: Some(retained),
                pause: Some(pause),
                holder,
            },
            PyDict::new(py).unbind(),
            call_options(true),
        ),
        true,
    )
}

#[rstest]
fn every_suspension_is_awaited_inline_on_the_callers_task(
    #[from(interpreter)] _python: Interpreter,
) {
    Python::attach(|py| {
        let module = install_lifecycle_module(py);
        let locals = PyDict::new(py);
        locals
            .set_item("drive", module.getattr("drive").unwrap())
            .unwrap();
        locals
            .set_item(
                "await_execution",
                wrap_pyfunction!(await_execution, py).unwrap(),
            )
            .unwrap();
        locals
            .set_item(
                "calling_execution",
                wrap_pyfunction!(calling_execution, py).unwrap(),
            )
            .unwrap();
        let probe = std::ffi::CString::new(include_str!("lifecycle.py")).unwrap();
        py.run(&probe, Some(&locals), Some(&locals)).unwrap();
    });
}

#[rstest]
fn retained_exception_frames_are_collectable(#[from(interpreter)] _python: Interpreter) {
    Python::attach(|py| {
        let locals = PyDict::new(py);
        locals
            .set_item(
                "error_execution",
                wrap_pyfunction!(error_execution, py).unwrap(),
            )
            .unwrap();
        py.run(
            pyo3::ffi::c_str!(
                r#"
import gc
import weakref

class Retained:
    pass

def cycle():
    retained = Retained()
    try:
        raise ValueError('retained traceback')
    except ValueError as error:
        retained.owner = error_execution(error)
    return weakref.ref(retained)

reference = cycle()
gc.collect()
assert reference() is None
"#
            ),
            Some(&locals),
            Some(&locals),
        )
        .unwrap();
    });
}

#[rstest]
fn cycles_through_a_retained_execution_are_collectable(#[from(interpreter)] _python: Interpreter) {
    Python::attach(|py| {
        let locals = PyDict::new(py);
        locals
            .set_item(
                "retaining_coroutine",
                wrap_pyfunction!(retaining_coroutine, py).unwrap(),
            )
            .unwrap();
        py.run(
            pyo3::ffi::c_str!(
                r#"
import gc
import weakref

class Retained:
    pass

def cycle():
    retained = Retained()
    coroutine = retaining_coroutine(retained)
    retained.coroutine = coroutine
    return weakref.ref(retained)

retained_ref = cycle()
gc.collect()
assert retained_ref() is None
"#
            ),
            Some(&locals),
            Some(&locals),
        )
        .unwrap();
    });
}

#[rstest]
#[case::held_by_the_hooks(Holder::Hooks)]
#[case::selected_as_the_response(Holder::Response)]
#[case::selected_as_the_exception(Holder::Exception)]
fn cycles_through_a_call_suspended_in_a_hook_are_collectable(
    #[from(interpreter)] _python: Interpreter,
    #[case] holder: Holder,
) {
    Python::attach(|py| {
        let locals = PyDict::new(py);
        let call = PyCFunction::new_closure(py, None, None, move |args, _| {
            let (retained, pause) = args.extract::<(Py<PyAny>, Py<PyAny>)>()?;
            paused_call(args.py(), retained, pause, holder)
        })
        .unwrap();
        locals.set_item("paused_call", call).unwrap();
        py.run(
            pyo3::ffi::c_str!(
                r#"
import gc
import weakref

class Retained(Exception):
    pass

class Pause:
    def __await__(self):
        yield

def cycle():
    retained = Retained()
    coroutine = paused_call(retained, Pause())
    coroutine.send(None)
    retained.coroutine = coroutine
    return weakref.ref(retained)

reference = cycle()
gc.collect()
assert reference() is None
"#
            ),
            Some(&locals),
            Some(&locals),
        )
        .unwrap();
    });
}