mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
fix(rust): preserve bridge behavior after core move
This commit is contained in:
parent
851e706cb4
commit
ecc03d81ae
1 changed files with 0 additions and 72 deletions
|
|
@ -1,18 +1,9 @@
|
|||
use std::num::NonZeroUsize;
|
||||
|
||||
use litellm_core::provider_callbacks::{
|
||||
CallbackDecision, ProviderAttemptObserver, ProviderError, ProviderPostCall, ProviderPreCall,
|
||||
};
|
||||
use litellm_python_interop::callback_runtime::{AsyncContext, CallbackRuntime, SyncContext};
|
||||
use pyo3::prelude::*;
|
||||
|
||||
use crate::constants::OCR_CALLBACK_CAPACITY;
|
||||
use crate::execution::PythonCallContext;
|
||||
|
||||
litellm_core::provider_attempt_observer_catalog!(crate::bind_python_hooks,
|
||||
pub(crate) struct PythonProviderSession;
|
||||
trait ProviderAttemptObserver;
|
||||
);
|
||||
litellm_core::streaming_observer_catalog!(crate::bind_python_hooks,
|
||||
pub struct PythonStreamingSession;
|
||||
trait litellm_core::provider_callbacks::StreamingObserver;
|
||||
|
|
@ -36,41 +27,6 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
|
|||
)
|
||||
}
|
||||
|
||||
pub(crate) enum PythonProviderObserver {
|
||||
Disabled,
|
||||
Sync(PythonProviderSession<SyncContext>),
|
||||
Async(PythonProviderSession<AsyncContext>),
|
||||
}
|
||||
|
||||
impl PythonProviderObserver {
|
||||
pub(crate) fn new(
|
||||
adapter: Option<Py<PyAny>>,
|
||||
context: PythonCallContext<'_>,
|
||||
) -> PyResult<Self> {
|
||||
let Some(adapter) = adapter else {
|
||||
return Ok(Self::Disabled);
|
||||
};
|
||||
let py = context.py;
|
||||
let module = py.import("litellm.rust_bridge._native")?;
|
||||
let runtime = module
|
||||
.getattr("__python_callback_runtime__")?
|
||||
.extract::<PyRef<'_, PythonCallbackRuntime>>()?
|
||||
.0
|
||||
.clone();
|
||||
if context.asynchronous {
|
||||
Ok(Self::Async(PythonProviderSession::new(
|
||||
adapter.bind(py),
|
||||
runtime.async_context(py)?,
|
||||
)?))
|
||||
} else {
|
||||
Ok(Self::Sync(PythonProviderSession::new(
|
||||
adapter.bind(py),
|
||||
runtime.sync_context(py)?,
|
||||
)?))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn python_async_session(
|
||||
adapter: Py<PyAny>,
|
||||
py: Python<'_>,
|
||||
|
|
@ -83,31 +39,3 @@ pub(crate) fn python_async_session(
|
|||
.clone();
|
||||
PythonSession::new(adapter.bind(py), runtime.async_context(py)?)
|
||||
}
|
||||
|
||||
impl ProviderAttemptObserver for PythonProviderObserver {
|
||||
type Error = PyErr;
|
||||
|
||||
async fn pre_call(&mut self, input: &ProviderPreCall) -> PyResult<CallbackDecision> {
|
||||
match self {
|
||||
Self::Disabled => Ok(CallbackDecision::Unchanged),
|
||||
Self::Sync(session) => session.pre_call(input).await,
|
||||
Self::Async(session) => session.pre_call(input).await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn post_call(&mut self, input: &ProviderPostCall) -> PyResult<CallbackDecision> {
|
||||
match self {
|
||||
Self::Disabled => Ok(CallbackDecision::Unchanged),
|
||||
Self::Sync(session) => session.post_call(input).await,
|
||||
Self::Async(session) => session.post_call(input).await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn error(&mut self, input: &ProviderError) -> PyResult<()> {
|
||||
match self {
|
||||
Self::Disabled => Ok(()),
|
||||
Self::Sync(session) => session.error(input).await,
|
||||
Self::Async(session) => session.error(input).await,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue