diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs index ed16b424604..07f2a5bed29 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs @@ -24,7 +24,14 @@ pub async fn ocr( context: &LiteLlmRequestContext, hooks: RequestHooks, ) -> Result { - ocr_with_observer(request, context, hooks, &mut NoopProviderAttemptObserver).await + ocr_with_observer( + request, + options, + context, + hooks, + &mut NoopProviderAttemptObserver, + ) + .await } #[tracing::instrument( @@ -35,6 +42,7 @@ pub async fn ocr( )] pub async fn ocr_with_observer( request: OcrRequest<'_>, + options: &RequestOptions, context: &LiteLlmRequestContext, hooks: RequestHooks, observer: &mut Observer, diff --git a/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs b/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs index eb72e765b91..2d9aea95484 100644 --- a/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs +++ b/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs @@ -97,18 +97,21 @@ impl ProviderAttemptObserver for ProviderObserver { } } -fn observer_request(api_base: &str) -> OcrRequest<'_> { +fn observer_request() -> OcrRequest<'static> { OcrRequest { model: "mistral/mistral-ocr-4-1", document: json!({"type":"document_url","document_url":"https://example.com/document.pdf"}), optional_params: Map::new(), - options: RequestOptions { - api_key: Some("test-key".into()), - api_base: Some(api_base.into()), - custom_llm_provider: Some("mistral".into()), - timeout: Some(Duration::from_secs(2)), - ..Default::default() - }, + } +} + +fn observer_options(api_base: &str) -> RequestOptions { + RequestOptions { + api_key: Some("test-key".into()), + api_base: Some(api_base.into()), + custom_llm_provider: Some("mistral".into()), + timeout: Some(Duration::from_secs(2)), + ..Default::default() } } @@ -146,7 +149,8 @@ async fn observer_case(status: u16, body: &'static str, decision: Option<&'stati decision, }; let result = ocr_with_observer( - observer_request(&url), + observer_request(), + &observer_options(&url), &observer_context(), RequestHooks::default(), &mut observer, @@ -210,7 +214,8 @@ async fn provider_callback_rejection_stops_before_http() { }; let result = ocr_with_observer( - observer_request(&url), + observer_request(), + &observer_options(&url), &observer_context(), RequestHooks::default(), &mut observer, @@ -241,11 +246,12 @@ async fn invalid_ocr_preparation_does_not_call_observers_or_provider() { }; let request = OcrRequest { document: json!(42), - ..observer_request(&url) + ..observer_request() }; assert!( ocr_with_observer( request, + &observer_options(&url), &observer_context(), RequestHooks::default(), &mut observer diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 6d31a7b937a..5b4f1b68e37 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -81,7 +81,9 @@ impl ResponsesWebSocketConnection { } } } - let inner = match RustResponsesWebSocketConnection::connect(request, &options, &context).await { + let inner = match RustResponsesWebSocketConnection::connect(request, &options, &context) + .await + { Ok(inner) => inner, Err(error) => { if let Some(observer) = observer.as_mut() { diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs index 5c4bf1ba2f7..c9ff25d8965 100644 --- a/litellm-rust/crates/python-bridge/src/routes/definition.rs +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -159,6 +159,8 @@ mod tests { inputs: EchoInputs, _options: crate::marshal::NativeRequestOptions, _context: crate::marshal::NativeRequestContext, + _callback_adapter: Option>, + _python_context: crate::execution::PythonCallContext<'_>, ) -> PyResult> + Send + 'static> { FUTURE_DROPPED.store(false, Ordering::SeqCst); let drop_guard = (inputs.value == "pending").then_some(DropGuard); @@ -199,13 +201,21 @@ mod tests { let module = PyModule::new(py, "routes").expect("module should be created"); crate::routes::register(&module).expect("routes should register"); let routes = [ - ("ocr", "aocr", "(request, *, options, context, callback_adapter=None)"), + ( + "ocr", + "aocr", + "(request, *, options, context, callback_adapter=None)", + ), ( "transcription", "atranscription", "(request, *, options, context, callback_adapter=None)", ), - ("messages", "amessages", "(request, *, options, context, callback_adapter=None)"), + ( + "messages", + "amessages", + "(request, *, options, context, callback_adapter=None)", + ), ( "chat_completions", "achat_completions",