diff --git a/litellm-rust/crates/python-bridge/tests/callbacks/mod.rs b/litellm-rust/crates/python-bridge/tests/callbacks/mod.rs index 296a902f8fc..57ed5314760 100644 --- a/litellm-rust/crates/python-bridge/tests/callbacks/mod.rs +++ b/litellm-rust/crates/python-bridge/tests/callbacks/mod.rs @@ -1,8 +1,13 @@ +use std::collections::BTreeMap; use std::ffi::{CStr, CString}; use std::num::NonZeroUsize; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; +use litellm_core::provider_callbacks::{ + ProviderError, ProviderPostCall, ProviderPreCall, ProviderStreamClose, ProviderStreamEvent, + SessionEvent, SessionObserver, StreamingObserver, +}; use litellm_python_interop::callback_runtime::CallbackRuntime; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; @@ -235,6 +240,151 @@ impl Harness { fn calls(&self) -> usize { self.calls.load(Ordering::SeqCst) } + + fn streaming(&self, py: Python<'_>, adapter: &Bound<'_, PyAny>) -> PyResult> { + let mut session = crate::callback_bindings::PythonStreamingSession::new( + adapter, + self.runtime.sync_context(py)?, + )?; + crate::execution::run_sync( + py, + async move { + let pre = provider_pre_call(); + assert!(matches!( + session.pre_call(&pre).await?, + litellm_core::provider_callbacks::CallbackDecision::Unchanged + )); + assert!(matches!( + session.post_call(&provider_post_call()).await?, + litellm_core::provider_callbacks::CallbackDecision::Unchanged + )); + assert!(matches!( + session.stream_event(&provider_stream_event()).await?, + litellm_core::provider_callbacks::CallbackDecision::Unchanged + )); + session.stream_close(&provider_stream_close()).await?; + session.error(&provider_error()).await + }, + std::convert::identity, + ) + } + + fn session<'py>( + &self, + py: Python<'py>, + adapter: &Bound<'py, PyAny>, + ) -> PyResult> { + let mut session = + crate::callback_bindings::PythonSession::new(adapter, self.runtime.async_context(py)?)?; + crate::execution::run_async( + py, + async move { + let event = session_event(); + assert!(matches!( + session.before_connect(&event).await?, + litellm_core::provider_callbacks::CallbackDecision::Unchanged + )); + session.connected(&event).await?; + assert!(matches!( + session.before_send(&event).await?, + litellm_core::provider_callbacks::CallbackDecision::Unchanged + )); + assert!(matches!( + session.after_receive(&event).await?, + litellm_core::provider_callbacks::CallbackDecision::Unchanged + )); + session.response_complete(&event).await?; + session.response_error(&event).await?; + session.error(&event).await?; + session.close(&event).await + }, + std::convert::identity, + ) + } +} + +fn provider_pre_call() -> ProviderPreCall { + ProviderPreCall { + provider: "test".to_string(), + model: "model".to_string(), + call_id: "call".to_string(), + trace_id: Some("trace".to_string()), + attempt: 2, + started_at: 1.0, + request: BTreeMap::new(), + api_base: "https://provider.test".to_string(), + headers: BTreeMap::new(), + } +} + +fn provider_post_call() -> ProviderPostCall { + ProviderPostCall { + provider: "test".to_string(), + model: "model".to_string(), + call_id: "call".to_string(), + trace_id: Some("trace".to_string()), + attempt: 2, + started_at: 1.0, + response: serde_json::json!({}), + status_code: 200, + headers: BTreeMap::new(), + ended_at: 2.0, + } +} + +fn provider_stream_event() -> ProviderStreamEvent { + ProviderStreamEvent { + provider: "test".to_string(), + model: "model".to_string(), + call_id: "call".to_string(), + trace_id: Some("trace".to_string()), + attempt: 2, + started_at: 1.0, + event: serde_json::json!({"type": "delta"}), + sequence: 1, + } +} + +fn provider_stream_close() -> ProviderStreamClose { + ProviderStreamClose { + provider: "test".to_string(), + model: "model".to_string(), + call_id: "call".to_string(), + trace_id: Some("trace".to_string()), + attempt: 2, + started_at: 1.0, + outcome: "completed".to_string(), + ended_at: 2.0, + } +} + +fn provider_error() -> ProviderError { + ProviderError { + provider: "test".to_string(), + model: "model".to_string(), + call_id: "call".to_string(), + trace_id: Some("trace".to_string()), + attempt: 2, + started_at: 1.0, + message: "retrying".to_string(), + stage: "provider_response", + committed: true, + status_code: Some(429), + will_retry: true, + ended_at: 2.0, + } +} + +fn session_event() -> SessionEvent { + SessionEvent { + session_id: "session".to_string(), + call_id: "call".to_string(), + trace_id: Some("trace".to_string()), + event: Some(serde_json::json!({"type": "response.create"})), + response_id: Some("response".to_string()), + sequence: Some(1), + message: None, + } } fn run_python_test(name: &str, capacity: usize) { @@ -279,4 +429,5 @@ python_tests! { cancellation_and_admission: 1, interrupted_session: 1, synchronous_callbacks: 1, + callback_catalogs: 64, } diff --git a/litellm-rust/crates/python-bridge/tests/callbacks/test_callbacks.py b/litellm-rust/crates/python-bridge/tests/callbacks/test_callbacks.py index 451ca9dc945..5783468b382 100644 --- a/litellm-rust/crates/python-bridge/tests/callbacks/test_callbacks.py +++ b/litellm-rust/crates/python-bridge/tests/callbacks/test_callbacks.py @@ -17,6 +17,8 @@ class Harness(Protocol): def sync(self, adapter: object) -> Payload: ... def interrupt(self, adapter: object, stop: Awaitable[object]) -> Awaitable[Payload]: ... def calls(self) -> int: ... + def streaming(self, adapter: object) -> None: ... + def session(self, adapter: object) -> Awaitable[None]: ... REQUEST_CONTEXT: Final = contextvars.ContextVar("request_context", default="default") @@ -328,6 +330,80 @@ async def synchronous_callbacks(harness: Harness) -> None: assert inspect.getcoroutinestate(coroutine) == inspect.CORO_CLOSED +class CatalogAdapter: + def __init__(self) -> None: + self.events: tuple[str, ...] = () + + def _record(self, name: str) -> None: + self.events += (name,) + + def pre_call(self, payload: object) -> dict[str, str]: + self._record("pre_call") + return {"action": "unchanged"} + + def post_call(self, payload: object) -> dict[str, str]: + self._record("post_call") + return {"action": "unchanged"} + + def error(self, payload: object) -> None: + self._record("error") + + def stream_event(self, payload: object) -> dict[str, str]: + self._record("stream_event") + return {"action": "unchanged"} + + def stream_close(self, payload: object) -> None: + self._record("stream_close") + + async def before_connect(self, payload: object) -> dict[str, str]: + self._record("before_connect") + return {"action": "unchanged"} + + async def connected(self, payload: object) -> None: + self._record("connected") + + async def before_send(self, payload: object) -> dict[str, str]: + self._record("before_send") + return {"action": "unchanged"} + + async def after_receive(self, payload: object) -> dict[str, str]: + self._record("after_receive") + return {"action": "unchanged"} + + async def response_complete(self, payload: object) -> None: + self._record("response_complete") + + async def response_error(self, payload: object) -> None: + self._record("response_error") + + async def close(self, payload: object) -> None: + self._record("close") + + +class SessionCatalogAdapter(CatalogAdapter): + async def error(self, payload: object) -> None: + self._record("error") + + +async def callback_catalogs(harness: Harness) -> None: + streaming: Final = CatalogAdapter() + harness.streaming(streaming) + assert streaming.events == ("pre_call", "post_call", "stream_event", "stream_close", "error") + + session: Final = SessionCatalogAdapter() + await harness.session(session) + assert session.events == ( + "before_connect", + "connected", + "before_send", + "after_receive", + "response_complete", + "response_error", + "error", + "close", + ) + + TESTS: Final = ( transforms_and_context, callback_errors, @@ -337,6 +413,7 @@ TESTS: Final = ( cancellation_and_admission, interrupted_session, synchronous_callbacks, + callback_catalogs, )