mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
fix(rust): initialize callback lifecycle runtime
This commit is contained in:
parent
e4f8aa57af
commit
7780d877bb
4 changed files with 7 additions and 1 deletions
|
|
@ -162,6 +162,8 @@ mod _native {
|
|||
#[pymodule_init]
|
||||
fn init(module: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
super::errors::register(module)?;
|
||||
litellm_python_interop::callback_runtime::register(module)?;
|
||||
super::callback_bindings::register(module)?;
|
||||
super::routes::register(module)?;
|
||||
module.add_class::<super::ResponsesWebSocketConnection>()?;
|
||||
module.add_function(wrap_pyfunction!(
|
||||
|
|
|
|||
|
|
@ -398,6 +398,7 @@ def test_load_rust_ocr_uses_compiled_extension(monkeypatch):
|
|||
fake_module = types.ModuleType("litellm.rust_bridge._native")
|
||||
fake_module.ocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined]
|
||||
fake_module.aocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined]
|
||||
fake_module.ready_endpoints = {"ocr": {"callbacks"}} # type: ignore[attr-defined]
|
||||
monkeypatch.setattr(
|
||||
importlib.import_module("litellm.rust_bridge.bindings"),
|
||||
"get_native_bridge",
|
||||
|
|
|
|||
|
|
@ -113,6 +113,7 @@ async def test_connection_forwards_session_callback_adapter() -> None:
|
|||
cls,
|
||||
request: NativeResponsesWebSocketRequest,
|
||||
*,
|
||||
options: object,
|
||||
context: NativeRequestContext,
|
||||
callback_adapter: object | None = None,
|
||||
) -> _FakeNativeConnection:
|
||||
|
|
@ -167,7 +168,7 @@ async def test_connection_dispatch_cleans_up_without_reconnecting(native, sessio
|
|||
|
||||
class Native:
|
||||
@classmethod
|
||||
async def connect(cls, request, *, options, context):
|
||||
async def connect(cls, request, *, options, context, callback_adapter=None):
|
||||
connections.append("native")
|
||||
assert options.custom_llm_provider == "azure"
|
||||
return native_socket
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ class SyncBridge:
|
|||
*,
|
||||
options: NativeRequestOptions,
|
||||
context: NativeRequestContext,
|
||||
callback_adapter: object | None = None,
|
||||
) -> dict[str, object]:
|
||||
self.calls.append(
|
||||
{
|
||||
|
|
@ -53,6 +54,7 @@ class AsyncBridge:
|
|||
*,
|
||||
options: NativeRequestOptions,
|
||||
context: NativeRequestContext,
|
||||
callback_adapter: object | None = None,
|
||||
) -> dict[str, object]:
|
||||
return {"text": "async"}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue