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

use std::panic::{AssertUnwindSafe, catch_unwind};

use crate::{Error, panic_to_pyerr};
use pyo3::exceptions::PyBaseException;
use pyo3::gc::{PyTraverseError, PyVisit};
use pyo3::prelude::*;

pub type StepTypes = for<'py> fn(Python<'py>) -> PyResult<Bound<'py, PyModule>>;

/// What a suspended execution waits on: an awaitable the caller awaits, or a stream
/// boundary the caller's stream reads past.
pub enum Suspension {
    Await(Py<PyAny>),
    /// The call streams: the caller gets a stream over this execution carrying this head,
    /// and the execution stays suspended until the stream asks for a chunk.
    StreamOpen(Py<PyAny>),
    StreamChunk(Py<PyAny>),
}

enum Resume {
    Send(Py<PyAny>),
    Throw(PyErr),
}

pub enum ExecutionStep {
    Suspended(Suspension),
    Return(Py<PyAny>),
}

pub trait ExecutionBody: Send + Sync {
    fn resume(&mut self, result: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep>;
    fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>;
}

enum ExecutionState {
    Created(Box<dyn ExecutionBody>),
    Running,
    Suspended(Box<dyn ExecutionBody>),
    Closed,
}

#[pyclass]
pub struct Execution {
    state: ExecutionState,
    steps: StepTypes,
}

impl Execution {
    pub fn new(body: impl ExecutionBody + 'static, steps: StepTypes) -> Self {
        Self {
            state: ExecutionState::Created(Box::new(body)),
            steps,
        }
    }

    fn advance(slf: &Bound<'_, Self>, py: Python<'_>, input: Resume) -> PyResult<Py<PyAny>> {
        let (mut body, started) = {
            let mut execution = slf.borrow_mut();
            let started = match (&execution.state, &input) {
                (ExecutionState::Created(_), Resume::Send(value)) if value.is_none(py) => false,
                (ExecutionState::Created(_), Resume::Send(_)) => {
                    return Err(Error::SendBeforeStart.into());
                }
                (ExecutionState::Suspended(_), _) => true,
                (ExecutionState::Running, _) => {
                    return Err(Error::ExecutionRunning.into());
                }
                (ExecutionState::Closed, _) => {
                    return Err(Error::ExecutionClosed.into());
                }
                (ExecutionState::Created(_), Resume::Throw(_)) => {
                    return Err(Error::ExecutionOutOfOrder.into());
                }
            };
            match std::mem::replace(&mut execution.state, ExecutionState::Running) {
                ExecutionState::Created(body) | ExecutionState::Suspended(body) => (body, started),
                _ => unreachable!(),
            }
        };
        let result = match input {
            Resume::Send(value) => started.then_some(Ok(value)),
            Resume::Throw(error) => Some(Err(error)),
        };
        let steps = slf.borrow().steps;
        let outcome = catch_unwind(AssertUnwindSafe(|| {
            let step = body.resume(result)?;
            let (tag, value, suspended) = match step {
                ExecutionStep::Suspended(Suspension::Await(value)) => ("Await", value, true),
                ExecutionStep::Suspended(Suspension::StreamOpen(head)) => {
                    ("StreamOpen", head, true)
                }
                ExecutionStep::Suspended(Suspension::StreamChunk(chunk)) => {
                    ("StreamChunk", chunk, true)
                }
                ExecutionStep::Return(value) => ("Complete", value, false),
            };
            let step = steps(py)?.getattr(tag)?.call1((value,))?.unbind();
            Ok((step, suspended))
        }))
        .map_err(panic_to_pyerr)
        .and_then(|result| result);
        match outcome {
            Ok((step, true)) if matches!(slf.borrow().state, ExecutionState::Running) => {
                slf.borrow_mut().state = ExecutionState::Suspended(body);
                Ok(step)
            }
            outcome => {
                slf.borrow_mut().state = ExecutionState::Closed;
                drop(body);
                outcome.and_then(|(step, suspended)| {
                    if suspended {
                        Err(Error::ExecutionClosedWhileRunning.into())
                    } else {
                        Ok(step)
                    }
                })
            }
        }
    }
}

#[pymethods]
impl Execution {
    fn send(slf: &Bound<'_, Self>, py: Python<'_>, value: Py<PyAny>) -> PyResult<Py<PyAny>> {
        Self::advance(slf, py, Resume::Send(value))
    }

    fn throw(
        slf: &Bound<'_, Self>,
        py: Python<'_>,
        error: Bound<'_, PyBaseException>,
    ) -> PyResult<Py<PyAny>> {
        Self::advance(slf, py, Resume::Throw(PyErr::from_value(error.into_any())))
    }

    fn close(slf: &Bound<'_, Self>) {
        let state = std::mem::replace(&mut slf.borrow_mut().state, ExecutionState::Closed);
        drop(state);
    }

    fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
        match &self.state {
            ExecutionState::Created(body) | ExecutionState::Suspended(body) => {
                body.traverse(&visit)
            }
            _ => Ok(()),
        }
    }

    fn __clear__(slf: &Bound<'_, Self>) {
        Self::close(slf);
    }
}