diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs index b9dfe8dc4f8..40f6c10e5af 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs @@ -54,19 +54,19 @@ pub async fn run( let request = MessagesRequest { model: provider_model, body, - options: RequestOptions { - api_key: (deployment.litellm_params.api_key.as_deref()).map(|value| value.to_string()), - api_base: (deployment.litellm_params.api_base.as_deref()) - .map(|value| value.to_string()), - custom_llm_provider: (custom_llm_provider).map(|value| value.to_string()), - extra_headers, - timeout: None, - ..Default::default() - }, + }; + let options = RequestOptions { + api_key: (deployment.litellm_params.api_key.as_deref()).map(|value| value.to_string()), + api_base: (deployment.litellm_params.api_base.as_deref()).map(|value| value.to_string()), + custom_llm_provider: (custom_llm_provider).map(|value| value.to_string()), + extra_headers, + timeout: None, + ..Default::default() }; if request.body.get("stream").and_then(Value::as_bool) == Some(true) { return messages_stream( request, + &options, &LiteLlmRequestContext { ..Default::default() }, @@ -77,6 +77,7 @@ pub async fn run( let response = messages( request, + &options, &LiteLlmRequestContext { ..Default::default() }, diff --git a/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs b/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs index 2d9aea95484..d3836a8c8c6 100644 --- a/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs +++ b/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs @@ -86,7 +86,6 @@ impl ProviderAttemptObserver for ProviderObserver { async fn error(&mut self, input: &ProviderError) -> Result<(), Self::Error> { assert!(input.committed); - assert!(!input.will_retry); assert!(!input.message.is_empty()); self.events.lock().unwrap().push("error"); if self.rejected_callback == Some("error") { diff --git a/litellm-rust/crates/core/src/provider_callbacks/handler.rs b/litellm-rust/crates/core/src/provider_callbacks/handler.rs index facf31dae09..08706a244d2 100644 --- a/litellm-rust/crates/core/src/provider_callbacks/handler.rs +++ b/litellm-rust/crates/core/src/provider_callbacks/handler.rs @@ -143,7 +143,6 @@ where Error::Http { status, .. } => Some(*status), _ => None, }, - will_retry: false, ended_at: epoch_seconds(), }; observer.error(&event).await.map_err(callback_error) @@ -240,7 +239,6 @@ mod tests { assert_eq!(event.attempt, 3); assert_eq!(event.trace_id.as_deref(), Some("trace-1")); assert!(event.committed); - assert!(!event.will_retry); self.events.push("error"); Ok(()) } diff --git a/litellm-rust/crates/core/src/provider_callbacks/mod.rs b/litellm-rust/crates/core/src/provider_callbacks/mod.rs index e2e902dc654..f13e70da2ec 100644 --- a/litellm-rust/crates/core/src/provider_callbacks/mod.rs +++ b/litellm-rust/crates/core/src/provider_callbacks/mod.rs @@ -58,7 +58,6 @@ pub struct ProviderError { pub stage: &'static str, pub committed: bool, pub status_code: Option, - pub will_retry: bool, pub ended_at: f64, } diff --git a/litellm-rust/crates/python-bridge/tests/callbacks/mod.rs b/litellm-rust/crates/python-bridge/tests/callbacks/mod.rs index 57ed5314760..ebed1db4b10 100644 --- a/litellm-rust/crates/python-bridge/tests/callbacks/mod.rs +++ b/litellm-rust/crates/python-bridge/tests/callbacks/mod.rs @@ -370,7 +370,6 @@ fn provider_error() -> ProviderError { stage: "provider_response", committed: true, status_code: Some(429), - will_retry: true, ended_at: 2.0, } } diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 65f6328aafe..695b2420f52 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -30,6 +30,7 @@ from litellm.llms.base_llm.ocr.transformation import ( ) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.rust_bridge import ocr as rust_ocr_bridge +from litellm.rust_bridge.callback_adapters import ProviderLoggingAdapter from litellm.rust_bridge.request import ( NativeRequestCapabilities, NativeRequestOptions, @@ -240,27 +241,8 @@ def _prepare_rust_ocr_call( api_base=prepared_request.api_base, litellm_params=prepared_request.litellm_params, ) - resolved_complete_url: Final = provider_config.get_complete_url( - api_base=prepared_request.api_base, - model=prepared_request.model, - optional_params=prepared_request.optional_params, - litellm_params=prepared_request.litellm_params, - ) rust_api_base: Final = _rust_bridge_api_base(prepared_request, resolve_api_key) rust_optional_params: Final = _rust_bridge_optional_params(prepared_request, resolve_api_key) - prepared_request.litellm_logging_obj.pre_call( - input="OCR document processing", - api_key=resolved_api_key, - additional_args={ - "complete_input_dict": { - "model": prepared_request.model, - "document": prepared_request.document, - **rust_optional_params, - }, - "api_base": resolved_complete_url, - "headers": resolved_headers, - }, - ) return PreparedNativeCall( request=rust_ocr_bridge.NativeOCRRequest( model=prepared_request.model, @@ -292,6 +274,11 @@ def _prepare_rust_ocr_call( native_response_format=(prepared_request.optional_params.get(OCR_REQUEST_FORMAT_PARAM) == "native"), ), ), + callback_adapter=ProviderLoggingAdapter( + prepared_request.litellm_logging_obj, + "OCR document processing", + resolved_api_key, + ), ) diff --git a/litellm/rust_bridge/callback_adapters.py b/litellm/rust_bridge/callback_adapters.py index cfb35c87cea..1e08404e14c 100644 --- a/litellm/rust_bridge/callback_adapters.py +++ b/litellm/rust_bridge/callback_adapters.py @@ -53,7 +53,6 @@ class ProviderError(ProviderEvent): stage: str committed: bool status_code: int | None - will_retry: bool ended_at: float diff --git a/litellm/rust_bridge/callbacks.py b/litellm/rust_bridge/callbacks.py index c21b75c3a1d..b17edabf8b4 100644 --- a/litellm/rust_bridge/callbacks.py +++ b/litellm/rust_bridge/callbacks.py @@ -22,12 +22,11 @@ CallbackDecision: TypeAlias = CallbackUnchanged | CallbackReplace | CallbackReje class ProviderAttemptCallbackHandle(Protocol): - """Observe one provider attempt. + """Observe the provider operation inside one native call. - Retries repeat ``pre_call`` and ``error`` with incremented attempt metadata. - Only the successful attempt receives ``post_call``. These callbacks observe - the provider operation; outer SDK success and failure callbacks remain owned - by Python after endpoint dispatch completes. + Successful operations receive ``pre_call`` and ``post_call`` once. Failed + operations receive ``pre_call`` and ``error`` once. Outer SDK success and + failure callbacks remain owned by Python after endpoint dispatch completes. """ def pre_call(self, payload: object, /) -> CallbackDecision: ... diff --git a/litellm/rust_bridge/request.py b/litellm/rust_bridge/request.py index daf991037b5..27c72023ce9 100644 --- a/litellm/rust_bridge/request.py +++ b/litellm/rust_bridge/request.py @@ -74,6 +74,7 @@ def vertex_options(params: Mapping[str, object]) -> NativeVertexOptions: location=location if isinstance(location, str) else None, ) + from typing_extensions import ReadOnly, TypedDict, TypeVar @@ -192,6 +193,12 @@ def call_native( native: NativeFunction[RequestT, ResultT, CallbackT], prepared: PreparedNativeCall[RequestT, CallbackT], ) -> ResultT: + if prepared.callback_adapter is None: + return native( + prepared.request, + options=prepared.options, + context=prepared.context, + ) return native( prepared.request, options=prepared.options, diff --git a/tests/test_litellm/ocr/live_callback_smoke.py b/tests/test_litellm/ocr/live_callback_smoke.py index e7d9270af63..ffa8ce1b114 100644 --- a/tests/test_litellm/ocr/live_callback_smoke.py +++ b/tests/test_litellm/ocr/live_callback_smoke.py @@ -62,5 +62,5 @@ if __name__ == "__main__": else: native: Final = get_native_bridge() assert native is not None - with patch.object(native, "ready_endpoints", {"ocr": frozenset({"callbacks"})}): + with patch.object(native, "ready_endpoints", {"ocr": frozenset({"callbacks"})}, create=True): asyncio.run(smoke(baseline=False)) diff --git a/tests/test_litellm/ocr/sdk_callback_contract.py b/tests/test_litellm/ocr/sdk_callback_contract.py index 9fed80557de..e5e02de17aa 100644 --- a/tests/test_litellm/ocr/sdk_callback_contract.py +++ b/tests/test_litellm/ocr/sdk_callback_contract.py @@ -79,7 +79,7 @@ async def exercise(asynchronous: bool, rust: bool, case: str, *, native_expected following: Final = tuple(event for event in await follower.wait() if native or event.name != "post") names: Final = tuple(event.name for event in events) prefix: Final = ("pre", "post") if native and status == 200 and case != "timeout" else ("pre",) - assert names[: len(prefix)] == prefix + assert names[: len(prefix)] == prefix, (asynchronous, rust, case, native, names, prefix) terminal: Final = "success" if successful else "failure" expected: Final = ( ("async_success",) @@ -267,7 +267,7 @@ async def verify_foundation() -> None: assert_native_unavailable() native: Final = get_native_bridge() assert native is not None - with patch.object(native, "ready_endpoints", {"ocr": frozenset({"callbacks"})}): + with patch.object(native, "ready_endpoints", {"ocr": frozenset({"callbacks"})}, create=True): await verify_parity() assert_native_unavailable() diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 29c24f21c16..3dd5f672ba9 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -11,6 +11,7 @@ import pytest import litellm from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.rust_bridge.callback_adapters import ProviderLoggingAdapter from litellm.rust_bridge import configuration from litellm.rust_bridge.request import ( NativeOCRRequest, @@ -52,6 +53,7 @@ class RecordingBridge: def __init__(self) -> None: self.calls: list[dict[str, object]] = [] + self.callback_adapter: object | None = None def __call__( self, @@ -61,6 +63,7 @@ class RecordingBridge: context: NativeRequestContext, callback_adapter: object | None = None, ) -> dict[str, object]: + self.callback_adapter = callback_adapter self.calls.append( { "model": request.model, @@ -82,6 +85,7 @@ class RecordingAsyncBridge: def __init__(self) -> None: self.calls: list[dict[str, object]] = [] + self.callback_adapter: object | None = None async def __call__( self, @@ -91,6 +95,7 @@ class RecordingAsyncBridge: context: NativeRequestContext, callback_adapter: object | None = None, ) -> dict[str, object]: + self.callback_adapter = callback_adapter self.calls.append( { "model": request.model, @@ -691,7 +696,7 @@ def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): assert bridge.calls[0]["api_base"] == "https://document-intelligence.example.com" -def test_run_rust_ocr_runs_pre_call_logging(): +def test_run_rust_ocr_passes_provider_logging_adapter(): logging_obj = RecordingLogging() bridge = RecordingBridge() litellm.rust(True) @@ -709,17 +714,11 @@ def test_run_rust_ocr_runs_pre_call_logging(): resolve_api_key=lambda _name: None, ) - assert logging_obj.pre_call_kwargs is not None - assert logging_obj.pre_call_kwargs["input"] == "OCR document processing" - additional_args = logging_obj.pre_call_kwargs["additional_args"] - complete_input = additional_args["complete_input_dict"] - assert complete_input["document"] == DOCUMENT - assert complete_input["include_image_base64"] is True - assert additional_args["api_base"] == "https://api.mistral.ai/v1/ocr" - assert additional_args["headers"] == { - "Authorization": "Bearer sk-test", - "x-trace-id": "trace-1", - } + adapter = bridge.callback_adapter + assert isinstance(adapter, ProviderLoggingAdapter) + assert adapter.logging_obj is logging_obj + assert adapter.input == "OCR document processing" + assert adapter.api_key == "sk-test" def test_ocr_routes_azure_ai_to_rust_when_enabled(fake_bridge): diff --git a/tests/test_litellm/rust_bridge/test_callback_adapters.py b/tests/test_litellm/rust_bridge/test_callback_adapters.py index d97098eff0d..f5cdb6cf128 100644 --- a/tests/test_litellm/rust_bridge/test_callback_adapters.py +++ b/tests/test_litellm/rust_bridge/test_callback_adapters.py @@ -45,7 +45,6 @@ def test_provider_logging_adapter_preserves_provider_lifecycle() -> None: stage="provider_response", committed=True, status_code=429, - will_retry=True, ended_at=11.0, ) ) @@ -69,7 +68,6 @@ def test_provider_logging_adapter_preserves_provider_lifecycle() -> None: "stage": "provider_response", "committed": True, "status_code": 429, - "will_retry": True, "ended_at": 11.0, }