litellm/litellm-rust/crates/python-bridge/src/cache/python/host.rs

use litellm_cache::Error;
use litellm_host::protocol::Reply;
use litellm_host_python::{from_py, to_py};
use pyo3::{
    gc::{PyTraverseError, PyVisit},
    prelude::*,
    types::PyDict,
};
use serde_json::Value;

use super::service::CacheCall;
use crate::cache::HandleError;

enum Pending {
    Lookup(Reply<Result<Option<Value>, Error>>),
    Store(Reply<Result<(), Error>>),
}

pub(crate) struct PythonCache {
    cache: Option<Py<PyAny>>,
    arguments: Option<Py<PyDict>>,
    pending: Option<Pending>,
    asynchronous: bool,
}

impl PythonCache {
    pub fn new(asynchronous: bool) -> Self {
        Self {
            cache: None,
            arguments: None,
            pending: None,
            asynchronous,
        }
    }

    pub(in crate::cache) fn bind(
        &mut self,
        cache: Bound<'_, PyAny>,
        arguments: &Bound<'_, PyDict>,
    ) {
        self.cache = Some(cache.unbind());
        self.arguments = Some(arguments.clone().unbind());
    }

    pub fn begin(&mut self, py: Python<'_>, call: CacheCall) -> PyResult<Option<Py<PyAny>>> {
        let Some(cache) = self.cache.as_ref() else {
            return Err(HandleError::Unconfigured.into());
        };
        let arguments = self
            .arguments
            .as_ref()
            .ok_or(HandleError::MissingArguments)?
            .bind(py)
            .copy()?;
        let (method, result) = match call {
            CacheCall::Lookup { reply } => {
                self.pending = Some(Pending::Lookup(reply));
                (
                    if self.asynchronous {
                        "async_get_cache"
                    } else {
                        "get_cache"
                    },
                    None,
                )
            }
            CacheCall::Store { value, reply } => {
                self.pending = Some(Pending::Store(reply));
                (
                    if self.asynchronous {
                        "async_add_cache"
                    } else {
                        "add_cache"
                    },
                    Some(to_py(py, &value)?),
                )
            }
        };
        let result = match result {
            Some(value) => cache
                .bind(py)
                .call_method(method, (value,), Some(&arguments)),
            None => cache.bind(py).call_method(method, (), Some(&arguments)),
        }
        .map(Bound::unbind);
        if self.asynchronous && result.is_ok() {
            return result.map(Some);
        }
        self.resume(py, result)
    }

    pub fn resume(
        &mut self,
        py: Python<'_>,
        result: PyResult<Py<PyAny>>,
    ) -> PyResult<Option<Py<PyAny>>> {
        let result = match result {
            Err(error) if !error.is_instance_of::<pyo3::exceptions::PyException>(py) => {
                self.pending = None;
                return Err(error);
            }
            result => result,
        };
        match self.pending.take() {
            Some(Pending::Lookup(reply)) => {
                let value = result.map_err(Error::unavailable).and_then(|value| {
                    if value.bind(py).is_none() {
                        Ok(None)
                    } else {
                        from_py(value.bind(py))
                            .map(Some)
                            .map_err(Error::invalid_entry)
                    }
                });
                reply.send(value);
            }
            Some(Pending::Store(reply)) => {
                reply.send(result.map(|_| ()).map_err(Error::unavailable));
            }
            None => return Err(HandleError::NothingPending.into()),
        }
        Ok(None)
    }

    pub fn close(&mut self) {
        self.pending = None;
        self.cache = None;
        self.arguments = None;
    }

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