From 678e59c6b4415fef77b23373e647f843733a08ff Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 15 Sep 2026 20:47:52 -0700 Subject: [PATCH] refactor(python-bridge): drop OCR callback fast path and body re-aliasing The bridge kept a Rust-side shadow of Logging's callback registries and elided pre_call/post_call/success_handler when it thought nothing was listening. That forked the logging contract and hid a bug: with callbacks present, the during-call hook re-inserted the caller's original document and unmapped optional params over the provider-transformed body. Call the real Logging handlers unconditionally, like the Python wrapper does, and hand callbacks a fresh dict built from core's composed body instead of re-aliasing caller objects into it. --- .../python-bridge/src/lifecycle/bindings.rs | 57 ---------- .../crates/python-bridge/src/lifecycle/mod.rs | 27 +---- .../python-bridge/src/routes/ocr/callbacks.rs | 28 +---- .../python-bridge/src/routes/ocr/lifecycle.rs | 36 ++---- litellm/litellm_core_utils/litellm_logging.py | 1 - litellm/rust_bridge/lifecycle.py | 85 -------------- tests/test_litellm_rust/ocr/test_callbacks.py | 63 ----------- tests/test_litellm_rust/ocr/test_lifecycle.py | 104 +----------------- tests/test_litellm_rust/ocr/test_requests.py | 4 +- 9 files changed, 24 insertions(+), 381 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs b/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs index 06b32b67fd5..23ac0283646 100644 --- a/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs +++ b/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs @@ -19,34 +19,6 @@ impl PythonLogger { visit.call(&self.0) } - pub(crate) fn callbacks_needed(&self, py: Python<'_>, phase: &str) -> PyResult { - if !self - .object(py) - .getattr("_native_callback_fast_path") - .is_ok_and(|value| value.is_truthy().unwrap_or(false)) - { - return Ok(true); - } - py.import("litellm.rust_bridge.lifecycle")? - .getattr("callbacks_needed")? - .call1((self.object(py), phase))? - .extract() - } - - pub(super) fn success_bookkeeping( - &self, - py: Python<'_>, - response: &Option>, - start: &Py, - end: &Option>, - asynchronous: bool, - ) -> PyResult<()> { - py.import("litellm.rust_bridge.lifecycle")? - .getattr("success_bookkeeping")? - .call1((self.object(py), response, start, end, asynchronous))?; - Ok(()) - } - pub(super) fn defers_async_logging(&self, py: Python<'_>) -> bool { self.object(py) .getattr("_defer_async_logging") @@ -68,9 +40,6 @@ impl PythonLogger { start: &Py, end: &Option>, ) -> PyResult<()> { - if !self.callbacks_needed(py, "sync_success_async")? { - return Ok(()); - } self.object(py).call_method1( "handle_sync_success_callbacks_for_async_calls", (response, start, end), @@ -86,19 +55,6 @@ impl PythonLogger { end: &Option>, asynchronous: bool, ) -> PyResult>> { - if !self.callbacks_needed( - py, - if asynchronous { - "async_failure" - } else { - "sync_failure" - }, - )? { - py.import("litellm.rust_bridge.lifecycle")? - .getattr("failure_bookkeeping")? - .call1((self.object(py), error, start, end, asynchronous))?; - return Ok(None); - } let trace = py .import("traceback")? .getattr("format_exception")? @@ -129,9 +85,6 @@ impl PythonLogger { start: &Py, end: &Option>, ) -> PyResult<()> { - if !self.callbacks_needed(py, "sync_success")? { - return self.success_bookkeeping(py, response, start, end, false); - } let context = py.import("contextvars")?.call_method0("copy_context")?; py.import("litellm.litellm_core_utils.litellm_logging")? .getattr("executor")? @@ -155,9 +108,6 @@ impl PythonLogger { start: &Py, end: &Option>, ) -> PyResult<()> { - if !self.callbacks_needed(py, "async_success")? { - return self.success_bookkeeping(py, response, start, end, true); - } let context = py.import("contextvars")?.call_method0("copy_context")?; let worker = py .import("litellm.litellm_core_utils.logging_worker")? @@ -226,13 +176,6 @@ pub(super) fn is_internal_call(py: Python<'_>) -> PyResult { pub(super) struct DeploymentHooks; impl DeploymentHooks { - pub(super) fn needed(py: Python<'_>) -> PyResult { - py.import("litellm.rust_bridge.lifecycle")? - .getattr("deployment_callbacks_needed")? - .call0()? - .extract() - } - pub(super) fn before_call( py: Python<'_>, kwargs: &Py, diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs index ddc3ad1ce12..7ae22b2a5be 100644 --- a/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs +++ b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs @@ -312,9 +312,6 @@ impl PythonCallState { match phase { HostPhase::Setup => self.setup(py)?, HostPhase::DeploymentPreCall => { - if !DeploymentHooks::needed(py)? { - return Ok(HostStep::Ready(self.kwargs.clone_ref(py).into_any())); - } return Ok(HostStep::Suspend(DeploymentHooks::before_call( py, &self.kwargs, @@ -323,13 +320,6 @@ impl PythonCallState { } HostPhase::Prepare => self.prepare(py)?, HostPhase::DeploymentPostCall => { - if !DeploymentHooks::needed(py)? { - return self - .response - .as_ref() - .map(|value| HostStep::Ready(value.clone_ref(py))) - .ok_or_else(missing_state); - } return Ok(HostStep::Suspend(DeploymentHooks::after_success( py, &self.kwargs, @@ -340,9 +330,7 @@ impl PythonCallState { HostPhase::Finalize => self.finalize(py)?, HostPhase::Success => self.dispatch_success(py)?, HostPhase::DeploymentFailure => { - if let Some(error) = &self.error - && DeploymentHooks::needed(py)? - { + if let Some(error) = &self.error { return Ok(HostStep::Suspend(DeploymentHooks::after_failure( py, &self.kwargs, @@ -455,15 +443,6 @@ impl PythonCallState { end: self.end.as_ref().map(|value| value.clone_ref(py)), }; if !self.asynchronous { - if !logger.callbacks_needed(py, "sync_success")? { - return logger.success_bookkeeping( - py, - &self.response, - &self.start, - &self.end, - false, - ); - } pending().sync(py) } else { if !self.internal @@ -473,9 +452,7 @@ impl PythonCallState { .get_item("fallbacks")? .is_none_or(|value| value.is_none()) { - if !logger.callbacks_needed(py, "async_success")? { - logger.success_bookkeeping(py, &self.response, &self.start, &self.end, true)?; - } else if logger.defers_async_logging(py) { + if logger.defers_async_logging(py) { logger.defer_success( py, Py::new( diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs index e3fc344b20a..2dbb968b764 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs @@ -88,13 +88,7 @@ impl PythonLogger { kwargs.set_item("input", "OCR document processing")?; kwargs.set_item("api_key", api_key)?; kwargs.set_item("additional_args", &additional)?; - if self.callbacks_needed(py, "input")? { - self.object(py).call_method("pre_call", (), Some(&kwargs))?; - } else { - self.object(py) - .call_method("_pre_call", (), Some(&kwargs))?; - self.object(py).call_method0("record_api_call_start_time")?; - } + self.object(py).call_method("pre_call", (), Some(&kwargs))?; Ok(()) } @@ -108,21 +102,11 @@ impl PythonLogger { let additional = PyDict::new(py); additional.set_item("complete_input_dict", body)?; additional.set_item("headers", headers)?; - if self.callbacks_needed(py, "input")? { - let kwargs = PyDict::new(py); - kwargs.set_item("original_response", to_py(py, original_response)?)?; - kwargs.set_item("additional_args", &additional)?; - self.object(py) - .call_method("post_call", (), Some(&kwargs))?; - } else { - let response = py - .import("json")? - .call_method1("dumps", (to_py(py, original_response)?,))?; - self.object(py).call_method1( - "record_post_call", - (response, py.None(), py.None(), additional), - )?; - } + let kwargs = PyDict::new(py); + kwargs.set_item("original_response", to_py(py, original_response)?)?; + kwargs.set_item("additional_args", &additional)?; + self.object(py) + .call_method("post_call", (), Some(&kwargs))?; Ok(()) } } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs index 684bc751262..e74228778eb 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs @@ -87,28 +87,9 @@ impl PythonOcrHost { &projected.fields.secret_fields, &request.url, )?; - if !self.state.logger()?.callbacks_needed(py, "payload")? { - self.state - .logger()? - .object(py) - .call_method0("record_api_call_start_time")?; - return Ok(request); - } - if let Some(body) = request.body.as_object_mut() { - for name in &request.retained_fields { - body.remove(name); - } - } let body = to_py(py, &request.body)? .into_bound(py) .cast_into::()?; - if let Some(retained) = &self.projected()?.callback_inputs { - for name in &request.retained_fields { - if let Some(value) = retained.bind(py).get_item(name)? { - body.set_item(name, value)?; - } - } - } let headers = PyDict::new(py); for (name, value) in &request.headers { headers.set_item(name, value)?; @@ -134,16 +115,13 @@ impl PythonOcrHost { py: Python<'_>, request: OcrPostCallRequest, ) -> PyResult { - let logger = self.state.logger()?; - if logger.callbacks_needed(py, "payload")? { - let projected = self.projected()?; - logger.post_ocr( - py, - &request.original_response, - projected.body.as_ref(), - projected.headers.as_ref(), - )?; - } + let projected = self.projected()?; + self.state.logger()?.post_ocr( + py, + &request.original_response, + projected.body.as_ref(), + projected.headers.as_ref(), + )?; Ok(request) } } diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 40621a2f68d..d3410e5d1e5 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -573,7 +573,6 @@ class Logging(LiteLLMLoggingBaseClass): self.streaming_chunks: list[Any] = [] # for generating complete stream response self.sync_streaming_chunks: list[Any] = [] # for generating complete stream response self.log_raw_request_response = log_raw_request_response - self._native_callback_fast_path: bool = False # Initialize dynamic callbacks self.dynamic_input_callbacks: list[str | Callable | CustomLogger] | None = dynamic_input_callbacks diff --git a/litellm/rust_bridge/lifecycle.py b/litellm/rust_bridge/lifecycle.py index f1cc912129d..35429ea7ccc 100644 --- a/litellm/rust_bridge/lifecycle.py +++ b/litellm/rust_bridge/lifecycle.py @@ -1,7 +1,6 @@ from __future__ import annotations import datetime -import os import uuid from collections.abc import Awaitable, Mapping from dataclasses import dataclass @@ -87,13 +86,10 @@ def setup( } supplied: Final = arguments.get("litellm_logging_obj") if isinstance(supplied, Logging): - supplied._native_callback_fast_path = False # pyright: ignore[reportPrivateUsage] # supplied loggers retain all dispatch contracts return CallSetup(supplied, arguments) logger, prepared = utils.function_setup( call_type, utils.Rules(), start_time, *args, is_async_call=asynchronous, **arguments ) - if type(logger) is Logging and call_type in ("ocr", "aocr"): - logger._native_callback_fast_path = True # pyright: ignore[reportPrivateUsage] # only bridge-created OCR loggers opt into callback elision return CallSetup(logger, prepared) @@ -122,84 +118,3 @@ def finalize( MetadataUpdater, response_metadata.update_response_metadata ) update(response, logger, model if isinstance(model, str) else None, kwargs, start_time, end_time) - - -def deployment_callbacks_needed() -> bool: - import litellm - from litellm.integrations.custom_logger import CustomLogger - - return any(isinstance(callback, CustomLogger) for callback in litellm.callbacks) - - -def callbacks_needed(logger: Logging, phase: str) -> bool: - import litellm - from litellm._logging import ( - _is_debugging_on, # pyright: ignore[reportPrivateUsage] # use the same debug gate as Logging - ) - - if ( - _is_debugging_on() - or getattr(logger, "litellm_request_debug", False) - or os.getenv("LITELLM_PRINT_STANDARD_LOGGING_PAYLOAD") - ): - return True - input_needed: Final = bool( - litellm.input_callback - or litellm._async_input_callback # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - or logger.dynamic_input_callbacks - or callable(getattr(logger, "logger_fn", None)) - or logger.log_raw_request_response - or litellm.log_raw_request_response - ) - match phase: - case "input": - return input_needed - case "sync_success": - return bool(litellm.success_callback or logger.dynamic_success_callbacks) - case "sync_success_async": - return bool( - (litellm.success_callback or logger.dynamic_success_callbacks) - and logger._should_run_sync_callbacks_for_async_calls() # pyright: ignore[reportPrivateUsage] # preserve async call filtering of sync callbacks - ) - case "async_success": - return bool(litellm._async_success_callback or logger.dynamic_async_success_callbacks) # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - case "sync_failure": - return bool(litellm.failure_callback or logger.dynamic_failure_callbacks) - case "async_failure": - return bool(litellm._async_failure_callback or logger.dynamic_async_failure_callbacks) # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - case "payload": - return bool( - input_needed - or litellm.success_callback - or litellm.failure_callback - or litellm._async_success_callback # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - or litellm._async_failure_callback # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - or logger.dynamic_success_callbacks - or logger.dynamic_async_success_callbacks - or logger.dynamic_failure_callbacks - or logger.dynamic_async_failure_callbacks - ) - case _: - return True - - -def success_bookkeeping( - logger: Logging, response: object, start: datetime.datetime, end: datetime.datetime, asynchronous: bool -) -> None: - phase: Final = "async_success" if asynchronous else "sync_success" - if logger.should_run_logging(phase): - logger._success_handler_helper_fn( # pyright: ignore[reportPrivateUsage] # retain success bookkeeping without constructing a callback payload - result=response, start_time=start, end_time=end, build_logging_payload=False - ) - logger.has_run_logging(phase) - - -def failure_bookkeeping( - logger: Logging, error: BaseException, start: datetime.datetime, end: datetime.datetime, asynchronous: bool -) -> None: - phase: Final = "async_failure" if asynchronous else "sync_failure" - if logger.should_run_logging(phase): - logger._failure_handler_helper_fn( # pyright: ignore[reportPrivateUsage] # retain failure accounting without formatting an unused traceback or payload - error, "", start, end, build_logging_payload=False - ) - logger.has_run_logging(phase) diff --git a/tests/test_litellm_rust/ocr/test_callbacks.py b/tests/test_litellm_rust/ocr/test_callbacks.py index 1cfd04b1bff..66e2bea0186 100644 --- a/tests/test_litellm_rust/ocr/test_callbacks.py +++ b/tests/test_litellm_rust/ocr/test_callbacks.py @@ -119,69 +119,6 @@ def test_native_ocr_pre_call_header_rebinding_does_not_replace_execution_root(oc assert "x-rebound" not in ocr_server.requests[0].headers -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -async def test_native_ocr_pre_call_nested_document_edit_updates_caller_callback_and_provider_references( - ocr_server: RecordingServer, asynchronous: bool -) -> None: - original: Final = dict(OCR_DOCUMENT) - replacement_url: Final = "data:application/pdf;base64,ZGVm" - retained: Final = [] - aliases: Final = [] - - class Retain(CustomLogger): - def log_pre_api_call(self, model, messages, kwargs): - aliases.append(request_body(kwargs)["document"] is original) - retained.append(request_body(kwargs)["document"]) - - class Edit(CustomLogger): - def log_pre_api_call(self, model, messages, kwargs): - original["document_url"] = replacement_url - - arguments: Final = { - "model": "mistral/mistral-ocr-latest", - "document": original, - "api_key": "test-key", - "api_base": ocr_server.base_url, - "callbacks": [Retain(), Edit()], - } - response: Final = ( - await call_native_aocr(ocr_server, **arguments) - if asynchronous - else call_native_ocr(ocr_server, **arguments) - ) - - assert aliases == [True] - assert retained[0]["document_url"] == replacement_url - assert original["document_url"] == replacement_url - assert ocr_server.requests[0].body["document"]["document_url"] == replacement_url - assert response.pages[0].markdown == "native OCR response" - - -def test_native_ocr_pre_call_document_replacement_does_not_mutate_original_document( - ocr_server: RecordingServer, -) -> None: - original: Final = dict(OCR_DOCUMENT) - replacement: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,ZGVm"} - retained: Final = [] - - class RetainAndReplace(CustomLogger): - def log_pre_api_call(self, model, messages, kwargs): - body = request_body(kwargs) - retained.append(body["document"]) - body["document"] = replacement - - call_native_ocr( - ocr_server, - document=original, - callbacks=[RetainAndReplace()], - ) - - assert retained[0] is original - assert original["document_url"] == OCR_DOCUMENT["document_url"] - assert ocr_server.requests[0].body["document"] == replacement - - def test_native_ocr_pre_call_body_rebinding_is_visible_to_callbacks_but_not_provider( ocr_server: RecordingServer, ) -> None: diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index b0df9168719..2ca9e77db4f 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -79,7 +79,7 @@ async def test_native_bindings_run_callbacks_and_send_their_mutations( class Observe(Logging): def pre_call(self, input, api_key, additional_args): body: Final = additional_args["complete_input_dict"] - assert body["pages"] is pages + assert body["pages"] == pages and body["pages"] is not pages observed.append(self.model_call_details["litellm_params"]["metadata"]["marker"]) body["pages"].append(2) additional_args["headers"]["x-callback"] = "native" @@ -98,6 +98,7 @@ async def test_native_bindings_run_callbacks_and_send_their_mutations( assert response.pages[0].markdown == "native OCR response" assert observed == [marker] and observed[0] is marker + assert pages == [0] assert len(ocr_server.requests) == 1 assert ocr_server.requests[0].body["pages"] == [0, 2] assert ocr_server.requests[0].headers["x-callback"] == "native" @@ -263,7 +264,6 @@ async def test_deployment_hook_replaces_complete_routing_request(ocr_server: Rec assert response.pages[0].markdown == "native OCR response" assert observed == [(replacement, "replacement-key")] - assert observed[0][0] is replacement assert replacement == original assert replacement is not original assert original == {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"} @@ -621,7 +621,7 @@ def test_sync_pre_call_can_make_nested_native_request(ocr_server: RecordingServe @pytest.mark.asyncio -async def test_retained_argument_aliases_and_body_roots_survive_envelope_replacement( +async def test_callback_body_roots_survive_envelope_replacement( ocr_server: RecordingServer, ) -> None: pages: Final = [0] @@ -633,8 +633,8 @@ async def test_retained_argument_aliases_and_body_roots_survive_envelope_replace def pre_call(self, input, api_key, additional_args): body: Final = additional_args["complete_input_dict"] headers: Final = additional_args["headers"] - observed.append((body["document"] is document, body["pages"] is pages)) - pages.append(2) + observed.append((body["document"] == document, body["pages"] == pages)) + body["pages"].append(2) headers["x-retained"] = "yes" additional_args["complete_input_dict"] = {"discarded": True} additional_args["headers"] = {} @@ -672,6 +672,7 @@ async def test_retained_argument_aliases_and_body_roots_survive_envelope_replace assert observed[0] == (False, False, True) assert observed[1] == (True, True) assert observed[3] == (True, True) + assert pages == [0] assert ocr_server.requests[0].body["pages"] == [0, 2] assert ocr_server.requests[0].headers["x-retained"] == "yes" @@ -933,72 +934,6 @@ async def test_response_limit_is_enforced_at_the_public_boundary( assert "max_response_bytes" not in body -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", [False, True]) -@pytest.mark.parametrize("failure", [False, True]) -async def test_empty_callbacks_keep_bookkeeping_without_optional_dispatch( - ocr_server: RecordingServer, - monkeypatch: pytest.MonkeyPatch, - asynchronous: bool, - failure: bool, - created_loggers: list[Logging], -) -> None: - from litellm import utils - from litellm.litellm_core_utils import litellm_logging, logging_worker - - class DispatchProbe: - deployments = 0 - submissions = 0 - enqueues = 0 - - def deployment(self, *args: object, **kwargs: object) -> None: - self.deployments += 1 - - def submit(self, *args: object, **kwargs: object) -> None: - self.submissions += 1 - - def ensure_initialized_and_enqueue(self, coroutine: Coroutine[object, object, object]) -> None: - self.enqueues += 1 - coroutine.close() - - probe: Final = DispatchProbe() - for name in ( - "async_pre_call_deployment_hook", - "async_post_call_success_deployment_hook", - "async_post_call_failure_deployment_hook", - ): - monkeypatch.setattr(utils, name, probe.deployment) - monkeypatch.setattr(litellm_logging, "executor", probe) - monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", probe) - if failure: - ocr_server.enqueue(ResponseSpec(body={"message": "provider failed"}, status=500)) - trace_id_var.set("callback-free-parent") - arguments: Final = {"litellm_trace_id": "callback-free-call", "litellm_call_id": "callback-free-id"} - if failure: - with pytest.raises(litellm.InternalServerError): - await call_aocr(ocr_server, **arguments) if asynchronous else call_ocr(ocr_server, **arguments) - else: - response: Final = ( - await call_aocr(ocr_server, **arguments) if asynchronous else call_ocr(ocr_server, **arguments) - ) - assert response.pages[0].markdown == "native OCR response" - assert response._hidden_params["litellm_call_id"] == "callback-free-id" - assert response._hidden_params["response_cost"] is not None - assert response._hidden_params["_response_ms"] > 0 - assert trace_id_var.get() == "callback-free-parent" - assert probe.deployments == probe.submissions == probe.enqueues == 0 - assert len(created_loggers) == 1 - logger: Final = created_loggers[0] - assert not hasattr(logger, "_native_pending_logging") - assert logger.model_call_details["first_api_call_start_time"] <= logger.model_call_details["end_time"] - assert "standard_logging_object" not in logger.model_call_details - assert ( - "original_response" not in logger.model_call_details or logger.model_call_details["original_response"] is None - ) - assert "complete_input_dict" not in logger.model_call_details.get("additional_args", {}) - assert logger.model_call_details["response_cost"] == (0 if failure else response._hidden_params["response_cost"]) - - @pytest.mark.asyncio @pytest.mark.parametrize( "registration", ["success_callback", "_async_success_callback", "failure_callback", "_async_failure_callback"] @@ -1076,30 +1011,3 @@ async def test_explicit_logging_consumers_keep_request_and_response_payloads( assert details["raw_request_typed_dict"]["raw_request_body"]["model"] == "mistral-ocr-latest" if consumer == "logger_fn": assert [item["log_event_type"] for item in snapshots] == ["pre_api_call", "post_api_call"] - - -@pytest.mark.asyncio -async def test_registration_removed_before_deferred_release_skips_queue( - ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, created_loggers: list[Logging] -) -> None: - from litellm.litellm_core_utils import logging_worker - - class QueueProbe: - enqueues = 0 - - def ensure_initialized_and_enqueue(self, coroutine: Coroutine[object, object, object]) -> None: - self.enqueues += 1 - coroutine.close() - - observer: Final = RecordingLogger() - litellm._async_success_callback.append(observer) - await call_aocr(ocr_server) - logger: Final = created_loggers[0] - assert hasattr(logger, "_native_pending_logging") - litellm._async_success_callback.clear() - probe: Final = QueueProbe() - monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", probe) - ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, False) - assert probe.enqueues == 0 - assert not observer.names - assert logger.model_call_details["response_cost"] is not None diff --git a/tests/test_litellm_rust/ocr/test_requests.py b/tests/test_litellm_rust/ocr/test_requests.py index 5bcc65d75af..47eb4f0f577 100644 --- a/tests/test_litellm_rust/ocr/test_requests.py +++ b/tests/test_litellm_rust/ocr/test_requests.py @@ -12,6 +12,7 @@ from tests.test_litellm_rust.support.requests import ( OCR_RESPONSE, call_native_aocr, call_native_ocr, + request_body, ) pytestmark = pytest.mark.requires_rust_extension @@ -535,7 +536,7 @@ async def test_native_ocr_inherits_named_credentials_without_overwriting_argumen class Observer(RecordingLogger): def log_pre_api_call(self, model, messages, kwargs): super().log_pre_api_call(model, messages, kwargs) - pages.append(2) + request_body(kwargs)["pages"].append(2) arguments: Final = { "model": "mistral/mistral-ocr-latest", @@ -548,3 +549,4 @@ async def test_native_ocr_inherits_named_credentials_without_overwriting_argumen assert response.pages[0].markdown == "native OCR response" assert ocr_server.requests[0].headers["authorization"] == f"Bearer {expected_key}" assert ocr_server.requests[0].body["pages"] == [0, 2] + assert pages == [0]