From e2781c47132a129b67667ed824d668899b4c3661 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 20:19:00 +0000 Subject: [PATCH] refactor(rust): move tests.rs files inline or under tests/ and drop autotests = false (#43028) * refactor(rust): move tests.rs files inline or under tests/ Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(rust): move tests.rs files inline or under tests/ Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(rust): inline path-included test files into their owning src files Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(rust): drop stray proptest regression file Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(rust): cover lowercase, empty and non-authorization headers in bearer detection Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Yujong Lee Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/AGENTS.md | 10 + .../crates/callbacks-legacy-python/Cargo.toml | 1 - .../callbacks-legacy-python/src/adapter.rs | 1116 ++++- .../callbacks-legacy-python/src/deferred.rs | 150 +- .../crates/callbacks-legacy-python/src/lib.rs | 213 +- .../callbacks-legacy-python/tests/deferred.rs | 146 - .../tests/deployment_hooks.rs | 282 -- .../callbacks-legacy-python/tests/payload.rs | 523 --- .../callbacks-legacy-python/tests/support.rs | 205 - .../callbacks-legacy-python/tests/terminal.rs | 291 -- litellm-rust/crates/core/Cargo.toml | 1 - .../core/src/audio_transcription/mod.rs | 3 - .../crates/core/src/chat_completions/mod.rs | 3 - .../core/src/chat_completions/prepare.rs | 838 ++++ .../crates/core/src/chat_completions/tests.rs | 833 ---- .../crates/core/src/messages/common_utils.rs | 183 + litellm-rust/crates/core/src/messages/mod.rs | 3 - litellm-rust/crates/core/src/ocr/document.rs | 157 + litellm-rust/crates/core/src/ocr/mod.rs | 238 +- litellm-rust/crates/core/src/ocr/route.rs | 3586 +++++++++++++++++ .../tests.rs => tests/audio_transcription.rs} | 4 +- .../crates/core/tests/aws_textract_ocr.rs | 193 - .../crates/core/tests/azure_ai_ocr.rs | 293 -- .../tests/azure_document_intelligence_ocr.rs | 712 ---- litellm-rust/crates/core/tests/cohere_ocr.rs | 136 - .../crates/core/tests/deepseek_ocr.rs | 133 - .../messages/tests.rs => tests/messages.rs} | 132 +- litellm-rust/crates/core/tests/ocr.rs | 1033 ----- .../crates/core/tests/ocr/document.rs | 152 - litellm-rust/crates/core/tests/ocr/support.rs | 203 - litellm-rust/crates/core/tests/reducto_ocr.rs | 584 --- .../core/tests/vertex_ai_deepseek_ocr.rs | 143 - .../crates/core/tests/vertex_ai_ocr.rs | 293 -- litellm-rust/crates/http/src/request.rs | 12 + .../llms/src/anthropic/chat/transformation.rs | 4 - .../bedrock/chat/converse_transformation.rs | 4 - .../anthropic_chat_transformation.rs} | 18 +- .../bedrock_converse_transformation.rs} | 12 +- litellm-rust/crates/types/src/utils.rs | 2 +- 39 files changed, 6486 insertions(+), 6359 deletions(-) create mode 100644 litellm-rust/AGENTS.md delete mode 100644 litellm-rust/crates/callbacks-legacy-python/tests/deferred.rs delete mode 100644 litellm-rust/crates/callbacks-legacy-python/tests/deployment_hooks.rs delete mode 100644 litellm-rust/crates/callbacks-legacy-python/tests/payload.rs delete mode 100644 litellm-rust/crates/callbacks-legacy-python/tests/support.rs delete mode 100644 litellm-rust/crates/callbacks-legacy-python/tests/terminal.rs delete mode 100644 litellm-rust/crates/core/src/chat_completions/tests.rs rename litellm-rust/crates/core/{src/audio_transcription/tests.rs => tests/audio_transcription.rs} (95%) delete mode 100644 litellm-rust/crates/core/tests/aws_textract_ocr.rs delete mode 100644 litellm-rust/crates/core/tests/azure_ai_ocr.rs delete mode 100644 litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs delete mode 100644 litellm-rust/crates/core/tests/cohere_ocr.rs delete mode 100644 litellm-rust/crates/core/tests/deepseek_ocr.rs rename litellm-rust/crates/core/{src/messages/tests.rs => tests/messages.rs} (79%) delete mode 100644 litellm-rust/crates/core/tests/ocr.rs delete mode 100644 litellm-rust/crates/core/tests/ocr/document.rs delete mode 100644 litellm-rust/crates/core/tests/ocr/support.rs delete mode 100644 litellm-rust/crates/core/tests/reducto_ocr.rs delete mode 100644 litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs delete mode 100644 litellm-rust/crates/core/tests/vertex_ai_ocr.rs rename litellm-rust/crates/llms/{src/anthropic/chat/tests.rs => tests/anthropic_chat_transformation.rs} (96%) rename litellm-rust/crates/llms/{src/bedrock/chat/tests.rs => tests/bedrock_converse_transformation.rs} (98%) diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md new file mode 100644 index 00000000000..e5ffcd1c57a --- /dev/null +++ b/litellm-rust/AGENTS.md @@ -0,0 +1,10 @@ +# Rust workspace rules + +## Test placement + +- Never create a `tests.rs` (or `test.rs`) file under `src/`, and never `#[path = "tests.rs"] mod tests;` +- A test that reaches private items lives inline, in a `#[cfg(test)] mod tests { ... }` at the bottom of the file that owns those items +- A test that only uses the crate's public API lives in `crates//tests/.rs`, next to `src/` +- Split a mixed test file along that line instead of widening visibility to move it +- A test for another crate's item belongs in that crate, not in a downstream one +- Never set `autotests = false` or hand-list `[[test]]` targets; every file directly under `tests/` is discovered by cargo, and a shared helper goes in `tests//mod.rs` or `tests//support.rs` so it is not picked up as a test crate of its own diff --git a/litellm-rust/crates/callbacks-legacy-python/Cargo.toml b/litellm-rust/crates/callbacks-legacy-python/Cargo.toml index fe19578e04d..ed5e0fb9691 100644 --- a/litellm-rust/crates/callbacks-legacy-python/Cargo.toml +++ b/litellm-rust/crates/callbacks-legacy-python/Cargo.toml @@ -4,7 +4,6 @@ version = "0.1.0" edition.workspace = true license.workspace = true repository.workspace = true -autotests = false [dependencies] litellm-host.workspace = true diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index e1742190205..75a635e9c63 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -487,11 +487,1115 @@ impl PythonLifecycle for LegacyLogging { } #[cfg(test)] -#[path = "../tests/deployment_hooks.rs"] -mod deployment_hooks_tests; +mod deployment_hooks_tests { + use std::ffi::CStr; + + use litellm_host::event::{FailureOrigin, Timing}; + use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; + use pyo3::exceptions::asyncio::CancelledError; + use pyo3::prelude::*; + use pyo3::types::PyDict; + use rstest::rstest; + + use super::LegacyLogging; + use crate::test_support::{legacy_call, local, namespace, run}; + + const CALL: &CStr = c" +document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} +kwargs = {'logger': logger, 'document': document} +"; + + const TIMING: Timing = Timing { + start_time: 0.0, + end_time: 1.0, + }; + + fn begin<'py>( + py: Python<'py>, + locals: &Bound<'py, PyDict>, + asynchronous: bool, + ) -> (LegacyLogging, LifecycleStep) { + let mut logging = legacy_call(py, locals, asynchronous); + let kwargs = local(locals, "kwargs") + .cast_into::() + .unwrap() + .unbind(); + let step = logging.begin(py, kwargs, 0.0).unwrap(); + (logging, step) + } + + fn arguments<'py>(py: Python<'py>, step: LifecycleStep) -> Bound<'py, PyDict> { + let LifecycleStep::Arguments(arguments) = step else { + panic!("expected the prepared arguments"); + }; + arguments.into_bound(py) + } + + fn awaits_deployment_hook(step: &LifecycleStep) -> bool { + matches!(step, LifecycleStep::Await(_)) + } + + #[rstest] + #[case::synchronous(false)] + #[case::asynchronous(true)] + fn deployment_pre_call_hook_runs_only_for_asynchronous_calls(#[case] asynchronous: bool) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, CALL); + let (_, step) = begin(py, &locals, asynchronous); + assert_eq!(awaits_deployment_hook(&step), asynchronous); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names.contains(&"pre_hook".to_string()), asynchronous); + }); + } + + #[test] + fn kwargs_returned_by_the_pre_call_hook_are_what_the_call_prepares() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c" +document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} +replacement = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,ZWRpdGVk'} +kwargs = {'logger': logger, 'document': document} +replaced_kwargs = {'logger': logger, 'document': replacement, 'pages': [0]} +", + ); + let (mut logging, step) = begin(py, &locals, true); + assert!(awaits_deployment_hook(&step)); + let step = logging + .resume(py, Ok(local(&locals, "replaced_kwargs").unbind())) + .unwrap(); + locals.set_item("prepared", arguments(py, step)).unwrap(); + run( + py, + &locals, + c" +assert prepared['document'] is replacement +assert prepared['pages'] is replaced_kwargs['pages'] +assert prepared['litellm_logging_obj'] is logger +assert 'litellm_logging_obj' not in replaced_kwargs +[checked] = [value for name, value in logger.calls if name == 'check_limits'] +assert checked is prepared +", + ); + }); + } + + #[rstest] + #[case::synchronous(false)] + #[case::asynchronous(true)] + fn a_keyword_the_bridge_never_reads_reaches_every_reader_as_the_callers_object( + #[case] asynchronous: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c" +opaque = object() +hooked = [] +logger.hooks = {'pre': lambda kwargs: hooked.append(kwargs['vendor_extension']) or kwargs} +kwargs = {'logger': logger, 'vendor_extension': opaque} +", + ); + let (mut logging, step) = begin(py, &locals, asynchronous); + let step = match step { + LifecycleStep::Await(hook_result) => logging.resume(py, Ok(hook_result)).unwrap(), + step => step, + }; + locals.set_item("prepared", arguments(py, step)).unwrap(); + locals.set_item("asynchronous", asynchronous).unwrap(); + run( + py, + &locals, + c" +assert prepared['vendor_extension'] is opaque +[checked] = [value for name, value in logger.calls if name == 'check_limits'] +assert checked['vendor_extension'] is opaque +assert hooked == ([opaque] if asynchronous else []), hooked +", + ); + }); + } + + #[test] + fn response_returned_by_the_post_call_hook_is_finalized_and_returned() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c" +kwargs = {'logger': logger} +response = object() +replacement = object() +logger.hooks = {'pre': lambda kwargs: kwargs} +", + ); + let (mut logging, _) = begin(py, &locals, true); + logging + .resume(py, Ok(local(&locals, "kwargs").unbind())) + .unwrap(); + let step = logging + .after_success(py, local(&locals, "response").unbind(), TIMING) + .unwrap(); + assert!(awaits_deployment_hook(&step)); + let step = logging + .resume(py, Ok(local(&locals, "replacement").unbind())) + .unwrap(); + let LifecycleStep::Response(returned) = step else { + panic!("expected the finalized response"); + }; + assert!(returned.bind(py).is(local(&locals, "replacement"))); + run( + py, + &locals, + c" +[finalized] = [value for name, value in logger.calls if name == 'finalize'] +assert finalized is replacement +", + ); + }); + } + + #[rstest] + #[case::pre_call(false)] + #[case::post_call(true)] + fn cancelling_a_deployment_hook_ends_the_call_with_that_cancellation(#[case] post_call: bool) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c"kwargs = {'logger': logger}\nresponse = object()"); + let (mut logging, _) = begin(py, &locals, true); + if post_call { + logging + .resume(py, Ok(local(&locals, "kwargs").unbind())) + .unwrap(); + logging + .after_success(py, local(&locals, "response").unbind(), TIMING) + .unwrap(); + } + let cancellation = CancelledError::new_err("cancelled"); + let cancelled = cancellation.value(py).clone(); + let error = logging.resume(py, Err(cancellation)).err().unwrap(); + assert!(error.value(py).is(&cancelled)); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert!(!names.iter().any(|name| name.contains("handler"))); + }); + } + + #[rstest] + #[case::hook_completed(false)] + #[case::hook_cancelled(true)] + fn failure_callbacks_run_after_the_failure_hook_however_it_ends(#[case] cancelled: bool) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c"kwargs = {'logger': logger}\nfailure = ValueError('provider')", + ); + let (mut logging, _) = begin(py, &locals, true); + logging + .resume(py, Ok(local(&locals, "kwargs").unbind())) + .unwrap(); + let failure = PyErr::from_value(local(&locals, "failure")); + let failed = LifecycleEvent::Failed { + timing: TIMING, + origin: FailureOrigin::Call, + error: &failure, + }; + let step = logging.emit(py, failed).unwrap(); + assert!(awaits_deployment_hook(&step)); + let hook_result = if cancelled { + Err(CancelledError::new_err("cancelled")) + } else { + Ok(py.None()) + }; + assert!(matches!( + logging.resume(py, hook_result).unwrap(), + LifecycleStep::Await(_) + )); + run( + py, + &locals, + c" +assert logger.names()[-3:] == ['failure_hook', 'failure_handler', 'async_failure_handler'], logger.calls +assert all(value is failure for name, value in logger.calls if name.endswith('_handler')) +", + ); + }); + } + + #[rstest] + #[case::synchronous(false)] + #[case::asynchronous(true)] + fn a_limit_rejected_before_the_call_surfaces_as_the_callers_error(#[case] asynchronous: bool) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c" +class BudgetExceeded(Exception): + pass + +rejection = BudgetExceeded('over budget') + +class LimitedLogger(StubLogger): + def check_limits(self, arguments): + raise rejection + +logger = LimitedLogger() +logger.hooks = {'pre': lambda kwargs: kwargs} +kwargs = {'logger': logger} +", + ); + let mut logging = legacy_call(py, &locals, asynchronous); + let kwargs = local(&locals, "kwargs") + .cast_into::() + .unwrap() + .unbind(); + let result = logging.begin(py, kwargs, 0.0).and_then(|step| match step { + LifecycleStep::Await(_) => { + logging.resume(py, Ok(local(&locals, "kwargs").unbind())) + } + step => Ok(step), + }); + let error = result.err().unwrap(); + assert!(error.value(py).is(local(&locals, "rejection"))); + }); + } +} + #[cfg(test)] -#[path = "../tests/payload.rs"] -mod payload_tests; +mod payload_tests { + use std::ffi::CStr; + + use litellm_auth::SecretValue; + use litellm_host::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; + use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle, to_py}; + use proptest::prelude::*; + use pyo3::prelude::*; + use rstest::rstest; + use serde_json::{Map, Value, json}; + + use super::LegacyLogging; + use crate::PythonLogger; + use crate::test_support::{legacy_call, local, namespace, run}; + + /// The payload phases of `Logging` on top of `StubLogger`, with `pre_call` handing the + /// payload to the case's `on_pre_call`. + const PAYLOAD_LOGGER: &CStr = c" +class Request: + pass + +class PayloadLogger(StubLogger): + def update_from_kwargs(self, **update): + self.update = update + + def pre_call(self, input, api_key, additional_args): + self.record('pre_call', None) + self.pre = additional_args + self.pre_api_key = api_key + on_pre_call(additional_args) + + def post_call(self, original_response, api_key, additional_args): + self.record('post_call', None) + self.post = (original_response, api_key, additional_args) + +request = Request() +kwargs = {} +logger = PayloadLogger() +on_pre_call = lambda additional_args: None +check = lambda: None +"; + + const DOCUMENT: &str = "data:application/pdf;base64,YWJj"; + const EDITED: &str = "data:application/pdf;base64,ZWRpdGVk"; + + fn document(source: &str) -> Value { + json!({"type": "document_url", "document_url": source}) + } + + fn before_send(script: &CStr, body: Value) -> WireRequest { + before_send_with_secrets(script, json!({}), body, &[]) + } + + /// Runs `before_send` over `body` for a route whose parameters are `optional_params`, with + /// the Python objects `script` binds, then delivers the provider's raw response the way the + /// driver does and runs the script's `check()`. + fn before_send_with_secrets( + script: &CStr, + optional_params: Value, + body: Value, + secret_fields: &[&str], + ) -> WireRequest { + before_send_bound(&[], script, optional_params, body, secret_fields) + } + + /// [`before_send_with_secrets`] with `bindings` placed in the namespace before `script` runs. + fn before_send_bound( + bindings: &[(&str, &Value)], + script: &CStr, + optional_params: Value, + body: Value, + secret_fields: &[&str], + ) -> WireRequest { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, PAYLOAD_LOGGER); + for &(name, value) in bindings { + locals.set_item(name, to_py(py, value).unwrap()).unwrap(); + } + run(py, &locals, script); + let mut logging = LegacyLogging { + logger: Some(PythonLogger::new(local(&locals, "logger").unbind())), + ..legacy_call(py, &locals, false) + }; + let context = RequestContext { + model: "model".into(), + custom_llm_provider: "provider".into(), + optional_params, + secret_fields: secret_fields.iter().map(|name| name.to_string()).collect(), + api_key: Some(SecretValue::new("route-key")), + }; + let wire = WireRequest { + url: "https://provider.invalid/ocr".into(), + headers: vec![("x-route".into(), "route".into())], + body, + }; + let step = logging.before_send(py, Box::new(wire), &context).unwrap(); + let raw = MachineEvent::ResponseReceived { + raw: RawResponse { + body: "raw response".into(), + }, + }; + assert!(matches!( + logging.emit(py, LifecycleEvent::Machine(&raw)).unwrap(), + LifecycleStep::Done + )); + run(py, &locals, c"check()"); + let LifecycleStep::Wire(wire) = step else { + panic!("before_send did not hand back the wire request"); + }; + *wire + }) + } + + #[rstest] + #[case::caller_keyword(c" +document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} +pages = [0] +kwargs = {'document': document, 'pages': pages} +observed = [] +on_pre_call = lambda args: observed.append( + (args['complete_input_dict']['document'] is document, args['complete_input_dict']['pages'] is pages) +) +def check(): + assert observed == [(True, True)], observed +")] + #[case::request_attribute_behind_an_omitted_keyword(c" +document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} +pages = [0] +request.document = document +kwargs = {'pages': pages} +observed = [] +on_pre_call = lambda args: observed.append( + (args['complete_input_dict']['document'] is document, args['complete_input_dict']['pages'] is pages) +) +def check(): + assert observed == [(True, True)], observed +")] + fn passthrough_keys_reach_pre_call_as_the_callers_own_objects(#[case] script: &CStr) { + let body = json!({"model": "model", "document": document(DOCUMENT), "pages": [0]}); + let wire = before_send(script, body.clone()); + assert_eq!(wire.body, body); + } + + #[test] + fn pre_call_edit_of_a_passthrough_object_reaches_the_caller_and_the_wire() { + let wire = before_send( + c" +document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} +kwargs = {'document': document} +def on_pre_call(args): + args['complete_input_dict']['document']['document_url'] = 'data:application/pdf;base64,ZWRpdGVk' +def check(): + assert document['document_url'] == 'data:application/pdf;base64,ZWRpdGVk' +", + json!({"document": document(DOCUMENT)}), + ); + assert_eq!(wire.body["document"], document(EDITED)); + } + + #[test] + fn a_body_key_the_route_rewrote_is_not_the_callers_object() { + let wire = before_send( + c" +document = {'type': 'document_url', 'document_url': 'https://example.invalid/scan.pdf'} +kwargs = {'document': document} +observed = [] +def on_pre_call(args): + observed.append(args['complete_input_dict']['document'] is document) + args['complete_input_dict']['document']['document_name'] = 'edited.pdf' +def check(): + assert observed == [False], observed + assert document == {'type': 'document_url', 'document_url': 'https://example.invalid/scan.pdf'} +", + json!({"document": document(DOCUMENT)}), + ); + assert_eq!( + wire.body["document"], + json!({"type": "document_url", "document_url": DOCUMENT, "document_name": "edited.pdf"}) + ); + } + + #[test] + fn a_caller_value_with_no_json_form_is_left_out_of_realiasing() { + let body = json!({"pages": [0]}); + let wire = before_send( + c" +opaque = object() +kwargs = {'pages': opaque} +observed = [] +on_pre_call = lambda args: observed.append(args['complete_input_dict']['pages']) +def check(): + assert observed == [[0]], observed +", + body.clone(), + ); + assert_eq!(wire.body, body); + } + + #[rstest] + #[case::body( + c" +def on_pre_call(args): + args['complete_input_dict'] = {'replacement': True} +" + )] + #[case::headers( + c" +def on_pre_call(args): + args['headers'] = {'x-replacement': 'yes'} +" + )] + fn rebinding_the_payload_envelope_does_not_reach_the_wire(#[case] script: &CStr) { + let body = json!({"document": document(DOCUMENT)}); + let wire = before_send(script, body.clone()); + assert_eq!(wire.body, body); + assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]); + } + + #[test] + fn pre_call_header_edit_reaches_the_wire() { + let wire = before_send( + c" +def on_pre_call(args): + args['headers']['x-callback'] = 'edited' +", + json!({}), + ); + assert_eq!( + wire.headers, + [ + ("x-route".to_string(), "route".to_string()), + ("x-callback".to_string(), "edited".to_string()), + ] + ); + } + + #[test] + fn pre_call_receives_the_wire_request_and_the_logger_its_redacted_request() { + let body = json!({"model": "model", "document": document(DOCUMENT)}); + before_send_with_secrets( + c" +logger_fn = lambda *args: None +kwargs = { + 'litellm_call_id': 'call-1', + 'client_secret': 'shh', + 'proxy_server_request': {'body': {}}, + 'logger_fn': logger_fn, + 'litellm_request_debug': True, + 'ocr_cost_per_page': 0.05, +} +observed = [] +on_pre_call = observed.append +def check(): + [args] = observed + assert args['api_base'] == 'https://provider.invalid/ocr', args + assert args['complete_input_dict'] == { + 'model': 'model', + 'document': {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'}, + }, args + update = logger.update + assert update['model'] == 'model' and update['custom_llm_provider'] == 'provider', update + assert update['litellm_params']['litellm_call_id'] == 'call-1', update + assert update['litellm_params']['api_base'] == 'https://provider.invalid/ocr', update + assert update['litellm_params']['logger_fn'] is logger_fn, update + assert update['litellm_params']['litellm_request_debug'] is True, update + assert update['litellm_params']['ocr_cost_per_page'] == 0.05, update + assert update['kwargs']['client_secret'] == '****', update + assert 'proxy_server_request' not in update['kwargs'], update + assert update['optional_params']['client_secret'] == '****', update +", + json!({"client_secret": "shh"}), + body, + &["client_secret"], + ); + } + + #[rstest] + #[case::added_key( + c" +def on_pre_call(args): + args['complete_input_dict']['include_image_base64'] = True +", + json!({"document": document(DOCUMENT), "include_image_base64": true}) + )] + #[case::replaced_document( + c" +document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} +kwargs = {'document': document} +def on_pre_call(args): + args['complete_input_dict']['document'] = { + 'type': 'document_url', 'document_url': 'data:application/pdf;base64,ZWRpdGVk' + } +def check(): + assert document['document_url'] == 'data:application/pdf;base64,YWJj', document +", + json!({"document": document(EDITED)}) + )] + #[case::retained_body_edited_after_rebinding( + c" +def on_pre_call(args): + retained = args['complete_input_dict'] + args['complete_input_dict'] = {'rebound': True} + retained['include_image_base64'] = True +", + json!({"document": document(DOCUMENT), "include_image_base64": true}) + )] + fn pre_call_body_edits_reach_the_wire(#[case] script: &CStr, #[case] expected: Value) { + let body = json!({"document": document(DOCUMENT)}); + let wire = before_send(script, body); + assert_eq!(wire.body, expected); + } + + #[test] + fn retained_headers_edited_after_rebinding_reach_the_wire() { + let wire = before_send( + c" +def on_pre_call(args): + retained = args['headers'] + args['headers'] = {'x-rebound': 'rebound'} + retained['x-retained'] = 'sent' +", + json!({}), + ); + assert_eq!( + wire.headers, + [ + ("x-route".to_string(), "route".to_string()), + ("x-retained".to_string(), "sent".to_string()), + ] + ); + } + + #[test] + fn post_call_receives_the_raw_response_the_route_key_and_the_body_and_headers_pre_call_saw() { + before_send( + c" +def check(): + original_response, api_key, additional_args = logger.post + assert original_response == 'raw response', original_response + assert api_key == logger.pre_api_key == 'route-key', (api_key, logger.pre_api_key) + assert additional_args == { + 'complete_input_dict': logger.pre['complete_input_dict'], + 'headers': logger.pre['headers'], + }, additional_args + assert additional_args['complete_input_dict'] is logger.pre['complete_input_dict'] + assert additional_args['headers'] is logger.pre['headers'] +", + json!({"document": document(DOCUMENT)}), + ); + } + + #[test] + fn every_request_runs_the_full_pre_call_and_post_call() { + let wire = before_send( + c" +def on_pre_call(args): + args['complete_input_dict']['include_image_base64'] = True +def check(): + assert logger.names() == ['pre_call', 'post_call'], logger.calls +", + json!({"document": document(DOCUMENT)}), + ); + assert_eq!( + wire.body, + json!({"document": document(DOCUMENT), "include_image_base64": true}) + ); + } + + /// What one pre-call callback does to the payload it is handed. + #[derive(Clone, Debug)] + enum Edit { + Nothing, + Set(String, Value), + Remove(String), + Rebind(Value), + RebindThenSetRetained(String, Value), + } + + impl Edit { + fn script(&self) -> Value { + match self { + Self::Nothing => json!({"kind": "nothing"}), + Self::Set(key, value) => json!({"kind": "set", "key": key, "value": value}), + Self::Remove(key) => json!({"kind": "remove", "key": key}), + Self::Rebind(value) => json!({"kind": "rebind", "value": value}), + Self::RebindThenSetRetained(key, value) => { + json!({"kind": "rebind_then_set_retained", "key": key, "value": value}) + } + } + } + + /// The legacy contract: the provider is sent the body object `pre_call` received, as + /// the callback left it. Rebinding the envelope's key points the envelope elsewhere and + /// leaves that object alone. + fn sent(&self, body: &Map) -> Value { + let mut sent = body.clone(); + match self { + Self::Nothing | Self::Rebind(_) => {} + Self::Set(key, value) | Self::RebindThenSetRetained(key, value) => { + sent.insert(key.clone(), value.clone()); + } + Self::Remove(key) => { + sent.remove(key); + } + } + Value::Object(sent) + } + } + + /// How the caller's keyword for a body key relates to what the route sends under it. + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + enum Caller { + PassedUnchanged, + RewrittenByTheRoute, + NotPassed, + } + + const MODEL: &CStr = c" +aliased = {} +def on_pre_call(args): + body = args['complete_input_dict'] + aliased.update({name: body[name] is kwargs[name] for name in unchanged}) + kind = edit['kind'] + if kind == 'set': + body[edit['key']] = edit['value'] + elif kind == 'remove': + body.pop(edit['key'], None) + elif kind == 'rebind': + args['complete_input_dict'] = edit['value'] + elif kind == 'rebind_then_set_retained': + args['complete_input_dict'] = {} + body[edit['key']] = edit['value'] +def check(): + assert aliased == {name: True for name in unchanged}, aliased + assert logger.names() == ['pre_call', 'post_call'], logger.calls +"; + + fn json_value() -> impl Strategy { + let leaf = prop_oneof![ + Just(Value::Null), + any::().prop_map(Value::from), + any::().prop_map(Value::from), + any::() + .prop_filter("JSON has no NaN or infinity", |number| number.is_finite()) + .prop_map(Value::from), + ".{0,8}".prop_map(Value::from), + ]; + leaf.prop_recursive(3, 24, 4, |inner| { + prop_oneof![ + prop::collection::vec(inner.clone(), 0..4).prop_map(Value::from), + prop::collection::btree_map(key(), inner, 0..4) + .prop_map(|fields| Value::Object(fields.into_iter().collect())), + ] + }) + } + + fn key() -> impl Strategy { + "[a-z]{1,6}" + } + + fn caller() -> impl Strategy { + prop_oneof![ + Just(Caller::PassedUnchanged), + Just(Caller::RewrittenByTheRoute), + Just(Caller::NotPassed), + ] + } + + fn edit() -> impl Strategy { + prop_oneof![ + Just(Edit::Nothing), + (key(), json_value()).prop_map(|(key, value)| Edit::Set(key, value)), + key().prop_map(Edit::Remove), + json_value().prop_map(Edit::Rebind), + (key(), json_value()).prop_map(|(key, value)| Edit::RebindThenSetRetained(key, value)), + ] + } + + proptest! { + #![proptest_config(ProptestConfig::with_cases(128))] + + /// For any body, any caller keywords and any callback edit: every keyword the route + /// sends unchanged reaches `pre_call` as the caller's own object, and the provider is + /// sent exactly what the model says, so a callback that edits nothing changes nothing. + #[test] + fn the_wire_is_the_body_pre_call_received_as_the_callback_left_it( + fields in prop::collection::btree_map(key(), (json_value(), caller()), 0..5), + edit in edit(), + ) { + let body: Map = fields + .iter() + .map(|(name, (value, _))| (name.clone(), value.clone())) + .collect(); + let kwargs: Map = fields + .iter() + .filter_map(|(name, (value, caller))| match caller { + Caller::PassedUnchanged => Some((name.clone(), value.clone())), + Caller::RewrittenByTheRoute => Some((name.clone(), json!([value]))), + Caller::NotPassed => None, + }) + .collect(); + let unchanged: Value = fields + .iter() + .filter(|(_, (_, caller))| *caller == Caller::PassedUnchanged) + .map(|(name, _)| Value::from(name.clone())) + .collect(); + + let wire = before_send_bound( + &[ + ("kwargs", &Value::Object(kwargs)), + ("unchanged", &unchanged), + ("edit", &edit.script()), + ], + MODEL, + json!({}), + Value::Object(body.clone()), + &[], + ); + + prop_assert_eq!(wire.body, edit.sent(&body)); + prop_assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]); + } + } +} + #[cfg(test)] -#[path = "../tests/terminal.rs"] -mod terminal_tests; +mod terminal_tests { + use std::ffi::CStr; + + use litellm_host::event::{FailureOrigin, Timing}; + use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; + use pyo3::exceptions::PyRuntimeError; + use pyo3::exceptions::asyncio::CancelledError; + use pyo3::prelude::*; + use pyo3::types::PyDict; + use rstest::rstest; + + use super::LegacyLogging; + use crate::PythonLogger; + use crate::test_support::{legacy_call, local, namespace, run}; + + const TIMING: Timing = Timing { + start_time: 0.0, + end_time: 1.0, + }; + + fn logged(py: Python<'_>, locals: &Bound<'_, PyDict>, asynchronous: bool) -> LegacyLogging { + LegacyLogging { + logger: Some(PythonLogger::new(local(locals, "logger").unbind())), + ..legacy_call(py, locals, asynchronous) + } + } + + fn succeed( + py: Python<'_>, + locals: &Bound<'_, PyDict>, + logging: &mut LegacyLogging, + ) -> LifecycleStep { + let response = local(locals, "response").unbind(); + logging + .emit( + py, + LifecycleEvent::Succeeded { + timing: TIMING, + response: &response, + }, + ) + .unwrap() + } + + fn fail( + py: Python<'_>, + locals: &Bound<'_, PyDict>, + logging: &mut LegacyLogging, + ) -> LifecycleStep { + let failure = PyErr::from_value(local(locals, "failure")); + logging + .emit( + py, + LifecycleEvent::Failed { + timing: TIMING, + origin: FailureOrigin::Host, + error: &failure, + }, + ) + .unwrap() + } + + #[rstest] + #[case::sync_listened(false, c"", &["submit"])] + #[case::async_listened( + true, + c"", + &["async_success_handler", "enqueued", "sync_success_for_async_call"] + )] + #[case::async_deferred(true, c"logger._defer_async_logging = True", &["sync_success_for_async_call"])] + #[case::async_with_fallbacks(true, c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])] + fn success_reaches_the_logging_handlers( + #[case] asynchronous: bool, + #[case] script: &CStr, + #[case] expected: &[&str], + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c"response = object()"); + run(py, &locals, script); + let mut logging = logged(py, &locals, asynchronous); + assert!(matches!( + succeed(py, &locals, &mut logging), + LifecycleStep::Done + )); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names, expected); + run( + py, + &locals, + c" +assert all(value is response for name, value in logger.calls if name.endswith('_handler')) +assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_async_logging', False) +", + ); + }); + } + + #[rstest] + #[case::synchronous(false, &["failure_handler"])] + #[case::asynchronous(true, &[])] + fn internal_calls_skip_failure_callbacks_only_when_asynchronous( + #[case] asynchronous: bool, + #[case] expected: &[&str], + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c"failure = ValueError('provider')"); + let mut logging = LegacyLogging { + internal: true, + ..logged(py, &locals, asynchronous) + }; + assert!(matches!( + fail(py, &locals, &mut logging), + LifecycleStep::Done + )); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names, expected); + }); + } + + #[test] + fn internal_async_calls_skip_the_async_success_fan_out() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c"response = object()"); + let mut logging = LegacyLogging { + internal: true, + ..logged(py, &locals, true) + }; + succeed(py, &locals, &mut logging); + run( + py, + &locals, + c"assert logger.names() == ['sync_success_for_async_call'], logger.calls", + ); + }); + } + + #[test] + fn a_failing_success_callback_is_reported_without_replacing_the_response() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c" +response = object() +failure = ValueError('terminal diagnostic') + +class FailingLogger(StubLogger): + def handle_sync_success_callbacks_for_async_calls(self, *args): + raise failure + +logger = FailingLogger() +", + ); + let mut logging = logged(py, &locals, true); + assert!(matches!( + succeed(py, &locals, &mut logging), + LifecycleStep::Done + )); + assert!( + logging + .response + .as_ref() + .unwrap() + .bind(py) + .is(local(&locals, "response")) + ); + run(py, &locals, c"assert unraisable_from(logger) == [failure]"); + }); + } + + #[rstest] + #[case::sync_listened(false, c"", &["failure_handler"])] + #[case::async_listened(true, c"", &["failure_handler", "async_failure_handler"])] + fn failure_reaches_the_logging_handlers( + #[case] asynchronous: bool, + #[case] script: &CStr, + #[case] expected: &[&str], + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c"failure = ValueError('provider')"); + run(py, &locals, script); + let mut logging = logged(py, &locals, asynchronous); + let step = fail(py, &locals, &mut logging); + let awaits_async_handler = expected.contains(&"async_failure_handler"); + assert_eq!( + matches!(step, LifecycleStep::Await(_)), + awaits_async_handler + ); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names, expected); + run( + py, + &locals, + c"assert all(value is failure for name, value in logger.calls if name.endswith('_handler'))", + ); + }); + } + + #[test] + fn a_failing_sync_failure_callback_keeps_the_error_and_still_runs_the_async_family() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c" +failure = ValueError('selected') + +class FailingLogger(StubLogger): + def failure_handler(self, error, trace, start, end): + self.record('failure_handler', error) + raise RuntimeError('handler failed') + +logger = FailingLogger() +", + ); + let mut logging = logged(py, &locals, true); + assert!(matches!( + fail(py, &locals, &mut logging), + LifecycleStep::Await(_) + )); + assert!( + logging + .error + .as_ref() + .unwrap() + .bind(py) + .is(local(&locals, "failure")) + ); + run( + py, + &locals, + c"assert logger.names() == ['failure_handler', 'async_failure_handler'], logger.calls", + ); + }); + } + + #[rstest] + #[case::completed(None, true)] + #[case::handler_error(Some(false), true)] + #[case::cancelled(Some(true), false)] + fn the_async_failure_handler_ends_the_call_unless_it_was_cancelled( + #[case] error: Option, + #[case] done: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c"failure = ValueError('provider')"); + let mut logging = logged(py, &locals, true); + fail(py, &locals, &mut logging); + let result = match error { + None => Ok(py.None()), + Some(false) => Err(PyRuntimeError::new_err("handler failed")), + Some(true) => Err(CancelledError::new_err("cancelled")), + }; + let expected = result.as_ref().err().map(|error| error.value(py).clone()); + match logging.resume(py, result) { + Ok(step) => assert!(done && matches!(step, LifecycleStep::Done)), + Err(propagated) => { + assert!(!done); + assert!(propagated.value(py).is(expected.unwrap())); + } + } + }); + } + + #[test] + fn closing_restores_the_correlation_context_once() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c""); + let mut logging = logged(py, &locals, true); + logging.close(py); + logging.close(py); + run( + py, + &locals, + c"assert logger.names() == ['restore'], logger.calls", + ); + }); + } +} diff --git a/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs b/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs index b18012f926e..292648ba69c 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs @@ -63,5 +63,151 @@ impl PendingLogging { } #[cfg(test)] -#[path = "../tests/deferred.rs"] -mod tests; +mod tests { + use std::ffi::CStr; + + use pyo3::prelude::*; + use pyo3::types::PyDict; + use rstest::rstest; + + use super::{PendingLogging, PendingSuccess}; + use crate::PythonLogger; + use crate::test_support::{local, namespace, run}; + + /// A deferred success for the namespace's `logger` and `response`, bound as `pending`. + fn defer<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { + let locals = namespace(py, c"response = object()"); + run(py, &locals, script); + let pending = Py::new( + py, + PendingLogging { + pending: Some(PendingSuccess { + logger: PythonLogger::new(local(&locals, "logger").unbind()), + response: Some(local(&locals, "response").unbind()), + start: py.None(), + end: Some(py.None()), + }), + }, + ) + .unwrap(); + locals.set_item("pending", pending).unwrap(); + locals + } + + #[test] + fn release_enqueues_the_success_once_in_the_releasing_context() { + Python::initialize(); + Python::attach(|py| { + let locals = defer( + py, + c" +from contextvars import ContextVar + +marker = ContextVar('marker', default='unset') +observed = [] + +def on_enqueue(coroutine): + observed.append(marker.get()) + pending.release(True) + +logger.on_enqueue = on_enqueue +", + ); + run( + py, + &locals, + c" +marker.set('release') +pending.release(True) +pending.release(True) +assert observed == ['release'], observed +assert logger.names() == ['async_success_handler', 'enqueued'], logger.calls +assert logger.calls[0][1] is response +", + ); + }); + } + + #[test] + fn a_blocked_release_drops_the_success_for_good() { + Python::initialize(); + Python::attach(|py| { + let locals = defer(py, c""); + run( + py, + &locals, + c" +pending.release(False) +pending.release(True) +assert logger.calls == [], logger.calls +", + ); + }); + } + + #[rstest] + #[case::ordinary_error(c"RuntimeError('queue full')", false)] + #[case::cancellation(c"asyncio.CancelledError()", true)] + fn a_failed_enqueue_closes_the_coroutine_and_is_never_replayed( + #[case] failure: &CStr, + #[case] propagates: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let locals = defer( + py, + c" +import asyncio + +def on_enqueue(coroutine): + raise failure + +logger.on_enqueue = on_enqueue +", + ); + locals + .set_item("failure", py.eval(failure, None, Some(&locals)).unwrap()) + .unwrap(); + let released = local(&locals, "pending").call_method1("release", (true,)); + match released { + Ok(_) => assert!(!propagates), + Err(error) => { + assert!(propagates); + assert!(error.value(py).is(local(&locals, "failure"))); + } + } + locals.set_item("propagates", propagates).unwrap(); + run( + py, + &locals, + c" +pending.release(True) +assert logger.names() == ['async_success_handler', 'enqueued', 'closed'], logger.calls +assert unraisable_from(logger) == ([] if propagates else [failure]) +", + ); + }); + } + + #[test] + fn an_unreleased_success_does_not_keep_its_logger_alive() { + Python::initialize(); + Python::attach(|py| { + let locals = defer(py, c""); + run( + py, + &locals, + c" +import gc +import weakref + +logger.pending = pending +reference = weakref.ref(logger) +del logger, pending +gc.collect() +assert reference() is None +", + ); + }); + } +} diff --git a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs index 44393792d1f..030bf03d4ba 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs @@ -16,13 +16,218 @@ mod deferred; mod logger; mod preparation; mod python; -#[cfg(test)] -#[path = "../tests/support.rs"] -mod test_support; - pub(crate) use adapter::LegacyLogging; pub use adapter::{LegacySurface, PassThroughStream}; pub use call::{PublicCall, run_legacy_call}; pub(crate) use callbacks::{LegacyCallbacks, is_internal_call}; pub(crate) use logger::{DeploymentHooks, PythonLogger, finalize, setup}; pub(crate) use preparation::prepare; + +#[cfg(test)] +mod test_support { + use std::ffi::CStr; + + use pyo3::prelude::*; + use pyo3::types::{PyDict, PyTuple}; + + use crate::{LegacyLogging, LegacySurface, PublicCall}; + + /// The parameters of every `callbacks_legacy_python` function, as the real module declares them. + /// `tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python + /// signatures, and [`namespace`] binds every fake call against it. + pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json"); + + /// Stand-ins for `callbacks_legacy_python`, the only Python module the crate calls. Tests + /// share one interpreter and run concurrently, so each fake is installed idempotently and + /// forwards to the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`). + /// Every fake is bound against the contract first, so a call the real module would reject + /// fails here too. + const STUBS: &CStr = c" +import contextvars +import inspect +import json +import sys +import traceback +import types + +for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.callbacks_legacy_python'): + sys.modules.setdefault(name, types.ModuleType(name)) + +legacy = sys.modules['litellm.rust_bridge.callbacks_legacy_python'] +CONTRACT = json.loads(python_contract) + + +def contracted(name, fake): + signature = inspect.Signature( + [inspect.Parameter(parameter, inspect.Parameter.POSITIONAL_OR_KEYWORD) for parameter in CONTRACT[name]] + ) + + def checked(*args, **kwargs): + signature.bind(*args, **kwargs) + return fake(*args, **kwargs) + + return checked + + +if not hasattr(legacy, 'is_internal'): + legacy.is_internal = contextvars.ContextVar('is_internal_call', default=False) + +FAKES = { + 'setup': lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace( + logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'], + kwargs=kwargs, + ), + 'check_limits': lambda arguments: arguments['logger'].check_limits(arguments), + 'finalize': lambda response, logger, kwargs, start, end: logger.record('finalize', response), + 'update_logging': lambda logger, kwargs, model, optional_params, litellm_params, provider: logger.update_from_kwargs( + kwargs=kwargs, + model=model, + optional_params=optional_params, + litellm_params=litellm_params, + custom_llm_provider=provider, + ), + 'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args), + 'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call( + original_response, api_key, additional_args + ), + 'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)), + 'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending), + 'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls( + response, start, end + ), + 'failure_handler': lambda logger, error, start, end, asynchronous: ( + logger.async_failure_handler if asynchronous else logger.failure_handler + )(error, ''.join(traceback.format_exception(error)), start, end), + 'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)), + 'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end), + 'enqueue_logging': lambda coroutine: coroutine.enqueue(), + 'restore_context': lambda logger: logger.record('restore', None), + 'custom_pricing_fields': lambda: ('ocr_cost_per_page',), + 'is_internal_call': lambda: legacy.is_internal.get(), + 'credential_list': lambda: [], + 'warn_unknown_credential': lambda name, loaded: None, + 'before_deployment_call': lambda kwargs, call_type: kwargs['logger'].hook('pre', kwargs, call_type), + 'after_deployment_success': lambda kwargs, response, call_type: kwargs['logger'].hook( + 'success', response, call_type + ), + 'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type), + 'stream_opened': lambda logger: logger.record('stream_opened', None), + 'stream_success': lambda logger, request_body, chunks, start, end, first_chunk: logger.record( + 'stream_success', list(chunks) + ), + 'stream_failure': lambda logger, request_body, chunks, error: logger.record('stream_failure', error), +} +assert FAKES.keys() == CONTRACT.keys(), sorted(FAKES.keys() ^ CONTRACT.keys()) +for name, fake in FAKES.items(): + setattr(legacy, name, contracted(name, fake)) + + +unraisable = sys.modules.setdefault( + 'litellm_test_unraisable', types.ModuleType('litellm_test_unraisable') +) +if not hasattr(unraisable, 'events'): + unraisable.events = [] + sys.unraisablehook = lambda event: unraisable.events.append((event.object, event.exc_value)) + + +def unraisable_from(owner): + return [error for source, error in unraisable.events if source is owner] + + +class StubCoroutine: + def __init__(self, logger): + self.logger = logger + + def enqueue(self): + self.logger.record('enqueued', None) + self.logger.on_enqueue(self) + + def close(self): + self.logger.record('closed', None) + + +class StubLogger: + def __init__(self): + self.calls = [] + self.hooks = {} + self.on_enqueue = lambda coroutine: None + + def record(self, name, value): + self.calls.append((name, value)) + + def names(self): + return [name for name, _ in self.calls] + + def hook(self, phase, value, call_type): + self.record(phase + '_hook', call_type) + return self.hooks.get(phase, lambda value: 'awaitable')(value) + + def check_limits(self, arguments): + self.record('check_limits', arguments) + + def failure_handler(self, error, trace, start, end): + self.record('failure_handler', error) + + def async_failure_handler(self, error, trace, start, end): + self.record('async_failure_handler', error) + return 'awaitable' + + def success_handler(self, response, start, end): + self.record('success_handler', response) + + def async_success_handler(self, response, start, end): + self.record('async_success_handler', response) + return StubCoroutine(self) + + def handle_sync_success_callbacks_for_async_calls(self, response, start, end): + self.record('sync_success_for_async_call', response) + + +logger = StubLogger() +"; + + /// A namespace with the stubs, `StubLogger` and a fresh `logger`, after `script` ran in it. + pub(crate) fn namespace<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { + let locals = PyDict::new(py); + locals.set_item("python_contract", PYTHON_CONTRACT).unwrap(); + py.run(STUBS, Some(&locals), Some(&locals)).unwrap(); + py.run(script, Some(&locals), Some(&locals)).unwrap(); + locals + } + + pub(crate) fn run(py: Python<'_>, locals: &Bound<'_, PyDict>, code: &CStr) { + py.run(code, Some(locals), Some(locals)).unwrap(); + } + + pub(crate) fn local<'py>(locals: &Bound<'py, PyDict>, name: &str) -> Bound<'py, PyAny> { + locals.get_item(name).unwrap().unwrap() + } + + /// A legacy call over the namespace's `kwargs` (or none) and `request` (or `None`). + pub(crate) fn legacy_call( + py: Python<'_>, + locals: &Bound<'_, PyDict>, + asynchronous: bool, + ) -> LegacyLogging { + let request = locals + .get_item("request") + .unwrap() + .unwrap_or_else(|| py.None().into_bound(py)); + let kwargs = locals + .get_item("kwargs") + .unwrap() + .map(|kwargs| kwargs.cast_into::().unwrap()) + .unwrap_or_else(|| PyDict::new(py)); + let call = PublicCall::capture(&request, &PyTuple::empty(py), &kwargs).unwrap(); + LegacyLogging::new( + py, + LegacySurface { + call_type: "test", + input_description: "test input", + stream: None, + }, + call, + asynchronous, + ) + } +} diff --git a/litellm-rust/crates/callbacks-legacy-python/tests/deferred.rs b/litellm-rust/crates/callbacks-legacy-python/tests/deferred.rs deleted file mode 100644 index 289ea1b2e7f..00000000000 --- a/litellm-rust/crates/callbacks-legacy-python/tests/deferred.rs +++ /dev/null @@ -1,146 +0,0 @@ -use std::ffi::CStr; - -use pyo3::prelude::*; -use pyo3::types::PyDict; -use rstest::rstest; - -use super::{PendingLogging, PendingSuccess}; -use crate::PythonLogger; -use crate::test_support::{local, namespace, run}; - -/// A deferred success for the namespace's `logger` and `response`, bound as `pending`. -fn defer<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { - let locals = namespace(py, c"response = object()"); - run(py, &locals, script); - let pending = Py::new( - py, - PendingLogging { - pending: Some(PendingSuccess { - logger: PythonLogger::new(local(&locals, "logger").unbind()), - response: Some(local(&locals, "response").unbind()), - start: py.None(), - end: Some(py.None()), - }), - }, - ) - .unwrap(); - locals.set_item("pending", pending).unwrap(); - locals -} - -#[test] -fn release_enqueues_the_success_once_in_the_releasing_context() { - Python::initialize(); - Python::attach(|py| { - let locals = defer( - py, - c" -from contextvars import ContextVar - -marker = ContextVar('marker', default='unset') -observed = [] - -def on_enqueue(coroutine): - observed.append(marker.get()) - pending.release(True) - -logger.on_enqueue = on_enqueue -", - ); - run( - py, - &locals, - c" -marker.set('release') -pending.release(True) -pending.release(True) -assert observed == ['release'], observed -assert logger.names() == ['async_success_handler', 'enqueued'], logger.calls -assert logger.calls[0][1] is response -", - ); - }); -} - -#[test] -fn a_blocked_release_drops_the_success_for_good() { - Python::initialize(); - Python::attach(|py| { - let locals = defer(py, c""); - run( - py, - &locals, - c" -pending.release(False) -pending.release(True) -assert logger.calls == [], logger.calls -", - ); - }); -} - -#[rstest] -#[case::ordinary_error(c"RuntimeError('queue full')", false)] -#[case::cancellation(c"asyncio.CancelledError()", true)] -fn a_failed_enqueue_closes_the_coroutine_and_is_never_replayed( - #[case] failure: &CStr, - #[case] propagates: bool, -) { - Python::initialize(); - Python::attach(|py| { - let locals = defer( - py, - c" -import asyncio - -def on_enqueue(coroutine): - raise failure - -logger.on_enqueue = on_enqueue -", - ); - locals - .set_item("failure", py.eval(failure, None, Some(&locals)).unwrap()) - .unwrap(); - let released = local(&locals, "pending").call_method1("release", (true,)); - match released { - Ok(_) => assert!(!propagates), - Err(error) => { - assert!(propagates); - assert!(error.value(py).is(local(&locals, "failure"))); - } - } - locals.set_item("propagates", propagates).unwrap(); - run( - py, - &locals, - c" -pending.release(True) -assert logger.names() == ['async_success_handler', 'enqueued', 'closed'], logger.calls -assert unraisable_from(logger) == ([] if propagates else [failure]) -", - ); - }); -} - -#[test] -fn an_unreleased_success_does_not_keep_its_logger_alive() { - Python::initialize(); - Python::attach(|py| { - let locals = defer(py, c""); - run( - py, - &locals, - c" -import gc -import weakref - -logger.pending = pending -reference = weakref.ref(logger) -del logger, pending -gc.collect() -assert reference() is None -", - ); - }); -} diff --git a/litellm-rust/crates/callbacks-legacy-python/tests/deployment_hooks.rs b/litellm-rust/crates/callbacks-legacy-python/tests/deployment_hooks.rs deleted file mode 100644 index 52c5e47f83f..00000000000 --- a/litellm-rust/crates/callbacks-legacy-python/tests/deployment_hooks.rs +++ /dev/null @@ -1,282 +0,0 @@ -use std::ffi::CStr; - -use litellm_host::event::{FailureOrigin, Timing}; -use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; -use pyo3::exceptions::asyncio::CancelledError; -use pyo3::prelude::*; -use pyo3::types::PyDict; -use rstest::rstest; - -use super::LegacyLogging; -use crate::test_support::{legacy_call, local, namespace, run}; - -const CALL: &CStr = c" -document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} -kwargs = {'logger': logger, 'document': document} -"; - -const TIMING: Timing = Timing { - start_time: 0.0, - end_time: 1.0, -}; - -fn begin<'py>( - py: Python<'py>, - locals: &Bound<'py, PyDict>, - asynchronous: bool, -) -> (LegacyLogging, LifecycleStep) { - let mut logging = legacy_call(py, locals, asynchronous); - let kwargs = local(locals, "kwargs") - .cast_into::() - .unwrap() - .unbind(); - let step = logging.begin(py, kwargs, 0.0).unwrap(); - (logging, step) -} - -fn arguments<'py>(py: Python<'py>, step: LifecycleStep) -> Bound<'py, PyDict> { - let LifecycleStep::Arguments(arguments) = step else { - panic!("expected the prepared arguments"); - }; - arguments.into_bound(py) -} - -fn awaits_deployment_hook(step: &LifecycleStep) -> bool { - matches!(step, LifecycleStep::Await(_)) -} - -#[rstest] -#[case::synchronous(false)] -#[case::asynchronous(true)] -fn deployment_pre_call_hook_runs_only_for_asynchronous_calls(#[case] asynchronous: bool) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, CALL); - let (_, step) = begin(py, &locals, asynchronous); - assert_eq!(awaits_deployment_hook(&step), asynchronous); - let names: Vec = local(&locals, "logger") - .call_method0("names") - .unwrap() - .extract() - .unwrap(); - assert_eq!(names.contains(&"pre_hook".to_string()), asynchronous); - }); -} - -#[test] -fn kwargs_returned_by_the_pre_call_hook_are_what_the_call_prepares() { - Python::initialize(); - Python::attach(|py| { - let locals = namespace( - py, - c" -document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} -replacement = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,ZWRpdGVk'} -kwargs = {'logger': logger, 'document': document} -replaced_kwargs = {'logger': logger, 'document': replacement, 'pages': [0]} -", - ); - let (mut logging, step) = begin(py, &locals, true); - assert!(awaits_deployment_hook(&step)); - let step = logging - .resume(py, Ok(local(&locals, "replaced_kwargs").unbind())) - .unwrap(); - locals.set_item("prepared", arguments(py, step)).unwrap(); - run( - py, - &locals, - c" -assert prepared['document'] is replacement -assert prepared['pages'] is replaced_kwargs['pages'] -assert prepared['litellm_logging_obj'] is logger -assert 'litellm_logging_obj' not in replaced_kwargs -[checked] = [value for name, value in logger.calls if name == 'check_limits'] -assert checked is prepared -", - ); - }); -} - -#[rstest] -#[case::synchronous(false)] -#[case::asynchronous(true)] -fn a_keyword_the_bridge_never_reads_reaches_every_reader_as_the_callers_object( - #[case] asynchronous: bool, -) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace( - py, - c" -opaque = object() -hooked = [] -logger.hooks = {'pre': lambda kwargs: hooked.append(kwargs['vendor_extension']) or kwargs} -kwargs = {'logger': logger, 'vendor_extension': opaque} -", - ); - let (mut logging, step) = begin(py, &locals, asynchronous); - let step = match step { - LifecycleStep::Await(hook_result) => logging.resume(py, Ok(hook_result)).unwrap(), - step => step, - }; - locals.set_item("prepared", arguments(py, step)).unwrap(); - locals.set_item("asynchronous", asynchronous).unwrap(); - run( - py, - &locals, - c" -assert prepared['vendor_extension'] is opaque -[checked] = [value for name, value in logger.calls if name == 'check_limits'] -assert checked['vendor_extension'] is opaque -assert hooked == ([opaque] if asynchronous else []), hooked -", - ); - }); -} - -#[test] -fn response_returned_by_the_post_call_hook_is_finalized_and_returned() { - Python::initialize(); - Python::attach(|py| { - let locals = namespace( - py, - c" -kwargs = {'logger': logger} -response = object() -replacement = object() -logger.hooks = {'pre': lambda kwargs: kwargs} -", - ); - let (mut logging, _) = begin(py, &locals, true); - logging - .resume(py, Ok(local(&locals, "kwargs").unbind())) - .unwrap(); - let step = logging - .after_success(py, local(&locals, "response").unbind(), TIMING) - .unwrap(); - assert!(awaits_deployment_hook(&step)); - let step = logging - .resume(py, Ok(local(&locals, "replacement").unbind())) - .unwrap(); - let LifecycleStep::Response(returned) = step else { - panic!("expected the finalized response"); - }; - assert!(returned.bind(py).is(local(&locals, "replacement"))); - run( - py, - &locals, - c" -[finalized] = [value for name, value in logger.calls if name == 'finalize'] -assert finalized is replacement -", - ); - }); -} - -#[rstest] -#[case::pre_call(false)] -#[case::post_call(true)] -fn cancelling_a_deployment_hook_ends_the_call_with_that_cancellation(#[case] post_call: bool) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, c"kwargs = {'logger': logger}\nresponse = object()"); - let (mut logging, _) = begin(py, &locals, true); - if post_call { - logging - .resume(py, Ok(local(&locals, "kwargs").unbind())) - .unwrap(); - logging - .after_success(py, local(&locals, "response").unbind(), TIMING) - .unwrap(); - } - let cancellation = CancelledError::new_err("cancelled"); - let cancelled = cancellation.value(py).clone(); - let error = logging.resume(py, Err(cancellation)).err().unwrap(); - assert!(error.value(py).is(&cancelled)); - let names: Vec = local(&locals, "logger") - .call_method0("names") - .unwrap() - .extract() - .unwrap(); - assert!(!names.iter().any(|name| name.contains("handler"))); - }); -} - -#[rstest] -#[case::hook_completed(false)] -#[case::hook_cancelled(true)] -fn failure_callbacks_run_after_the_failure_hook_however_it_ends(#[case] cancelled: bool) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace( - py, - c"kwargs = {'logger': logger}\nfailure = ValueError('provider')", - ); - let (mut logging, _) = begin(py, &locals, true); - logging - .resume(py, Ok(local(&locals, "kwargs").unbind())) - .unwrap(); - let failure = PyErr::from_value(local(&locals, "failure")); - let failed = LifecycleEvent::Failed { - timing: TIMING, - origin: FailureOrigin::Call, - error: &failure, - }; - let step = logging.emit(py, failed).unwrap(); - assert!(awaits_deployment_hook(&step)); - let hook_result = if cancelled { - Err(CancelledError::new_err("cancelled")) - } else { - Ok(py.None()) - }; - assert!(matches!( - logging.resume(py, hook_result).unwrap(), - LifecycleStep::Await(_) - )); - run( - py, - &locals, - c" -assert logger.names()[-3:] == ['failure_hook', 'failure_handler', 'async_failure_handler'], logger.calls -assert all(value is failure for name, value in logger.calls if name.endswith('_handler')) -", - ); - }); -} - -#[rstest] -#[case::synchronous(false)] -#[case::asynchronous(true)] -fn a_limit_rejected_before_the_call_surfaces_as_the_callers_error(#[case] asynchronous: bool) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace( - py, - c" -class BudgetExceeded(Exception): - pass - -rejection = BudgetExceeded('over budget') - -class LimitedLogger(StubLogger): - def check_limits(self, arguments): - raise rejection - -logger = LimitedLogger() -logger.hooks = {'pre': lambda kwargs: kwargs} -kwargs = {'logger': logger} -", - ); - let mut logging = legacy_call(py, &locals, asynchronous); - let kwargs = local(&locals, "kwargs") - .cast_into::() - .unwrap() - .unbind(); - let result = logging.begin(py, kwargs, 0.0).and_then(|step| match step { - LifecycleStep::Await(_) => logging.resume(py, Ok(local(&locals, "kwargs").unbind())), - step => Ok(step), - }); - let error = result.err().unwrap(); - assert!(error.value(py).is(local(&locals, "rejection"))); - }); -} diff --git a/litellm-rust/crates/callbacks-legacy-python/tests/payload.rs b/litellm-rust/crates/callbacks-legacy-python/tests/payload.rs deleted file mode 100644 index 5459b36af27..00000000000 --- a/litellm-rust/crates/callbacks-legacy-python/tests/payload.rs +++ /dev/null @@ -1,523 +0,0 @@ -use std::ffi::CStr; - -use litellm_auth::SecretValue; -use litellm_host::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; -use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle, to_py}; -use proptest::prelude::*; -use pyo3::prelude::*; -use rstest::rstest; -use serde_json::{Map, Value, json}; - -use super::LegacyLogging; -use crate::PythonLogger; -use crate::test_support::{legacy_call, local, namespace, run}; - -/// The payload phases of `Logging` on top of `StubLogger`, with `pre_call` handing the -/// payload to the case's `on_pre_call`. -const PAYLOAD_LOGGER: &CStr = c" -class Request: - pass - -class PayloadLogger(StubLogger): - def update_from_kwargs(self, **update): - self.update = update - - def pre_call(self, input, api_key, additional_args): - self.record('pre_call', None) - self.pre = additional_args - self.pre_api_key = api_key - on_pre_call(additional_args) - - def post_call(self, original_response, api_key, additional_args): - self.record('post_call', None) - self.post = (original_response, api_key, additional_args) - -request = Request() -kwargs = {} -logger = PayloadLogger() -on_pre_call = lambda additional_args: None -check = lambda: None -"; - -const DOCUMENT: &str = "data:application/pdf;base64,YWJj"; -const EDITED: &str = "data:application/pdf;base64,ZWRpdGVk"; - -fn document(source: &str) -> Value { - json!({"type": "document_url", "document_url": source}) -} - -fn before_send(script: &CStr, body: Value) -> WireRequest { - before_send_with_secrets(script, json!({}), body, &[]) -} - -/// Runs `before_send` over `body` for a route whose parameters are `optional_params`, with -/// the Python objects `script` binds, then delivers the provider's raw response the way the -/// driver does and runs the script's `check()`. -fn before_send_with_secrets( - script: &CStr, - optional_params: Value, - body: Value, - secret_fields: &[&str], -) -> WireRequest { - before_send_bound(&[], script, optional_params, body, secret_fields) -} - -/// [`before_send_with_secrets`] with `bindings` placed in the namespace before `script` runs. -fn before_send_bound( - bindings: &[(&str, &Value)], - script: &CStr, - optional_params: Value, - body: Value, - secret_fields: &[&str], -) -> WireRequest { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, PAYLOAD_LOGGER); - for &(name, value) in bindings { - locals.set_item(name, to_py(py, value).unwrap()).unwrap(); - } - run(py, &locals, script); - let mut logging = LegacyLogging { - logger: Some(PythonLogger::new(local(&locals, "logger").unbind())), - ..legacy_call(py, &locals, false) - }; - let context = RequestContext { - model: "model".into(), - custom_llm_provider: "provider".into(), - optional_params, - secret_fields: secret_fields.iter().map(|name| name.to_string()).collect(), - api_key: Some(SecretValue::new("route-key")), - }; - let wire = WireRequest { - url: "https://provider.invalid/ocr".into(), - headers: vec![("x-route".into(), "route".into())], - body, - }; - let step = logging.before_send(py, Box::new(wire), &context).unwrap(); - let raw = MachineEvent::ResponseReceived { - raw: RawResponse { - body: "raw response".into(), - }, - }; - assert!(matches!( - logging.emit(py, LifecycleEvent::Machine(&raw)).unwrap(), - LifecycleStep::Done - )); - run(py, &locals, c"check()"); - let LifecycleStep::Wire(wire) = step else { - panic!("before_send did not hand back the wire request"); - }; - *wire - }) -} - -#[rstest] -#[case::caller_keyword(c" -document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} -pages = [0] -kwargs = {'document': document, 'pages': pages} -observed = [] -on_pre_call = lambda args: observed.append( - (args['complete_input_dict']['document'] is document, args['complete_input_dict']['pages'] is pages) -) -def check(): - assert observed == [(True, True)], observed -")] -#[case::request_attribute_behind_an_omitted_keyword(c" -document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} -pages = [0] -request.document = document -kwargs = {'pages': pages} -observed = [] -on_pre_call = lambda args: observed.append( - (args['complete_input_dict']['document'] is document, args['complete_input_dict']['pages'] is pages) -) -def check(): - assert observed == [(True, True)], observed -")] -fn passthrough_keys_reach_pre_call_as_the_callers_own_objects(#[case] script: &CStr) { - let body = json!({"model": "model", "document": document(DOCUMENT), "pages": [0]}); - let wire = before_send(script, body.clone()); - assert_eq!(wire.body, body); -} - -#[test] -fn pre_call_edit_of_a_passthrough_object_reaches_the_caller_and_the_wire() { - let wire = before_send( - c" -document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} -kwargs = {'document': document} -def on_pre_call(args): - args['complete_input_dict']['document']['document_url'] = 'data:application/pdf;base64,ZWRpdGVk' -def check(): - assert document['document_url'] == 'data:application/pdf;base64,ZWRpdGVk' -", - json!({"document": document(DOCUMENT)}), - ); - assert_eq!(wire.body["document"], document(EDITED)); -} - -#[test] -fn a_body_key_the_route_rewrote_is_not_the_callers_object() { - let wire = before_send( - c" -document = {'type': 'document_url', 'document_url': 'https://example.invalid/scan.pdf'} -kwargs = {'document': document} -observed = [] -def on_pre_call(args): - observed.append(args['complete_input_dict']['document'] is document) - args['complete_input_dict']['document']['document_name'] = 'edited.pdf' -def check(): - assert observed == [False], observed - assert document == {'type': 'document_url', 'document_url': 'https://example.invalid/scan.pdf'} -", - json!({"document": document(DOCUMENT)}), - ); - assert_eq!( - wire.body["document"], - json!({"type": "document_url", "document_url": DOCUMENT, "document_name": "edited.pdf"}) - ); -} - -#[test] -fn a_caller_value_with_no_json_form_is_left_out_of_realiasing() { - let body = json!({"pages": [0]}); - let wire = before_send( - c" -opaque = object() -kwargs = {'pages': opaque} -observed = [] -on_pre_call = lambda args: observed.append(args['complete_input_dict']['pages']) -def check(): - assert observed == [[0]], observed -", - body.clone(), - ); - assert_eq!(wire.body, body); -} - -#[rstest] -#[case::body( - c" -def on_pre_call(args): - args['complete_input_dict'] = {'replacement': True} -" -)] -#[case::headers( - c" -def on_pre_call(args): - args['headers'] = {'x-replacement': 'yes'} -" -)] -fn rebinding_the_payload_envelope_does_not_reach_the_wire(#[case] script: &CStr) { - let body = json!({"document": document(DOCUMENT)}); - let wire = before_send(script, body.clone()); - assert_eq!(wire.body, body); - assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]); -} - -#[test] -fn pre_call_header_edit_reaches_the_wire() { - let wire = before_send( - c" -def on_pre_call(args): - args['headers']['x-callback'] = 'edited' -", - json!({}), - ); - assert_eq!( - wire.headers, - [ - ("x-route".to_string(), "route".to_string()), - ("x-callback".to_string(), "edited".to_string()), - ] - ); -} - -#[test] -fn pre_call_receives_the_wire_request_and_the_logger_its_redacted_request() { - let body = json!({"model": "model", "document": document(DOCUMENT)}); - before_send_with_secrets( - c" -logger_fn = lambda *args: None -kwargs = { - 'litellm_call_id': 'call-1', - 'client_secret': 'shh', - 'proxy_server_request': {'body': {}}, - 'logger_fn': logger_fn, - 'litellm_request_debug': True, - 'ocr_cost_per_page': 0.05, -} -observed = [] -on_pre_call = observed.append -def check(): - [args] = observed - assert args['api_base'] == 'https://provider.invalid/ocr', args - assert args['complete_input_dict'] == { - 'model': 'model', - 'document': {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'}, - }, args - update = logger.update - assert update['model'] == 'model' and update['custom_llm_provider'] == 'provider', update - assert update['litellm_params']['litellm_call_id'] == 'call-1', update - assert update['litellm_params']['api_base'] == 'https://provider.invalid/ocr', update - assert update['litellm_params']['logger_fn'] is logger_fn, update - assert update['litellm_params']['litellm_request_debug'] is True, update - assert update['litellm_params']['ocr_cost_per_page'] == 0.05, update - assert update['kwargs']['client_secret'] == '****', update - assert 'proxy_server_request' not in update['kwargs'], update - assert update['optional_params']['client_secret'] == '****', update -", - json!({"client_secret": "shh"}), - body, - &["client_secret"], - ); -} - -#[rstest] -#[case::added_key( - c" -def on_pre_call(args): - args['complete_input_dict']['include_image_base64'] = True -", - json!({"document": document(DOCUMENT), "include_image_base64": true}) -)] -#[case::replaced_document( - c" -document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} -kwargs = {'document': document} -def on_pre_call(args): - args['complete_input_dict']['document'] = { - 'type': 'document_url', 'document_url': 'data:application/pdf;base64,ZWRpdGVk' - } -def check(): - assert document['document_url'] == 'data:application/pdf;base64,YWJj', document -", - json!({"document": document(EDITED)}) -)] -#[case::retained_body_edited_after_rebinding( - c" -def on_pre_call(args): - retained = args['complete_input_dict'] - args['complete_input_dict'] = {'rebound': True} - retained['include_image_base64'] = True -", - json!({"document": document(DOCUMENT), "include_image_base64": true}) -)] -fn pre_call_body_edits_reach_the_wire(#[case] script: &CStr, #[case] expected: Value) { - let body = json!({"document": document(DOCUMENT)}); - let wire = before_send(script, body); - assert_eq!(wire.body, expected); -} - -#[test] -fn retained_headers_edited_after_rebinding_reach_the_wire() { - let wire = before_send( - c" -def on_pre_call(args): - retained = args['headers'] - args['headers'] = {'x-rebound': 'rebound'} - retained['x-retained'] = 'sent' -", - json!({}), - ); - assert_eq!( - wire.headers, - [ - ("x-route".to_string(), "route".to_string()), - ("x-retained".to_string(), "sent".to_string()), - ] - ); -} - -#[test] -fn post_call_receives_the_raw_response_the_route_key_and_the_body_and_headers_pre_call_saw() { - before_send( - c" -def check(): - original_response, api_key, additional_args = logger.post - assert original_response == 'raw response', original_response - assert api_key == logger.pre_api_key == 'route-key', (api_key, logger.pre_api_key) - assert additional_args == { - 'complete_input_dict': logger.pre['complete_input_dict'], - 'headers': logger.pre['headers'], - }, additional_args - assert additional_args['complete_input_dict'] is logger.pre['complete_input_dict'] - assert additional_args['headers'] is logger.pre['headers'] -", - json!({"document": document(DOCUMENT)}), - ); -} - -#[test] -fn every_request_runs_the_full_pre_call_and_post_call() { - let wire = before_send( - c" -def on_pre_call(args): - args['complete_input_dict']['include_image_base64'] = True -def check(): - assert logger.names() == ['pre_call', 'post_call'], logger.calls -", - json!({"document": document(DOCUMENT)}), - ); - assert_eq!( - wire.body, - json!({"document": document(DOCUMENT), "include_image_base64": true}) - ); -} - -/// What one pre-call callback does to the payload it is handed. -#[derive(Clone, Debug)] -enum Edit { - Nothing, - Set(String, Value), - Remove(String), - Rebind(Value), - RebindThenSetRetained(String, Value), -} - -impl Edit { - fn script(&self) -> Value { - match self { - Self::Nothing => json!({"kind": "nothing"}), - Self::Set(key, value) => json!({"kind": "set", "key": key, "value": value}), - Self::Remove(key) => json!({"kind": "remove", "key": key}), - Self::Rebind(value) => json!({"kind": "rebind", "value": value}), - Self::RebindThenSetRetained(key, value) => { - json!({"kind": "rebind_then_set_retained", "key": key, "value": value}) - } - } - } - - /// The legacy contract: the provider is sent the body object `pre_call` received, as - /// the callback left it. Rebinding the envelope's key points the envelope elsewhere and - /// leaves that object alone. - fn sent(&self, body: &Map) -> Value { - let mut sent = body.clone(); - match self { - Self::Nothing | Self::Rebind(_) => {} - Self::Set(key, value) | Self::RebindThenSetRetained(key, value) => { - sent.insert(key.clone(), value.clone()); - } - Self::Remove(key) => { - sent.remove(key); - } - } - Value::Object(sent) - } -} - -/// How the caller's keyword for a body key relates to what the route sends under it. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum Caller { - PassedUnchanged, - RewrittenByTheRoute, - NotPassed, -} - -const MODEL: &CStr = c" -aliased = {} -def on_pre_call(args): - body = args['complete_input_dict'] - aliased.update({name: body[name] is kwargs[name] for name in unchanged}) - kind = edit['kind'] - if kind == 'set': - body[edit['key']] = edit['value'] - elif kind == 'remove': - body.pop(edit['key'], None) - elif kind == 'rebind': - args['complete_input_dict'] = edit['value'] - elif kind == 'rebind_then_set_retained': - args['complete_input_dict'] = {} - body[edit['key']] = edit['value'] -def check(): - assert aliased == {name: True for name in unchanged}, aliased - assert logger.names() == ['pre_call', 'post_call'], logger.calls -"; - -fn json_value() -> impl Strategy { - let leaf = prop_oneof![ - Just(Value::Null), - any::().prop_map(Value::from), - any::().prop_map(Value::from), - any::() - .prop_filter("JSON has no NaN or infinity", |number| number.is_finite()) - .prop_map(Value::from), - ".{0,8}".prop_map(Value::from), - ]; - leaf.prop_recursive(3, 24, 4, |inner| { - prop_oneof![ - prop::collection::vec(inner.clone(), 0..4).prop_map(Value::from), - prop::collection::btree_map(key(), inner, 0..4) - .prop_map(|fields| Value::Object(fields.into_iter().collect())), - ] - }) -} - -fn key() -> impl Strategy { - "[a-z]{1,6}" -} - -fn caller() -> impl Strategy { - prop_oneof![ - Just(Caller::PassedUnchanged), - Just(Caller::RewrittenByTheRoute), - Just(Caller::NotPassed), - ] -} - -fn edit() -> impl Strategy { - prop_oneof![ - Just(Edit::Nothing), - (key(), json_value()).prop_map(|(key, value)| Edit::Set(key, value)), - key().prop_map(Edit::Remove), - json_value().prop_map(Edit::Rebind), - (key(), json_value()).prop_map(|(key, value)| Edit::RebindThenSetRetained(key, value)), - ] -} - -proptest! { - #![proptest_config(ProptestConfig::with_cases(128))] - - /// For any body, any caller keywords and any callback edit: every keyword the route - /// sends unchanged reaches `pre_call` as the caller's own object, and the provider is - /// sent exactly what the model says, so a callback that edits nothing changes nothing. - #[test] - fn the_wire_is_the_body_pre_call_received_as_the_callback_left_it( - fields in prop::collection::btree_map(key(), (json_value(), caller()), 0..5), - edit in edit(), - ) { - let body: Map = fields - .iter() - .map(|(name, (value, _))| (name.clone(), value.clone())) - .collect(); - let kwargs: Map = fields - .iter() - .filter_map(|(name, (value, caller))| match caller { - Caller::PassedUnchanged => Some((name.clone(), value.clone())), - Caller::RewrittenByTheRoute => Some((name.clone(), json!([value]))), - Caller::NotPassed => None, - }) - .collect(); - let unchanged: Value = fields - .iter() - .filter(|(_, (_, caller))| *caller == Caller::PassedUnchanged) - .map(|(name, _)| Value::from(name.clone())) - .collect(); - - let wire = before_send_bound( - &[ - ("kwargs", &Value::Object(kwargs)), - ("unchanged", &unchanged), - ("edit", &edit.script()), - ], - MODEL, - json!({}), - Value::Object(body.clone()), - &[], - ); - - prop_assert_eq!(wire.body, edit.sent(&body)); - prop_assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]); - } -} diff --git a/litellm-rust/crates/callbacks-legacy-python/tests/support.rs b/litellm-rust/crates/callbacks-legacy-python/tests/support.rs deleted file mode 100644 index d0c02fa6da5..00000000000 --- a/litellm-rust/crates/callbacks-legacy-python/tests/support.rs +++ /dev/null @@ -1,205 +0,0 @@ -use std::ffi::CStr; - -use pyo3::prelude::*; -use pyo3::types::{PyDict, PyTuple}; - -use crate::{LegacyLogging, LegacySurface, PublicCall}; - -/// The parameters of every `callbacks_legacy_python` function, as the real module declares them. -/// `tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python -/// signatures, and [`namespace`] binds every fake call against it. -pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json"); - -/// Stand-ins for `callbacks_legacy_python`, the only Python module the crate calls. Tests -/// share one interpreter and run concurrently, so each fake is installed idempotently and -/// forwards to the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`). -/// Every fake is bound against the contract first, so a call the real module would reject -/// fails here too. -const STUBS: &CStr = c" -import contextvars -import inspect -import json -import sys -import traceback -import types - -for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.callbacks_legacy_python'): - sys.modules.setdefault(name, types.ModuleType(name)) - -legacy = sys.modules['litellm.rust_bridge.callbacks_legacy_python'] -CONTRACT = json.loads(python_contract) - - -def contracted(name, fake): - signature = inspect.Signature( - [inspect.Parameter(parameter, inspect.Parameter.POSITIONAL_OR_KEYWORD) for parameter in CONTRACT[name]] - ) - - def checked(*args, **kwargs): - signature.bind(*args, **kwargs) - return fake(*args, **kwargs) - - return checked - - -if not hasattr(legacy, 'is_internal'): - legacy.is_internal = contextvars.ContextVar('is_internal_call', default=False) - -FAKES = { - 'setup': lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace( - logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'], - kwargs=kwargs, - ), - 'check_limits': lambda arguments: arguments['logger'].check_limits(arguments), - 'finalize': lambda response, logger, kwargs, start, end: logger.record('finalize', response), - 'update_logging': lambda logger, kwargs, model, optional_params, litellm_params, provider: logger.update_from_kwargs( - kwargs=kwargs, - model=model, - optional_params=optional_params, - litellm_params=litellm_params, - custom_llm_provider=provider, - ), - 'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args), - 'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call( - original_response, api_key, additional_args - ), - 'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)), - 'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending), - 'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls( - response, start, end - ), - 'failure_handler': lambda logger, error, start, end, asynchronous: ( - logger.async_failure_handler if asynchronous else logger.failure_handler - )(error, ''.join(traceback.format_exception(error)), start, end), - 'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)), - 'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end), - 'enqueue_logging': lambda coroutine: coroutine.enqueue(), - 'restore_context': lambda logger: logger.record('restore', None), - 'custom_pricing_fields': lambda: ('ocr_cost_per_page',), - 'is_internal_call': lambda: legacy.is_internal.get(), - 'credential_list': lambda: [], - 'warn_unknown_credential': lambda name, loaded: None, - 'before_deployment_call': lambda kwargs, call_type: kwargs['logger'].hook('pre', kwargs, call_type), - 'after_deployment_success': lambda kwargs, response, call_type: kwargs['logger'].hook( - 'success', response, call_type - ), - 'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type), - 'stream_opened': lambda logger: logger.record('stream_opened', None), - 'stream_success': lambda logger, request_body, chunks, start, end, first_chunk: logger.record( - 'stream_success', list(chunks) - ), - 'stream_failure': lambda logger, request_body, chunks, error: logger.record('stream_failure', error), -} -assert FAKES.keys() == CONTRACT.keys(), sorted(FAKES.keys() ^ CONTRACT.keys()) -for name, fake in FAKES.items(): - setattr(legacy, name, contracted(name, fake)) - - -unraisable = sys.modules.setdefault( - 'litellm_test_unraisable', types.ModuleType('litellm_test_unraisable') -) -if not hasattr(unraisable, 'events'): - unraisable.events = [] - sys.unraisablehook = lambda event: unraisable.events.append((event.object, event.exc_value)) - - -def unraisable_from(owner): - return [error for source, error in unraisable.events if source is owner] - - -class StubCoroutine: - def __init__(self, logger): - self.logger = logger - - def enqueue(self): - self.logger.record('enqueued', None) - self.logger.on_enqueue(self) - - def close(self): - self.logger.record('closed', None) - - -class StubLogger: - def __init__(self): - self.calls = [] - self.hooks = {} - self.on_enqueue = lambda coroutine: None - - def record(self, name, value): - self.calls.append((name, value)) - - def names(self): - return [name for name, _ in self.calls] - - def hook(self, phase, value, call_type): - self.record(phase + '_hook', call_type) - return self.hooks.get(phase, lambda value: 'awaitable')(value) - - def check_limits(self, arguments): - self.record('check_limits', arguments) - - def failure_handler(self, error, trace, start, end): - self.record('failure_handler', error) - - def async_failure_handler(self, error, trace, start, end): - self.record('async_failure_handler', error) - return 'awaitable' - - def success_handler(self, response, start, end): - self.record('success_handler', response) - - def async_success_handler(self, response, start, end): - self.record('async_success_handler', response) - return StubCoroutine(self) - - def handle_sync_success_callbacks_for_async_calls(self, response, start, end): - self.record('sync_success_for_async_call', response) - - -logger = StubLogger() -"; - -/// A namespace with the stubs, `StubLogger` and a fresh `logger`, after `script` ran in it. -pub(crate) fn namespace<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { - let locals = PyDict::new(py); - locals.set_item("python_contract", PYTHON_CONTRACT).unwrap(); - py.run(STUBS, Some(&locals), Some(&locals)).unwrap(); - py.run(script, Some(&locals), Some(&locals)).unwrap(); - locals -} - -pub(crate) fn run(py: Python<'_>, locals: &Bound<'_, PyDict>, code: &CStr) { - py.run(code, Some(locals), Some(locals)).unwrap(); -} - -pub(crate) fn local<'py>(locals: &Bound<'py, PyDict>, name: &str) -> Bound<'py, PyAny> { - locals.get_item(name).unwrap().unwrap() -} - -/// A legacy call over the namespace's `kwargs` (or none) and `request` (or `None`). -pub(crate) fn legacy_call( - py: Python<'_>, - locals: &Bound<'_, PyDict>, - asynchronous: bool, -) -> LegacyLogging { - let request = locals - .get_item("request") - .unwrap() - .unwrap_or_else(|| py.None().into_bound(py)); - let kwargs = locals - .get_item("kwargs") - .unwrap() - .map(|kwargs| kwargs.cast_into::().unwrap()) - .unwrap_or_else(|| PyDict::new(py)); - let call = PublicCall::capture(&request, &PyTuple::empty(py), &kwargs).unwrap(); - LegacyLogging::new( - py, - LegacySurface { - call_type: "test", - input_description: "test input", - stream: None, - }, - call, - asynchronous, - ) -} diff --git a/litellm-rust/crates/callbacks-legacy-python/tests/terminal.rs b/litellm-rust/crates/callbacks-legacy-python/tests/terminal.rs deleted file mode 100644 index f68209233f2..00000000000 --- a/litellm-rust/crates/callbacks-legacy-python/tests/terminal.rs +++ /dev/null @@ -1,291 +0,0 @@ -use std::ffi::CStr; - -use litellm_host::event::{FailureOrigin, Timing}; -use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; -use pyo3::exceptions::PyRuntimeError; -use pyo3::exceptions::asyncio::CancelledError; -use pyo3::prelude::*; -use pyo3::types::PyDict; -use rstest::rstest; - -use super::LegacyLogging; -use crate::PythonLogger; -use crate::test_support::{legacy_call, local, namespace, run}; - -const TIMING: Timing = Timing { - start_time: 0.0, - end_time: 1.0, -}; - -fn logged(py: Python<'_>, locals: &Bound<'_, PyDict>, asynchronous: bool) -> LegacyLogging { - LegacyLogging { - logger: Some(PythonLogger::new(local(locals, "logger").unbind())), - ..legacy_call(py, locals, asynchronous) - } -} - -fn succeed( - py: Python<'_>, - locals: &Bound<'_, PyDict>, - logging: &mut LegacyLogging, -) -> LifecycleStep { - let response = local(locals, "response").unbind(); - logging - .emit( - py, - LifecycleEvent::Succeeded { - timing: TIMING, - response: &response, - }, - ) - .unwrap() -} - -fn fail(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) -> LifecycleStep { - let failure = PyErr::from_value(local(locals, "failure")); - logging - .emit( - py, - LifecycleEvent::Failed { - timing: TIMING, - origin: FailureOrigin::Host, - error: &failure, - }, - ) - .unwrap() -} - -#[rstest] -#[case::sync_listened(false, c"", &["submit"])] -#[case::async_listened( - true, - c"", - &["async_success_handler", "enqueued", "sync_success_for_async_call"] -)] -#[case::async_deferred(true, c"logger._defer_async_logging = True", &["sync_success_for_async_call"])] -#[case::async_with_fallbacks(true, c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])] -fn success_reaches_the_logging_handlers( - #[case] asynchronous: bool, - #[case] script: &CStr, - #[case] expected: &[&str], -) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, c"response = object()"); - run(py, &locals, script); - let mut logging = logged(py, &locals, asynchronous); - assert!(matches!( - succeed(py, &locals, &mut logging), - LifecycleStep::Done - )); - let names: Vec = local(&locals, "logger") - .call_method0("names") - .unwrap() - .extract() - .unwrap(); - assert_eq!(names, expected); - run( - py, - &locals, - c" -assert all(value is response for name, value in logger.calls if name.endswith('_handler')) -assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_async_logging', False) -", - ); - }); -} - -#[rstest] -#[case::synchronous(false, &["failure_handler"])] -#[case::asynchronous(true, &[])] -fn internal_calls_skip_failure_callbacks_only_when_asynchronous( - #[case] asynchronous: bool, - #[case] expected: &[&str], -) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, c"failure = ValueError('provider')"); - let mut logging = LegacyLogging { - internal: true, - ..logged(py, &locals, asynchronous) - }; - assert!(matches!( - fail(py, &locals, &mut logging), - LifecycleStep::Done - )); - let names: Vec = local(&locals, "logger") - .call_method0("names") - .unwrap() - .extract() - .unwrap(); - assert_eq!(names, expected); - }); -} - -#[test] -fn internal_async_calls_skip_the_async_success_fan_out() { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, c"response = object()"); - let mut logging = LegacyLogging { - internal: true, - ..logged(py, &locals, true) - }; - succeed(py, &locals, &mut logging); - run( - py, - &locals, - c"assert logger.names() == ['sync_success_for_async_call'], logger.calls", - ); - }); -} - -#[test] -fn a_failing_success_callback_is_reported_without_replacing_the_response() { - Python::initialize(); - Python::attach(|py| { - let locals = namespace( - py, - c" -response = object() -failure = ValueError('terminal diagnostic') - -class FailingLogger(StubLogger): - def handle_sync_success_callbacks_for_async_calls(self, *args): - raise failure - -logger = FailingLogger() -", - ); - let mut logging = logged(py, &locals, true); - assert!(matches!( - succeed(py, &locals, &mut logging), - LifecycleStep::Done - )); - assert!( - logging - .response - .as_ref() - .unwrap() - .bind(py) - .is(local(&locals, "response")) - ); - run(py, &locals, c"assert unraisable_from(logger) == [failure]"); - }); -} - -#[rstest] -#[case::sync_listened(false, c"", &["failure_handler"])] -#[case::async_listened(true, c"", &["failure_handler", "async_failure_handler"])] -fn failure_reaches_the_logging_handlers( - #[case] asynchronous: bool, - #[case] script: &CStr, - #[case] expected: &[&str], -) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, c"failure = ValueError('provider')"); - run(py, &locals, script); - let mut logging = logged(py, &locals, asynchronous); - let step = fail(py, &locals, &mut logging); - let awaits_async_handler = expected.contains(&"async_failure_handler"); - assert_eq!( - matches!(step, LifecycleStep::Await(_)), - awaits_async_handler - ); - let names: Vec = local(&locals, "logger") - .call_method0("names") - .unwrap() - .extract() - .unwrap(); - assert_eq!(names, expected); - run( - py, - &locals, - c"assert all(value is failure for name, value in logger.calls if name.endswith('_handler'))", - ); - }); -} - -#[test] -fn a_failing_sync_failure_callback_keeps_the_error_and_still_runs_the_async_family() { - Python::initialize(); - Python::attach(|py| { - let locals = namespace( - py, - c" -failure = ValueError('selected') - -class FailingLogger(StubLogger): - def failure_handler(self, error, trace, start, end): - self.record('failure_handler', error) - raise RuntimeError('handler failed') - -logger = FailingLogger() -", - ); - let mut logging = logged(py, &locals, true); - assert!(matches!( - fail(py, &locals, &mut logging), - LifecycleStep::Await(_) - )); - assert!( - logging - .error - .as_ref() - .unwrap() - .bind(py) - .is(local(&locals, "failure")) - ); - run( - py, - &locals, - c"assert logger.names() == ['failure_handler', 'async_failure_handler'], logger.calls", - ); - }); -} - -#[rstest] -#[case::completed(None, true)] -#[case::handler_error(Some(false), true)] -#[case::cancelled(Some(true), false)] -fn the_async_failure_handler_ends_the_call_unless_it_was_cancelled( - #[case] error: Option, - #[case] done: bool, -) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, c"failure = ValueError('provider')"); - let mut logging = logged(py, &locals, true); - fail(py, &locals, &mut logging); - let result = match error { - None => Ok(py.None()), - Some(false) => Err(PyRuntimeError::new_err("handler failed")), - Some(true) => Err(CancelledError::new_err("cancelled")), - }; - let expected = result.as_ref().err().map(|error| error.value(py).clone()); - match logging.resume(py, result) { - Ok(step) => assert!(done && matches!(step, LifecycleStep::Done)), - Err(propagated) => { - assert!(!done); - assert!(propagated.value(py).is(expected.unwrap())); - } - } - }); -} - -#[test] -fn closing_restores_the_correlation_context_once() { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, c""); - let mut logging = logged(py, &locals, true); - logging.close(py); - logging.close(py); - run( - py, - &locals, - c"assert logger.names() == ['restore'], logger.calls", - ); - }); -} diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 4626781f3e3..d7096cdd774 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -4,7 +4,6 @@ version = "0.1.0" edition.workspace = true license.workspace = true repository.workspace = true -autotests = false [dependencies] litellm-secrets.workspace = true diff --git a/litellm-rust/crates/core/src/audio_transcription/mod.rs b/litellm-rust/crates/core/src/audio_transcription/mod.rs index af9c398c065..801fd5e9673 100644 --- a/litellm-rust/crates/core/src/audio_transcription/mod.rs +++ b/litellm-rust/crates/core/src/audio_transcription/mod.rs @@ -14,6 +14,3 @@ pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Resu execute_audio_transcription_provider_call(prepare_audio_transcription_provider_call(request)?) .await } - -#[cfg(test)] -mod tests; diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs index 81d35044d08..224c9d8cfed 100644 --- a/litellm-rust/crates/core/src/chat_completions/mod.rs +++ b/litellm-rust/crates/core/src/chat_completions/mod.rs @@ -51,6 +51,3 @@ pub fn chat_completions_decline_reason( .unsupported_reason(&messages, optional_params) .map(|reason| reason.0) } - -#[cfg(test)] -mod tests; diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index c8e6365121e..afea46221f5 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -143,3 +143,841 @@ pub(super) fn prepare_provider_request( timeout: request.timeout, }) } + +#[cfg(test)] +mod tests { + use litellm_llms::base_llm::chat::transformation::RequestAuth; + use serde_json::{Map, Value, json}; + + use super::{prepare_provider_request, resolve_request}; + use crate::chat_completions::{ + Error, + types::{ChatCompletionsRequest, ProviderChatCompletionsRequest}, + }; + + fn prepare_chat_completions_call( + request: ChatCompletionsRequest<'_>, + ) -> Result { + prepare_provider_request(resolve_request(request)?) + } + + fn request<'a>( + model: &'a str, + provider: Option<&'a str>, + messages: Value, + optional_params: Value, + ) -> ChatCompletionsRequest<'a> { + ChatCompletionsRequest { + model, + messages, + optional_params: match optional_params { + Value::Object(map) => map, + other => panic!("params must be an object, got {other}"), + }, + api_key: Some("sk-test"), + api_base: None, + custom_llm_provider: provider, + extra_headers: None, + timeout: None, + } + } + + /// `ProviderChatCompletionsRequest` deliberately has no `Debug` (its headers + /// carry resolved credentials), so unwrap the failure case by hand. + fn decline(request: ChatCompletionsRequest<'_>) -> Error { + match prepare_chat_completions_call(request) { + Err(error) => error, + Ok(prepared) => panic!("expected a decline, prepared a call to {}", prepared.url), + } + } + + #[test] + fn resolves_the_provider_from_the_model_prefix() { + let prepared = prepare_chat_completions_call(request( + "anthropic/claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + )) + .expect("prepares"); + assert_eq!(prepared.model, "claude-sonnet-4-5"); + assert_eq!(prepared.url, "https://api.anthropic.com/v1/messages"); + assert_eq!(prepared.body["model"], json!("claude-sonnet-4-5")); + } + + #[test] + fn strips_an_explicit_provider_prefix_from_the_model() { + let prepared = prepare_chat_completions_call(request( + "anthropic/claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + )) + .expect("prepares"); + assert_eq!(prepared.model, "claude-sonnet-4-5"); + } + + #[test] + fn adds_the_auth_and_default_headers() { + let prepared = prepare_chat_completions_call(request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + )) + .expect("prepares"); + assert!( + prepared + .upstream_headers + .contains(&("x-api-key".to_string(), "sk-test".to_string())) + ); + assert!( + prepared + .upstream_headers + .contains(&("anthropic-version".to_string(), "2023-06-01".to_string())) + ); + assert!(matches!( + prepared.auth, + RequestAuth::Header { + name: "x-api-key", + .. + } + )); + } + + #[test] + fn the_deployment_credential_replaces_a_caller_supplied_auth_header() { + // Python builds `{**headers, **anthropic_headers}`, so the deployment's key + // overwrites a forwarded one. Honouring the caller's would let whoever sends + // the request choose the Anthropic principal it bills to. + let mut call = request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + ); + call.extra_headers = Some(Map::from_iter([( + "X-Api-Key".to_string(), + json!("sk-caller"), + )])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let keys: Vec<_> = prepared + .upstream_headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) + .collect(); + assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers); + assert_eq!(keys[0].1, "sk-test"); + } + + #[test] + fn a_forwarded_authorization_header_suppresses_the_resolved_api_key_header() { + // Anthropic's `validate_environment` pops `x-api-key` and sets `authorization` + // for an OAuth token, so re-adding the key here would put the credential into + // a header the host removed on purpose. + let mut call = request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + ); + call.extra_headers = Some(Map::from_iter([ + ( + "Authorization".to_string(), + json!("Bearer sk-ant-oat01-token"), + ), + ("X-Api-Key".to_string(), json!("sk-caller")), + ])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + assert!( + !prepared + .upstream_headers + .iter() + .any(|(name, value)| name.eq_ignore_ascii_case("x-api-key") && value == "sk-test"), + "the resolved key must not be applied over an OAuth bearer, got {:?}", + prepared.upstream_headers + ); + assert!( + prepared + .upstream_headers + .iter() + .any(|(name, value)| name.eq_ignore_ascii_case("authorization") + && value == "Bearer sk-ant-oat01-token") + ); + } + + #[test] + fn an_unrelated_forwarded_authorization_does_not_defer_the_resolved_key() { + // Only an OAuth bearer replaces the credential. Python sends the deployment's + // `x-api-key` alongside any other forwarded `authorization`, so deferring on + // the mere presence of that header would drop the deployment's auth. + let mut call = request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + ); + call.extra_headers = Some(Map::from_iter([ + ("Authorization".to_string(), json!("Bearer unrelated")), + ("X-Api-Key".to_string(), json!("sk-caller")), + ])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let keys: Vec<_> = prepared + .upstream_headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) + .collect(); + assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers); + assert_eq!(keys[0].1, "sk-test"); + assert!( + prepared + .upstream_headers + .iter() + .any(|(name, value)| name.eq_ignore_ascii_case("authorization") + && value == "Bearer unrelated"), + "the unrelated authorization must survive, got {:?}", + prepared.upstream_headers + ); + } + + #[test] + fn declines_an_unsupported_request_before_resolving_credentials() { + let mut call = request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({"stream": true}), + ); + call.api_key = None; + // No api_key is set and no env is consulted: the gate must run first, so the + // error is the decline rather than a missing-credential error. + assert_eq!(decline(call), Error::Unsupported("streaming")); + } + + #[test] + fn rejects_an_unknown_provider() { + assert_eq!( + decline(request( + "openai/gpt-4o", + None, + json!([{"role": "user", "content": "hi"}]), + json!({}), + )), + Error::InvalidProvider("openai".to_string()) + ); + } + + #[test] + fn rejects_a_model_with_no_resolvable_provider() { + assert!(matches!( + decline(request( + "claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({}), + )), + Error::InvalidProvider(_) + )); + } + + #[test] + fn rejects_an_empty_or_malformed_message_list() { + assert_eq!( + decline(request( + "anthropic/claude-sonnet-4-5", + None, + json!([]), + json!({}), + )), + Error::InvalidRequest("chat completions requires at least one message".to_string()) + ); + assert!(matches!( + decline(request( + "anthropic/claude-sonnet-4-5", + None, + json!("not a list"), + json!({}), + )), + Error::InvalidRequest(_) + )); + } + + #[test] + fn rejects_non_string_extra_headers() { + let mut call = request( + "anthropic/claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({}), + ); + call.extra_headers = Some(Map::from_iter([("x-trace".to_string(), json!(7))])); + assert_eq!( + decline(call), + Error::Headers(litellm_http::request::HeaderError { + context: "chat completions", + name: "x-trace".to_string(), + actual: "number", + }) + ); + } + + #[test] + fn prepares_a_bedrock_call_without_resolving_credentials() { + let mut call = request( + "bedrock/us-east-1/anthropic.claude-v2", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"maxTokens": 16}), + ); + call.api_key = None; + let prepared = prepare_chat_completions_call(call).expect("prepares"); + assert_eq!( + prepared.url, + "https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/converse" + ); + assert_eq!( + prepared.auth, + RequestAuth::AwsSigV4 { + region: "us-east-1".to_string(), + service: "bedrock", + } + ); + // SigV4 signs the serialized body, so prepare must not have added an + // Authorization header; the handler does it. + assert!( + !prepared + .upstream_headers + .iter() + .any(|(name, _)| name.eq_ignore_ascii_case("authorization")) + ); + assert_eq!(prepared.body["inferenceConfig"], json!({"maxTokens": 16})); + } + + #[tokio::test] + async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() { + // Python signs only the AWS header set and reattaches the rest, so a header + // the caller forwarded rides along without joining the canonical request. + // Signing it makes Converse 403 on a deployment that works on Python. + let mut call = request( + "bedrock/us-east-1/anthropic.claude-v2", + None, + json!([{"role": "user", "content": "hi"}]), + json!({ + "maxTokens": 16, + "aws_access_key_id": "AKIDEXAMPLE", + "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY" + }), + ); + // A key would resolve to a bearer token and never reach the signer. + call.api_key = None; + call.extra_headers = Some(Map::from_iter([( + "x-request-id".to_string(), + json!("abc-123"), + )])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let signed = crate::chat_completions::handler::outbound_request(&prepared) + .await + .expect("signs"); + + let authorization = signed + .header("authorization") + .expect("carries an authorization header") + .to_string(); + assert!( + authorization.starts_with("AWS4-HMAC-SHA256"), + "expected a SigV4 signature, got {authorization}" + ); + assert!( + !authorization.contains("x-request-id"), + "forwarded header reached SignedHeaders: {authorization}" + ); + // It still goes on the wire, it is just not part of the signature. + assert!( + signed + .headers() + .iter() + .any(|(name, value)| name == "x-request-id" && value == "abc-123"), + "forwarded header was dropped instead of reattached" + ); + } + + #[tokio::test] + async fn a_forwarded_header_the_signer_computes_declines_to_python() { + // Reattaching the caller's copy next to the computed one puts the name on + // the wire twice and Bedrock rejects the pair, so a request carrying one + // has to go to Python instead of being signed here. + for forwarded in [ + "Authorization", + "x-amz-date", + "x-amz-security-token", + "Date", + ] { + let mut call = request( + "bedrock/us-east-1/anthropic.claude-v2", + None, + json!([{"role": "user", "content": "hi"}]), + json!({ + "maxTokens": 16, + "aws_access_key_id": "AKIDEXAMPLE", + "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY" + }), + ); + call.api_key = None; + call.extra_headers = Some(Map::from_iter([(forwarded.to_string(), json!("forged"))])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let error = crate::chat_completions::handler::outbound_request(&prepared) + .await + .expect_err("{forwarded} should decline instead of being signed"); + assert!( + matches!(error, Error::Unsupported(_)), + "{forwarded} declined as {error:?}, which the host would not fall back on" + ); + } + } + + #[test] + fn a_bedrock_deployment_bearer_outranks_a_forwarded_authorization() { + // `get_request_headers` assigns `headers["Authorization"]` unconditionally + // once a bearer token resolves, so the deployment's identity wins on + // Python. Keeping the caller's would authorize and bill the call as a + // different principal, and only when the deployment carries `rust: true`. + let mut call = request( + "bedrock/us-east-1/anthropic.claude-v2", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"maxTokens": 16}), + ); + call.extra_headers = Some(Map::from_iter([( + "Authorization".to_string(), + json!("Bearer caller-supplied"), + )])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let authorizations: Vec<_> = prepared + .upstream_headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("authorization")) + .map(|(_, value)| value.as_str()) + .collect(); + assert_eq!( + authorizations, + vec!["Bearer sk-test"], + "the deployment token must be the only authorization on the wire" + ); + } + + #[test] + fn an_anthropic_forwarded_oauth_bearer_still_outranks_the_resolved_key() { + // The opposite precedence, and deliberate: Anthropic's own transform + // honours a forwarded OAuth bearer, so the Bedrock fix above must not be + // generalized into a rule that the configured key always wins. + // + // An OAuth bearer is the whole of that exception. This forwarded a plain + // `x-api-key` until round 17, which read as the same claim and was not: + // Python overwrites a forwarded `x-api-key` with the deployment's. + let mut call = request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + ); + call.extra_headers = Some(Map::from_iter([( + "authorization".to_string(), + json!("Bearer sk-ant-oat01-forwarded"), + )])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let keys: Vec<_> = prepared + .upstream_headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) + .map(|(_, value)| value.as_str()) + .collect(); + assert!(keys.is_empty(), "got {:?}", prepared.upstream_headers); + assert!( + prepared + .upstream_headers + .iter() + .any(|(name, value)| name.eq_ignore_ascii_case("authorization") + && value == "Bearer sk-ant-oat01-forwarded") + ); + } + + #[test] + fn a_bedrock_api_key_is_sent_as_a_bearer_token_instead_of_being_signed() { + // The configured bearer identity has its own account and quota boundary, + // so a request carrying one must not be signed as whatever principal the + // host's AWS credentials resolve to. + let prepared = prepare_chat_completions_call(request( + "bedrock/us-east-1/anthropic.claude-v2", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"maxTokens": 16}), + )) + .expect("prepares"); + assert_eq!( + prepared.auth, + RequestAuth::Bearer { + token: "sk-test".to_string() + } + ); + assert!( + prepared + .upstream_headers + .iter() + .any(|(name, value)| name.eq_ignore_ascii_case("authorization") + && value == "Bearer sk-test"), + "prepare did not carry the bearer token" + ); + } + + fn decline_reason( + model: &str, + provider: Option<&str>, + messages: Value, + params: Value, + ) -> Option<&'static str> { + let params = match params { + Value::Object(map) => map, + other => panic!("params must be an object, got {other}"), + }; + crate::chat_completions::chat_completions_decline_reason(model, provider, messages, ¶ms) + } + + #[test] + fn the_gate_accepts_what_prepare_accepts() { + assert_eq!( + decline_reason( + "anthropic/claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + ), + None + ); + } + + #[test] + fn the_gate_declines_without_resolving_credentials_or_calling_out() { + assert_eq!( + decline_reason( + "anthropic/claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"stream": true}), + ), + Some("streaming") + ); + assert_eq!( + decline_reason( + "openai/gpt-4o", + None, + json!([{"role": "user", "content": "hi"}]), + json!({}), + ), + Some("provider is not on the rust chat completions path") + ); + assert_eq!( + decline_reason( + "claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({}), + ), + Some("provider is not on the rust chat completions path") + ); + assert_eq!( + decline_reason( + "anthropic/claude-sonnet-4-5", + None, + json!("nope"), + json!({}) + ), + Some("unreadable message list") + ); + assert_eq!( + decline_reason("anthropic/claude-sonnet-4-5", None, json!([]), json!({})), + Some("empty message list") + ); + } + + #[test] + fn the_gate_agrees_with_prepare_on_every_case_it_accepts() { + // A gate that accepts what prepare then declines would make the host emit + // its pre-call logging on a path that falls back, so pin the agreement. + for (messages, params) in [ + ( + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 8}), + ), + ( + json!([{"role": "system", "content": "s"}, {"role": "user", "content": "hi"}]), + json!({"temperature": 0.1}), + ), + ( + json!([{"role": "user", "content": "hi"}, {"role": "assistant", "content": "yo"}]), + json!({}), + ), + ] { + assert_eq!( + decline_reason( + "anthropic/claude-sonnet-4-5", + None, + messages.clone(), + params.clone() + ), + None, + "gate declined {messages}" + ); + prepare_chat_completions_call(request( + "anthropic/claude-sonnet-4-5", + None, + messages.clone(), + params, + )) + .unwrap_or_else(|error| panic!("prepare declined {messages}: {error}")); + } + } + + mod round_trip { + use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::{TcpListener, TcpStream}, + }; + + use super::*; + use crate::chat_completions::chat_completions; + + async fn read_http_request(socket: &mut TcpStream) -> String { + let mut request = Vec::new(); + let mut buffer = [0_u8; 1024]; + let header_end = loop { + let n = socket.read(&mut buffer).await.expect("reads request"); + if n == 0 { + break request.len(); + } + request.extend_from_slice(&buffer[..n]); + if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") + { + break position + 4; + } + }; + let headers = String::from_utf8_lossy(&request[..header_end]); + let content_length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) + .unwrap_or(0); + while request.len().saturating_sub(header_end) < content_length { + let n = socket.read(&mut buffer).await.expect("reads body"); + if n == 0 { + break; + } + request.extend_from_slice(&buffer[..n]); + } + String::from_utf8(request).expect("request is utf8") + } + + fn http_response(status: &str, body: &str) -> String { + format!( + "HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + body.len(), + body + ) + } + + /// Serve one request from a stub upstream and hand back what it received. + async fn serve_once( + status: &'static str, + body: &'static str, + ) -> (String, tokio::task::JoinHandle) { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let port = listener.local_addr().expect("addr").port(); + let handle = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts"); + let received = read_http_request(&mut socket).await; + socket + .write_all(http_response(status, body).as_bytes()) + .await + .expect("writes response"); + socket.flush().await.expect("flushes"); + received + }); + (format!("http://127.0.0.1:{port}/v1/messages"), handle) + } + + fn call(api_base: &str, messages: Value, params: Value) -> ChatCompletionsRequest<'_> { + ChatCompletionsRequest { + model: "anthropic/claude-sonnet-4-5", + messages, + optional_params: match params { + Value::Object(map) => map, + other => panic!("params must be an object, got {other}"), + }, + api_key: Some("sk-test"), + api_base: Some(api_base), + custom_llm_provider: None, + extra_headers: None, + timeout: Some(std::time::Duration::from_secs(10)), + } + } + + const GOOD_BODY: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5-20260101","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#; + + #[tokio::test] + async fn round_trip_sends_the_translated_body_and_normalizes_the_response() { + let (api_base, handle) = serve_once("200 OK", GOOD_BODY).await; + let response = chat_completions(call( + &api_base, + json!([ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "hi"} + ]), + json!({"max_tokens": 16}), + )) + .await + .expect("call succeeds"); + + let received = handle.await.expect("server task"); + let sent: Value = serde_json::from_str( + received + .split_once("\r\n\r\n") + .expect("request has a body") + .1, + ) + .expect("body is json"); + assert_eq!( + sent["messages"], + json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]) + ); + assert_eq!( + sent["system"], + json!([{"type": "text", "text": "be terse"}]) + ); + assert_eq!(sent["max_tokens"], json!(16)); + assert!(received.to_lowercase().contains("x-api-key: sk-test")); + + assert_eq!( + response.choices[0].message.content.as_deref(), + Some("hello") + ); + assert_eq!(response.usage.total_tokens, 15); + } + + #[tokio::test] + async fn a_response_it_cannot_normalize_is_reported_as_already_sent() { + // The provider was called and billed, so the host must not retry this + // on its own path. `MissingField` here would read as a pre-send + // decline and be retried; `InvalidResponse` cannot. + const NO_USAGE: &str = + r#"{"model":"m","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}"#; + let (api_base, handle) = serve_once("200 OK", NO_USAGE).await; + let err = chat_completions(call( + &api_base, + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + )) + .await + .expect_err("response cannot be normalized"); + handle.await.expect("server task"); + assert!( + matches!(err, Error::InvalidResponse(_)), + "expected a post-send error, got {err:?}" + ); + } + + #[tokio::test] + async fn a_tool_use_block_in_the_response_is_also_reported_as_already_sent() { + const TOOL_USE: &str = r#"{"model":"m","content":[{"type":"tool_use","id":"t","name":"f","input":{}}],"stop_reason":"tool_use","usage":{"input_tokens":1,"output_tokens":1}}"#; + let (api_base, handle) = serve_once("200 OK", TOOL_USE).await; + let err = chat_completions(call( + &api_base, + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + )) + .await + .expect_err("response cannot be normalized"); + handle.await.expect("server task"); + assert!( + matches!(err, Error::InvalidResponse(_)), + "expected a post-send error, got {err:?}" + ); + } + + #[tokio::test] + async fn an_upstream_error_status_keeps_its_code() { + let (api_base, handle) = + serve_once("429 Too Many Requests", r#"{"error":"slow down"}"#).await; + let err = chat_completions(call( + &api_base, + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + )) + .await + .expect_err("upstream rejects"); + handle.await.expect("server task"); + assert!( + matches!( + err, + Error::Transport(litellm_http::transport::Error::Http { status: 429, .. }) + ), + "expected a 429, got {err:?}" + ); + } + + #[tokio::test] + async fn a_connection_that_is_never_established_declines_instead_of_failing() { + // Nothing was sent, so nothing was billed and the host can still serve + // the request. Classing this with the post-send failures would turn a + // recoverable fallback into a user-facing error on exactly the + // deployments whose transport is configured only on the Python client. + let port = { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + listener.local_addr().expect("has an address").port() + // Dropped here, so the port is closed and the connect is refused. + }; + let err = chat_completions(call( + &format!("http://127.0.0.1:{port}/v1/messages"), + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + )) + .await + .expect_err("nothing is listening"); + assert!( + matches!( + err, + Error::Transport(litellm_http::transport::Error::Connect(_)) + ), + "expected a pre-send connect failure, got {err:?}" + ); + } + + #[test] + fn response_errors_collapse_to_one_variant_that_can_only_mean_already_sent() { + use crate::chat_completions::handler::as_response_error; + + for original in [ + Error::MissingField("usage"), + Error::Unsupported("non-text response content block"), + Error::InvalidRequest("whatever".to_string()), + Error::Auth(litellm_auth::Error::InvalidHeader), + ] { + let label = format!("{original:?}"); + assert!( + matches!(as_response_error(original), Error::InvalidResponse(_)), + "{label} must not stay retryable once the provider has answered" + ); + } + // An upstream status is already unambiguous, so it survives intact. + assert!(matches!( + as_response_error(Error::Transport(litellm_http::transport::Error::Http { + status: 500, + body: "boom".to_string() + })), + Error::Transport(litellm_http::transport::Error::Http { status: 500, .. }) + )); + } + } +} diff --git a/litellm-rust/crates/core/src/chat_completions/tests.rs b/litellm-rust/crates/core/src/chat_completions/tests.rs deleted file mode 100644 index dd5938cf168..00000000000 --- a/litellm-rust/crates/core/src/chat_completions/tests.rs +++ /dev/null @@ -1,833 +0,0 @@ -use litellm_llms::base_llm::chat::transformation::RequestAuth; -use serde_json::{Map, Value, json}; - -use super::{ - Error, - prepare::{prepare_provider_request, resolve_request}, -}; -use crate::chat_completions::types::{ChatCompletionsRequest, ProviderChatCompletionsRequest}; - -fn prepare_chat_completions_call( - request: ChatCompletionsRequest<'_>, -) -> Result { - prepare_provider_request(resolve_request(request)?) -} - -fn request<'a>( - model: &'a str, - provider: Option<&'a str>, - messages: Value, - optional_params: Value, -) -> ChatCompletionsRequest<'a> { - ChatCompletionsRequest { - model, - messages, - optional_params: match optional_params { - Value::Object(map) => map, - other => panic!("params must be an object, got {other}"), - }, - api_key: Some("sk-test"), - api_base: None, - custom_llm_provider: provider, - extra_headers: None, - timeout: None, - } -} - -/// `ProviderChatCompletionsRequest` deliberately has no `Debug` (its headers -/// carry resolved credentials), so unwrap the failure case by hand. -fn decline(request: ChatCompletionsRequest<'_>) -> Error { - match prepare_chat_completions_call(request) { - Err(error) => error, - Ok(prepared) => panic!("expected a decline, prepared a call to {}", prepared.url), - } -} - -#[test] -fn resolves_the_provider_from_the_model_prefix() { - let prepared = prepare_chat_completions_call(request( - "anthropic/claude-sonnet-4-5", - None, - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - )) - .expect("prepares"); - assert_eq!(prepared.model, "claude-sonnet-4-5"); - assert_eq!(prepared.url, "https://api.anthropic.com/v1/messages"); - assert_eq!(prepared.body["model"], json!("claude-sonnet-4-5")); -} - -#[test] -fn strips_an_explicit_provider_prefix_from_the_model() { - let prepared = prepare_chat_completions_call(request( - "anthropic/claude-sonnet-4-5", - Some("anthropic"), - json!([{"role": "user", "content": "hi"}]), - json!({}), - )) - .expect("prepares"); - assert_eq!(prepared.model, "claude-sonnet-4-5"); -} - -#[test] -fn adds_the_auth_and_default_headers() { - let prepared = prepare_chat_completions_call(request( - "claude-sonnet-4-5", - Some("anthropic"), - json!([{"role": "user", "content": "hi"}]), - json!({}), - )) - .expect("prepares"); - assert!( - prepared - .upstream_headers - .contains(&("x-api-key".to_string(), "sk-test".to_string())) - ); - assert!( - prepared - .upstream_headers - .contains(&("anthropic-version".to_string(), "2023-06-01".to_string())) - ); - assert!(matches!( - prepared.auth, - RequestAuth::Header { - name: "x-api-key", - .. - } - )); -} - -#[test] -fn the_deployment_credential_replaces_a_caller_supplied_auth_header() { - // Python builds `{**headers, **anthropic_headers}`, so the deployment's key - // overwrites a forwarded one. Honouring the caller's would let whoever sends - // the request choose the Anthropic principal it bills to. - let mut call = request( - "claude-sonnet-4-5", - Some("anthropic"), - json!([{"role": "user", "content": "hi"}]), - json!({}), - ); - call.extra_headers = Some(Map::from_iter([( - "X-Api-Key".to_string(), - json!("sk-caller"), - )])); - let prepared = prepare_chat_completions_call(call).expect("prepares"); - let keys: Vec<_> = prepared - .upstream_headers - .iter() - .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) - .collect(); - assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers); - assert_eq!(keys[0].1, "sk-test"); -} - -#[test] -fn a_forwarded_authorization_header_suppresses_the_resolved_api_key_header() { - // Anthropic's `validate_environment` pops `x-api-key` and sets `authorization` - // for an OAuth token, so re-adding the key here would put the credential into - // a header the host removed on purpose. - let mut call = request( - "claude-sonnet-4-5", - Some("anthropic"), - json!([{"role": "user", "content": "hi"}]), - json!({}), - ); - call.extra_headers = Some(Map::from_iter([ - ( - "Authorization".to_string(), - json!("Bearer sk-ant-oat01-token"), - ), - ("X-Api-Key".to_string(), json!("sk-caller")), - ])); - let prepared = prepare_chat_completions_call(call).expect("prepares"); - assert!( - !prepared - .upstream_headers - .iter() - .any(|(name, value)| name.eq_ignore_ascii_case("x-api-key") && value == "sk-test"), - "the resolved key must not be applied over an OAuth bearer, got {:?}", - prepared.upstream_headers - ); - assert!( - prepared - .upstream_headers - .iter() - .any(|(name, value)| name.eq_ignore_ascii_case("authorization") - && value == "Bearer sk-ant-oat01-token") - ); -} - -#[test] -fn an_unrelated_forwarded_authorization_does_not_defer_the_resolved_key() { - // Only an OAuth bearer replaces the credential. Python sends the deployment's - // `x-api-key` alongside any other forwarded `authorization`, so deferring on - // the mere presence of that header would drop the deployment's auth. - let mut call = request( - "claude-sonnet-4-5", - Some("anthropic"), - json!([{"role": "user", "content": "hi"}]), - json!({}), - ); - call.extra_headers = Some(Map::from_iter([ - ("Authorization".to_string(), json!("Bearer unrelated")), - ("X-Api-Key".to_string(), json!("sk-caller")), - ])); - let prepared = prepare_chat_completions_call(call).expect("prepares"); - let keys: Vec<_> = prepared - .upstream_headers - .iter() - .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) - .collect(); - assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers); - assert_eq!(keys[0].1, "sk-test"); - assert!( - prepared - .upstream_headers - .iter() - .any(|(name, value)| name.eq_ignore_ascii_case("authorization") - && value == "Bearer unrelated"), - "the unrelated authorization must survive, got {:?}", - prepared.upstream_headers - ); -} - -#[test] -fn declines_an_unsupported_request_before_resolving_credentials() { - let mut call = request( - "claude-sonnet-4-5", - Some("anthropic"), - json!([{"role": "user", "content": "hi"}]), - json!({"stream": true}), - ); - call.api_key = None; - // No api_key is set and no env is consulted: the gate must run first, so the - // error is the decline rather than a missing-credential error. - assert_eq!(decline(call), Error::Unsupported("streaming")); -} - -#[test] -fn rejects_an_unknown_provider() { - assert_eq!( - decline(request( - "openai/gpt-4o", - None, - json!([{"role": "user", "content": "hi"}]), - json!({}), - )), - Error::InvalidProvider("openai".to_string()) - ); -} - -#[test] -fn rejects_a_model_with_no_resolvable_provider() { - assert!(matches!( - decline(request( - "claude-sonnet-4-5", - None, - json!([{"role": "user", "content": "hi"}]), - json!({}), - )), - Error::InvalidProvider(_) - )); -} - -#[test] -fn rejects_an_empty_or_malformed_message_list() { - assert_eq!( - decline(request( - "anthropic/claude-sonnet-4-5", - None, - json!([]), - json!({}), - )), - Error::InvalidRequest("chat completions requires at least one message".to_string()) - ); - assert!(matches!( - decline(request( - "anthropic/claude-sonnet-4-5", - None, - json!("not a list"), - json!({}), - )), - Error::InvalidRequest(_) - )); -} - -#[test] -fn rejects_non_string_extra_headers() { - let mut call = request( - "anthropic/claude-sonnet-4-5", - None, - json!([{"role": "user", "content": "hi"}]), - json!({}), - ); - call.extra_headers = Some(Map::from_iter([("x-trace".to_string(), json!(7))])); - assert_eq!( - decline(call), - Error::Headers(litellm_http::request::HeaderError { - context: "chat completions", - name: "x-trace".to_string(), - actual: "number", - }) - ); -} - -#[test] -fn prepares_a_bedrock_call_without_resolving_credentials() { - let mut call = request( - "bedrock/us-east-1/anthropic.claude-v2", - None, - json!([{"role": "user", "content": "hi"}]), - json!({"maxTokens": 16}), - ); - call.api_key = None; - let prepared = prepare_chat_completions_call(call).expect("prepares"); - assert_eq!( - prepared.url, - "https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/converse" - ); - assert_eq!( - prepared.auth, - RequestAuth::AwsSigV4 { - region: "us-east-1".to_string(), - service: "bedrock", - } - ); - // SigV4 signs the serialized body, so prepare must not have added an - // Authorization header; the handler does it. - assert!( - !prepared - .upstream_headers - .iter() - .any(|(name, _)| name.eq_ignore_ascii_case("authorization")) - ); - assert_eq!(prepared.body["inferenceConfig"], json!({"maxTokens": 16})); -} - -#[tokio::test] -async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() { - // Python signs only the AWS header set and reattaches the rest, so a header - // the caller forwarded rides along without joining the canonical request. - // Signing it makes Converse 403 on a deployment that works on Python. - let mut call = request( - "bedrock/us-east-1/anthropic.claude-v2", - None, - json!([{"role": "user", "content": "hi"}]), - json!({ - "maxTokens": 16, - "aws_access_key_id": "AKIDEXAMPLE", - "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY" - }), - ); - // A key would resolve to a bearer token and never reach the signer. - call.api_key = None; - call.extra_headers = Some(Map::from_iter([( - "x-request-id".to_string(), - json!("abc-123"), - )])); - let prepared = prepare_chat_completions_call(call).expect("prepares"); - let signed = super::handler::outbound_request(&prepared) - .await - .expect("signs"); - - let authorization = signed - .header("authorization") - .expect("carries an authorization header") - .to_string(); - assert!( - authorization.starts_with("AWS4-HMAC-SHA256"), - "expected a SigV4 signature, got {authorization}" - ); - assert!( - !authorization.contains("x-request-id"), - "forwarded header reached SignedHeaders: {authorization}" - ); - // It still goes on the wire, it is just not part of the signature. - assert!( - signed - .headers() - .iter() - .any(|(name, value)| name == "x-request-id" && value == "abc-123"), - "forwarded header was dropped instead of reattached" - ); -} - -#[tokio::test] -async fn a_forwarded_header_the_signer_computes_declines_to_python() { - // Reattaching the caller's copy next to the computed one puts the name on - // the wire twice and Bedrock rejects the pair, so a request carrying one - // has to go to Python instead of being signed here. - for forwarded in [ - "Authorization", - "x-amz-date", - "x-amz-security-token", - "Date", - ] { - let mut call = request( - "bedrock/us-east-1/anthropic.claude-v2", - None, - json!([{"role": "user", "content": "hi"}]), - json!({ - "maxTokens": 16, - "aws_access_key_id": "AKIDEXAMPLE", - "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY" - }), - ); - call.api_key = None; - call.extra_headers = Some(Map::from_iter([(forwarded.to_string(), json!("forged"))])); - let prepared = prepare_chat_completions_call(call).expect("prepares"); - let error = super::handler::outbound_request(&prepared) - .await - .expect_err("{forwarded} should decline instead of being signed"); - assert!( - matches!(error, Error::Unsupported(_)), - "{forwarded} declined as {error:?}, which the host would not fall back on" - ); - } -} - -#[test] -fn a_bedrock_deployment_bearer_outranks_a_forwarded_authorization() { - // `get_request_headers` assigns `headers["Authorization"]` unconditionally - // once a bearer token resolves, so the deployment's identity wins on - // Python. Keeping the caller's would authorize and bill the call as a - // different principal, and only when the deployment carries `rust: true`. - let mut call = request( - "bedrock/us-east-1/anthropic.claude-v2", - None, - json!([{"role": "user", "content": "hi"}]), - json!({"maxTokens": 16}), - ); - call.extra_headers = Some(Map::from_iter([( - "Authorization".to_string(), - json!("Bearer caller-supplied"), - )])); - let prepared = prepare_chat_completions_call(call).expect("prepares"); - let authorizations: Vec<_> = prepared - .upstream_headers - .iter() - .filter(|(name, _)| name.eq_ignore_ascii_case("authorization")) - .map(|(_, value)| value.as_str()) - .collect(); - assert_eq!( - authorizations, - vec!["Bearer sk-test"], - "the deployment token must be the only authorization on the wire" - ); -} - -#[test] -fn an_anthropic_forwarded_oauth_bearer_still_outranks_the_resolved_key() { - // The opposite precedence, and deliberate: Anthropic's own transform - // honours a forwarded OAuth bearer, so the Bedrock fix above must not be - // generalized into a rule that the configured key always wins. - // - // An OAuth bearer is the whole of that exception. This forwarded a plain - // `x-api-key` until round 17, which read as the same claim and was not: - // Python overwrites a forwarded `x-api-key` with the deployment's. - let mut call = request( - "claude-sonnet-4-5", - Some("anthropic"), - json!([{"role": "user", "content": "hi"}]), - json!({}), - ); - call.extra_headers = Some(Map::from_iter([( - "authorization".to_string(), - json!("Bearer sk-ant-oat01-forwarded"), - )])); - let prepared = prepare_chat_completions_call(call).expect("prepares"); - let keys: Vec<_> = prepared - .upstream_headers - .iter() - .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) - .map(|(_, value)| value.as_str()) - .collect(); - assert!(keys.is_empty(), "got {:?}", prepared.upstream_headers); - assert!( - prepared - .upstream_headers - .iter() - .any(|(name, value)| name.eq_ignore_ascii_case("authorization") - && value == "Bearer sk-ant-oat01-forwarded") - ); -} - -#[test] -fn a_bedrock_api_key_is_sent_as_a_bearer_token_instead_of_being_signed() { - // The configured bearer identity has its own account and quota boundary, - // so a request carrying one must not be signed as whatever principal the - // host's AWS credentials resolve to. - let prepared = prepare_chat_completions_call(request( - "bedrock/us-east-1/anthropic.claude-v2", - None, - json!([{"role": "user", "content": "hi"}]), - json!({"maxTokens": 16}), - )) - .expect("prepares"); - assert_eq!( - prepared.auth, - RequestAuth::Bearer { - token: "sk-test".to_string() - } - ); - assert!( - prepared - .upstream_headers - .iter() - .any(|(name, value)| name.eq_ignore_ascii_case("authorization") - && value == "Bearer sk-test"), - "prepare did not carry the bearer token" - ); -} - -fn decline_reason( - model: &str, - provider: Option<&str>, - messages: Value, - params: Value, -) -> Option<&'static str> { - let params = match params { - Value::Object(map) => map, - other => panic!("params must be an object, got {other}"), - }; - super::chat_completions_decline_reason(model, provider, messages, ¶ms) -} - -#[test] -fn the_gate_accepts_what_prepare_accepts() { - assert_eq!( - decline_reason( - "anthropic/claude-sonnet-4-5", - None, - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - ), - None - ); -} - -#[test] -fn the_gate_declines_without_resolving_credentials_or_calling_out() { - assert_eq!( - decline_reason( - "anthropic/claude-sonnet-4-5", - None, - json!([{"role": "user", "content": "hi"}]), - json!({"stream": true}), - ), - Some("streaming") - ); - assert_eq!( - decline_reason( - "openai/gpt-4o", - None, - json!([{"role": "user", "content": "hi"}]), - json!({}), - ), - Some("provider is not on the rust chat completions path") - ); - assert_eq!( - decline_reason( - "claude-sonnet-4-5", - None, - json!([{"role": "user", "content": "hi"}]), - json!({}), - ), - Some("provider is not on the rust chat completions path") - ); - assert_eq!( - decline_reason( - "anthropic/claude-sonnet-4-5", - None, - json!("nope"), - json!({}) - ), - Some("unreadable message list") - ); - assert_eq!( - decline_reason("anthropic/claude-sonnet-4-5", None, json!([]), json!({})), - Some("empty message list") - ); -} - -#[test] -fn the_gate_agrees_with_prepare_on_every_case_it_accepts() { - // A gate that accepts what prepare then declines would make the host emit - // its pre-call logging on a path that falls back, so pin the agreement. - for (messages, params) in [ - ( - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 8}), - ), - ( - json!([{"role": "system", "content": "s"}, {"role": "user", "content": "hi"}]), - json!({"temperature": 0.1}), - ), - ( - json!([{"role": "user", "content": "hi"}, {"role": "assistant", "content": "yo"}]), - json!({}), - ), - ] { - assert_eq!( - decline_reason( - "anthropic/claude-sonnet-4-5", - None, - messages.clone(), - params.clone() - ), - None, - "gate declined {messages}" - ); - prepare_chat_completions_call(request( - "anthropic/claude-sonnet-4-5", - None, - messages.clone(), - params, - )) - .unwrap_or_else(|error| panic!("prepare declined {messages}: {error}")); - } -} - -mod round_trip { - use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::{TcpListener, TcpStream}, - }; - - use super::*; - use crate::chat_completions::chat_completions; - - async fn read_http_request(socket: &mut TcpStream) -> String { - let mut request = Vec::new(); - let mut buffer = [0_u8; 1024]; - let header_end = loop { - let n = socket.read(&mut buffer).await.expect("reads request"); - if n == 0 { - break request.len(); - } - request.extend_from_slice(&buffer[..n]); - if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") { - break position + 4; - } - }; - let headers = String::from_utf8_lossy(&request[..header_end]); - let content_length = headers - .lines() - .find_map(|line| { - let (name, value) = line.split_once(':')?; - name.eq_ignore_ascii_case("content-length") - .then(|| value.trim().parse::().ok()) - .flatten() - }) - .unwrap_or(0); - while request.len().saturating_sub(header_end) < content_length { - let n = socket.read(&mut buffer).await.expect("reads body"); - if n == 0 { - break; - } - request.extend_from_slice(&buffer[..n]); - } - String::from_utf8(request).expect("request is utf8") - } - - fn http_response(status: &str, body: &str) -> String { - format!( - "HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - body.len(), - body - ) - } - - /// Serve one request from a stub upstream and hand back what it received. - async fn serve_once( - status: &'static str, - body: &'static str, - ) -> (String, tokio::task::JoinHandle) { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - let port = listener.local_addr().expect("addr").port(); - let handle = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts"); - let received = read_http_request(&mut socket).await; - socket - .write_all(http_response(status, body).as_bytes()) - .await - .expect("writes response"); - socket.flush().await.expect("flushes"); - received - }); - (format!("http://127.0.0.1:{port}/v1/messages"), handle) - } - - fn call(api_base: &str, messages: Value, params: Value) -> ChatCompletionsRequest<'_> { - ChatCompletionsRequest { - model: "anthropic/claude-sonnet-4-5", - messages, - optional_params: match params { - Value::Object(map) => map, - other => panic!("params must be an object, got {other}"), - }, - api_key: Some("sk-test"), - api_base: Some(api_base), - custom_llm_provider: None, - extra_headers: None, - timeout: Some(std::time::Duration::from_secs(10)), - } - } - - const GOOD_BODY: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5-20260101","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#; - - #[tokio::test] - async fn round_trip_sends_the_translated_body_and_normalizes_the_response() { - let (api_base, handle) = serve_once("200 OK", GOOD_BODY).await; - let response = chat_completions(call( - &api_base, - json!([ - {"role": "system", "content": "be terse"}, - {"role": "user", "content": "hi"} - ]), - json!({"max_tokens": 16}), - )) - .await - .expect("call succeeds"); - - let received = handle.await.expect("server task"); - let sent: Value = serde_json::from_str( - received - .split_once("\r\n\r\n") - .expect("request has a body") - .1, - ) - .expect("body is json"); - assert_eq!( - sent["messages"], - json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]) - ); - assert_eq!( - sent["system"], - json!([{"type": "text", "text": "be terse"}]) - ); - assert_eq!(sent["max_tokens"], json!(16)); - assert!(received.to_lowercase().contains("x-api-key: sk-test")); - - assert_eq!( - response.choices[0].message.content.as_deref(), - Some("hello") - ); - assert_eq!(response.usage.total_tokens, 15); - } - - #[tokio::test] - async fn a_response_it_cannot_normalize_is_reported_as_already_sent() { - // The provider was called and billed, so the host must not retry this - // on its own path. `MissingField` here would read as a pre-send - // decline and be retried; `InvalidResponse` cannot. - const NO_USAGE: &str = - r#"{"model":"m","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}"#; - let (api_base, handle) = serve_once("200 OK", NO_USAGE).await; - let err = chat_completions(call( - &api_base, - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - )) - .await - .expect_err("response cannot be normalized"); - handle.await.expect("server task"); - assert!( - matches!(err, Error::InvalidResponse(_)), - "expected a post-send error, got {err:?}" - ); - } - - #[tokio::test] - async fn a_tool_use_block_in_the_response_is_also_reported_as_already_sent() { - const TOOL_USE: &str = r#"{"model":"m","content":[{"type":"tool_use","id":"t","name":"f","input":{}}],"stop_reason":"tool_use","usage":{"input_tokens":1,"output_tokens":1}}"#; - let (api_base, handle) = serve_once("200 OK", TOOL_USE).await; - let err = chat_completions(call( - &api_base, - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - )) - .await - .expect_err("response cannot be normalized"); - handle.await.expect("server task"); - assert!( - matches!(err, Error::InvalidResponse(_)), - "expected a post-send error, got {err:?}" - ); - } - - #[tokio::test] - async fn an_upstream_error_status_keeps_its_code() { - let (api_base, handle) = - serve_once("429 Too Many Requests", r#"{"error":"slow down"}"#).await; - let err = chat_completions(call( - &api_base, - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - )) - .await - .expect_err("upstream rejects"); - handle.await.expect("server task"); - assert!( - matches!( - err, - Error::Transport(litellm_http::transport::Error::Http { status: 429, .. }) - ), - "expected a 429, got {err:?}" - ); - } - - #[tokio::test] - async fn a_connection_that_is_never_established_declines_instead_of_failing() { - // Nothing was sent, so nothing was billed and the host can still serve - // the request. Classing this with the post-send failures would turn a - // recoverable fallback into a user-facing error on exactly the - // deployments whose transport is configured only on the Python client. - let port = { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - listener.local_addr().expect("has an address").port() - // Dropped here, so the port is closed and the connect is refused. - }; - let err = chat_completions(call( - &format!("http://127.0.0.1:{port}/v1/messages"), - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - )) - .await - .expect_err("nothing is listening"); - assert!( - matches!( - err, - Error::Transport(litellm_http::transport::Error::Connect(_)) - ), - "expected a pre-send connect failure, got {err:?}" - ); - } - - #[test] - fn response_errors_collapse_to_one_variant_that_can_only_mean_already_sent() { - use crate::chat_completions::handler::as_response_error; - - for original in [ - Error::MissingField("usage"), - Error::Unsupported("non-text response content block"), - Error::InvalidRequest("whatever".to_string()), - Error::Auth(litellm_auth::Error::InvalidHeader), - ] { - let label = format!("{original:?}"); - assert!( - matches!(as_response_error(original), Error::InvalidResponse(_)), - "{label} must not stay retryable once the provider has answered" - ); - } - // An upstream status is already unambiguous, so it survives intact. - assert!(matches!( - as_response_error(Error::Transport(litellm_http::transport::Error::Http { - status: 500, - body: "boom".to_string() - })), - Error::Transport(litellm_http::transport::Error::Http { status: 500, .. }) - )); - } -} diff --git a/litellm-rust/crates/core/src/messages/common_utils.rs b/litellm-rust/crates/core/src/messages/common_utils.rs index 015e026f6da..4327754ed05 100644 --- a/litellm-rust/crates/core/src/messages/common_utils.rs +++ b/litellm-rust/crates/core/src/messages/common_utils.rs @@ -26,3 +26,186 @@ pub(super) fn string_headers( ) -> Result, Error> { shared_string_headers(HEADER_CONTEXT, extra_headers).map_err(Error::from) } + +#[cfg(test)] +mod tests { + use std::{sync::Arc, time::Duration}; + + use futures_util::future::BoxFuture; + use litellm_secrets::{SecretValue, source::SecretSource}; + use serde_json::{Value, json}; + use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::{TcpListener, TcpStream}, + }; + + use super::{messages_provider_config, string_headers, truncate_error_body}; + use crate::messages::{ + Error, + route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine}, + types::MessagesShaping, + }; + + struct RecordingSecrets { + values: Vec<(&'static str, String)>, + requested: std::sync::Mutex>, + } + + impl SecretSource for RecordingSecrets { + fn get_secret_str<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result, litellm_secrets::Error>> { + Box::pin(async move { + self.requested.lock().unwrap().push(name.to_string()); + Ok(self + .values + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| SecretValue::new(value.clone()))) + }) + } + } + + fn secrets_call() -> MessagesCall { + let Value::Object(body) = json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}] + }) else { + unreachable!("literal object") + }; + MessagesCall { + model: "claude-sonnet-4-5".into(), + body, + api_key: None, + api_base: None, + custom_llm_provider: Some("anthropic".into()), + extra_headers: None, + provider_specific_header: None, + timeout: Some(Duration::from_secs(5)), + shaping: MessagesShaping::default(), + } + } + + async fn read_http_request(socket: &mut TcpStream) -> String { + let mut request = Vec::new(); + let mut buffer = [0_u8; 1024]; + let header_end = loop { + let n = socket.read(&mut buffer).await.expect("reads request"); + if n == 0 { + break request.len(); + } + request.extend_from_slice(&buffer[..n]); + if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") { + break position + 4; + } + }; + let headers = String::from_utf8_lossy(&request[..header_end]); + let content_length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) + .unwrap_or(0); + while request.len().saturating_sub(header_end) < content_length { + let n = socket.read(&mut buffer).await.expect("reads body"); + if n == 0 { + break; + } + request.extend_from_slice(&buffer[..n]); + } + String::from_utf8(request).expect("request is utf8") + } + + #[tokio::test] + async fn route_reads_the_provider_credential_and_base_from_the_secret_source() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let addr = listener.local_addr().expect("addr"); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let request = read_http_request(&mut socket).await; + let response_body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4-5","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1}}"#; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + response_body.len(), + response_body + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes response"); + request + }); + let secrets = Arc::new(RecordingSecrets { + values: vec![ + ("ANTHROPIC_API_KEY", "sk-from-manager".to_string()), + ("ANTHROPIC_BASE_URL", format!("http://{addr}")), + ], + requested: std::sync::Mutex::new(Vec::new()), + }); + + let output = litellm_host::run::run( + messages_machine(secrets.clone()), + &LocalMessagesHost::new(secrets_call()), + ) + .await + .expect("messages request succeeds"); + + assert!(matches!(output, MessagesOutput::Message(_))); + let request = server.await.expect("server task completes"); + assert!( + request + .to_ascii_lowercase() + .contains("x-api-key: sk-from-manager"), + "{request}" + ); + let requested = secrets.requested.lock().unwrap().clone(); + assert_eq!( + requested, + messages_provider_config("anthropic") + .unwrap() + .secret_names() + .iter() + .map(ToString::to_string) + .collect::>() + ); + } + + #[test] + fn provider_config_resolves_anthropic_and_azure_ai() { + assert!(messages_provider_config("anthropic").is_some()); + assert!(messages_provider_config("azure_ai").is_some()); + assert!(messages_provider_config("openai").is_none()); + } + + #[test] + fn truncate_error_body_caps_long_payloads() { + let body = "x".repeat(400); + let truncated = truncate_error_body(&body); + assert!(truncated.ends_with("... (truncated)")); + let prefix_chars = truncated + .strip_suffix("... (truncated)") + .expect("truncated marker present") + .chars() + .count(); + assert_eq!(prefix_chars, 256); + } + + #[test] + fn string_headers_rejects_non_string_values() { + let headers = json!({"x-count": 3}).as_object().unwrap().clone(); + let err = string_headers(Some(headers)).expect_err("non-string header rejected"); + assert_eq!( + err, + Error::Headers(litellm_http::request::HeaderError { + context: "messages", + name: "x-count".to_string(), + actual: "number", + }) + ); + } +} diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index 8795d4f8507..180eb08810e 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -46,6 +46,3 @@ pub async fn messages(request: MessagesRequest<'_>) -> Result &'static str { + match self { + Self::Mistral => "mistral/model", + Self::AzureAi => "azure_ai/model", + Self::VertexMistral => "vertex_ai/mistral-ocr-maas", + Self::AzureCohereParse => "azure_ai/cohere-parse", + Self::Cohere => "cohere/model", + } + } + + fn document_type(self) -> &'static str { + match self { + Self::Mistral | Self::AzureAi | Self::VertexMistral => "document_url", + Self::AzureCohereParse | Self::Cohere => "image_url", + } + } + + fn options(self) -> Value { + match self { + Self::Mistral | Self::AzureAi => json!({"pages": [0]}), + Self::VertexMistral => json!({"pages": [0], "vertex_project": "project-1"}), + Self::AzureCohereParse | Self::Cohere => json!({"output_format": "markdown"}), + } + } + } + + /// What the host does to the wire request in `before_send`. + #[derive(Clone, Copy, Debug)] + enum Host { + Detached, + ReplacesDocument, + } + + const REPLACED_DOCUMENT: &str = "data:image/png;base64,cmVwbGFjZWQ="; + + impl Host { + fn before_send(self, wire: WireRequest) -> WireRequest { + let Value::Object(fields) = wire.body else { + return wire; + }; + let body = fields + .into_iter() + .map(|(name, value)| match self { + Self::Detached => (name, value), + Self::ReplacesDocument if name == "document" => { + let document_type = value["type"].clone(); + let key = document_type.as_str().unwrap_or_default().to_string(); + (name, json!({"type": document_type, key: REPLACED_DOCUMENT})) + } + Self::ReplacesDocument => (name, value), + }) + .collect(); + WireRequest { + body: Value::Object(body), + ..wire + } + } + } + + struct Sent { + result: Result<(), Error>, + provider_body: Option, + } + + async fn send(route: Route, host: Host, document_base: &str) -> Sent { + let (base, seen, provider) = + mock_server(vec![MockResponse::json(json!({"pages": []}))]).await; + let document_type = route.document_type(); + let document = + json!({"type": document_type, document_type: format!("{document_base}/scan.png")}); + let request = wire_request_with_document(route.model(), &base, document, route.options()); + let local = + LocalOcrHost::new(request).with_before_send(move |wire, _| Ok(host.before_send(wire))); + let result = perform_ocr_with(local).await.map(|_| ()); + match result { + Ok(()) => provider.await.unwrap(), + Err(_) => provider.abort(), + } + let provider_body = seen + .lock() + .unwrap() + .first() + .map(|request| request_body(request)); + Sent { + result, + provider_body, + } + } + + fn served_document_uri() -> String { + use base64::Engine; + format!( + "data:image/png;base64,{}", + base64::engine::general_purpose::STANDARD.encode(SERVED_DOCUMENT) + ) + } + + #[rstest] + #[case::azure_ai(Route::AzureAi)] + #[case::vertex_mistral(Route::VertexMistral)] + #[case::azure_cohere_parse(Route::AzureCohereParse)] + #[tokio::test] + async fn inlining_routes_send_the_downloaded_document(#[case] route: Route) { + let (document_base, _documents) = document_server().await; + let sent = send(route, Host::Detached, &document_base).await; + sent.result.unwrap(); + assert_eq!( + sent.provider_body.unwrap()["document"][route.document_type()], + json!(served_document_uri()) + ); + } + + #[rstest] + #[tokio::test] + async fn document_replaced_by_the_host_reaches_the_provider( + #[values( + Route::Mistral, + Route::AzureAi, + Route::VertexMistral, + Route::AzureCohereParse, + Route::Cohere + )] + route: Route, + ) { + let (document_base, _documents) = document_server().await; + let sent = send(route, Host::ReplacesDocument, &document_base).await; + sent.result.unwrap(); + assert_eq!( + sent.provider_body.unwrap()["document"][route.document_type()], + json!(REPLACED_DOCUMENT) + ); + } +} diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index f298f106a5f..270a402c9fa 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -9,36 +9,210 @@ pub mod types; pub mod wire; #[cfg(test)] -#[path = "../../tests/aws_textract_ocr.rs"] -mod aws_textract_tests; +pub(crate) mod test_support { + use std::sync::{Arc, Mutex}; -#[cfg(test)] -#[path = "../../tests/azure_ai_ocr.rs"] -mod azure_ai_tests; -#[cfg(test)] -#[path = "../../tests/azure_document_intelligence_ocr.rs"] -mod azure_document_intelligence_tests; -#[cfg(test)] -#[path = "../../tests/cohere_ocr.rs"] -mod cohere_tests; -#[cfg(test)] -#[path = "../../tests/deepseek_ocr.rs"] -mod deepseek_tests; -#[cfg(test)] -#[path = "../../tests/ocr/document.rs"] -mod document_tests; -#[cfg(test)] -#[path = "../../tests/reducto_ocr.rs"] -mod reducto_tests; -#[cfg(test)] -#[path = "../../tests/ocr/support.rs"] -pub(crate) mod test_support; -#[cfg(test)] -#[path = "../../tests/ocr.rs"] -pub(crate) mod tests; -#[cfg(test)] -#[path = "../../tests/vertex_ai_deepseek_ocr.rs"] -mod vertex_ai_deepseek_tests; -#[cfg(test)] -#[path = "../../tests/vertex_ai_ocr.rs"] -mod vertex_ai_tests; + use futures_util::future::BoxFuture; + use litellm_host::event::WireRequest; + use litellm_llms::base_llm::ocr::{ + error::Error, + handler::{CallHooks, OcrClient}, + transformation::LiteLLMOcrResponse, + }; + use serde_json::{Value, json}; + use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpListener, + }; + + use crate::ocr::{ + route::{LocalOcrHost, ocr_machine}, + types::LiteLLMOcrRequest, + wire::{OcrWireRequest, decode_request}, + }; + + /// Stands in for a host with no hooks registered: the wire request goes out unchanged + /// and response events go nowhere. + pub(crate) struct NoHooks; + + impl CallHooks for NoHooks { + fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result> { + Box::pin(async move { Ok(wire) }) + } + + fn response_received<'a>(&'a self, _body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> { + Box::pin(async { Ok(()) }) + } + } + + pub(crate) fn ocr_client() -> OcrClient { + let document_http = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .expect("test document client builds"); + OcrClient::for_test(reqwest::Client::new(), document_http) + } + + pub(crate) async fn perform_ocr( + request: LiteLLMOcrRequest, + ) -> Result { + crate::ocr::client::perform(&ocr_client(), request).await + } + + pub(crate) async fn perform_ocr_with(host: LocalOcrHost) -> Result { + litellm_host::run::run(ocr_machine(ocr_client()), &host).await + } + + pub(crate) fn wire_request(model: &str, base: &str, options: Value) -> LiteLLMOcrRequest { + wire_request_with_document( + model, + base, + json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), + options, + ) + } + + pub(crate) fn wire_request_with_document( + model: &str, + base: &str, + document: Value, + options: Value, + ) -> LiteLLMOcrRequest { + decode_request(OcrWireRequest { + model: model.into(), + document, + api_key: Some(litellm_auth::SecretValue::new("test-key")), + api_base: Some(base.into()), + custom_llm_provider: None, + extra_headers: None, + optional_params: options.as_object().unwrap().clone(), + input_sources: Default::default(), + timeout_seconds: Some(2.0), + }) + .unwrap() + } + + pub(crate) fn resolved_request( + request: LiteLLMOcrRequest, + ) -> crate::ocr::types::ResolvedOcrRequest { + request + .map_document(crate::ocr::document::prepare_document) + .unwrap() + } + + pub(crate) fn with_source(request: LiteLLMOcrRequest, source: &str) -> LiteLLMOcrRequest { + let request = resolved_request(request); + let document = request.document.clone().with_source(source.into()); + request.with_document(document.into()) + } + + pub(crate) fn request_body(request: &str) -> Value { + serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() + } + + pub(crate) const SERVED_DOCUMENT: &[u8] = b"\x89PNG served document"; + + /// Serves [`SERVED_DOCUMENT`] as `image/png` to every connection until aborted. + pub(crate) async fn document_server() -> (String, tokio::task::JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let task = tokio::spawn(async move { + loop { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut buffer = [0u8; 4096]; + let _ = socket.read(&mut buffer).await.unwrap(); + let head = format!( + "HTTP/1.1 200 OK\r\nContent-Type: image/png\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + SERVED_DOCUMENT.len() + ); + socket.write_all(head.as_bytes()).await.unwrap(); + socket.write_all(SERVED_DOCUMENT).await.unwrap(); + } + }); + (base, task) + } + + pub(crate) struct MockResponse { + pub status: u16, + pub headers: Vec<(&'static str, String)>, + pub body: Value, + } + + impl MockResponse { + pub fn json(body: Value) -> Self { + Self { + status: 200, + headers: vec![], + body, + } + } + } + + pub(crate) async fn mock_server( + responses: Vec, + ) -> (String, Arc>>, tokio::task::JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let requests = Arc::new(Mutex::new(Vec::new())); + let seen = requests.clone(); + let server_base = base.clone(); + let task = tokio::spawn(async move { + for response in responses { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut bytes = Vec::new(); + let mut buffer = [0u8; 4096]; + let header_end = loop { + let n = socket.read(&mut buffer).await.unwrap(); + assert!(n > 0); + bytes.extend_from_slice(&buffer[..n]); + if let Some(index) = bytes.windows(4).position(|s| s == b"\r\n\r\n") { + break index + 4; + } + }; + let length = String::from_utf8_lossy(&bytes[..header_end]) + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().unwrap()) + }) + .unwrap_or(0); + while bytes.len() < header_end + length { + let n = socket.read(&mut buffer).await.unwrap(); + assert!(n > 0); + bytes.extend_from_slice(&buffer[..n]); + } + seen.lock() + .unwrap() + .push(String::from_utf8_lossy(&bytes).into_owned()); + let body = serde_json::to_vec(&response.body).unwrap(); + let headers = response + .headers + .into_iter() + .map(|(name, value)| { + format!("{name}: {}\r\n", value.replace("{base}", &server_base)) + }) + .collect::(); + let head = format!( + "HTTP/1.1 {} OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n{}\r\n", + response.status, + body.len(), + headers + ); + socket.write_all(head.as_bytes()).await.unwrap(); + socket.write_all(&body).await.unwrap(); + } + }); + (base, requests, task) + } + + pub(crate) fn header<'a>(request: &'a str, name: &str) -> Option<&'a str> { + request + .lines() + .take_while(|line| !line.is_empty()) + .find_map(|line| { + let (key, value) = line.split_once(':')?; + key.eq_ignore_ascii_case(name).then(|| value.trim()) + }) + } +} diff --git a/litellm-rust/crates/core/src/ocr/route.rs b/litellm-rust/crates/core/src/ocr/route.rs index 26c9ac27102..7f83291bdab 100644 --- a/litellm-rust/crates/core/src/ocr/route.rs +++ b/litellm-rust/crates/core/src/ocr/route.rs @@ -207,3 +207,3589 @@ impl litellm_host::host::Host for LocalOcrHost { Ok(()) } } + +#[cfg(test)] +mod aws_textract_tests { + use std::{collections::BTreeMap, time::SystemTime}; + + use litellm_auth_aws::{Credentials, aws_signature_headers, sign_post}; + use litellm_llms::base_llm::ocr::error::Error; + use serde_json::{Value, json}; + use time::{PrimitiveDateTime, format_description}; + + use crate::ocr::{ + route::LocalOcrHost, + test_support::{ + MockResponse, header, mock_server, perform_ocr_with, request_body, + wire_request_with_document, + }, + types::LiteLLMOcrRequest, + }; + + const ACCESS_KEY_ID: &str = "AKIDEXAMPLE"; + const SECRET_ACCESS_KEY: &str = "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"; + + fn textract_request(base: &str) -> LiteLLMOcrRequest { + textract_request_for("aws_textract/detect-document-text", base) + } + + fn textract_request_for(model: &str, base: &str) -> LiteLLMOcrRequest { + wire_request_with_document( + model, + &format!("{base}/"), + json!({"type": "image_url", "image_url": "data:image/png;base64,b3JpZ2luYWw="}), + json!({ + "aws_access_key_id": ACCESS_KEY_ID, + "aws_secret_access_key": SECRET_ACCESS_KEY, + "aws_region_name": "eu-west-1" + }), + ) + } + + fn textract_response() -> MockResponse { + MockResponse::json(json!({ + "DocumentMetadata": {"Pages": 1}, + "Blocks": [{"BlockType": "PAGE"}, {"BlockType": "LINE", "Text": "Invoice 12345"}] + })) + } + + /// Recomputes SigV4 over the bytes the server received, at the time the client claimed. + fn expected_authorization(url: &str, raw_request: &str) -> String { + let format = + format_description::parse_borrowed::<2>("[year][month][day]T[hour][minute][second]Z") + .unwrap(); + let signed_at: SystemTime = + PrimitiveDateTime::parse(header(raw_request, "x-amz-date").unwrap(), &format) + .unwrap() + .assume_utc() + .into(); + let headers: BTreeMap = ["content-type", "x-amz-target"] + .into_iter() + .map(|name| { + ( + name.to_string(), + header(raw_request, name).unwrap().to_string(), + ) + }) + .collect(); + let body = raw_request.split_once("\r\n\r\n").unwrap().1; + sign_post( + url, + body.as_bytes(), + &aws_signature_headers(&headers), + "eu-west-1", + "textract", + &Credentials::new(ACCESS_KEY_ID, SECRET_ACCESS_KEY, None, None, "test"), + signed_at, + ) + .unwrap()["Authorization"] + .clone() + } + + #[tokio::test] + async fn the_request_is_signed_for_textract_and_lines_become_the_page() { + let (base, seen, server) = mock_server(vec![textract_response()]).await; + + let response = perform_ocr_with(LocalOcrHost::new(textract_request(&base))) + .await + .unwrap(); + server.await.unwrap(); + + let raw = seen.lock().unwrap()[0].clone(); + assert_eq!( + header(&raw, "x-amz-target"), + Some("Textract.DetectDocumentText") + ); + assert_eq!( + header(&raw, "content-type"), + Some("application/x-amz-json-1.1") + ); + assert_eq!( + request_body(&raw), + json!({"Document": {"Bytes": "b3JpZ2luYWw="}}) + ); + assert_eq!( + header(&raw, "authorization"), + Some(expected_authorization(&format!("{base}/"), &raw).as_str()) + ); + assert_eq!(response.pages[0].markdown, "Invoice 12345"); + assert_eq!(response.usage_info.unwrap().pages_processed, Some(1)); + } + + #[tokio::test] + async fn a_body_rewritten_by_before_send_is_what_gets_signed_and_sent() { + let (base, seen, server) = mock_server(vec![textract_response()]).await; + let host = LocalOcrHost::new(textract_request(&base)).with_before_send(|mut wire, _| { + assert!( + !wire + .headers + .iter() + .any(|(name, _)| name.eq_ignore_ascii_case("authorization")), + "the hook ran after signing" + ); + wire.body["Document"]["Bytes"] = Value::from("cmVkYWN0ZWQ="); + Ok(wire) + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + + let raw = seen.lock().unwrap()[0].clone(); + assert_eq!( + request_body(&raw), + json!({"Document": {"Bytes": "cmVkYWN0ZWQ="}}) + ); + assert_eq!( + header(&raw, "authorization"), + Some(expected_authorization(&format!("{base}/"), &raw).as_str()) + ); + } + + #[tokio::test] + async fn a_multi_page_rejection_reaches_the_caller_with_the_single_page_limit() { + let (base, _, server) = mock_server(vec![MockResponse { + status: 400, + headers: vec![], + body: json!({ + "__type": "UnsupportedDocumentException", + "Message": "Request has unsupported document format" + }), + }]) + .await; + + let error = perform_ocr_with(LocalOcrHost::new(textract_request(&base))) + .await + .unwrap_err(); + server.await.unwrap(); + + let Error::Provider { status, body, .. } = error else { + panic!("expected a provider error, got {error:?}"); + }; + assert_eq!(status, 400); + assert!( + body.contains("multi-page documents are not supported"), + "{body}" + ); + } + + #[tokio::test] + async fn analyze_document_asks_for_layout_and_tables_and_returns_markdown() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "DocumentMetadata": {"Pages": 1}, + "Blocks": [ + {"Id": "l1", "BlockType": "LINE", "Text": "Quarterly Report"}, + {"Id": "t", "BlockType": "LAYOUT_TITLE", + "Relationships": [{"Type": "CHILD", "Ids": ["l1"]}]} + ] + }))]) + .await; + let request = textract_request_for("aws_textract/analyze-document", &base); + + let response = perform_ocr_with(LocalOcrHost::new(request)).await.unwrap(); + server.await.unwrap(); + + let raw = seen.lock().unwrap()[0].clone(); + assert_eq!( + header(&raw, "x-amz-target"), + Some("Textract.AnalyzeDocument") + ); + assert_eq!( + request_body(&raw)["FeatureTypes"], + json!(["LAYOUT", "TABLES"]) + ); + assert_eq!( + header(&raw, "authorization"), + Some(expected_authorization(&format!("{base}/"), &raw).as_str()) + ); + assert_eq!(response.pages[0].markdown, "# Quarterly Report"); + } +} + +#[cfg(test)] +mod azure_ai_tests { + use litellm_llms::base_llm::ocr::error::Error; + use serde_json::{Value, json}; + + use crate::ocr::route::LocalOcrHost; + use crate::ocr::test_support::{ + MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request, + }; + + #[tokio::test] + async fn facade_executes_azure_mistral_with_prepared_auth() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"hello"}], + "usage_info":{"pages_processed":1} + }))]) + .await; + let mut request = wire_request( + "azure_ai/model", + &base, + json!({"include_image_base64":true}), + ); + request.credentials.api_key = None; + request.transport.extra_headers = vec![( + "Authorization".into(), + "Bearer python-prepared-token".into(), + )]; + + let result = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(result.pages[0].markdown, "hello"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /providers/mistral/azure/ocr ")); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer python-prepared-token\r\n") + ); + let body: Value = + serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!( + body, + json!({ + "model":"model", + "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, + "include_image_base64":true + }) + ); + } + + #[tokio::test] + async fn facade_acquires_supplied_entra_token_for_final_request() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let mut request = wire_request( + "azure_ai/model", + &base, + json!({"azure_ad_token":"rust-owned-token"}), + ); + request.credentials.api_key = None; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer rust-owned-token\r\n") + ); + } + + #[tokio::test] + async fn rejects_non_inline_body_after_guardrails() { + let request = wire_request("azure_ai/model", "http://127.0.0.1:1", json!({})); + let host = LocalOcrHost::new(request).with_before_send(|mut wire, _| { + wire.body["document"] = json!({ + "type":"document_url", + "document_url":"https://example.com/not-inline.pdf" + }); + Ok(wire) + }); + let error = perform_ocr_with(host).await.unwrap_err(); + assert!(error.to_string().contains("data URI")); + } + + mod transformation { + use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }; + + use litellm_auth::{ + ResolvedCredential, SecretValue, TokenFuture, TokenProvider, TokenProviderHandle, + }; + use rstest::rstest; + use serde_json::json; + + use super::*; + use crate::ocr::{ + test_support::{MockResponse, header, mock_server, perform_ocr}, + types::LiteLLMOcrRequest, + wire::decode_request, + }; + + #[derive(Debug)] + struct CountingToken { + token: fn(usize) -> String, + calls: AtomicUsize, + } + + impl CountingToken { + fn new(token: fn(usize) -> String) -> Arc { + Arc::new(Self { + token, + calls: AtomicUsize::new(0), + }) + } + + fn calls(&self) -> usize { + self.calls.load(Ordering::SeqCst) + } + } + + impl TokenProvider for CountingToken { + fn acquire(&self) -> TokenFuture<'_> { + let call = self.calls.fetch_add(1, Ordering::SeqCst) + 1; + let token = SecretValue::new((self.token)(call)); + Box::pin(async move { + Ok(ResolvedCredential::AccessToken { + token, + expires_on: None, + }) + }) + } + } + + fn numbered_token(call: usize) -> String { + format!("callback-{call}") + } + + fn azure_request( + provider: &Arc, + api_base: Option<&str>, + api_key: Option<&str>, + extra_headers: Value, + optional_params: Value, + ) -> LiteLLMOcrRequest { + let wire = serde_json::from_value(json!({ + "model": "azure_ai/mistral-ocr-latest", + "document": {"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, + "api_key": api_key, + "api_base": api_base, + "custom_llm_provider": null, + "extra_headers": extra_headers, + "optional_params": optional_params, + "timeout_seconds": 2.0 + })) + .unwrap(); + LiteLLMOcrRequest { + azure_ad_token_provider: Some(TokenProviderHandle::new(provider.clone())), + ..decode_request(wire).unwrap() + } + } + + fn ocr_page() -> MockResponse { + MockResponse::json(json!({"pages":[{"index":0,"markdown":"hello"}]})) + } + + #[tokio::test] + async fn token_provider_result_is_the_bearer_and_is_acquired_for_each_request() { + let provider = CountingToken::new(numbered_token); + let (base, seen, server) = mock_server(vec![ocr_page(), ocr_page()]).await; + + for _ in 0..2 { + perform_ocr(azure_request( + &provider, + Some(&base), + None, + Value::Null, + json!({}), + )) + .await + .unwrap(); + } + server.await.unwrap(); + + assert_eq!(provider.calls(), 2); + let requests = seen.lock().unwrap(); + assert_eq!( + requests + .iter() + .map(|request| header(request, "authorization")) + .collect::>(), + [Some("Bearer callback-1"), Some("Bearer callback-2")] + ); + } + + #[rstest] + #[case::api_key_skips_provider(Some("resource-key"), Value::Null, json!({}), "Bearer resource-key", 0)] + #[case::provider_beats_static_token( + None, + Value::Null, + json!({"azure_ad_token":"static-token"}), + "Bearer callback-1", + 1 + )] + #[case::header_wins_on_the_wire_but_provider_still_runs( + None, + json!({"Authorization":"Bearer override"}), + json!({}), + "Bearer override", + 1 + )] + #[tokio::test] + async fn credential_precedence( + #[case] api_key: Option<&str>, + #[case] extra_headers: Value, + #[case] optional_params: Value, + #[case] expected_authorization: &str, + #[case] expected_calls: usize, + ) { + let provider = CountingToken::new(numbered_token); + let (base, seen, server) = mock_server(vec![ocr_page()]).await; + + perform_ocr(azure_request( + &provider, + Some(&base), + api_key, + extra_headers, + optional_params, + )) + .await + .unwrap(); + server.await.unwrap(); + + assert_eq!(provider.calls(), expected_calls); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!( + header(&requests[0], "authorization"), + Some(expected_authorization) + ); + } + + #[rstest] + #[case::missing_api_base( + false, + json!({}), + numbered_token, + |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::MissingApiBase { + provider: "Azure AI", + environment_variable: "AZURE_AI_API_BASE", + })), + 0 + )] + #[case::unsupported_oidc_reference( + true, + json!({"azure_ad_token":"oidc/assertion","client_id":"client","tenant_id":"tenant"}), + numbered_token, + |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::UnsupportedOidcReference)), + 0 + )] + #[case::empty_provider_token_ignores_static_token( + true, + json!({"azure_ad_token":"static-token"}), + |_| String::new(), + |error: &Error| matches!(error, Error::MissingAzureAiCredentials), + 1 + )] + #[tokio::test] + async fn credential_failures_send_no_provider_request( + #[case] with_api_base: bool, + #[case] optional_params: Value, + #[case] token: fn(usize) -> String, + #[case] expected: fn(&Error) -> bool, + #[case] expected_calls: usize, + ) { + let provider = CountingToken::new(token); + let (base, seen, server) = mock_server(vec![ocr_page()]).await; + + let error = perform_ocr(azure_request( + &provider, + with_api_base.then_some(base.as_str()), + None, + Value::Null, + optional_params, + )) + .await + .unwrap_err(); + server.abort(); + + assert!(expected(&error), "unexpected error: {error:?}"); + assert_eq!(provider.calls(), expected_calls); + assert!(seen.lock().unwrap().is_empty()); + } + } +} + +#[cfg(test)] +mod azure_document_intelligence_tests { + use litellm_host::event::{CallEvent, MachineEvent}; + use litellm_llms::base_llm::ocr::{error::Error, settings::OcrSettings}; + use rstest::rstest; + use serde_json::{Value, json}; + + use crate::ocr::route::LocalOcrHost; + use crate::ocr::{ + test_support::{ + MockResponse, mock_server, ocr_client, perform_ocr, perform_ocr_with, wire_request, + }, + wire::{OcrWireRequest, decode_request}, + }; + + fn query_value(url: &str, key: &str) -> Option { + url::Url::parse(url) + .unwrap() + .query_pairs() + .find_map(|(name, value)| (name == key).then(|| value.into_owned())) + } + + #[tokio::test] + async fn facade_maps_pages_features_and_url_document() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded", + "analyzeResult":{"pages":[]} + }))]) + .await; + let mut request = wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"]}), + ); + request.document = serde_json::from_value::< + litellm_llms::base_llm::ocr::transformation::OcrDocument, + >(json!({ + "type":"document_url", + "document_url":"https://example.com/document.pdf" + })) + .unwrap() + .into(); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let request = &seen.lock().unwrap()[0]; + let target = request.split_whitespace().nth(1).unwrap(); + let url = format!("{base}{target}"); + assert_eq!(query_value(&url, "pages").as_deref(), Some("1,2,3")); + assert_eq!( + query_value(&url, "features").as_deref(), + Some("keyValuePairs,languages") + ); + let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!( + body, + json!({"urlSource":"https://example.com/document.pdf"}) + ); + } + + #[rstest] + #[case(json!({"pages":[true]}), Error::Pages("expected only integers or only strings".into()))] + #[case(json!({"pages":[1,"2"]}), Error::Pages("expected only integers or only strings".into()))] + #[case(json!({"pages":[-1]}), Error::Pages("negative page index".into()))] + #[case(json!({"pages":"1&&features=bad"}), Error::Pages("invalid native page range".into()))] + #[case(json!({"features":"languages&pages=1"}), Error::Features)] + #[case(json!({"req_format":"azure"}), Error::RequestFormat)] + #[tokio::test] + async fn rejects_invalid_pages_features_and_format( + #[case] options: Value, + #[case] expected: Error, + ) { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({}))]).await; + let result = decode_request(OcrWireRequest { + model: "azure_ai/doc-intelligence/prebuilt-read".into(), + document: json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), + api_key: Some(litellm_auth::SecretValue::new("key")), + api_base: Some(base), + custom_llm_provider: None, + extra_headers: None, + optional_params: options.as_object().unwrap().clone(), + input_sources: Default::default(), + timeout_seconds: Some(2.0), + }); + let result = match result { + Ok(request) => perform_ocr(request).await, + Err(error) => Err(error), + }; + server.abort(); + let _ = server.await; + assert!( + seen.lock().unwrap().is_empty(), + "sent invalid options: {options}" + ); + let error = result.unwrap_err(); + assert_eq!( + std::mem::discriminant(&error), + std::mem::discriminant(&expected) + ); + assert_eq!(error.http_status_code(), Some(400)); + assert_eq!(error.to_string(), expected.to_string()); + } + + #[rstest] + #[case(json!({}))] + #[case(json!({"req_format":"litellm"}))] + #[tokio::test] + async fn missing_native_fields_keep_page_text_without_retaining_raw_response( + #[case] options: Value, + ) { + let operation = json!({ + "status":"succeeded", + "analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"hello"}]}]} + }); + let (base, seen, server) = mock_server(vec![MockResponse::json(operation)]).await; + let response = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + options, + )) + .await + .unwrap(); + server.await.unwrap(); + + assert_eq!(response.pages.len(), 1); + assert_eq!(response.pages[0].index, 0); + assert_eq!(response.pages[0].markdown, "hello"); + assert_eq!(response.provider_native_response, None); + let serialized = response.into_json(); + assert_eq!(serialized.get("content"), Some(&Value::Null)); + assert_eq!(serialized.get("tables"), Some(&Value::Null)); + assert_eq!(serialized.get("keyValuePairs"), Some(&Value::Null)); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + let target = requests[0].split_whitespace().nth(1).unwrap(); + let url = format!("{base}{target}"); + for field in ["pages", "features", "req_format"] { + assert_eq!(query_value(&url, field), None); + } + let body: Value = + serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!(body, json!({"base64Source":"YWJj"})); + } + + #[tokio::test] + async fn inline_document_decodes_to_base64_source() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded" + }))]) + .await; + let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let request = &seen.lock().unwrap()[0]; + let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!(body, json!({"base64Source":"YWJj"})); + } + + #[tokio::test] + async fn immediate_response_normalizes_pages_and_preserves_native() { + let operation = json!({ + "status":"succeeded", + "operationExtension":42, + "analyzeResult":{ + "content":"A\n\nB", + "tables":[{"cells":[]}], + "keyValuePairs":[{"key":{"content":"A"}}], + "pages":[{ + "pageNumber":"2", + "width":"8.5", + "height":11, + "unit":"inch", + "lines":[{"content":"A"},{"content":null},{"content":"B"}] + }] + } + }); + let (base, _, server) = mock_server(vec![MockResponse::json(operation.clone())]).await; + let result = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"req_format":"native"}), + )) + .await + .unwrap(); + server.await.unwrap(); + + assert_eq!(result.pages[0].index, 1); + assert_eq!(result.pages[0].markdown, "A\n\nB"); + assert_eq!( + serde_json::to_value(&result.pages[0].dimensions).unwrap(), + json!({"width":816,"height":1056,"dpi":96}) + ); + assert_eq!(result.usage_info.as_ref().unwrap().pages_processed, Some(1)); + let serialized = result.clone().into_json(); + assert_eq!(serialized["content"], "A\n\nB"); + assert_eq!(serialized["tables"], json!([{"cells":[]}])); + assert_eq!( + serialized["keyValuePairs"], + json!([{"key":{"content":"A"}}]) + ); + assert!(serialized.get("key_value_pairs").is_none()); + assert_eq!( + result.provider_native_response.map(Value::Object), + Some(operation) + ); + } + + #[tokio::test] + async fn client_settings_choose_the_api_version_and_the_inch_to_pixel_dpi() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded", + "analyzeResult":{"pages":[{"pageNumber":1,"width":8.5,"height":11,"unit":"inch"}]} + }))]) + .await; + let client = ocr_client().with_settings(OcrSettings { + document_intelligence_api_version: "2099-01-01".into(), + document_intelligence_dpi: 72, + ..OcrSettings::default() + }); + + let result = crate::ocr::client::perform( + &client, + wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})), + ) + .await + .unwrap(); + server.await.unwrap(); + + let target = seen.lock().unwrap()[0] + .split_whitespace() + .nth(1) + .unwrap() + .to_string(); + assert_eq!( + query_value(&format!("{base}{target}"), "api-version").as_deref(), + Some("2099-01-01") + ); + assert_eq!( + serde_json::to_value(&result.pages[0].dimensions).unwrap(), + json!({"width":612,"height":792,"dpi":72}) + ); + } + + #[tokio::test] + async fn accepted_response_polls_to_success_with_only_credentials() { + let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}}); + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse { + status: 200, + headers: vec![("Retry-After", "0".into())], + body: json!({"status":"running"}), + }, + MockResponse::json(operation.clone()), + ]) + .await; + let mut request = wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"req_format":"native"}), + ); + request + .transport + .extra_headers + .push(("X-Trace".into(), "initial-only".into())); + + let result = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!( + result.provider_native_response.map(Value::Object), + Some(operation) + ); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 3); + assert!(requests[0].to_ascii_lowercase().contains("x-trace:")); + for poll in &requests[1..] { + assert!(!poll.to_ascii_lowercase().contains("x-trace:")); + assert!( + poll.to_ascii_lowercase() + .contains("ocp-apim-subscription-key: test-key") + ); + } + } + + #[tokio::test] + async fn accepted_response_emits_response_received_before_polling() { + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({"submitted": true}), + }, + MockResponse::json(json!({"status":"succeeded"})), + ]) + .await; + let request_count = seen.clone(); + let host = LocalOcrHost::new(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .with_observer(move |event| { + let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event else { + return; + }; + match request_count.lock().unwrap().len() { + 1 => assert_eq!(raw.body, r#"{"submitted":true}"#), + 2 => assert!(raw.body.contains("succeeded")), + count => panic!("unexpected callback after {count} requests"), + } + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 2); + } + + #[tokio::test] + async fn polling_forwards_bearer_credentials() { + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse::json(json!({"status":"succeeded"})), + ]) + .await; + let mut request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); + request.credentials.api_key = None; + request.transport.extra_headers = vec![("Authorization".into(), "Bearer token".into())]; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert!( + requests[1] + .to_ascii_lowercase() + .contains("authorization: bearer token") + ); + } + + #[tokio::test] + async fn polling_does_not_follow_redirects() { + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse { + status: 302, + headers: vec![("Location", "{base}/redirected".into())], + body: json!({}), + }, + MockResponse::json(json!({"status":"succeeded"})), + ]) + .await; + + let error = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .await + .unwrap_err(); + + assert!(error.to_string().contains("status 302"), "{error}"); + assert_eq!(seen.lock().unwrap().len(), 2); + server.abort(); + } + + #[tokio::test] + async fn polling_rejects_terminal_failure() { + let (base, _, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse::json(json!({"status":"failed"})), + ]) + .await; + + let error = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .await + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains("status failed")); + } + + #[tokio::test] + async fn malformed_provider_pages_report_response_paths() { + for (analysis, path) in [ + (json!({"pages":null}), "pages"), + (json!({"pages":[null]}), "pages[0]"), + (json!({"pages":[{"lines":null}]}), "lines"), + (json!({"pages":[{"width":"bad"}]}), "width"), + ] { + let (base, _, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded", + "analyzeResult":analysis + }))]) + .await; + let error = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .await + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains(path), "{error}"); + } + } + + #[tokio::test] + async fn rejects_missing_invalid_and_cross_origin_operation_locations() { + for headers in [ + Vec::new(), + vec![("Operation-Location", "/relative".into())], + vec![("Operation-Location", "http://example.com/operation".into())], + vec![( + "Operation-Location", + "http://user:password@127.0.0.1/operation".into(), + )], + ] { + let (base, _, server) = mock_server(vec![MockResponse { + status: 202, + headers, + body: json!({}), + }]) + .await; + let error = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .await + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains("operation-location")); + } + } + + #[tokio::test] + async fn polling_deadline_bounds_retry_delay() { + let (base, _, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse { + status: 200, + headers: vec![("Retry-After", "9999".into())], + body: json!({"status":"notStarted"}), + }, + ]) + .await; + let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); + let client = ocr_client().with_settings(OcrSettings { + poll_timeout: std::time::Duration::from_millis(100), + ..OcrSettings::default() + }); + + let error = tokio::time::timeout( + std::time::Duration::from_secs(1), + crate::ocr::client::perform(&client, request), + ) + .await + .unwrap() + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains("timed out")); + } + + #[tokio::test] + async fn model_id_is_encoded_and_dot_segments_are_rejected() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded" + }))]) + .await; + perform_ocr(wire_request( + "azure_ai/doc-intelligence/a ?#é", + &base, + json!({}), + )) + .await + .unwrap(); + server.await.unwrap(); + assert!(seen.lock().unwrap()[0].contains("a%20%3F%23%C3%A9:analyze")); + + for model in [ + "azure_ai/doc-intelligence/.", + "azure_ai/doc-intelligence/..", + ] { + let error = perform_ocr(wire_request(model, "http://127.0.0.1:1", json!({}))) + .await + .unwrap_err(); + assert!(error.to_string().contains("dot segment")); + } + } + + mod transformation { + use std::sync::{Arc, Mutex}; + + use litellm_host::event::{CallEvent, MachineEvent}; + use litellm_llms::base_llm::ocr::transformation::OcrDocument; + use serde_json::{Value, json}; + + use super::*; + use crate::ocr::{ + route::LocalOcrHost, + test_support::{ + MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request, + }, + }; + + #[tokio::test] + async fn facade_maps_pages_features_and_url_document() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded", + "analyzeResult":{"pages":[]} + }))]) + .await; + let mut request = wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"], "future_option": {"nested":null}, "extra_body":{"provider_option":false}}), + ); + request.document = serde_json::from_value::(json!({ + "type":"document_url", + "document_url":"https://example.com/document.pdf" + })) + .unwrap() + .into(); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let request = &seen.lock().unwrap()[0]; + let target = request.split_whitespace().nth(1).unwrap(); + let url = format!("{base}{target}"); + assert_eq!(query_value(&url, "pages").as_deref(), Some("1,2,3")); + assert_eq!( + query_value(&url, "features").as_deref(), + Some("keyValuePairs,languages") + ); + let body: Value = + serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!( + body, + json!({"urlSource":"https://example.com/document.pdf", "future_option":{"nested":null}, "provider_option":false}) + ); + } + + #[tokio::test] + async fn rejects_invalid_pages_features_and_format() { + for options in [ + json!({"pages":[true]}), + json!({"pages":[1,"2"]}), + json!({"pages":[-1]}), + json!({"pages":"1&&features=bad"}), + json!({"features":"languages&pages=1"}), + json!({"req_format":"azure"}), + ] { + let request = wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + "http://127.0.0.1:1", + options.clone(), + ); + let rejected = perform_ocr(request).await.is_err(); + assert!(rejected, "accepted {options}"); + } + } + + #[tokio::test] + async fn immediate_response_normalizes_pages_and_preserves_native() { + let operation = json!({ + "status":"succeeded", + "operationExtension":42, + "analyzeResult":{ + "content":"A\n\nB", + "tables":[{"cells":[]}], + "keyValuePairs":[{"key":{"content":"A"}}], + "pages":[{ + "pageNumber":"2", + "width":"8.5", + "height":11, + "unit":"inch", + "lines":[{"content":"A"},{"content":null},{"content":"B"}] + }] + } + }); + let (base, _, server) = mock_server(vec![MockResponse::json(operation.clone())]).await; + let result = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"req_format":"native"}), + )) + .await + .unwrap(); + server.await.unwrap(); + + assert_eq!(result.pages[0].index, 1); + assert_eq!(result.pages[0].markdown, "A\n\nB"); + assert_eq!( + serde_json::to_value(&result.pages[0].dimensions).unwrap(), + json!({"width":816,"height":1056,"dpi":96}) + ); + assert_eq!(result.usage_info.as_ref().unwrap().pages_processed, Some(1)); + let serialized = result.clone().into_json(); + assert_eq!(serialized["content"], "A\n\nB"); + assert_eq!(serialized["tables"], json!([{"cells":[]}])); + assert_eq!( + serialized["keyValuePairs"], + json!([{"key":{"content":"A"}}]) + ); + assert!(serialized.get("key_value_pairs").is_none()); + assert_eq!( + result.provider_native_response.as_ref(), + operation.as_object() + ); + } + + #[tokio::test] + async fn accepted_response_polls_to_success_with_only_credentials() { + let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}}); + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse { + status: 200, + headers: vec![("Retry-After", "0".into())], + body: json!({"status":"running"}), + }, + MockResponse::json(operation.clone()), + ]) + .await; + let mut request = wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"req_format":"native"}), + ); + request + .transport + .extra_headers + .push(("X-Trace".into(), "initial-only".into())); + + let result = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!( + result.provider_native_response.as_ref(), + operation.as_object() + ); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 3); + assert!(requests[0].to_ascii_lowercase().contains("x-trace:")); + for poll in &requests[1..] { + assert!(!poll.to_ascii_lowercase().contains("x-trace:")); + assert!( + poll.to_ascii_lowercase() + .contains("ocp-apim-subscription-key: test-key") + ); + } + } + + #[tokio::test] + async fn accepted_response_emits_response_received_for_submission_and_completed_poll() { + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({"submitted": true}), + }, + MockResponse::json(json!({"status":"succeeded"})), + ]) + .await; + let responses_received = Arc::new(Mutex::new(Vec::new())); + let request_count = seen.clone(); + let observed = responses_received.clone(); + let host = LocalOcrHost::new(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .with_observer(move |event| { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { + observed + .lock() + .unwrap() + .push((request_count.lock().unwrap().len(), raw.body.clone())); + } + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 2); + assert_eq!( + *responses_received.lock().unwrap(), + [ + (1, r#"{"submitted":true}"#.to_string()), + (2, r#"{"status":"succeeded"}"#.to_string()), + ] + ); + } + } +} + +#[cfg(test)] +mod cohere_tests { + mod transformation { + use litellm_llms::{ + base_llm::ocr::{ + error::Error, + transformation::{BaseOcrConfig, OcrDocument, OcrResponseFormat}, + }, + cohere::ocr::transformation::*, + }; + use rstest::rstest; + use serde_json::{Value, json}; + + #[tokio::test] + async fn composed_body_preserves_native_document_fields_and_untyped_overrides() { + let request = crate::ocr::test_support::wire_request( + "cohere/parse", + "https://example.com", + json!({ + "output_format":"markdown", "timeout":30, + "extra_body":{ + "output_format": {"future":true}, + "document":{"type":"image_url","image_url":"https://example.com/a.png", + "provider_options":{"nested":[false,0,null]}} + } + }), + ); + let request = request.with_document( + serde_json::from_value(json!({ + "type":"image_url","image_url":"https://example.com/original.png" + })) + .unwrap(), + ); + let request = crate::ocr::prepare::prepare_request_for_test(request); + let http = CohereParseConfig + .prepare_request( + &request, + &crate::ocr::test_support::ocr_client(), + &crate::ocr::test_support::NoHooks, + ) + .await + .unwrap(); + let body: Value = serde_json::from_slice(http.body()).unwrap(); + assert_eq!( + body, + json!({ + "model":"parse", "output_format":{"future":true}, + "document":{"type":"image_url","image_url":"https://example.com/a.png", + "provider_options":{"nested":[false,0,null]}} + }) + ); + } + + #[tokio::test] + async fn explicit_null_options_use_defaults_before_http() { + let request = crate::ocr::test_support::wire_request( + "cohere/parse", + "https://example.com", + json!({"output_format":null,"req_format":null}), + ); + let request = request.with_document( + serde_json::from_value( + json!({"type":"image_url","image_url":"https://example.com/a.png"}), + ) + .unwrap(), + ); + assert_eq!( + request.response_format().unwrap(), + OcrResponseFormat::Litellm + ); + let request = crate::ocr::prepare::prepare_request_for_test(request); + let http = CohereParseConfig + .prepare_request( + &request, + &crate::ocr::test_support::ocr_client(), + &crate::ocr::test_support::NoHooks, + ) + .await + .unwrap(); + let body: Value = serde_json::from_slice(http.body()).unwrap(); + assert_eq!(body["output_format"], "markdown"); + assert!(body.get("req_format").is_none()); + } + + #[rstest] + #[case::cohere("cohere/parse-v5.0", "POST /v2/parse ")] + #[case::azure_ai("azure_ai/Cohere-parse-v5.0", "POST /providers/cohere/v2/parse ")] + #[tokio::test] + async fn route_sends_image_to_its_parse_endpoint_with_the_bearer_key( + #[case] model: &str, + #[case] request_line: &str, + ) { + use crate::ocr::test_support::{MockResponse, header, mock_server, perform_ocr}; + + let (base, seen, server) = + mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let request = crate::ocr::test_support::wire_request(model, &base, json!({})) + .with_document( + serde_json::from_value::( + json!({"type":"image_url","image_url":"data:image/png;base64,YWJj"}), + ) + .unwrap() + .into(), + ); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with(request_line), "{}", requests[0]); + assert_eq!( + header(&requests[0], "authorization"), + Some("Bearer test-key") + ); + } + + #[rstest] + #[tokio::test] + async fn route_rejects_non_image_document_without_a_request( + #[values("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0")] model: &str, + ) { + use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr}; + + let (base, seen, server) = + mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + + let error = perform_ocr(crate::ocr::test_support::wire_request( + model, + &base, + json!({}), + )) + .await + .unwrap_err(); + server.abort(); + + assert!(matches!(error, Error::CohereImageOnly), "{error:?}"); + assert!(seen.lock().unwrap().is_empty()); + } + } +} + +#[cfg(test)] +mod deepseek_tests { + use litellm_llms::{ + base_llm::ocr::transformation::{BaseOcrConfig, OcrDocument}, + vertex_ai::ocr::deepseek_transformation::{ + DeepSeekOcrParams, DeepSeekOcrResponse, VertexAIDeepSeekOCRConfig, + normalize_response as transform_ocr_response, + }, + }; + use rstest::rstest; + use serde_json::{Value, json}; + + fn document() -> OcrDocument { + serde_json::from_value(json!({"type":"image_url","image_url":"gs://bucket/a.png"})).unwrap() + } + + #[rstest] + #[case("stream", json!(true))] + #[case("temperature", json!(0.1))] + #[case("max_tokens", json!(1024))] + #[case("top_p", json!(0.9))] + #[case("n", json!(2))] + #[case("stop", json!("done"))] + #[case("stop", json!(["done", "stop"]))] + fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) { + let params: DeepSeekOcrParams = + serde_json::from_value(json!({name: value.clone(), "ignored": true})).unwrap(); + let result = serde_json::to_value( + VertexAIDeepSeekOCRConfig + .transform_ocr_request("deepseek-ai/deepseek-ocr-maas", document(), ¶ms, &[]) + .unwrap(), + ) + .unwrap(); + assert_eq!(result["model"], "deepseek-ai/deepseek-ocr-maas"); + assert_eq!( + result["messages"][0]["content"][0], + json!({"type":"image_url","image_url":"gs://bucket/a.png"}) + ); + assert_eq!(result[name], value); + assert!(result.get("ignored").is_none()); + } + + #[rstest] + #[case(json!({"type":"image_url","image_url":"data:image/png;base64,AA=="}))] + #[case(json!({"type":"document_url","document_url":"data:application/pdf;base64,AA=="}))] + fn request_maps_both_document_types_to_image_content(#[case] document: Value) { + let source = document + .get("image_url") + .or_else(|| document.get("document_url")) + .unwrap() + .clone(); + let request = VertexAIDeepSeekOCRConfig + .transform_ocr_request( + "deepseek-ai/deepseek-ocr-maas", + serde_json::from_value(document).unwrap(), + &DeepSeekOcrParams::default(), + &[], + ) + .unwrap(); + let result = serde_json::to_value(request).unwrap(); + assert_eq!( + result["messages"][0]["content"][0], + json!({"type":"image_url","image_url":source}) + ); + } + + #[rstest] + #[case(json!("# hello"), "# hello")] + #[case(json!("{broken"), "{broken")] + #[case(json!(" {\"pages\":[]} "), " {\"pages\":[]} ")] + #[case(json!({"pages":[]}), "")] + #[case(json!("[]"), "[]")] + #[case(json!("{\"pages\":[{\"markdown\":\"json text\"}]}"), "json text")] + #[case(json!({"pages":[{"markdown":"object"}]}), "object")] + fn response_codec_handles_text_json_and_objects( + #[case] content: Value, + #[case] expected: &str, + ) { + let structured = content + .as_object() + .is_some_and(|object| object.contains_key("pages")) + || content + .as_str() + .is_some_and(|text| text.contains("\"pages\"")); + let response: DeepSeekOcrResponse = serde_json::from_value( + json!({"choices":[{"message":{"content":content}}],"usage":{"prompt_tokens":1}}), + ) + .unwrap(); + let result = transform_ocr_response("model", response) + .unwrap() + .into_json(); + assert_eq!(result["pages"][0]["markdown"], expected); + assert_eq!(result["pages"][0]["index"], 0); + if structured { + assert!(result["usage_info"].is_null()); + } else { + assert_eq!(result["usage_info"]["prompt_tokens"], 1); + } + } + + #[test] + fn structured_result_maps_pages_usage_model_and_annotation() { + let response: DeepSeekOcrResponse = serde_json::from_value(json!({ + "choices":[{"message":{"content":{ + "pages":[{"index":2,"markdown":"page","images":[{"id":"one"}],"dimensions":{"width":10}}], + "model":"provider-model", + "usage_info":{"pages_processed":1}, + "document_annotation":{"language":"en"}, + "future":"kept" + }}}] + })) + .unwrap(); + let result = transform_ocr_response("requested", response) + .unwrap() + .into_json(); + assert_eq!(result["pages"][0]["index"], 2); + assert_eq!(result["pages"][0]["images"][0]["id"], "one"); + assert_eq!(result["model"], "provider-model"); + assert_eq!(result["usage_info"]["pages_processed"], 1); + assert_eq!(result["document_annotation"]["language"], "en"); + assert_eq!(result["future"], "kept"); + } + + #[test] + fn response_codec_rejects_missing_empty_and_malformed_content() { + for value in [ + json!({"choices":[{"message":{"content":{}}}]}), + json!({"choices":[]}), + json!({"choices":[{"message":{"content":""}}]}), + json!({"choices":[{"message":{"content":"{\"pages\":[{\"markdown\":42}]}"}}]}), + json!({"choices":[{"message":{"content":{"pages":[{"markdown":42}]}}}]}), + ] { + let result = serde_json::from_value::(value) + .map_err(|_| ()) + .and_then(|response| transform_ocr_response("model", response).map_err(|_| ())); + assert!(result.is_err()); + } + } +} + +#[cfg(test)] +mod reducto_tests { + use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; + use litellm_llms::base_llm::ocr::{error::Error, transformation::OcrDocument}; + use rstest::rstest; + use serde_json::{Value, json}; + + use crate::ocr::route::LocalOcrHost; + use crate::ocr::test_support::{ + MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request, + }; + + fn request_body(request: &str) -> Value { + serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() + } + + #[rstest] + #[case( + "reducto/parse-v3", + json!({ + "formatting":{"table_output_format":"html"}, + "retrieval":{"chunk_mode":"section"}, + "settings":{"ocr_system":"standard"}, + "future_ocr_option":true, + "extra_body":{"provider_option":"value"} + }), + "reducto://already.pdf", + json!({ + "input":"reducto://already.pdf", + "formatting":{"table_output_format":"html"}, + "retrieval":{"chunk_mode":"section"}, + "settings":{"ocr_system":"standard"}, + "future_ocr_option":true, + "provider_option":"value" + }) + )] + #[case( + "reducto/parse-legacy", + json!({ + "enhance":{"agentic":[{"type":"table"}]}, + "future_ocr_option":true, + "extra_body":{"provider_option":"value"} + }), + "reducto://legacy.pdf", + json!({ + "document_url":"reducto://legacy.pdf", + "options":{"enhance":{"agentic":[{"type":"table"}]}}, + "future_ocr_option":true, + "provider_option":"value" + }) + )] + #[tokio::test] + async fn request_mapping_matches_python( + #[case] model: &str, + #[case] options: Value, + #[case] source: &str, + #[case] expected: Value, + ) { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "result":{"chunks":[]} + }))]) + .await; + let request = + crate::ocr::test_support::with_source(wire_request(model, &base, options), source); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /parse ")); + assert_eq!(request_body(&requests[0]), expected); + } + + #[rstest] + #[case("parse-v3")] + #[case("parse-legacy")] + #[tokio::test] + async fn data_uri_upload_preserves_multipart_headers( + #[case] model: &str, + #[values("application/pdf", "image/png")] mime_type: &str, + ) { + let (base, seen, server) = mock_server(vec![ + MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), + MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})), + ]) + .await; + let document = if mime_type.starts_with("image/") { + json!({"type":"image_url","image_url":format!("data:{mime_type};base64,YWJj")}) + } else { + json!({"type":"document_url","document_url":format!("data:{mime_type};base64,YWJj")}) + }; + let mut request = crate::ocr::types::LiteLLMOcrRequest { + document: serde_json::from_value::(document) + .unwrap() + .into(), + ..wire_request(&format!("reducto/{model}"), &base, json!({})) + }; + request.transport.extra_headers = vec![ + ("Content-Type".into(), "application/json".into()), + ("X-Trace".into(), "upload-test".into()), + ]; + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.pages[0].markdown, "hello"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert!(requests[0].starts_with("POST /upload ")); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("content-type: multipart/form-data; boundary=") + ); + assert!(requests[0].contains("x-trace: upload-test")); + let multipart = requests[0].split_once("\r\n\r\n").unwrap().1; + assert!(multipart.contains(&format!("Content-Type: {mime_type}\r\n"))); + assert!(multipart.contains("\r\n\r\nabc\r\n--")); + assert!(requests[1].starts_with("POST /parse ")); + let source_field = if model == "parse-legacy" { + "document_url" + } else { + "input" + }; + assert_eq!( + request_body(&requests[1]), + json!({source_field:"reducto://uploaded.pdf"}) + ); + for request in requests.iter() { + assert!( + request + .to_ascii_lowercase() + .contains("authorization: bearer test-key\r\n") + ); + } + } + + #[tokio::test] + async fn response_received_stays_after_reducto_upload_and_parse() { + let (base, seen, server) = mock_server(vec![ + MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), + MockResponse::json(json!({"result":{"chunks":[]}})), + ]) + .await; + let request_count = seen.clone(); + let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) + .with_observer(move |event| { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { + assert_eq!(request_count.lock().unwrap().len(), 2); + assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); + } + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 2); + } + + #[rstest] + #[case(json!({"file_id":""}))] + #[case(json!({}))] + #[case(json!({"file_id":null}))] + #[tokio::test] + async fn invalid_upload_ids_stop_before_parse(#[case] response: Value) { + let (base, seen, server) = mock_server(vec![MockResponse::json(response)]).await; + let error = perform_ocr(wire_request("reducto/parse-v3", &base, json!({}))) + .await + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains("file_id")); + assert_eq!(seen.lock().unwrap().len(), 1); + } + + #[tokio::test] + async fn upload_failure_stops_before_parse() { + let (base, seen, server) = mock_server(vec![MockResponse { + status: 503, + headers: vec![], + body: json!({"error":"unavailable"}), + }]) + .await; + assert!( + perform_ocr(wire_request("reducto/parse-v3", &base, json!({}))) + .await + .is_err() + ); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 1); + } + + #[rstest] + #[case("https://example.com/a.pdf", Error::ReductoSource)] + #[case("reducto://", Error::RequestField { path: "document file id".into() })] + #[case("data:application/pdf;base64", Error::InvalidDataUri)] + #[case("data:application/pdf;base64,INVALID!", Error::InvalidDataUri)] + #[tokio::test] + async fn rejects_invalid_document_sources_before_network( + #[case] source: &str, + #[case] expected: Error, + ) { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({}))]).await; + let request = crate::ocr::test_support::with_source( + wire_request("reducto/parse-v3", &base, json!({})), + source, + ); + let result = perform_ocr(request).await; + server.abort(); + let _ = server.await; + assert!( + seen.lock().unwrap().is_empty(), + "sent invalid source: {source}" + ); + let error = result.unwrap_err(); + assert_eq!( + std::mem::discriminant(&error), + std::mem::discriminant(&expected) + ); + assert_eq!(error.http_status_code(), Some(400)); + assert_eq!(error.to_string(), expected.to_string()); + } + + #[test] + fn response_normalization_groups_blocks_and_distinguishes_null_result() { + use litellm_llms::reducto::ocr::transformation::{ + ReductoResponse, normalize_response as transform_ocr_response, + }; + + let raw = json!({"usage":{"num_pages":"2","credits":"3"},"result":{"type":"full","chunks":[ + {"blocks":[{ + "type":"Table", + "content":"B", + "bbox":{"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}, + "confidence":"high", + "granular_confidence":{"parse_confidence":0.95,"extract_confidence":null}, + "image_url":null + }]}, + {"blocks":[{"content":"A","bbox":{"page":1},"type":"Text"},{"content":"C","bbox":{"page":1}}]} + ]}}); + let response: ReductoResponse = serde_json::from_value(raw).unwrap(); + let normalized = transform_ocr_response("parse-v3", response) + .unwrap() + .into_json(); + assert_eq!(normalized["pages"][0]["markdown"], "A\n\nC"); + assert_eq!(normalized["pages"][1]["markdown"], "B"); + assert_eq!(normalized["pages"][1]["blocks"][0]["type"], "Table"); + assert_eq!( + normalized["pages"][1]["blocks"][0]["bbox"], + json!({"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}) + ); + assert_eq!(normalized["pages"][1]["blocks"][0]["confidence"], "high"); + assert_eq!( + normalized["pages"][1]["blocks"][0]["granular_confidence"]["parse_confidence"], + 0.95 + ); + assert!(normalized["pages"][1]["blocks"][0]["image_url"].is_null()); + assert_eq!(normalized["usage_info"]["pages_processed"], 2); + assert_eq!(normalized["usage_info"]["credits"], 3.0); + + let missing: ReductoResponse = + serde_json::from_value(json!({"chunks":[{"content":"text"}]})).unwrap(); + let missing = transform_ocr_response("parse-v3", missing).unwrap(); + assert_eq!(missing.pages[0].markdown, "text"); + let null: ReductoResponse = serde_json::from_value( + json!({"result":null,"chunks":[{"content":"ignored"}],"usage":null}), + ) + .unwrap(); + let null = transform_ocr_response("parse-v3", null).unwrap(); + assert!(null.pages.is_empty()); + } + + #[tokio::test] + async fn facade_omits_native_response_by_default_and_preserves_auth_priority() { + let raw = json!({"job_id":"job-1","result":{"chunks":[]}}); + let (base, seen, server) = mock_server(vec![MockResponse::json(raw)]).await; + let mut request = crate::ocr::test_support::with_source( + wire_request("reducto/parse-v3", &base, json!({})), + "reducto://ready.pdf", + ); + request.transport.extra_headers = vec![("authorization".into(), "Bearer existing".into())]; + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.provider_native_response, None); + assert!( + seen.lock().unwrap()[0] + .to_ascii_lowercase() + .contains("authorization: bearer existing") + ); + } + + #[tokio::test] + async fn native_format_retains_the_provider_response() { + let raw = json!({ + "result":{"chunks":[{"content":"native OCR response"}]}, + "usage":{"num_pages":1} + }); + let (base, _, server) = mock_server(vec![MockResponse::json(raw.clone())]).await; + let request = crate::ocr::test_support::with_source( + wire_request("reducto/parse-v3", &base, json!({"req_format":"native"})), + "reducto://ready.pdf", + ); + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + + assert_eq!(response.pages[0].markdown, "native OCR response"); + assert_eq!(response.provider_native_response.as_ref(), raw.as_object()); + } + + #[tokio::test] + async fn unknown_model_reaches_parse_and_keeps_its_name() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "result":{"chunks":[{"content":"future model response"}]} + }))]) + .await; + let request = crate::ocr::test_support::with_source( + wire_request("reducto/future-parse-model", &base, json!({})), + "reducto://ready.pdf", + ); + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + + assert_eq!(response.model, "future-parse-model"); + assert_eq!(response.pages[0].markdown, "future model response"); + let requests = seen.lock().unwrap(); + assert!(requests[0].starts_with("POST /parse ")); + assert_eq!( + request_body(&requests[0]), + json!({"input":"reducto://ready.pdf"}) + ); + } + + #[tokio::test] + async fn guardrail_rewrites_document_before_upload() { + let (base, seen, server) = + mock_server(vec![MockResponse::json(json!({"result":{"chunks":[]}}))]).await; + let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) + .with_before_send(|wire, _| { + assert_eq!( + wire.body["document_url"], + "data:application/pdf;base64,YWJj" + ); + Ok(WireRequest { + body: json!({"type":"document_url","document_url":"reducto://guarded.pdf"}), + ..wire + }) + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /parse ")); + assert!(requests[0].contains("reducto://guarded.pdf")); + } + + mod transformation { + use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; + use litellm_llms::{ + base_llm::ocr::transformation::{BaseOcrConfig, OcrConnection, OcrRequestContext}, + reducto::ocr::transformation::*, + }; + use rstest::rstest; + + use super::*; + use crate::ocr::{ + route::LocalOcrHost, + test_support::{ + MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request, + }, + }; + + #[tokio::test] + async fn v3_options_preserve_explicit_null() { + let overrides = + serde_json::from_value(json!({"formatting":null,"settings":{},"unknown":true})) + .unwrap(); + let params = ReductoParseV3Config + .map_ocr_params(&overrides, "parse-v3") + .unwrap(); + let client = crate::ocr::test_support::ocr_client(); + let connection = OcrConnection::default(); + let document = serde_json::from_value( + json!({"type":"document_url","document_url":"reducto://ready.pdf"}), + ) + .unwrap(); + let body = ReductoParseV3Config + .async_transform_ocr_request( + "parse-v3", + document, + ¶ms, + &[], + OcrRequestContext { + client: &client, + connection: &connection, + }, + ) + .await + .unwrap(); + assert_eq!( + serde_json::to_value(body).unwrap(), + json!({ + "input":"reducto://ready.pdf", "formatting":null, "settings":{} + }) + ); + let absent = ReductoParseV3Config + .map_ocr_params( + &litellm_core_utils::call_arguments::CallArguments::default(), + "parse-v3", + ) + .unwrap(); + assert_eq!(serde_json::to_value(absent).unwrap(), json!({})); + } + + #[rstest] + #[case( + "reducto/parse-v3", + json!({ + "formatting":{"table_output_format":"html"}, + "retrieval":{"chunk_mode":"section"}, + "settings":{"ocr_system":"standard"}, + "future_ocr_option":true, + "extra_body":{"provider_option":"value"} + }), + "reducto://already.pdf", + json!({ + "input":"reducto://already.pdf", + "formatting":{"table_output_format":"html"}, + "retrieval":{"chunk_mode":"section"}, + "settings":{"ocr_system":"standard"}, + "future_ocr_option":true, + "provider_option":"value" + }) + )] + #[case( + "reducto/parse-legacy", + json!({ + "enhance":{"agentic":[{"type":"table"}]}, + "future_ocr_option":true, + "extra_body":{"provider_option":"value"} + }), + "reducto://legacy.pdf", + json!({ + "document_url":"reducto://legacy.pdf", + "options":{"enhance":{"agentic":[{"type":"table"}]}}, + "future_ocr_option":true, + "provider_option":"value" + }) + )] + #[tokio::test] + async fn request_mapping_matches_python( + #[case] model: &str, + #[case] options: Value, + #[case] source: &str, + #[case] expected: Value, + ) { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "result":{"chunks":[]} + }))]) + .await; + let request = + crate::ocr::test_support::with_source(wire_request(model, &base, options), source); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /parse ")); + assert_eq!(request_body(&requests[0]), expected); + } + + #[rstest] + #[case("parse-v3")] + #[case("parse-legacy")] + #[tokio::test] + async fn data_uri_upload_preserves_multipart_headers(#[case] model: &str) { + let (base, seen, server) = mock_server(vec![ + MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), + MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})), + ]) + .await; + let mut request = wire_request(&format!("reducto/{model}"), &base, json!({})); + request.transport.extra_headers = vec![ + ("Content-Type".into(), "application/json".into()), + ("X-Trace".into(), "upload-test".into()), + ]; + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.pages[0].markdown, "hello"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert!(requests[0].starts_with("POST /upload ")); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("content-type: multipart/form-data; boundary=") + ); + assert!(requests[0].contains("x-trace: upload-test")); + assert!(requests[0].contains("application/pdf")); + assert!(requests[0].contains("abc")); + assert!(requests[1].starts_with("POST /parse ")); + } + + #[tokio::test] + async fn response_received_stays_after_reducto_upload_and_parse() { + let (base, seen, server) = mock_server(vec![ + MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), + MockResponse::json(json!({"result":{"chunks":[]}})), + ]) + .await; + let request_count = seen.clone(); + let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) + .with_observer(move |event| { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { + assert_eq!(request_count.lock().unwrap().len(), 2); + assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); + } + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 2); + } + + #[rstest] + #[case("https://example.com/a.pdf")] + #[case("reducto://")] + #[case("data:application/pdf;base64")] + #[case("data:application/pdf;base64,INVALID!")] + #[tokio::test] + async fn rejects_invalid_document_sources_before_network(#[case] source: &str) { + let request = crate::ocr::test_support::with_source( + wire_request("reducto/parse-v3", "http://127.0.0.1:1", json!({})), + source, + ); + assert!(perform_ocr(request).await.is_err()); + } + + #[tokio::test] + async fn facade_omits_native_response_by_default_and_preserves_auth_priority() { + let raw = json!({"job_id":"job-1","result":{"chunks":[]}}); + let (base, seen, server) = mock_server(vec![MockResponse::json(raw)]).await; + let mut request = crate::ocr::test_support::with_source( + wire_request("reducto/parse-v3", &base, json!({})), + "reducto://ready.pdf", + ); + request.transport.extra_headers = + vec![("authorization".into(), "Bearer existing".into())]; + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.provider_native_response, None); + assert!( + seen.lock().unwrap()[0] + .to_ascii_lowercase() + .contains("authorization: bearer existing") + ); + } + + #[rstest] + #[case("reducto/parse-v3")] + #[case("reducto/parse-legacy")] + #[tokio::test] + async fn guardrail_headers_reach_upload_and_parse(#[case] model: &str) { + let (base, seen, server) = mock_server(vec![ + MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), + MockResponse::json(json!({"result":{"chunks":[]}})), + ]) + .await; + let mut request = wire_request(model, &base, json!({})); + request.transport.extra_headers = + vec![("authorization".into(), "Bearer original".into())]; + let host = LocalOcrHost::new(request).with_before_send(|wire, _| { + Ok(WireRequest { + headers: vec![("authorization".into(), "Bearer guarded".into())], + ..wire + }) + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert!(requests[0].starts_with("POST /upload ")); + assert!(requests[1].starts_with("POST /parse ")); + for request in requests.iter() { + assert!(request.contains("authorization: Bearer guarded")); + assert!(!request.contains("Bearer original")); + } + } + } +} + +#[cfg(test)] +mod vertex_ai_tests { + use litellm_auth::InputSource; + use litellm_llms::base_llm::ocr::{settings::OcrSettings, transformation::OcrResponseFormat}; + use serde_json::{Value, json}; + + use crate::ocr::test_support::{ + MockResponse, mock_server, ocr_client, perform_ocr, wire_request, + }; + + fn request_body(request: &str) -> Value { + serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() + } + + #[tokio::test] + async fn facade_executes_vertex_mistral_with_resolved_project_and_location() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"hello"}], + "usage_info":{"pages_processed":1} + }))]) + .await; + let request = wire_request( + "vertex_ai/mistral-ocr-maas", + &base, + json!({ + "vertex_project":"project-1", + "vertex_location":"europe-west4", + "extract_footer":true + }), + ); + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.pages[0].markdown, "hello"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with( + "POST /v1/projects/project-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict " + )); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer test-key") + ); + assert_eq!( + request_body(&requests[0]), + json!({ + "model":"mistral-ocr-maas", + "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, + "extract_footer":true + }) + ); + } + + #[tokio::test] + async fn configured_project_and_location_apply_when_the_call_sets_neither() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let client = ocr_client().with_settings(OcrSettings { + vertex_project: Some("configured-project".into()), + vertex_location: Some("europe-west4".into()), + ..OcrSettings::default() + }); + + crate::ocr::client::perform( + &client, + wire_request("vertex_ai/mistral-ocr-maas", &base, json!({})), + ) + .await + .unwrap(); + server.await.unwrap(); + assert!(seen.lock().unwrap()[0].starts_with( + "POST /v1/projects/configured-project/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict " + )); + } + + #[tokio::test] + async fn supplied_authorization_is_forwarded_without_a_static_token() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let mut request = wire_request( + "vertex_ai/model", + &base, + json!({"vertex_project":"project-1"}), + ); + request.credentials.api_key = None; + request.transport.extra_headers = vec![("authorization".into(), "Bearer supplied".into())]; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert!( + seen.lock().unwrap()[0] + .to_ascii_lowercase() + .contains("authorization: bearer supplied") + ); + } + + #[tokio::test] + async fn invalid_credentials_fail_before_provider_http() { + let request = wire_request( + "vertex_ai/model", + "http://127.0.0.1:1", + json!({"vertex_credentials": true}), + ); + let error = perform_ocr(request).await.unwrap_err(); + assert!(error.to_string().contains("vertex_credentials")); + } + + #[tokio::test] + async fn request_controlled_api_base_is_rejected_before_vertex_auth() { + let mut request = wire_request( + "vertex_ai/mistral-ocr-maas", + "https://caller.example", + json!({"vertex_project":"project-1"}), + ); + request.credentials.api_base = Some(litellm_auth::Sourced::new( + "https://caller.example".into(), + InputSource::Request, + )); + + let error = perform_ocr(request).await.unwrap_err(); + assert!( + error + .to_string() + .contains("request-controlled Vertex AI endpoint") + ); + } + + #[tokio::test] + async fn adapters_build_complete_requests_and_share_mistral_normalization() { + use std::time::Duration; + + use litellm_llms::{ + base_llm::ocr::transformation::BaseOcrConfig, + mistral::ocr::transformation::MistralOcrConfig, + vertex_ai::ocr::transformation::VertexAiOcrConfig, + }; + + use crate::ocr::test_support::ocr_client; + + let client = ocr_client(); + let options = json!({ + "pages": [0, 2], + "include_image_base64": true, + "vertex_project": "project-1", + "vertex_location": "us-central1", + "unknown": "ignored" + }); + let direct = wire_request( + "mistral/mistral-ocr-maas", + "https://mistral.test", + options.clone(), + ); + let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options); + let direct = crate::ocr::prepare::prepare_request_for_test( + crate::ocr::test_support::resolved_request(direct), + ); + let vertex = crate::ocr::prepare::prepare_request_for_test( + crate::ocr::test_support::resolved_request(vertex), + ); + let direct_http = MistralOcrConfig + .prepare_request(&direct, &client, &crate::ocr::test_support::NoHooks) + .await + .unwrap(); + let vertex_http = VertexAiOcrConfig + .prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks) + .await + .unwrap(); + assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr"); + assert_eq!( + vertex_http.url(), + "https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict" + ); + for http in [&direct_http, &vertex_http] { + assert_eq!(http.header("authorization").unwrap(), "Bearer test-key"); + assert_eq!(http.header("content-type").unwrap(), "application/json"); + assert_eq!(http.timeout(), Some(Duration::from_secs(2))); + let body: Value = serde_json::from_slice(http.body()).unwrap(); + assert_eq!( + body, + json!({ + "model": "mistral-ocr-maas", + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "pages": [0, 2], + "include_image_base64": true, + "unknown": "ignored" + }) + ); + } + let payload = json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}); + let raw = serde_json::to_vec(&payload).unwrap(); + let direct_response = MistralOcrConfig + .transform_ocr_response(&direct.model, &raw, OcrResponseFormat::Litellm) + .unwrap() + .into_json(); + let vertex_response = VertexAiOcrConfig + .transform_ocr_response(&vertex.model, &raw, OcrResponseFormat::Litellm) + .unwrap() + .into_json(); + assert_eq!(direct_response, vertex_response); + assert_eq!(direct_response["model"], "mistral-ocr-maas"); + assert_eq!(direct_response["object"], "ocr"); + assert_eq!(direct_response["extra"], "preserved"); + } + + mod transformation { + + use rstest::rstest; + use serde_json::{Value, json}; + + use crate::ocr::test_support::wire_request; + + #[rstest] + #[case::mistral(false)] + #[case::vertex(true)] + #[tokio::test] + async fn configs_build_complete_requests_and_share_mistral_normalization( + #[case] use_vertex: bool, + ) { + use std::time::Duration; + + use litellm_llms::{ + base_llm::ocr::transformation::BaseOcrConfig, + mistral::ocr::transformation::MistralOcrConfig, + vertex_ai::ocr::transformation::VertexAiOcrConfig, + }; + + use crate::ocr::test_support::ocr_client; + + let client = ocr_client(); + let options = json!({ + "pages": [0, 2], + "include_image_base64": true, + "vertex_project": "project-1", + "vertex_location": "us-central1", + "unknown": "preserved" + }); + let direct = wire_request( + "mistral/mistral-ocr-maas", + "https://mistral.test", + options.clone(), + ); + let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options); + let direct = crate::ocr::prepare::prepare_request_for_test( + crate::ocr::test_support::resolved_request(direct), + ); + let vertex = crate::ocr::prepare::prepare_request_for_test( + crate::ocr::test_support::resolved_request(vertex), + ); + let direct_http = MistralOcrConfig + .prepare_request(&direct, &client, &crate::ocr::test_support::NoHooks) + .await + .unwrap(); + let vertex_http = VertexAiOcrConfig + .prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks) + .await + .unwrap(); + assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr"); + assert_eq!( + vertex_http.url(), + "https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict" + ); + let http = if use_vertex { + &vertex_http + } else { + &direct_http + }; + assert_eq!(http.header("authorization").unwrap(), "Bearer test-key"); + assert_eq!(http.header("content-type").unwrap(), "application/json"); + assert_eq!(http.timeout(), Some(Duration::from_secs(2))); + let body: Value = serde_json::from_slice(http.body()).unwrap(); + assert_eq!( + body, + json!({ + "model": "mistral-ocr-maas", + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "pages": [0, 2], + "include_image_base64": true, + "unknown": "preserved" + }) + ); + let payload = serde_json::to_vec( + &json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}), + ) + .unwrap(); + let direct_response = MistralOcrConfig + .transform_ocr_response(&direct.model, &payload, Default::default()) + .unwrap() + .into_json(); + let vertex_response = VertexAiOcrConfig + .transform_ocr_response(&vertex.model, &payload, Default::default()) + .unwrap() + .into_json(); + assert_eq!(direct_response, vertex_response); + assert_eq!(direct_response["model"], "mistral-ocr-maas"); + assert_eq!(direct_response["object"], "ocr"); + assert_eq!(direct_response["extra"], "preserved"); + } + } +} + +#[cfg(test)] +mod vertex_ai_deepseek_tests { + use litellm_auth::InputSource; + use serde_json::{Value, json}; + + use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; + + fn request_body(request: &str) -> Value { + serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() + } + + #[tokio::test] + async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "choices":[{"message":{"content":"recognized"}}], + "usage":{"prompt_tokens":1} + }))]) + .await; + let request = wire_request( + "vertex_ai/deepseek-ocr-maas", + &base, + json!({ + "vertex_project":"project-1", + "vertex_location":"europe-west4", + "temperature":0.1, + "future_ocr_option":true, + "extra_body":{"provider_option":"value"} + }), + ); + let request = crate::ocr::test_support::with_source(request, "gs://bucket/document.pdf"); + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.pages[0].markdown, "recognized"); + assert_eq!( + response.usage_info.unwrap().extra_fields["prompt_tokens"], + 1 + ); + let requests = seen.lock().unwrap(); + assert!(requests[0].starts_with( + "POST /v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions " + )); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer test-key") + ); + let body = request_body(&requests[0]); + assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); + assert_eq!(body["temperature"], 0.1); + assert_eq!(body["future_ocr_option"], true); + assert!(body.get("extra_body").is_none()); + assert_eq!( + body["messages"][0]["content"][0], + json!({"type":"image_url","image_url":"gs://bucket/document.pdf"}) + ); + } + + #[test] + fn host_registration_selects_deepseek_without_affecting_mistral() { + assert!(crate::ocr::arguments::is_supported_request( + "deepseek-ocr-maas", + Some("vertex_ai") + )); + assert!(crate::ocr::arguments::is_supported_request( + "mistral-ocr-maas", + Some("vertex_ai") + )); + } + + #[tokio::test] + async fn request_controlled_api_base_is_rejected_before_vertex_auth() { + let mut request = wire_request( + "vertex_ai/deepseek-ocr-maas", + "https://caller.example", + json!({"vertex_project":"project-1"}), + ); + request.credentials.api_base = Some(litellm_auth::Sourced::new( + "https://caller.example".into(), + InputSource::Request, + )); + + let error = perform_ocr(request).await.unwrap_err(); + assert!( + error + .to_string() + .contains("request-controlled Vertex AI endpoint") + ); + } + + mod deepseek_transformation { + use serde_json::json; + + use super::*; + use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; + + #[tokio::test] + async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "choices":[{"message":{"content":"recognized"}}], + "usage":{"prompt_tokens":1} + }))]) + .await; + let request = wire_request( + "vertex_ai/deepseek-ocr-maas", + &base, + json!({ + "vertex_project":"project-1", + "vertex_location":"europe-west4", + "temperature":0.1, + "future_ocr_option":true, + "extra_body":{"provider_option":"value"} + }), + ); + let request = + crate::ocr::test_support::with_source(request, "gs://bucket/document.pdf"); + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.pages[0].markdown, "recognized"); + assert_eq!( + response.usage_info.unwrap().extra_fields["prompt_tokens"], + 1 + ); + let requests = seen.lock().unwrap(); + assert!(requests[0].starts_with( + "POST /v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions " + )); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer test-key") + ); + let body = request_body(&requests[0]); + assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); + assert_eq!(body["temperature"], 0.1); + assert_eq!(body["future_ocr_option"], true); + assert_eq!(body["provider_option"], "value"); + assert!(body.get("vertex_project").is_none()); + assert!(body.get("extra_body").is_none()); + assert_eq!( + body["messages"][0]["content"][0], + json!({"type":"image_url","image_url":"gs://bucket/document.pdf"}) + ); + } + } +} + +#[cfg(test)] +pub(crate) mod tests { + use std::sync::{Arc, Mutex}; + + use futures_util::future::BoxFuture; + use litellm_auth_gcp::VertexAuth; + use litellm_host::{ + event::{CallEvent, MachineEvent, WireRequest}, + host::{Host, HostOp, HostResult}, + machine::{HostFailure, Machine, MachineStep}, + }; + use litellm_http::{ + HttpClientPool, HttpSettings, Resolution, + media::{PublicDnsResolver, UrlPolicy}, + }; + use litellm_llms::base_llm::ocr::{ + error::Error as OcrError, + handler::OcrClient, + settings::OcrSettings, + transformation::{ + BaseOcrConfig, LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig, + }, + }; + use litellm_secrets::source::SecretSource; + use rstest::rstest; + use serde_json::{Value, json}; + + use crate::ocr::route::{LocalOcrHost, OcrOp, OcrOpResult, ocr_machine}; + use crate::ocr::{ + test_support::{ + MockResponse, mock_server, ocr_client, perform_ocr, perform_ocr_with, wire_request, + }, + wire::{OcrWireRequest, decode_request}, + }; + + struct RecordingSecretSource { + names: Arc>>, + values: &'static [(&'static str, &'static str)], + api_base: String, + } + + impl SecretSource for RecordingSecretSource { + fn get_secret_str<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result, litellm_secrets::Error>> + { + self.names.lock().unwrap().push(name.to_owned()); + Box::pin(async move { + Ok(match name { + "MISTRAL_AZURE_API_BASE" => Some(self.api_base.clone()), + _ => self + .values + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| value.to_string()), + } + .map(litellm_secrets::SecretValue::new)) + }) + } + } + + #[rstest] + #[case::mistral("mistral/model", json!({}))] + #[case::vertex("vertex_ai/mistral-ocr-latest", json!({"vertex_project":"test-project", "vertex_location":"us-central1"}))] + #[tokio::test] + async fn ocr_contract_upstream_error_preserves_status_body_and_headers( + #[case] model: &str, + #[case] options: Value, + ) { + let payload = json!({"message": format!("{} END-OF-PROVIDER-BODY", "x".repeat(4096))}); + let expected_body = serde_json::to_string(&payload).unwrap(); + let (base, seen, server) = mock_server(vec![MockResponse { + status: 422, + headers: vec![ + ("Retry-After", "17".into()), + ("X-Request-ID", "request-123".into()), + ("X-Future-Header", "retained".into()), + ], + body: payload, + }]) + .await; + let error = perform_ocr(wire_request(model, &base, options)) + .await + .unwrap_err(); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 1); + let OcrError::Provider { + status, + body, + headers, + } = error + else { + panic!("expected provider error, got {error:?}"); + }; + assert_eq!(status, 422); + for (name, value) in [ + ("retry-after", "17"), + ("x-request-id", "request-123"), + ("x-future-header", "retained"), + ] { + assert!( + headers + .iter() + .any(|(key, actual)| key.eq_ignore_ascii_case(name) && actual == value) + ); + } + assert_eq!( + body.len(), + expected_body.len(), + "provider error body was truncated" + ); + assert_eq!(body, expected_body); + } + + #[test] + fn request_boundary_selects_mistral_and_rejects_unknown_providers() { + let request = OcrWireRequest { + model: "mistral/model".into(), + document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), + api_key: Some(litellm_auth::SecretValue::new("key")), + api_base: None, + custom_llm_provider: None, + extra_headers: None, + optional_params: json!({"extract_header":true,"unknown":42}) + .as_object() + .unwrap() + .clone(), + input_sources: Default::default(), + timeout_seconds: None, + }; + assert!(decode_request(request).is_ok()); + assert!( + decode_request(OcrWireRequest { + model: "model".into(), + document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), + api_key: Some(litellm_auth::SecretValue::new("key")), + api_base: None, + custom_llm_provider: Some("unknown".into()), + extra_headers: None, + optional_params: serde_json::Map::new(), + input_sources: Default::default(), + timeout_seconds: None, + }) + .is_err() + ); + } + + #[tokio::test] + async fn facade_executes_direct_mistral_once() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"hello","custom":"preserved"}], + "usage_info":{"pages_processed":1} + }))]) + .await; + let result = perform_ocr(wire_request( + "mistral/model", + &base, + json!({"pages":"0,2-4","extract_header":true,"unknown":"ignored"}), + )) + .await + .unwrap(); + server.await.unwrap(); + assert_eq!(result.pages[0].markdown, "hello"); + assert_eq!(result.pages[0].extra_fields["custom"], "preserved"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /v1/ocr ")); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer test-key\r\n") + ); + let body: Value = + serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!( + body, + json!({ + "model":"model", + "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, + "pages":"0,2-4", + "extract_header":true, + "unknown":"ignored" + }) + ); + } + + #[tokio::test] + async fn facade_retains_native_response_when_requested() { + let provider_response = json!({ + "pages":[{"index":0,"markdown":"hello"}], + "usage_info":{"pages_processed":1}, + "provider_only":"preserved" + }); + let (base, _, server) = + mock_server(vec![MockResponse::json(provider_response.clone())]).await; + let response = perform_ocr(wire_request( + "mistral/model", + &base, + json!({"req_format":"native"}), + )) + .await + .unwrap(); + + server.await.unwrap(); + assert_eq!( + response.provider_native_response.map(Value::Object), + Some(provider_response) + ); + } + + #[rstest] + #[case::plain_key(&[("MISTRAL_API_KEY", "plain")], "plain")] + #[case::azure_key_wins(&[("MISTRAL_AZURE_API_KEY", "azure"), ("MISTRAL_API_KEY", "plain")], "azure")] + #[case::empty_azure_key_falls_through(&[("MISTRAL_AZURE_API_KEY", ""), ("MISTRAL_API_KEY", "plain")], "plain")] + #[tokio::test] + async fn mistral_env_fallbacks_follow_python_through_the_injected_secret_source( + #[case] secrets: &'static [(&'static str, &'static str)], + #[case] expected_key: &str, + ) { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let names = Arc::new(Mutex::new(Vec::new())); + let client = ocr_client().with_secrets(Arc::new(RecordingSecretSource { + names: names.clone(), + values: secrets, + api_base: base.clone(), + })); + let request = decode_request(OcrWireRequest { + model: "mistral/model".into(), + document: json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), + api_key: None, + api_base: None, + custom_llm_provider: None, + extra_headers: None, + optional_params: Default::default(), + input_sources: Default::default(), + timeout_seconds: Some(2.0), + }) + .unwrap(); + + crate::ocr::client::perform(&client, request).await.unwrap(); + server.await.unwrap(); + assert_eq!( + *names.lock().unwrap(), + litellm_llms::mistral::ocr::transformation::MistralOcrConfig.secret_names() + ); + assert!(seen.lock().unwrap()[0].contains(&format!("authorization: Bearer {expected_key}"))); + } + + #[tokio::test] + async fn mistral_ocr_resolves_provider_secrets_before_transformation() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let names = Arc::new(Mutex::new(Vec::new())); + let client = ocr_client().with_secrets(Arc::new(RecordingSecretSource { + names: names.clone(), + values: &[("MISTRAL_API_KEY", "source-key")], + api_base: base.clone(), + })); + let request = decode_request(OcrWireRequest { + model: "mistral/mistral-ocr-latest".into(), + document: json!({ + "type":"document_url", + "document_url":"data:application/pdf;base64,YWJj" + }), + api_key: None, + api_base: None, + custom_llm_provider: None, + extra_headers: None, + optional_params: Default::default(), + input_sources: Default::default(), + timeout_seconds: Some(2.0), + }) + .unwrap(); + + crate::ocr::client::perform(&client, request).await.unwrap(); + server.await.unwrap(); + assert_eq!( + *names.lock().unwrap(), + litellm_llms::mistral::ocr::transformation::MistralOcrConfig.secret_names() + ); + assert!(seen.lock().unwrap()[0].contains("authorization: Bearer source-key")); + } + + #[tokio::test] + async fn ocr_client_uses_the_injected_http_pool_configuration() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let settings = HttpSettings { + user_agent: Some("host-owned/1".into()), + ..HttpSettings::default() + }; + let client = OcrClient::new( + &HttpClientPool::new(Arc::new(PublicDnsResolver)), + &Resolution::from(&settings).config, + UrlPolicy::default(), + VertexAuth::default(), + OcrSettings::default(), + Arc::new(litellm_secrets::source::EnvironmentSecrets::default()), + ) + .unwrap(); + crate::ocr::client::perform(&client, wire_request("mistral/model", &base, json!({}))) + .await + .unwrap(); + server.await.unwrap(); + assert!(seen.lock().unwrap()[0].contains("user-agent: host-owned/1")); + } + + fn event_name(event: &CallEvent) -> &'static str { + match event { + CallEvent::Started { .. } => "started", + CallEvent::Machine(MachineEvent::ResponseReceived { .. }) => "response", + CallEvent::Succeeded { .. } => "success", + CallEvent::Failed { .. } => "failure", + } + } + + fn recording_host( + request: crate::ocr::types::LiteLLMOcrRequest, + events: Arc>>, + block: bool, + ) -> LocalOcrHost { + let before_send_events = events.clone(); + LocalOcrHost::new(request) + .with_before_send(move |wire, _| { + before_send_events.lock().unwrap().push("before_send"); + if block { + return Err(OcrError::InvalidRequest("blocked".into())); + } + Ok(wire) + }) + .with_observer(move |event| events.lock().unwrap().push(event_name(event))) + } + + #[tokio::test] + async fn lifecycle_sends_headers_returned_by_the_before_send_operation() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))) + .with_before_send(|mut wire, _| { + wire.headers + .push(("x-core-callback".into(), "edited".into())); + Ok(wire) + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + + assert!(seen.lock().unwrap()[0].contains("x-core-callback: edited")); + } + + #[tokio::test] + async fn before_send_context_names_the_route_and_its_secrets() { + let (base, _, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let observed = Arc::new(Mutex::new(None)); + let captured = observed.clone(); + let host = LocalOcrHost::new(wire_request( + "mistral/model", + &base, + json!({"pages": [0], "req_format": "native"}), + )) + .with_before_send(move |wire, context| { + *captured.lock().unwrap() = Some((wire.clone(), context.clone())); + Ok(wire) + }); + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + let (wire, context) = observed.lock().unwrap().take().unwrap(); + assert_eq!(context.custom_llm_provider, "mistral"); + assert_eq!(context.model, "model"); + assert_eq!(wire.body["pages"], json!([0])); + assert!(context.secret_fields.is_empty()); + assert_eq!(context.optional_params["req_format"], "native"); + + let (base, _, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let observed = Arc::new(Mutex::new(None)); + let captured = observed.clone(); + let request = wire_request( + "azure_ai/model", + &base, + json!({"client_secret": "shh", "tenant_id": "t"}), + ); + let request = request.with_document(crate::ocr::types::OcrDocumentInput::Bytes { + bytes: b"abc".as_slice().into(), + file_name: None, + mime_type: Some("application/pdf".into()), + }); + let host = LocalOcrHost::new(request).with_before_send(move |wire, context| { + *captured.lock().unwrap() = Some(context.clone()); + Ok(wire) + }); + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + let context = observed.lock().unwrap().take().unwrap(); + assert_eq!(context.secret_fields, ["client_secret"]); + } + + #[tokio::test] + async fn lifecycle_orders_hooks_and_emits_one_success() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let events = Arc::new(Mutex::new(Vec::new())); + let host = recording_host( + wire_request("mistral/model", &base, json!({})), + events.clone(), + false, + ); + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + assert_eq!( + *events.lock().unwrap(), + ["started", "before_send", "response", "success"] + ); + assert_eq!(seen.lock().unwrap().len(), 1); + } + + #[tokio::test] + async fn lifecycle_blocking_prevents_execution_and_emits_one_failure() { + let events = Arc::new(Mutex::new(Vec::new())); + let host = recording_host( + wire_request("mistral/model", "http://127.0.0.1:1", json!({})), + events.clone(), + true, + ); + let error = perform_ocr_with(host).await.unwrap_err(); + assert!(matches!(error, OcrError::InvalidRequest(message) if message == "blocked")); + assert_eq!( + *events.lock().unwrap(), + ["started", "before_send", "failure"] + ); + } + + #[tokio::test] + async fn upstream_failure_emits_one_terminal_failure() { + let (base, seen, server) = mock_server(vec![MockResponse { + status: 500, + headers: vec![], + body: json!({"error":"failed"}), + }]) + .await; + let events = Arc::new(Mutex::new(Vec::new())); + let host = recording_host( + wire_request("mistral/model", &base, json!({})), + events.clone(), + false, + ); + assert!(perform_ocr_with(host).await.is_err()); + server.await.unwrap(); + assert_eq!( + *events.lock().unwrap(), + ["started", "before_send", "failure"] + ); + assert_eq!(seen.lock().unwrap().len(), 1); + } + + /// Drives the machine by hand, answering every op through `host` except `before_send`, + /// which `intercept` answers so a test can fail or cancel exactly there. + async fn drive_until( + client: OcrClient, + host: &LocalOcrHost, + mut intercept: impl FnMut(WireRequest) -> Result>, + ) -> ( + Result, + Vec<&'static str>, + crate::ocr::route::OcrMachine, + ) { + let mut machine = ocr_machine(client); + let mut result = None; + let mut ops = Vec::new(); + let outcome = loop { + let op = match machine.resume(result.take()).await { + Ok(MachineStep::Host(op)) => op, + Ok(MachineStep::Complete(response)) => break Ok(response), + Err(error) => break Err(error), + }; + let answer = match op { + HostOp::Route(op) => { + ops.push(match op { + OcrOp::ProjectRequest => "ProjectRequest", + OcrOp::ReadDocument => "ReadDocument", + OcrOp::AcquireAzureAdToken => "AcquireAzureAdToken", + }); + host.route(op) + .await + .map(HostResult::Route) + .map_err(HostFailure::Error) + } + HostOp::BeforeSend { wire, .. } => { + ops.push("BeforeSend"); + intercept(*wire).map(|wire| HostResult::BeforeSend(Box::new(wire))) + } + HostOp::Emit(event) => { + let event = CallEvent::Machine(event); + ops.push(event_name(&event)); + host.emit(&event) + .await + .map(|()| HostResult::Emitted) + .map_err(HostFailure::Error) + } + }; + match answer { + Ok(answer) => result = Some(answer), + Err(failure) => break machine.interrupt(failure).await, + } + }; + (outcome, ops, machine) + } + + #[tokio::test] + async fn failed_before_send_does_not_replay_or_reach_transport() { + let host = LocalOcrHost::new(wire_request( + "mistral/model", + "http://127.0.0.1:1", + json!({}), + )); + let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, |_| { + Err(HostFailure::Error(OcrError::InvalidRequest( + "before_send failed".into(), + ))) + }) + .await; + assert!( + matches!(outcome, Err(OcrError::InvalidRequest(message)) if message == "before_send failed") + ); + assert_eq!(ops, ["ProjectRequest", "BeforeSend"]); + assert!(machine.resume(None).await.is_err()); + } + + #[tokio::test] + async fn invalid_provider_response_emits_response_received_before_normalization_failure() { + let (base, seen, server) = + mock_server(vec![MockResponse::json(json!({"pages":"invalid"}))]).await; + let responses_received = Arc::new(Mutex::new(Vec::new())); + let observed = responses_received.clone(); + let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))) + .with_observer(move |event| { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { + observed.lock().unwrap().push(raw.body.clone()); + } + }); + let error = perform_ocr_with(host).await.unwrap_err(); + server.await.unwrap(); + assert!(matches!(error, OcrError::ResponseField { .. })); + assert_eq!(seen.lock().unwrap().len(), 1); + assert_eq!( + *responses_received.lock().unwrap(), + [r#"{"pages":"invalid"}"#] + ); + } + + #[tokio::test] + async fn direct_native_host_drives_the_same_state_machine() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"native"}] + }))]) + .await; + let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))); + let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, Ok).await; + server.await.unwrap(); + assert_eq!(outcome.unwrap().pages[0].markdown, "native"); + assert_eq!(seen.lock().unwrap().len(), 1); + assert_eq!(ops, ["ProjectRequest", "BeforeSend", "response"]); + assert!(matches!( + machine.resume(None).await, + Err(OcrError::InvalidRequest(_)) + )); + } + + async fn drive_native_file_call( + request: crate::ocr::types::LiteLLMOcrRequest, + content: Result, + ) -> (Result, usize) { + let reads = Arc::new(Mutex::new(0)); + let counted = reads.clone(); + let content = Mutex::new(Some(content)); + let host = LocalOcrHost::new(request).with_reader(move || { + *counted.lock().unwrap() += 1; + content.lock().unwrap().take().unwrap() + }); + let outcome = perform_ocr_with(host).await; + let reads = *reads.lock().unwrap(); + (outcome, reads) + } + + #[tokio::test] + async fn host_reader_documents_are_read_once_at_the_core_selected_point_and_encoded() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"file"}] + }))]) + .await; + let request = wire_request("mistral/model", &base, json!({})).with_document( + crate::ocr::types::OcrDocumentInput::HostReader { + mime_type: Some("application/pdf".into()), + }, + ); + let (response, reads) = drive_native_file_call( + request, + Ok(crate::ocr::types::OcrFileContent { + bytes: b"abc".as_slice().into(), + file_name: Some("scan.png".into()), + }), + ) + .await; + server.await.unwrap(); + assert_eq!(response.unwrap().pages[0].markdown, "file"); + assert_eq!(reads, 1); + assert!(seen.lock().unwrap()[0].contains("data:application/pdf;base64,YWJj")); + } + + #[tokio::test] + async fn host_reader_failures_and_empty_files_fail_before_the_provider_is_called() { + let (base, seen, _server) = mock_server(vec![]).await; + let request = wire_request("mistral/model", &base, json!({})); + let failure = OcrError::InvalidRequest("reader exploded".into()); + let (response, reads) = drive_native_file_call( + request + .with_document(crate::ocr::types::OcrDocumentInput::HostReader { mime_type: None }), + Err(failure.clone()), + ) + .await; + assert!( + matches!(response.unwrap_err(), OcrError::InvalidRequest(message) if message == "reader exploded") + ); + assert_eq!(reads, 1); + + let request = wire_request("mistral/model", &base, json!({})); + let (response, _) = drive_native_file_call( + request + .with_document(crate::ocr::types::OcrDocumentInput::HostReader { mime_type: None }), + Ok(crate::ocr::types::OcrFileContent { + bytes: Default::default(), + file_name: None, + }), + ) + .await; + assert!(matches!(response.unwrap_err(), OcrError::EmptyFile)); + assert!(seen.lock().unwrap().is_empty()); + } + + #[tokio::test] + async fn path_documents_are_read_by_core_without_a_host_operation() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"path"}] + }))]) + .await; + let dir = std::env::temp_dir().join(format!("litellm-ocr-{}", rand::random::())); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("scan.png"); + std::fs::write(&path, b"abc").unwrap(); + let request = wire_request("mistral/model", &base, json!({})).with_document( + crate::ocr::types::OcrDocumentInput::Path { + path: path.clone(), + mime_type: None, + }, + ); + let (response, reads) = + drive_native_file_call(request, Err(OcrError::InvalidRequest("unused".into()))).await; + server.await.unwrap(); + std::fs::remove_dir_all(&dir).unwrap(); + assert_eq!(response.unwrap().pages[0].markdown, "path"); + assert_eq!(reads, 0); + assert!(seen.lock().unwrap()[0].contains("data:image/png;base64,YWJj")); + + let (base, seen, _server) = mock_server(vec![]).await; + let request = wire_request("mistral/model", &base, json!({})); + let (response, _) = drive_native_file_call( + request.with_document(crate::ocr::types::OcrDocumentInput::Path { + path: path.clone(), + mime_type: None, + }), + Err(OcrError::InvalidRequest("unused".into())), + ) + .await; + assert!(matches!( + response.unwrap_err(), + OcrError::FileRead { path: failed, source } if failed == path && source.kind() == std::io::ErrorKind::NotFound + )); + assert!(seen.lock().unwrap().is_empty()); + } + + #[tokio::test] + async fn cancellation_at_before_send_prevents_execution_and_further_resumption() { + let host = LocalOcrHost::new(wire_request( + "mistral/model", + "http://127.0.0.1:1", + json!({}), + )); + let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, |_| { + Err(HostFailure::Cancelled(OcrError::InvalidRequest( + "cancelled".into(), + ))) + }) + .await; + assert!( + matches!(outcome, Err(OcrError::InvalidRequest(message)) if message == "cancelled") + ); + assert_eq!(ops, ["ProjectRequest", "BeforeSend"]); + assert!(machine.resume(Some(HostResult::Emitted)).await.is_err()); + } + + #[tokio::test] + async fn missing_host_result_preserves_pending_operation() { + let request = wire_request("mistral/model", "http://127.0.0.1:1", json!({})); + let mut machine = ocr_machine(ocr_client()); + assert!(matches!( + machine.resume(None).await.unwrap(), + MachineStep::Host(HostOp::Route(OcrOp::ProjectRequest)) + )); + assert!(machine.resume(None).await.is_err()); + assert!(matches!( + machine + .resume(Some(HostResult::Route(OcrOpResult::Request { + request: Box::new(request), + caller_token: false, + }))) + .await + .unwrap(), + MachineStep::Host(HostOp::BeforeSend { .. }) + )); + } + + async fn read_bounded_response( + response: Vec, + limit: usize, + ) -> Result { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0; 4096]; + assert!(socket.read(&mut request).await.unwrap() > 0); + socket.write_all(&response).await.unwrap(); + std::future::pending::<()>().await; + }); + let response = reqwest::Client::new() + .get(format!("http://{address}")) + .send() + .await + .unwrap(); + let result = tokio::time::timeout( + std::time::Duration::from_secs(2), + litellm_llms::base_llm::ocr::handler::read_response_bytes(response, limit), + ) + .await; + server.abort(); + let _ = server.await; + result.expect("bounded reads must finish without waiting for the rest of an oversized body") + } + + #[tokio::test] + async fn response_limit_accepts_exact_size_and_rejects_declared_and_chunked_overflow() { + use litellm_llms::base_llm::ocr::error::Error; + + for response in [ + "HTTP/1.1 200 OK\r\nContent-Length: 8\r\n\r\nabcdefgh", + "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n4\r\nefgh\r\n0\r\n\r\n", + ] { + assert_eq!( + read_bounded_response(response.as_bytes().to_vec(), 8) + .await + .unwrap(), + "abcdefgh" + ); + } + for response in [ + "HTTP/1.1 200 OK\r\nContent-Length: 9\r\n\r\n", + "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n5\r\nefghi\r\n", + ] { + assert!(matches!( + read_bounded_response(response.as_bytes().to_vec(), 8).await, + Err(Error::TooLarge { limit: 8 }) + )); + } + } + + #[rstest] + #[case::declared("Content-Length: 1000000")] + #[case::chunked("Transfer-Encoding: chunked")] + #[tokio::test] + async fn oversized_error_retains_http_status_and_bounded_diagnostics_without_draining( + #[case] headers: &str, + ) { + let prefix = "x".repeat(4096); + let body = if headers.starts_with("Transfer") { + format!("{:x}\r\n{prefix}\r\n", prefix.len()) + } else { + prefix.clone() + }; + let response = format!("HTTP/1.1 429 Too Many Requests\r\n{headers}\r\n\r\n{body}"); + let error = read_bounded_response(response.into_bytes(), prefix.len()) + .await + .unwrap_err(); + match error { + OcrError::Transport(litellm_http::transport::Error::Http { status, body }) => { + assert_eq!(status, 429); + assert_eq!(body, prefix); + } + error => panic!("unexpected error: {error}"), + } + } + + #[test] + fn response_limit_is_validated_and_not_forwarded_to_the_provider() { + let request = wire_request( + "mistral/model", + "http://localhost", + json!({"max_response_bytes": 123}), + ); + assert_eq!(request.transport.max_response_bytes, 123); + assert!(!request.optional_params.contains_key("max_response_bytes")); + for value in [ + json!(0), + json!(-1), + json!(true), + json!("123"), + json!(1.5), + json!(OCR_RESPONSE_MAX_BYTES + 1), + Value::Null, + ] { + let wire = serde_json::from_value(json!({ + "model": "mistral/model", "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "optional_params": {"max_response_bytes": value} + })).unwrap(); + let Err(error) = decode_request(wire) else { + panic!("invalid response limit accepted") + }; + assert!(error.to_string().contains("max_response_bytes")); + } + } + + #[derive(Debug)] + struct PendingToken { + entered: Arc, + dropped: Arc, + } + + struct TokenFutureDrop(Arc); + + impl Drop for TokenFutureDrop { + fn drop(&mut self) { + self.0.store(true, std::sync::atomic::Ordering::SeqCst); + } + } + + impl litellm_auth::TokenProvider for PendingToken { + fn acquire(&self) -> litellm_auth::TokenFuture<'_> { + Box::pin(async move { + let _guard = TokenFutureDrop(self.dropped.clone()); + self.entered.notify_one(); + std::future::pending().await + }) + } + } + + #[tokio::test] + async fn interrupt_drops_provider_captures_before_returning() { + use std::sync::atomic::{AtomicBool, Ordering}; + + let entered = Arc::new(tokio::sync::Notify::new()); + let dropped = Arc::new(AtomicBool::new(false)); + let request = wire_request("azure_ai/mistral-ocr", "https://example.invalid", json!({})); + let request = crate::ocr::types::LiteLLMOcrRequest { + transport: OcrTransportConfig { + extra_headers: vec![("authorization".into(), "Bearer test-key".into())], + ..request.transport + }, + azure_ad_token_provider: Some(litellm_auth::TokenProviderHandle::new(Arc::new( + PendingToken { + entered: entered.clone(), + dropped: dropped.clone(), + }, + ))), + ..request + }; + let host = LocalOcrHost::new(request); + let mut machine = ocr_machine(ocr_client()); + let mut result = None; + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + tokio::select! { + _ = entered.notified() => break, + step = machine.resume(result.take()) => { + result = Some(match step.unwrap() { + MachineStep::Host(HostOp::Route(op)) => HostResult::Route(host.route(op).await.unwrap()), + MachineStep::Host(HostOp::BeforeSend { wire, .. }) => { + HostResult::BeforeSend(wire) + } + MachineStep::Host(HostOp::Emit(_)) => HostResult::Emitted, + MachineStep::Complete(_) => panic!("pending provider completed"), + }); + } + } + } + }) + .await + .unwrap(); + assert!(!dropped.load(Ordering::SeqCst)); + let selected = OcrError::InvalidRequest("cancelled".into()); + let acknowledgement = machine.interrupt(HostFailure::Cancelled(selected.clone())); + assert!( + dropped.load(Ordering::SeqCst), + "interrupt returned while provider captures were still alive" + ); + assert!( + matches!(acknowledgement.await, Err(OcrError::InvalidRequest(message)) if message == "cancelled") + ); + } + + struct CallerTokenHost { + request: Mutex>, + trace: Mutex>, + } + + impl Host for CallerTokenHost { + async fn route(&self, op: OcrOp) -> Result { + match op { + OcrOp::ProjectRequest => { + self.trace.lock().unwrap().push("project".into()); + Ok(OcrOpResult::Request { + request: Box::new(self.request.lock().unwrap().take().unwrap()), + caller_token: true, + }) + } + OcrOp::AcquireAzureAdToken => { + self.trace.lock().unwrap().push("token".into()); + Ok(OcrOpResult::AzureAdToken( + litellm_auth::ResolvedCredential::Static(litellm_auth::SecretValue::new( + "caller-token", + )), + )) + } + OcrOp::ReadDocument => Err(OcrError::InvalidRequest("no reader".into())), + } + } + + async fn before_send( + &self, + wire: WireRequest, + _: &litellm_host::event::RequestContext, + ) -> Result { + let is_authorization = |name: &str| name.eq_ignore_ascii_case("authorization"); + let authorization = wire + .headers + .iter() + .find(|(name, _)| is_authorization(name)) + .map(|(_, value)| value.clone()) + .unwrap_or_default(); + self.trace + .lock() + .unwrap() + .push(format!("before_send:{authorization}")); + let headers = wire + .headers + .into_iter() + .map(|(name, value)| match is_authorization(&name) { + true => (name, "Bearer edited".to_string()), + false => (name, value), + }) + .collect(); + Ok(WireRequest { headers, ..wire }) + } + } + + #[tokio::test] + async fn the_callers_azure_token_is_acquired_before_before_send_which_can_still_replace_it() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let mut request = wire_request("azure_ai/model", &base, json!({})); + request.credentials.api_key = None; + let host = CallerTokenHost { + request: Mutex::new(Some(request)), + trace: Mutex::new(Vec::new()), + }; + + litellm_host::run::run(ocr_machine(ocr_client()), &host) + .await + .unwrap(); + server.await.unwrap(); + + assert_eq!( + *host.trace.lock().unwrap(), + ["project", "token", "before_send:Bearer caller-token"] + ); + assert!( + seen.lock().unwrap()[0] + .to_ascii_lowercase() + .contains("authorization: bearer edited\r\n") + ); + } + + #[tokio::test] + async fn interrupting_an_in_flight_provider_request_closes_its_connection() { + use tokio::io::AsyncReadExt; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let received = Arc::new(tokio::sync::Notify::new()); + let server_received = received.clone(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = Vec::new(); + let mut buffer = [0u8; 4096]; + while !request.windows(4).any(|window| window == b"\r\n\r\n") { + let read = socket.read(&mut buffer).await.unwrap(); + request.extend_from_slice(&buffer[..read]); + } + server_received.notify_one(); + loop { + if socket.read(&mut buffer).await.unwrap() == 0 { + break; + } + } + }); + let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))); + let mut machine = ocr_machine(ocr_client()); + let mut result = None; + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + tokio::select! { + _ = received.notified() => break, + step = machine.resume(result.take()) => { + result = Some(match step.unwrap() { + MachineStep::Host(HostOp::Route(op)) => HostResult::Route(host.route(op).await.unwrap()), + MachineStep::Host(HostOp::BeforeSend { wire, .. }) => HostResult::BeforeSend(wire), + MachineStep::Host(HostOp::Emit(_)) => HostResult::Emitted, + MachineStep::Complete(_) => panic!("the stalled provider completed"), + }); + } + } + } + }) + .await + .unwrap(); + + let cancelled = OcrError::InvalidRequest("cancelled".into()); + assert!( + machine + .interrupt(HostFailure::Cancelled(cancelled)) + .await + .is_err() + ); + tokio::time::timeout(std::time::Duration::from_secs(1), server) + .await + .expect("the provider connection stayed open after the interrupt") + .unwrap(); + } +} diff --git a/litellm-rust/crates/core/src/audio_transcription/tests.rs b/litellm-rust/crates/core/tests/audio_transcription.rs similarity index 95% rename from litellm-rust/crates/core/src/audio_transcription/tests.rs rename to litellm-rust/crates/core/tests/audio_transcription.rs index 8ccf7a07a0f..aa5aef0149b 100644 --- a/litellm-rust/crates/core/src/audio_transcription/tests.rs +++ b/litellm-rust/crates/core/tests/audio_transcription.rs @@ -4,11 +4,9 @@ use std::{ thread, }; +use litellm_core::audio_transcription::{audio_transcription, types::AudioTranscriptionRequest}; use serde_json::{Map, json}; -use super::audio_transcription; -use crate::audio_transcription::types::AudioTranscriptionRequest; - #[tokio::test] async fn bedrock_request_is_signed_and_contains_audio() { let listener = TcpListener::bind("127.0.0.1:0").expect("listener"); diff --git a/litellm-rust/crates/core/tests/aws_textract_ocr.rs b/litellm-rust/crates/core/tests/aws_textract_ocr.rs deleted file mode 100644 index c536317ad5c..00000000000 --- a/litellm-rust/crates/core/tests/aws_textract_ocr.rs +++ /dev/null @@ -1,193 +0,0 @@ -use std::{collections::BTreeMap, time::SystemTime}; - -use litellm_auth_aws::{Credentials, aws_signature_headers, sign_post}; -use litellm_llms::base_llm::ocr::error::Error; -use serde_json::{Value, json}; -use time::{PrimitiveDateTime, format_description}; - -use crate::ocr::{ - route::LocalOcrHost, - test_support::{ - MockResponse, header, mock_server, perform_ocr_with, request_body, - wire_request_with_document, - }, - types::LiteLLMOcrRequest, -}; - -const ACCESS_KEY_ID: &str = "AKIDEXAMPLE"; -const SECRET_ACCESS_KEY: &str = "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"; - -fn textract_request(base: &str) -> LiteLLMOcrRequest { - textract_request_for("aws_textract/detect-document-text", base) -} - -fn textract_request_for(model: &str, base: &str) -> LiteLLMOcrRequest { - wire_request_with_document( - model, - &format!("{base}/"), - json!({"type": "image_url", "image_url": "data:image/png;base64,b3JpZ2luYWw="}), - json!({ - "aws_access_key_id": ACCESS_KEY_ID, - "aws_secret_access_key": SECRET_ACCESS_KEY, - "aws_region_name": "eu-west-1" - }), - ) -} - -fn textract_response() -> MockResponse { - MockResponse::json(json!({ - "DocumentMetadata": {"Pages": 1}, - "Blocks": [{"BlockType": "PAGE"}, {"BlockType": "LINE", "Text": "Invoice 12345"}] - })) -} - -/// Recomputes SigV4 over the bytes the server received, at the time the client claimed. -fn expected_authorization(url: &str, raw_request: &str) -> String { - let format = - format_description::parse_borrowed::<2>("[year][month][day]T[hour][minute][second]Z") - .unwrap(); - let signed_at: SystemTime = - PrimitiveDateTime::parse(header(raw_request, "x-amz-date").unwrap(), &format) - .unwrap() - .assume_utc() - .into(); - let headers: BTreeMap = ["content-type", "x-amz-target"] - .into_iter() - .map(|name| { - ( - name.to_string(), - header(raw_request, name).unwrap().to_string(), - ) - }) - .collect(); - let body = raw_request.split_once("\r\n\r\n").unwrap().1; - sign_post( - url, - body.as_bytes(), - &aws_signature_headers(&headers), - "eu-west-1", - "textract", - &Credentials::new(ACCESS_KEY_ID, SECRET_ACCESS_KEY, None, None, "test"), - signed_at, - ) - .unwrap()["Authorization"] - .clone() -} - -#[tokio::test] -async fn the_request_is_signed_for_textract_and_lines_become_the_page() { - let (base, seen, server) = mock_server(vec![textract_response()]).await; - - let response = perform_ocr_with(LocalOcrHost::new(textract_request(&base))) - .await - .unwrap(); - server.await.unwrap(); - - let raw = seen.lock().unwrap()[0].clone(); - assert_eq!( - header(&raw, "x-amz-target"), - Some("Textract.DetectDocumentText") - ); - assert_eq!( - header(&raw, "content-type"), - Some("application/x-amz-json-1.1") - ); - assert_eq!( - request_body(&raw), - json!({"Document": {"Bytes": "b3JpZ2luYWw="}}) - ); - assert_eq!( - header(&raw, "authorization"), - Some(expected_authorization(&format!("{base}/"), &raw).as_str()) - ); - assert_eq!(response.pages[0].markdown, "Invoice 12345"); - assert_eq!(response.usage_info.unwrap().pages_processed, Some(1)); -} - -#[tokio::test] -async fn a_body_rewritten_by_before_send_is_what_gets_signed_and_sent() { - let (base, seen, server) = mock_server(vec![textract_response()]).await; - let host = LocalOcrHost::new(textract_request(&base)).with_before_send(|mut wire, _| { - assert!( - !wire - .headers - .iter() - .any(|(name, _)| name.eq_ignore_ascii_case("authorization")), - "the hook ran after signing" - ); - wire.body["Document"]["Bytes"] = Value::from("cmVkYWN0ZWQ="); - Ok(wire) - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - - let raw = seen.lock().unwrap()[0].clone(); - assert_eq!( - request_body(&raw), - json!({"Document": {"Bytes": "cmVkYWN0ZWQ="}}) - ); - assert_eq!( - header(&raw, "authorization"), - Some(expected_authorization(&format!("{base}/"), &raw).as_str()) - ); -} - -#[tokio::test] -async fn a_multi_page_rejection_reaches_the_caller_with_the_single_page_limit() { - let (base, _, server) = mock_server(vec![MockResponse { - status: 400, - headers: vec![], - body: json!({ - "__type": "UnsupportedDocumentException", - "Message": "Request has unsupported document format" - }), - }]) - .await; - - let error = perform_ocr_with(LocalOcrHost::new(textract_request(&base))) - .await - .unwrap_err(); - server.await.unwrap(); - - let Error::Provider { status, body, .. } = error else { - panic!("expected a provider error, got {error:?}"); - }; - assert_eq!(status, 400); - assert!( - body.contains("multi-page documents are not supported"), - "{body}" - ); -} - -#[tokio::test] -async fn analyze_document_asks_for_layout_and_tables_and_returns_markdown() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "DocumentMetadata": {"Pages": 1}, - "Blocks": [ - {"Id": "l1", "BlockType": "LINE", "Text": "Quarterly Report"}, - {"Id": "t", "BlockType": "LAYOUT_TITLE", - "Relationships": [{"Type": "CHILD", "Ids": ["l1"]}]} - ] - }))]) - .await; - let request = textract_request_for("aws_textract/analyze-document", &base); - - let response = perform_ocr_with(LocalOcrHost::new(request)).await.unwrap(); - server.await.unwrap(); - - let raw = seen.lock().unwrap()[0].clone(); - assert_eq!( - header(&raw, "x-amz-target"), - Some("Textract.AnalyzeDocument") - ); - assert_eq!( - request_body(&raw)["FeatureTypes"], - json!(["LAYOUT", "TABLES"]) - ); - assert_eq!( - header(&raw, "authorization"), - Some(expected_authorization(&format!("{base}/"), &raw).as_str()) - ); - assert_eq!(response.pages[0].markdown, "# Quarterly Report"); -} diff --git a/litellm-rust/crates/core/tests/azure_ai_ocr.rs b/litellm-rust/crates/core/tests/azure_ai_ocr.rs deleted file mode 100644 index 1492aaaeb11..00000000000 --- a/litellm-rust/crates/core/tests/azure_ai_ocr.rs +++ /dev/null @@ -1,293 +0,0 @@ -use litellm_llms::base_llm::ocr::error::Error; -use serde_json::{Value, json}; - -use super::test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request}; -use crate::ocr::route::LocalOcrHost; - -#[tokio::test] -async fn facade_executes_azure_mistral_with_prepared_auth() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"hello"}], - "usage_info":{"pages_processed":1} - }))]) - .await; - let mut request = wire_request( - "azure_ai/model", - &base, - json!({"include_image_base64":true}), - ); - request.credentials.api_key = None; - request.transport.extra_headers = vec![( - "Authorization".into(), - "Bearer python-prepared-token".into(), - )]; - - let result = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(result.pages[0].markdown, "hello"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /providers/mistral/azure/ocr ")); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer python-prepared-token\r\n") - ); - let body: Value = serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!( - body, - json!({ - "model":"model", - "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, - "include_image_base64":true - }) - ); -} - -#[tokio::test] -async fn facade_acquires_supplied_entra_token_for_final_request() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let mut request = wire_request( - "azure_ai/model", - &base, - json!({"azure_ad_token":"rust-owned-token"}), - ); - request.credentials.api_key = None; - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer rust-owned-token\r\n") - ); -} - -#[tokio::test] -async fn rejects_non_inline_body_after_guardrails() { - let request = wire_request("azure_ai/model", "http://127.0.0.1:1", json!({})); - let host = LocalOcrHost::new(request).with_before_send(|mut wire, _| { - wire.body["document"] = json!({ - "type":"document_url", - "document_url":"https://example.com/not-inline.pdf" - }); - Ok(wire) - }); - let error = perform_ocr_with(host).await.unwrap_err(); - assert!(error.to_string().contains("data URI")); -} - -mod transformation { - use std::sync::{ - Arc, - atomic::{AtomicUsize, Ordering}, - }; - - use litellm_auth::{ - ResolvedCredential, SecretValue, TokenFuture, TokenProvider, TokenProviderHandle, - }; - use rstest::rstest; - use serde_json::json; - - use super::*; - use crate::ocr::{ - test_support::{MockResponse, header, mock_server, perform_ocr}, - types::LiteLLMOcrRequest, - wire::decode_request, - }; - - #[derive(Debug)] - struct CountingToken { - token: fn(usize) -> String, - calls: AtomicUsize, - } - - impl CountingToken { - fn new(token: fn(usize) -> String) -> Arc { - Arc::new(Self { - token, - calls: AtomicUsize::new(0), - }) - } - - fn calls(&self) -> usize { - self.calls.load(Ordering::SeqCst) - } - } - - impl TokenProvider for CountingToken { - fn acquire(&self) -> TokenFuture<'_> { - let call = self.calls.fetch_add(1, Ordering::SeqCst) + 1; - let token = SecretValue::new((self.token)(call)); - Box::pin(async move { - Ok(ResolvedCredential::AccessToken { - token, - expires_on: None, - }) - }) - } - } - - fn numbered_token(call: usize) -> String { - format!("callback-{call}") - } - - fn azure_request( - provider: &Arc, - api_base: Option<&str>, - api_key: Option<&str>, - extra_headers: Value, - optional_params: Value, - ) -> LiteLLMOcrRequest { - let wire = serde_json::from_value(json!({ - "model": "azure_ai/mistral-ocr-latest", - "document": {"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, - "api_key": api_key, - "api_base": api_base, - "custom_llm_provider": null, - "extra_headers": extra_headers, - "optional_params": optional_params, - "timeout_seconds": 2.0 - })) - .unwrap(); - LiteLLMOcrRequest { - azure_ad_token_provider: Some(TokenProviderHandle::new(provider.clone())), - ..decode_request(wire).unwrap() - } - } - - fn ocr_page() -> MockResponse { - MockResponse::json(json!({"pages":[{"index":0,"markdown":"hello"}]})) - } - - #[tokio::test] - async fn token_provider_result_is_the_bearer_and_is_acquired_for_each_request() { - let provider = CountingToken::new(numbered_token); - let (base, seen, server) = mock_server(vec![ocr_page(), ocr_page()]).await; - - for _ in 0..2 { - perform_ocr(azure_request( - &provider, - Some(&base), - None, - Value::Null, - json!({}), - )) - .await - .unwrap(); - } - server.await.unwrap(); - - assert_eq!(provider.calls(), 2); - let requests = seen.lock().unwrap(); - assert_eq!( - requests - .iter() - .map(|request| header(request, "authorization")) - .collect::>(), - [Some("Bearer callback-1"), Some("Bearer callback-2")] - ); - } - - #[rstest] - #[case::api_key_skips_provider(Some("resource-key"), Value::Null, json!({}), "Bearer resource-key", 0)] - #[case::provider_beats_static_token( - None, - Value::Null, - json!({"azure_ad_token":"static-token"}), - "Bearer callback-1", - 1 - )] - #[case::header_wins_on_the_wire_but_provider_still_runs( - None, - json!({"Authorization":"Bearer override"}), - json!({}), - "Bearer override", - 1 - )] - #[tokio::test] - async fn credential_precedence( - #[case] api_key: Option<&str>, - #[case] extra_headers: Value, - #[case] optional_params: Value, - #[case] expected_authorization: &str, - #[case] expected_calls: usize, - ) { - let provider = CountingToken::new(numbered_token); - let (base, seen, server) = mock_server(vec![ocr_page()]).await; - - perform_ocr(azure_request( - &provider, - Some(&base), - api_key, - extra_headers, - optional_params, - )) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!(provider.calls(), expected_calls); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert_eq!( - header(&requests[0], "authorization"), - Some(expected_authorization) - ); - } - - #[rstest] - #[case::missing_api_base( - false, - json!({}), - numbered_token, - |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::MissingApiBase { - provider: "Azure AI", - environment_variable: "AZURE_AI_API_BASE", - })), - 0 - )] - #[case::unsupported_oidc_reference( - true, - json!({"azure_ad_token":"oidc/assertion","client_id":"client","tenant_id":"tenant"}), - numbered_token, - |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::UnsupportedOidcReference)), - 0 - )] - #[case::empty_provider_token_ignores_static_token( - true, - json!({"azure_ad_token":"static-token"}), - |_| String::new(), - |error: &Error| matches!(error, Error::MissingAzureAiCredentials), - 1 - )] - #[tokio::test] - async fn credential_failures_send_no_provider_request( - #[case] with_api_base: bool, - #[case] optional_params: Value, - #[case] token: fn(usize) -> String, - #[case] expected: fn(&Error) -> bool, - #[case] expected_calls: usize, - ) { - let provider = CountingToken::new(token); - let (base, seen, server) = mock_server(vec![ocr_page()]).await; - - let error = perform_ocr(azure_request( - &provider, - with_api_base.then_some(base.as_str()), - None, - Value::Null, - optional_params, - )) - .await - .unwrap_err(); - server.abort(); - - assert!(expected(&error), "unexpected error: {error:?}"); - assert_eq!(provider.calls(), expected_calls); - assert!(seen.lock().unwrap().is_empty()); - } -} diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs deleted file mode 100644 index 6dc9bfa5e7e..00000000000 --- a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs +++ /dev/null @@ -1,712 +0,0 @@ -use litellm_host::event::{CallEvent, MachineEvent}; -use litellm_llms::base_llm::ocr::{error::Error, settings::OcrSettings}; -use rstest::rstest; -use serde_json::{Value, json}; - -use super::{ - test_support::{ - MockResponse, mock_server, ocr_client, perform_ocr, perform_ocr_with, wire_request, - }, - wire::{OcrWireRequest, decode_request}, -}; -use crate::ocr::route::LocalOcrHost; - -fn query_value(url: &str, key: &str) -> Option { - url::Url::parse(url) - .unwrap() - .query_pairs() - .find_map(|(name, value)| (name == key).then(|| value.into_owned())) -} - -#[tokio::test] -async fn facade_maps_pages_features_and_url_document() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded", - "analyzeResult":{"pages":[]} - }))]) - .await; - let mut request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"]}), - ); - request.document = - serde_json::from_value::(json!({ - "type":"document_url", - "document_url":"https://example.com/document.pdf" - })) - .unwrap() - .into(); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let request = &seen.lock().unwrap()[0]; - let target = request.split_whitespace().nth(1).unwrap(); - let url = format!("{base}{target}"); - assert_eq!(query_value(&url, "pages").as_deref(), Some("1,2,3")); - assert_eq!( - query_value(&url, "features").as_deref(), - Some("keyValuePairs,languages") - ); - let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!( - body, - json!({"urlSource":"https://example.com/document.pdf"}) - ); -} - -#[rstest] -#[case(json!({"pages":[true]}), Error::Pages("expected only integers or only strings".into()))] -#[case(json!({"pages":[1,"2"]}), Error::Pages("expected only integers or only strings".into()))] -#[case(json!({"pages":[-1]}), Error::Pages("negative page index".into()))] -#[case(json!({"pages":"1&&features=bad"}), Error::Pages("invalid native page range".into()))] -#[case(json!({"features":"languages&pages=1"}), Error::Features)] -#[case(json!({"req_format":"azure"}), Error::RequestFormat)] -#[tokio::test] -async fn rejects_invalid_pages_features_and_format( - #[case] options: Value, - #[case] expected: Error, -) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({}))]).await; - let result = decode_request(OcrWireRequest { - model: "azure_ai/doc-intelligence/prebuilt-read".into(), - document: json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), - api_key: Some(litellm_auth::SecretValue::new("key")), - api_base: Some(base), - custom_llm_provider: None, - extra_headers: None, - optional_params: options.as_object().unwrap().clone(), - input_sources: Default::default(), - timeout_seconds: Some(2.0), - }); - let result = match result { - Ok(request) => perform_ocr(request).await, - Err(error) => Err(error), - }; - server.abort(); - let _ = server.await; - assert!( - seen.lock().unwrap().is_empty(), - "sent invalid options: {options}" - ); - let error = result.unwrap_err(); - assert_eq!( - std::mem::discriminant(&error), - std::mem::discriminant(&expected) - ); - assert_eq!(error.http_status_code(), Some(400)); - assert_eq!(error.to_string(), expected.to_string()); -} - -#[rstest] -#[case(json!({}))] -#[case(json!({"req_format":"litellm"}))] -#[tokio::test] -async fn missing_native_fields_keep_page_text_without_retaining_raw_response( - #[case] options: Value, -) { - let operation = json!({ - "status":"succeeded", - "analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"hello"}]}]} - }); - let (base, seen, server) = mock_server(vec![MockResponse::json(operation)]).await; - let response = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - options, - )) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!(response.pages.len(), 1); - assert_eq!(response.pages[0].index, 0); - assert_eq!(response.pages[0].markdown, "hello"); - assert_eq!(response.provider_native_response, None); - let serialized = response.into_json(); - assert_eq!(serialized.get("content"), Some(&Value::Null)); - assert_eq!(serialized.get("tables"), Some(&Value::Null)); - assert_eq!(serialized.get("keyValuePairs"), Some(&Value::Null)); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - let target = requests[0].split_whitespace().nth(1).unwrap(); - let url = format!("{base}{target}"); - for field in ["pages", "features", "req_format"] { - assert_eq!(query_value(&url, field), None); - } - let body: Value = serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!(body, json!({"base64Source":"YWJj"})); -} - -#[tokio::test] -async fn inline_document_decodes_to_base64_source() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded" - }))]) - .await; - let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let request = &seen.lock().unwrap()[0]; - let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!(body, json!({"base64Source":"YWJj"})); -} - -#[tokio::test] -async fn immediate_response_normalizes_pages_and_preserves_native() { - let operation = json!({ - "status":"succeeded", - "operationExtension":42, - "analyzeResult":{ - "content":"A\n\nB", - "tables":[{"cells":[]}], - "keyValuePairs":[{"key":{"content":"A"}}], - "pages":[{ - "pageNumber":"2", - "width":"8.5", - "height":11, - "unit":"inch", - "lines":[{"content":"A"},{"content":null},{"content":"B"}] - }] - } - }); - let (base, _, server) = mock_server(vec![MockResponse::json(operation.clone())]).await; - let result = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"req_format":"native"}), - )) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!(result.pages[0].index, 1); - assert_eq!(result.pages[0].markdown, "A\n\nB"); - assert_eq!( - serde_json::to_value(&result.pages[0].dimensions).unwrap(), - json!({"width":816,"height":1056,"dpi":96}) - ); - assert_eq!(result.usage_info.as_ref().unwrap().pages_processed, Some(1)); - let serialized = result.clone().into_json(); - assert_eq!(serialized["content"], "A\n\nB"); - assert_eq!(serialized["tables"], json!([{"cells":[]}])); - assert_eq!( - serialized["keyValuePairs"], - json!([{"key":{"content":"A"}}]) - ); - assert!(serialized.get("key_value_pairs").is_none()); - assert_eq!( - result.provider_native_response.map(Value::Object), - Some(operation) - ); -} - -#[tokio::test] -async fn client_settings_choose_the_api_version_and_the_inch_to_pixel_dpi() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded", - "analyzeResult":{"pages":[{"pageNumber":1,"width":8.5,"height":11,"unit":"inch"}]} - }))]) - .await; - let client = ocr_client().with_settings(OcrSettings { - document_intelligence_api_version: "2099-01-01".into(), - document_intelligence_dpi: 72, - ..OcrSettings::default() - }); - - let result = crate::ocr::client::perform( - &client, - wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})), - ) - .await - .unwrap(); - server.await.unwrap(); - - let target = seen.lock().unwrap()[0] - .split_whitespace() - .nth(1) - .unwrap() - .to_string(); - assert_eq!( - query_value(&format!("{base}{target}"), "api-version").as_deref(), - Some("2099-01-01") - ); - assert_eq!( - serde_json::to_value(&result.pages[0].dimensions).unwrap(), - json!({"width":612,"height":792,"dpi":72}) - ); -} - -#[tokio::test] -async fn accepted_response_polls_to_success_with_only_credentials() { - let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}}); - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse { - status: 200, - headers: vec![("Retry-After", "0".into())], - body: json!({"status":"running"}), - }, - MockResponse::json(operation.clone()), - ]) - .await; - let mut request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"req_format":"native"}), - ); - request - .transport - .extra_headers - .push(("X-Trace".into(), "initial-only".into())); - - let result = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!( - result.provider_native_response.map(Value::Object), - Some(operation) - ); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 3); - assert!(requests[0].to_ascii_lowercase().contains("x-trace:")); - for poll in &requests[1..] { - assert!(!poll.to_ascii_lowercase().contains("x-trace:")); - assert!( - poll.to_ascii_lowercase() - .contains("ocp-apim-subscription-key: test-key") - ); - } -} - -#[tokio::test] -async fn accepted_response_emits_response_received_before_polling() { - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({"submitted": true}), - }, - MockResponse::json(json!({"status":"succeeded"})), - ]) - .await; - let request_count = seen.clone(); - let host = LocalOcrHost::new(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .with_observer(move |event| { - let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event else { - return; - }; - match request_count.lock().unwrap().len() { - 1 => assert_eq!(raw.body, r#"{"submitted":true}"#), - 2 => assert!(raw.body.contains("succeeded")), - count => panic!("unexpected callback after {count} requests"), - } - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 2); -} - -#[tokio::test] -async fn polling_forwards_bearer_credentials() { - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse::json(json!({"status":"succeeded"})), - ]) - .await; - let mut request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); - request.credentials.api_key = None; - request.transport.extra_headers = vec![("Authorization".into(), "Bearer token".into())]; - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert!( - requests[1] - .to_ascii_lowercase() - .contains("authorization: bearer token") - ); -} - -#[tokio::test] -async fn polling_does_not_follow_redirects() { - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse { - status: 302, - headers: vec![("Location", "{base}/redirected".into())], - body: json!({}), - }, - MockResponse::json(json!({"status":"succeeded"})), - ]) - .await; - - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - - assert!(error.to_string().contains("status 302"), "{error}"); - assert_eq!(seen.lock().unwrap().len(), 2); - server.abort(); -} - -#[tokio::test] -async fn polling_rejects_terminal_failure() { - let (base, _, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse::json(json!({"status":"failed"})), - ]) - .await; - - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("status failed")); -} - -#[tokio::test] -async fn malformed_provider_pages_report_response_paths() { - for (analysis, path) in [ - (json!({"pages":null}), "pages"), - (json!({"pages":[null]}), "pages[0]"), - (json!({"pages":[{"lines":null}]}), "lines"), - (json!({"pages":[{"width":"bad"}]}), "width"), - ] { - let (base, _, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded", - "analyzeResult":analysis - }))]) - .await; - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains(path), "{error}"); - } -} - -#[tokio::test] -async fn rejects_missing_invalid_and_cross_origin_operation_locations() { - for headers in [ - Vec::new(), - vec![("Operation-Location", "/relative".into())], - vec![("Operation-Location", "http://example.com/operation".into())], - vec![( - "Operation-Location", - "http://user:password@127.0.0.1/operation".into(), - )], - ] { - let (base, _, server) = mock_server(vec![MockResponse { - status: 202, - headers, - body: json!({}), - }]) - .await; - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("operation-location")); - } -} - -#[tokio::test] -async fn polling_deadline_bounds_retry_delay() { - let (base, _, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse { - status: 200, - headers: vec![("Retry-After", "9999".into())], - body: json!({"status":"notStarted"}), - }, - ]) - .await; - let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); - let client = ocr_client().with_settings(OcrSettings { - poll_timeout: std::time::Duration::from_millis(100), - ..OcrSettings::default() - }); - - let error = tokio::time::timeout( - std::time::Duration::from_secs(1), - crate::ocr::client::perform(&client, request), - ) - .await - .unwrap() - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("timed out")); -} - -#[tokio::test] -async fn model_id_is_encoded_and_dot_segments_are_rejected() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded" - }))]) - .await; - perform_ocr(wire_request( - "azure_ai/doc-intelligence/a ?#é", - &base, - json!({}), - )) - .await - .unwrap(); - server.await.unwrap(); - assert!(seen.lock().unwrap()[0].contains("a%20%3F%23%C3%A9:analyze")); - - for model in [ - "azure_ai/doc-intelligence/.", - "azure_ai/doc-intelligence/..", - ] { - let error = perform_ocr(wire_request(model, "http://127.0.0.1:1", json!({}))) - .await - .unwrap_err(); - assert!(error.to_string().contains("dot segment")); - } -} - -mod transformation { - use std::sync::{Arc, Mutex}; - - use litellm_host::event::{CallEvent, MachineEvent}; - use litellm_llms::base_llm::ocr::transformation::OcrDocument; - use serde_json::{Value, json}; - - use super::*; - use crate::ocr::{ - route::LocalOcrHost, - test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request}, - }; - - #[tokio::test] - async fn facade_maps_pages_features_and_url_document() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded", - "analyzeResult":{"pages":[]} - }))]) - .await; - let mut request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"], "future_option": {"nested":null}, "extra_body":{"provider_option":false}}), - ); - request.document = serde_json::from_value::(json!({ - "type":"document_url", - "document_url":"https://example.com/document.pdf" - })) - .unwrap() - .into(); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let request = &seen.lock().unwrap()[0]; - let target = request.split_whitespace().nth(1).unwrap(); - let url = format!("{base}{target}"); - assert_eq!(query_value(&url, "pages").as_deref(), Some("1,2,3")); - assert_eq!( - query_value(&url, "features").as_deref(), - Some("keyValuePairs,languages") - ); - let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!( - body, - json!({"urlSource":"https://example.com/document.pdf", "future_option":{"nested":null}, "provider_option":false}) - ); - } - - #[tokio::test] - async fn rejects_invalid_pages_features_and_format() { - for options in [ - json!({"pages":[true]}), - json!({"pages":[1,"2"]}), - json!({"pages":[-1]}), - json!({"pages":"1&&features=bad"}), - json!({"features":"languages&pages=1"}), - json!({"req_format":"azure"}), - ] { - let request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - "http://127.0.0.1:1", - options.clone(), - ); - let rejected = perform_ocr(request).await.is_err(); - assert!(rejected, "accepted {options}"); - } - } - - #[tokio::test] - async fn immediate_response_normalizes_pages_and_preserves_native() { - let operation = json!({ - "status":"succeeded", - "operationExtension":42, - "analyzeResult":{ - "content":"A\n\nB", - "tables":[{"cells":[]}], - "keyValuePairs":[{"key":{"content":"A"}}], - "pages":[{ - "pageNumber":"2", - "width":"8.5", - "height":11, - "unit":"inch", - "lines":[{"content":"A"},{"content":null},{"content":"B"}] - }] - } - }); - let (base, _, server) = mock_server(vec![MockResponse::json(operation.clone())]).await; - let result = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"req_format":"native"}), - )) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!(result.pages[0].index, 1); - assert_eq!(result.pages[0].markdown, "A\n\nB"); - assert_eq!( - serde_json::to_value(&result.pages[0].dimensions).unwrap(), - json!({"width":816,"height":1056,"dpi":96}) - ); - assert_eq!(result.usage_info.as_ref().unwrap().pages_processed, Some(1)); - let serialized = result.clone().into_json(); - assert_eq!(serialized["content"], "A\n\nB"); - assert_eq!(serialized["tables"], json!([{"cells":[]}])); - assert_eq!( - serialized["keyValuePairs"], - json!([{"key":{"content":"A"}}]) - ); - assert!(serialized.get("key_value_pairs").is_none()); - assert_eq!( - result.provider_native_response.as_ref(), - operation.as_object() - ); - } - - #[tokio::test] - async fn accepted_response_polls_to_success_with_only_credentials() { - let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}}); - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse { - status: 200, - headers: vec![("Retry-After", "0".into())], - body: json!({"status":"running"}), - }, - MockResponse::json(operation.clone()), - ]) - .await; - let mut request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"req_format":"native"}), - ); - request - .transport - .extra_headers - .push(("X-Trace".into(), "initial-only".into())); - - let result = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!( - result.provider_native_response.as_ref(), - operation.as_object() - ); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 3); - assert!(requests[0].to_ascii_lowercase().contains("x-trace:")); - for poll in &requests[1..] { - assert!(!poll.to_ascii_lowercase().contains("x-trace:")); - assert!( - poll.to_ascii_lowercase() - .contains("ocp-apim-subscription-key: test-key") - ); - } - } - - #[tokio::test] - async fn accepted_response_emits_response_received_for_submission_and_completed_poll() { - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({"submitted": true}), - }, - MockResponse::json(json!({"status":"succeeded"})), - ]) - .await; - let responses_received = Arc::new(Mutex::new(Vec::new())); - let request_count = seen.clone(); - let observed = responses_received.clone(); - let host = LocalOcrHost::new(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .with_observer(move |event| { - if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { - observed - .lock() - .unwrap() - .push((request_count.lock().unwrap().len(), raw.body.clone())); - } - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 2); - assert_eq!( - *responses_received.lock().unwrap(), - [ - (1, r#"{"submitted":true}"#.to_string()), - (2, r#"{"status":"succeeded"}"#.to_string()), - ] - ); - } -} diff --git a/litellm-rust/crates/core/tests/cohere_ocr.rs b/litellm-rust/crates/core/tests/cohere_ocr.rs deleted file mode 100644 index 12824f58b1d..00000000000 --- a/litellm-rust/crates/core/tests/cohere_ocr.rs +++ /dev/null @@ -1,136 +0,0 @@ -mod transformation { - use litellm_llms::{ - base_llm::ocr::{ - error::Error, - transformation::{BaseOcrConfig, OcrDocument, OcrResponseFormat}, - }, - cohere::ocr::transformation::*, - }; - use rstest::rstest; - use serde_json::{Value, json}; - - #[tokio::test] - async fn composed_body_preserves_native_document_fields_and_untyped_overrides() { - let request = crate::ocr::test_support::wire_request( - "cohere/parse", - "https://example.com", - json!({ - "output_format":"markdown", "timeout":30, - "extra_body":{ - "output_format": {"future":true}, - "document":{"type":"image_url","image_url":"https://example.com/a.png", - "provider_options":{"nested":[false,0,null]}} - } - }), - ); - let request = request.with_document( - serde_json::from_value(json!({ - "type":"image_url","image_url":"https://example.com/original.png" - })) - .unwrap(), - ); - let request = crate::ocr::prepare::prepare_request_for_test(request); - let http = CohereParseConfig - .prepare_request( - &request, - &crate::ocr::test_support::ocr_client(), - &crate::ocr::test_support::NoHooks, - ) - .await - .unwrap(); - let body: Value = serde_json::from_slice(http.body()).unwrap(); - assert_eq!( - body, - json!({ - "model":"parse", "output_format":{"future":true}, - "document":{"type":"image_url","image_url":"https://example.com/a.png", - "provider_options":{"nested":[false,0,null]}} - }) - ); - } - - #[tokio::test] - async fn explicit_null_options_use_defaults_before_http() { - let request = crate::ocr::test_support::wire_request( - "cohere/parse", - "https://example.com", - json!({"output_format":null,"req_format":null}), - ); - let request = request.with_document( - serde_json::from_value( - json!({"type":"image_url","image_url":"https://example.com/a.png"}), - ) - .unwrap(), - ); - assert_eq!( - request.response_format().unwrap(), - OcrResponseFormat::Litellm - ); - let request = crate::ocr::prepare::prepare_request_for_test(request); - let http = CohereParseConfig - .prepare_request( - &request, - &crate::ocr::test_support::ocr_client(), - &crate::ocr::test_support::NoHooks, - ) - .await - .unwrap(); - let body: Value = serde_json::from_slice(http.body()).unwrap(); - assert_eq!(body["output_format"], "markdown"); - assert!(body.get("req_format").is_none()); - } - - #[rstest] - #[case::cohere("cohere/parse-v5.0", "POST /v2/parse ")] - #[case::azure_ai("azure_ai/Cohere-parse-v5.0", "POST /providers/cohere/v2/parse ")] - #[tokio::test] - async fn route_sends_image_to_its_parse_endpoint_with_the_bearer_key( - #[case] model: &str, - #[case] request_line: &str, - ) { - use crate::ocr::test_support::{MockResponse, header, mock_server, perform_ocr}; - - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let request = crate::ocr::test_support::wire_request(model, &base, json!({})) - .with_document( - serde_json::from_value::( - json!({"type":"image_url","image_url":"data:image/png;base64,YWJj"}), - ) - .unwrap() - .into(), - ); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with(request_line), "{}", requests[0]); - assert_eq!( - header(&requests[0], "authorization"), - Some("Bearer test-key") - ); - } - - #[rstest] - #[tokio::test] - async fn route_rejects_non_image_document_without_a_request( - #[values("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0")] model: &str, - ) { - use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr}; - - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - - let error = perform_ocr(crate::ocr::test_support::wire_request( - model, - &base, - json!({}), - )) - .await - .unwrap_err(); - server.abort(); - - assert!(matches!(error, Error::CohereImageOnly), "{error:?}"); - assert!(seen.lock().unwrap().is_empty()); - } -} diff --git a/litellm-rust/crates/core/tests/deepseek_ocr.rs b/litellm-rust/crates/core/tests/deepseek_ocr.rs deleted file mode 100644 index 96e7451769d..00000000000 --- a/litellm-rust/crates/core/tests/deepseek_ocr.rs +++ /dev/null @@ -1,133 +0,0 @@ -use litellm_llms::{ - base_llm::ocr::transformation::{BaseOcrConfig, OcrDocument}, - vertex_ai::ocr::deepseek_transformation::{ - DeepSeekOcrParams, DeepSeekOcrResponse, VertexAIDeepSeekOCRConfig, - normalize_response as transform_ocr_response, - }, -}; -use rstest::rstest; -use serde_json::{Value, json}; - -fn document() -> OcrDocument { - serde_json::from_value(json!({"type":"image_url","image_url":"gs://bucket/a.png"})).unwrap() -} - -#[rstest] -#[case("stream", json!(true))] -#[case("temperature", json!(0.1))] -#[case("max_tokens", json!(1024))] -#[case("top_p", json!(0.9))] -#[case("n", json!(2))] -#[case("stop", json!("done"))] -#[case("stop", json!(["done", "stop"]))] -fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) { - let params: DeepSeekOcrParams = - serde_json::from_value(json!({name: value.clone(), "ignored": true})).unwrap(); - let result = serde_json::to_value( - VertexAIDeepSeekOCRConfig - .transform_ocr_request("deepseek-ai/deepseek-ocr-maas", document(), ¶ms, &[]) - .unwrap(), - ) - .unwrap(); - assert_eq!(result["model"], "deepseek-ai/deepseek-ocr-maas"); - assert_eq!( - result["messages"][0]["content"][0], - json!({"type":"image_url","image_url":"gs://bucket/a.png"}) - ); - assert_eq!(result[name], value); - assert!(result.get("ignored").is_none()); -} - -#[rstest] -#[case(json!({"type":"image_url","image_url":"data:image/png;base64,AA=="}))] -#[case(json!({"type":"document_url","document_url":"data:application/pdf;base64,AA=="}))] -fn request_maps_both_document_types_to_image_content(#[case] document: Value) { - let source = document - .get("image_url") - .or_else(|| document.get("document_url")) - .unwrap() - .clone(); - let request = VertexAIDeepSeekOCRConfig - .transform_ocr_request( - "deepseek-ai/deepseek-ocr-maas", - serde_json::from_value(document).unwrap(), - &DeepSeekOcrParams::default(), - &[], - ) - .unwrap(); - let result = serde_json::to_value(request).unwrap(); - assert_eq!( - result["messages"][0]["content"][0], - json!({"type":"image_url","image_url":source}) - ); -} - -#[rstest] -#[case(json!("# hello"), "# hello")] -#[case(json!("{broken"), "{broken")] -#[case(json!(" {\"pages\":[]} "), " {\"pages\":[]} ")] -#[case(json!({"pages":[]}), "")] -#[case(json!("[]"), "[]")] -#[case(json!("{\"pages\":[{\"markdown\":\"json text\"}]}"), "json text")] -#[case(json!({"pages":[{"markdown":"object"}]}), "object")] -fn response_codec_handles_text_json_and_objects(#[case] content: Value, #[case] expected: &str) { - let structured = content - .as_object() - .is_some_and(|object| object.contains_key("pages")) - || content - .as_str() - .is_some_and(|text| text.contains("\"pages\"")); - let response: DeepSeekOcrResponse = serde_json::from_value( - json!({"choices":[{"message":{"content":content}}],"usage":{"prompt_tokens":1}}), - ) - .unwrap(); - let result = transform_ocr_response("model", response) - .unwrap() - .into_json(); - assert_eq!(result["pages"][0]["markdown"], expected); - assert_eq!(result["pages"][0]["index"], 0); - if structured { - assert!(result["usage_info"].is_null()); - } else { - assert_eq!(result["usage_info"]["prompt_tokens"], 1); - } -} - -#[test] -fn structured_result_maps_pages_usage_model_and_annotation() { - let response: DeepSeekOcrResponse = serde_json::from_value(json!({ - "choices":[{"message":{"content":{ - "pages":[{"index":2,"markdown":"page","images":[{"id":"one"}],"dimensions":{"width":10}}], - "model":"provider-model", - "usage_info":{"pages_processed":1}, - "document_annotation":{"language":"en"}, - "future":"kept" - }}}] - })) - .unwrap(); - let result = transform_ocr_response("requested", response) - .unwrap() - .into_json(); - assert_eq!(result["pages"][0]["index"], 2); - assert_eq!(result["pages"][0]["images"][0]["id"], "one"); - assert_eq!(result["model"], "provider-model"); - assert_eq!(result["usage_info"]["pages_processed"], 1); - assert_eq!(result["document_annotation"]["language"], "en"); - assert_eq!(result["future"], "kept"); -} - -#[test] -fn response_codec_rejects_missing_empty_and_malformed_content() { - for value in [ - json!({"choices":[{"message":{"content":{}}}]}), - json!({"choices":[]}), - json!({"choices":[{"message":{"content":""}}]}), - json!({"choices":[{"message":{"content":"{\"pages\":[{\"markdown\":42}]}"}}]}), - json!({"choices":[{"message":{"content":{"pages":[{"markdown":42}]}}}]}), - ] { - let result = serde_json::from_value::(value) - .map_err(|_| ()) - .and_then(|response| transform_ocr_response("model", response).map_err(|_| ())); - assert!(result.is_err()); - } -} diff --git a/litellm-rust/crates/core/src/messages/tests.rs b/litellm-rust/crates/core/tests/messages.rs similarity index 79% rename from litellm-rust/crates/core/src/messages/tests.rs rename to litellm-rust/crates/core/tests/messages.rs index ce48752864a..18af8a7d619 100644 --- a/litellm-rust/crates/core/src/messages/tests.rs +++ b/litellm-rust/crates/core/tests/messages.rs @@ -1,7 +1,11 @@ use std::{sync::Arc, time::Duration}; use futures_util::future::BoxFuture; -use litellm_http::request::{has_bearer_auth, has_header}; +use litellm_core::messages::{ + Error, messages, + route::{LocalMessagesHost, MessagesCall, messages_machine}, + types::{MessagesRequest, MessagesShaping}, +}; use litellm_secrets::{SecretValue, source::SecretSource}; use serde_json::{Map, Value, json}; use tokio::{ @@ -9,14 +13,6 @@ use tokio::{ net::{TcpListener, TcpStream}, }; -use super::{ - Error, - common_utils::{messages_provider_config, string_headers, truncate_error_body}, - messages, - route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine}, -}; -use crate::messages::types::{MessagesRequest, MessagesShaping}; - struct RecordingSecrets { values: Vec<(&'static str, String)>, fails: bool, @@ -73,55 +69,6 @@ fn secrets_call() -> MessagesCall { } } -#[tokio::test] -async fn route_reads_the_provider_credential_and_base_from_the_secret_source() { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - let addr = listener.local_addr().expect("addr"); - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts request"); - let request = read_http_request(&mut socket).await; - let response_body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4-5","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1}}"#; - socket - .write_all(write_response(response_body).as_bytes()) - .await - .expect("writes response"); - request - }); - let secrets = Arc::new(RecordingSecrets::new( - vec![ - ("ANTHROPIC_API_KEY", "sk-from-manager".to_string()), - ("ANTHROPIC_BASE_URL", format!("http://{addr}")), - ], - false, - )); - - let output = litellm_host::run::run( - messages_machine(secrets.clone()), - &LocalMessagesHost::new(secrets_call()), - ) - .await - .expect("messages request succeeds"); - - assert!(matches!(output, MessagesOutput::Message(_))); - let request = server.await.expect("server task completes"); - assert!( - request - .to_ascii_lowercase() - .contains("x-api-key: sk-from-manager"), - "{request}" - ); - let requested = secrets.requested.lock().unwrap().clone(); - assert_eq!( - requested, - messages_provider_config("anthropic") - .unwrap() - .secret_names() - .iter() - .map(ToString::to_string) - .collect::>() - ); -} - #[tokio::test] async fn route_surfaces_a_secret_manager_failure_before_the_call() { let Err(error) = litellm_host::run::run( @@ -179,75 +126,6 @@ fn write_response(body: &str) -> String { ) } -#[test] -fn provider_config_resolves_anthropic_and_azure_ai() { - assert!(messages_provider_config("anthropic").is_some()); - assert!(messages_provider_config("azure_ai").is_some()); - assert!(messages_provider_config("openai").is_none()); -} - -#[test] -fn truncate_error_body_caps_long_payloads() { - let body = "x".repeat(400); - let truncated = truncate_error_body(&body); - assert!(truncated.ends_with("... (truncated)")); - let prefix_chars = truncated - .strip_suffix("... (truncated)") - .expect("truncated marker present") - .chars() - .count(); - assert_eq!(prefix_chars, 256); -} - -#[test] -fn string_headers_rejects_non_string_values() { - let headers = json!({"x-count": 3}).as_object().unwrap().clone(); - let err = string_headers(Some(headers)).expect_err("non-string header rejected"); - assert_eq!( - err, - Error::Headers(litellm_http::request::HeaderError { - context: "messages", - name: "x-count".to_string(), - actual: "number", - }) - ); -} - -#[test] -fn has_header_is_case_insensitive() { - let headers = vec![("X-Api-Key".to_string(), "secret".to_string())]; - assert!(has_header(&headers, "x-api-key")); - assert!(!has_header(&headers, "authorization")); -} - -#[test] -fn has_bearer_auth_requires_a_nonempty_bearer_token() { - assert!(has_bearer_auth(&[( - "Authorization".to_string(), - "Bearer tok".to_string() - )])); - assert!(has_bearer_auth(&[( - "authorization".to_string(), - "bearer tok".to_string() - )])); - assert!(!has_bearer_auth(&[( - "authorization".to_string(), - "Bearer ".to_string() - )])); - assert!(!has_bearer_auth(&[( - "authorization".to_string(), - String::new() - )])); - assert!(!has_bearer_auth(&[( - "authorization".to_string(), - "Basic abc".to_string() - )])); - assert!(!has_bearer_auth(&[( - "x-api-key".to_string(), - "sk".to_string() - )])); -} - #[tokio::test] async fn messages_round_trip_builds_azure_request_and_passes_response_through() { let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs deleted file mode 100644 index 26cd2153e8d..00000000000 --- a/litellm-rust/crates/core/tests/ocr.rs +++ /dev/null @@ -1,1033 +0,0 @@ -use std::sync::{Arc, Mutex}; - -use futures_util::future::BoxFuture; -use litellm_auth_gcp::VertexAuth; -use litellm_host::{ - event::{CallEvent, MachineEvent, WireRequest}, - host::{Host, HostOp, HostResult}, - machine::{HostFailure, Machine, MachineStep}, -}; -use litellm_http::{ - HttpClientPool, HttpSettings, Resolution, - media::{PublicDnsResolver, UrlPolicy}, -}; -use litellm_llms::base_llm::ocr::{ - error::Error as OcrError, - handler::OcrClient, - settings::OcrSettings, - transformation::{ - BaseOcrConfig, LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig, - }, -}; -use litellm_secrets::source::SecretSource; -use rstest::rstest; -use serde_json::{Value, json}; - -use super::{ - test_support::{ - MockResponse, mock_server, ocr_client, perform_ocr, perform_ocr_with, wire_request, - }, - wire::{OcrWireRequest, decode_request}, -}; -use crate::ocr::route::{LocalOcrHost, OcrOp, OcrOpResult, ocr_machine}; - -struct RecordingSecretSource { - names: Arc>>, - values: &'static [(&'static str, &'static str)], - api_base: String, -} - -impl SecretSource for RecordingSecretSource { - fn get_secret_str<'a>( - &'a self, - name: &'a str, - ) -> BoxFuture<'a, Result, litellm_secrets::Error>> { - self.names.lock().unwrap().push(name.to_owned()); - Box::pin(async move { - Ok(match name { - "MISTRAL_AZURE_API_BASE" => Some(self.api_base.clone()), - _ => self - .values - .iter() - .find(|(key, _)| *key == name) - .map(|(_, value)| value.to_string()), - } - .map(litellm_secrets::SecretValue::new)) - }) - } -} - -#[rstest] -#[case::mistral("mistral/model", json!({}))] -#[case::vertex("vertex_ai/mistral-ocr-latest", json!({"vertex_project":"test-project", "vertex_location":"us-central1"}))] -#[tokio::test] -async fn ocr_contract_upstream_error_preserves_status_body_and_headers( - #[case] model: &str, - #[case] options: Value, -) { - let payload = json!({"message": format!("{} END-OF-PROVIDER-BODY", "x".repeat(4096))}); - let expected_body = serde_json::to_string(&payload).unwrap(); - let (base, seen, server) = mock_server(vec![MockResponse { - status: 422, - headers: vec![ - ("Retry-After", "17".into()), - ("X-Request-ID", "request-123".into()), - ("X-Future-Header", "retained".into()), - ], - body: payload, - }]) - .await; - let error = perform_ocr(wire_request(model, &base, options)) - .await - .unwrap_err(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 1); - let OcrError::Provider { - status, - body, - headers, - } = error - else { - panic!("expected provider error, got {error:?}"); - }; - assert_eq!(status, 422); - for (name, value) in [ - ("retry-after", "17"), - ("x-request-id", "request-123"), - ("x-future-header", "retained"), - ] { - assert!( - headers - .iter() - .any(|(key, actual)| key.eq_ignore_ascii_case(name) && actual == value) - ); - } - assert_eq!( - body.len(), - expected_body.len(), - "provider error body was truncated" - ); - assert_eq!(body, expected_body); -} - -#[test] -fn request_boundary_selects_mistral_and_rejects_unknown_providers() { - let request = OcrWireRequest { - model: "mistral/model".into(), - document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), - api_key: Some(litellm_auth::SecretValue::new("key")), - api_base: None, - custom_llm_provider: None, - extra_headers: None, - optional_params: json!({"extract_header":true,"unknown":42}) - .as_object() - .unwrap() - .clone(), - input_sources: Default::default(), - timeout_seconds: None, - }; - assert!(decode_request(request).is_ok()); - assert!( - decode_request(OcrWireRequest { - model: "model".into(), - document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), - api_key: Some(litellm_auth::SecretValue::new("key")), - api_base: None, - custom_llm_provider: Some("unknown".into()), - extra_headers: None, - optional_params: serde_json::Map::new(), - input_sources: Default::default(), - timeout_seconds: None, - }) - .is_err() - ); -} - -#[tokio::test] -async fn facade_executes_direct_mistral_once() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"hello","custom":"preserved"}], - "usage_info":{"pages_processed":1} - }))]) - .await; - let result = perform_ocr(wire_request( - "mistral/model", - &base, - json!({"pages":"0,2-4","extract_header":true,"unknown":"ignored"}), - )) - .await - .unwrap(); - server.await.unwrap(); - assert_eq!(result.pages[0].markdown, "hello"); - assert_eq!(result.pages[0].extra_fields["custom"], "preserved"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /v1/ocr ")); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer test-key\r\n") - ); - let body: Value = serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!( - body, - json!({ - "model":"model", - "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, - "pages":"0,2-4", - "extract_header":true, - "unknown":"ignored" - }) - ); -} - -#[tokio::test] -async fn facade_retains_native_response_when_requested() { - let provider_response = json!({ - "pages":[{"index":0,"markdown":"hello"}], - "usage_info":{"pages_processed":1}, - "provider_only":"preserved" - }); - let (base, _, server) = mock_server(vec![MockResponse::json(provider_response.clone())]).await; - let response = perform_ocr(wire_request( - "mistral/model", - &base, - json!({"req_format":"native"}), - )) - .await - .unwrap(); - - server.await.unwrap(); - assert_eq!( - response.provider_native_response.map(Value::Object), - Some(provider_response) - ); -} - -#[rstest] -#[case::plain_key(&[("MISTRAL_API_KEY", "plain")], "plain")] -#[case::azure_key_wins(&[("MISTRAL_AZURE_API_KEY", "azure"), ("MISTRAL_API_KEY", "plain")], "azure")] -#[case::empty_azure_key_falls_through(&[("MISTRAL_AZURE_API_KEY", ""), ("MISTRAL_API_KEY", "plain")], "plain")] -#[tokio::test] -async fn mistral_env_fallbacks_follow_python_through_the_injected_secret_source( - #[case] secrets: &'static [(&'static str, &'static str)], - #[case] expected_key: &str, -) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let names = Arc::new(Mutex::new(Vec::new())); - let client = ocr_client().with_secrets(Arc::new(RecordingSecretSource { - names: names.clone(), - values: secrets, - api_base: base.clone(), - })); - let request = decode_request(OcrWireRequest { - model: "mistral/model".into(), - document: json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), - api_key: None, - api_base: None, - custom_llm_provider: None, - extra_headers: None, - optional_params: Default::default(), - input_sources: Default::default(), - timeout_seconds: Some(2.0), - }) - .unwrap(); - - crate::ocr::client::perform(&client, request).await.unwrap(); - server.await.unwrap(); - assert_eq!( - *names.lock().unwrap(), - litellm_llms::mistral::ocr::transformation::MistralOcrConfig.secret_names() - ); - assert!(seen.lock().unwrap()[0].contains(&format!("authorization: Bearer {expected_key}"))); -} - -#[tokio::test] -async fn mistral_ocr_resolves_provider_secrets_before_transformation() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let names = Arc::new(Mutex::new(Vec::new())); - let client = ocr_client().with_secrets(Arc::new(RecordingSecretSource { - names: names.clone(), - values: &[("MISTRAL_API_KEY", "source-key")], - api_base: base.clone(), - })); - let request = decode_request(OcrWireRequest { - model: "mistral/mistral-ocr-latest".into(), - document: json!({ - "type":"document_url", - "document_url":"data:application/pdf;base64,YWJj" - }), - api_key: None, - api_base: None, - custom_llm_provider: None, - extra_headers: None, - optional_params: Default::default(), - input_sources: Default::default(), - timeout_seconds: Some(2.0), - }) - .unwrap(); - - crate::ocr::client::perform(&client, request).await.unwrap(); - server.await.unwrap(); - assert_eq!( - *names.lock().unwrap(), - litellm_llms::mistral::ocr::transformation::MistralOcrConfig.secret_names() - ); - assert!(seen.lock().unwrap()[0].contains("authorization: Bearer source-key")); -} - -#[tokio::test] -async fn ocr_client_uses_the_injected_http_pool_configuration() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let settings = HttpSettings { - user_agent: Some("host-owned/1".into()), - ..HttpSettings::default() - }; - let client = OcrClient::new( - &HttpClientPool::new(Arc::new(PublicDnsResolver)), - &Resolution::from(&settings).config, - UrlPolicy::default(), - VertexAuth::default(), - OcrSettings::default(), - Arc::new(litellm_secrets::source::EnvironmentSecrets::default()), - ) - .unwrap(); - crate::ocr::client::perform(&client, wire_request("mistral/model", &base, json!({}))) - .await - .unwrap(); - server.await.unwrap(); - assert!(seen.lock().unwrap()[0].contains("user-agent: host-owned/1")); -} - -fn event_name(event: &CallEvent) -> &'static str { - match event { - CallEvent::Started { .. } => "started", - CallEvent::Machine(MachineEvent::ResponseReceived { .. }) => "response", - CallEvent::Succeeded { .. } => "success", - CallEvent::Failed { .. } => "failure", - } -} - -fn recording_host( - request: crate::ocr::types::LiteLLMOcrRequest, - events: Arc>>, - block: bool, -) -> LocalOcrHost { - let before_send_events = events.clone(); - LocalOcrHost::new(request) - .with_before_send(move |wire, _| { - before_send_events.lock().unwrap().push("before_send"); - if block { - return Err(OcrError::InvalidRequest("blocked".into())); - } - Ok(wire) - }) - .with_observer(move |event| events.lock().unwrap().push(event_name(event))) -} - -#[tokio::test] -async fn lifecycle_sends_headers_returned_by_the_before_send_operation() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))).with_before_send( - |mut wire, _| { - wire.headers - .push(("x-core-callback".into(), "edited".into())); - Ok(wire) - }, - ); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - - assert!(seen.lock().unwrap()[0].contains("x-core-callback: edited")); -} - -#[tokio::test] -async fn before_send_context_names_the_route_and_its_secrets() { - let (base, _, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let observed = Arc::new(Mutex::new(None)); - let captured = observed.clone(); - let host = LocalOcrHost::new(wire_request( - "mistral/model", - &base, - json!({"pages": [0], "req_format": "native"}), - )) - .with_before_send(move |wire, context| { - *captured.lock().unwrap() = Some((wire.clone(), context.clone())); - Ok(wire) - }); - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - let (wire, context) = observed.lock().unwrap().take().unwrap(); - assert_eq!(context.custom_llm_provider, "mistral"); - assert_eq!(context.model, "model"); - assert_eq!(wire.body["pages"], json!([0])); - assert!(context.secret_fields.is_empty()); - assert_eq!(context.optional_params["req_format"], "native"); - - let (base, _, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let observed = Arc::new(Mutex::new(None)); - let captured = observed.clone(); - let request = wire_request( - "azure_ai/model", - &base, - json!({"client_secret": "shh", "tenant_id": "t"}), - ); - let request = request.with_document(crate::ocr::types::OcrDocumentInput::Bytes { - bytes: b"abc".as_slice().into(), - file_name: None, - mime_type: Some("application/pdf".into()), - }); - let host = LocalOcrHost::new(request).with_before_send(move |wire, context| { - *captured.lock().unwrap() = Some(context.clone()); - Ok(wire) - }); - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - let context = observed.lock().unwrap().take().unwrap(); - assert_eq!(context.secret_fields, ["client_secret"]); -} - -#[tokio::test] -async fn lifecycle_orders_hooks_and_emits_one_success() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let events = Arc::new(Mutex::new(Vec::new())); - let host = recording_host( - wire_request("mistral/model", &base, json!({})), - events.clone(), - false, - ); - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!( - *events.lock().unwrap(), - ["started", "before_send", "response", "success"] - ); - assert_eq!(seen.lock().unwrap().len(), 1); -} - -#[tokio::test] -async fn lifecycle_blocking_prevents_execution_and_emits_one_failure() { - let events = Arc::new(Mutex::new(Vec::new())); - let host = recording_host( - wire_request("mistral/model", "http://127.0.0.1:1", json!({})), - events.clone(), - true, - ); - let error = perform_ocr_with(host).await.unwrap_err(); - assert!(matches!(error, OcrError::InvalidRequest(message) if message == "blocked")); - assert_eq!( - *events.lock().unwrap(), - ["started", "before_send", "failure"] - ); -} - -#[tokio::test] -async fn upstream_failure_emits_one_terminal_failure() { - let (base, seen, server) = mock_server(vec![MockResponse { - status: 500, - headers: vec![], - body: json!({"error":"failed"}), - }]) - .await; - let events = Arc::new(Mutex::new(Vec::new())); - let host = recording_host( - wire_request("mistral/model", &base, json!({})), - events.clone(), - false, - ); - assert!(perform_ocr_with(host).await.is_err()); - server.await.unwrap(); - assert_eq!( - *events.lock().unwrap(), - ["started", "before_send", "failure"] - ); - assert_eq!(seen.lock().unwrap().len(), 1); -} - -/// Drives the machine by hand, answering every op through `host` except `before_send`, -/// which `intercept` answers so a test can fail or cancel exactly there. -async fn drive_until( - client: OcrClient, - host: &LocalOcrHost, - mut intercept: impl FnMut(WireRequest) -> Result>, -) -> ( - Result, - Vec<&'static str>, - crate::ocr::route::OcrMachine, -) { - let mut machine = ocr_machine(client); - let mut result = None; - let mut ops = Vec::new(); - let outcome = loop { - let op = match machine.resume(result.take()).await { - Ok(MachineStep::Host(op)) => op, - Ok(MachineStep::Complete(response)) => break Ok(response), - Err(error) => break Err(error), - }; - let answer = match op { - HostOp::Route(op) => { - ops.push(match op { - OcrOp::ProjectRequest => "ProjectRequest", - OcrOp::ReadDocument => "ReadDocument", - OcrOp::AcquireAzureAdToken => "AcquireAzureAdToken", - }); - host.route(op) - .await - .map(HostResult::Route) - .map_err(HostFailure::Error) - } - HostOp::BeforeSend { wire, .. } => { - ops.push("BeforeSend"); - intercept(*wire).map(|wire| HostResult::BeforeSend(Box::new(wire))) - } - HostOp::Emit(event) => { - let event = CallEvent::Machine(event); - ops.push(event_name(&event)); - host.emit(&event) - .await - .map(|()| HostResult::Emitted) - .map_err(HostFailure::Error) - } - }; - match answer { - Ok(answer) => result = Some(answer), - Err(failure) => break machine.interrupt(failure).await, - } - }; - (outcome, ops, machine) -} - -#[tokio::test] -async fn failed_before_send_does_not_replay_or_reach_transport() { - let host = LocalOcrHost::new(wire_request( - "mistral/model", - "http://127.0.0.1:1", - json!({}), - )); - let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, |_| { - Err(HostFailure::Error(OcrError::InvalidRequest( - "before_send failed".into(), - ))) - }) - .await; - assert!( - matches!(outcome, Err(OcrError::InvalidRequest(message)) if message == "before_send failed") - ); - assert_eq!(ops, ["ProjectRequest", "BeforeSend"]); - assert!(machine.resume(None).await.is_err()); -} - -#[tokio::test] -async fn invalid_provider_response_emits_response_received_before_normalization_failure() { - let (base, seen, server) = - mock_server(vec![MockResponse::json(json!({"pages":"invalid"}))]).await; - let responses_received = Arc::new(Mutex::new(Vec::new())); - let observed = responses_received.clone(); - let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))).with_observer( - move |event| { - if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { - observed.lock().unwrap().push(raw.body.clone()); - } - }, - ); - let error = perform_ocr_with(host).await.unwrap_err(); - server.await.unwrap(); - assert!(matches!(error, OcrError::ResponseField { .. })); - assert_eq!(seen.lock().unwrap().len(), 1); - assert_eq!( - *responses_received.lock().unwrap(), - [r#"{"pages":"invalid"}"#] - ); -} - -#[tokio::test] -async fn direct_native_host_drives_the_same_state_machine() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"native"}] - }))]) - .await; - let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))); - let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, Ok).await; - server.await.unwrap(); - assert_eq!(outcome.unwrap().pages[0].markdown, "native"); - assert_eq!(seen.lock().unwrap().len(), 1); - assert_eq!(ops, ["ProjectRequest", "BeforeSend", "response"]); - assert!(matches!( - machine.resume(None).await, - Err(OcrError::InvalidRequest(_)) - )); -} - -async fn drive_native_file_call( - request: crate::ocr::types::LiteLLMOcrRequest, - content: Result, -) -> (Result, usize) { - let reads = Arc::new(Mutex::new(0)); - let counted = reads.clone(); - let content = Mutex::new(Some(content)); - let host = LocalOcrHost::new(request).with_reader(move || { - *counted.lock().unwrap() += 1; - content.lock().unwrap().take().unwrap() - }); - let outcome = perform_ocr_with(host).await; - let reads = *reads.lock().unwrap(); - (outcome, reads) -} - -#[tokio::test] -async fn host_reader_documents_are_read_once_at_the_core_selected_point_and_encoded() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"file"}] - }))]) - .await; - let request = wire_request("mistral/model", &base, json!({})).with_document( - crate::ocr::types::OcrDocumentInput::HostReader { - mime_type: Some("application/pdf".into()), - }, - ); - let (response, reads) = drive_native_file_call( - request, - Ok(crate::ocr::types::OcrFileContent { - bytes: b"abc".as_slice().into(), - file_name: Some("scan.png".into()), - }), - ) - .await; - server.await.unwrap(); - assert_eq!(response.unwrap().pages[0].markdown, "file"); - assert_eq!(reads, 1); - assert!(seen.lock().unwrap()[0].contains("data:application/pdf;base64,YWJj")); -} - -#[tokio::test] -async fn host_reader_failures_and_empty_files_fail_before_the_provider_is_called() { - let (base, seen, _server) = mock_server(vec![]).await; - let request = wire_request("mistral/model", &base, json!({})); - let failure = OcrError::InvalidRequest("reader exploded".into()); - let (response, reads) = drive_native_file_call( - request.with_document(crate::ocr::types::OcrDocumentInput::HostReader { mime_type: None }), - Err(failure.clone()), - ) - .await; - assert!( - matches!(response.unwrap_err(), OcrError::InvalidRequest(message) if message == "reader exploded") - ); - assert_eq!(reads, 1); - - let request = wire_request("mistral/model", &base, json!({})); - let (response, _) = drive_native_file_call( - request.with_document(crate::ocr::types::OcrDocumentInput::HostReader { mime_type: None }), - Ok(crate::ocr::types::OcrFileContent { - bytes: Default::default(), - file_name: None, - }), - ) - .await; - assert!(matches!(response.unwrap_err(), OcrError::EmptyFile)); - assert!(seen.lock().unwrap().is_empty()); -} - -#[tokio::test] -async fn path_documents_are_read_by_core_without_a_host_operation() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"path"}] - }))]) - .await; - let dir = std::env::temp_dir().join(format!("litellm-ocr-{}", rand::random::())); - std::fs::create_dir_all(&dir).unwrap(); - let path = dir.join("scan.png"); - std::fs::write(&path, b"abc").unwrap(); - let request = wire_request("mistral/model", &base, json!({})).with_document( - crate::ocr::types::OcrDocumentInput::Path { - path: path.clone(), - mime_type: None, - }, - ); - let (response, reads) = - drive_native_file_call(request, Err(OcrError::InvalidRequest("unused".into()))).await; - server.await.unwrap(); - std::fs::remove_dir_all(&dir).unwrap(); - assert_eq!(response.unwrap().pages[0].markdown, "path"); - assert_eq!(reads, 0); - assert!(seen.lock().unwrap()[0].contains("data:image/png;base64,YWJj")); - - let (base, seen, _server) = mock_server(vec![]).await; - let request = wire_request("mistral/model", &base, json!({})); - let (response, _) = drive_native_file_call( - request.with_document(crate::ocr::types::OcrDocumentInput::Path { - path: path.clone(), - mime_type: None, - }), - Err(OcrError::InvalidRequest("unused".into())), - ) - .await; - assert!(matches!( - response.unwrap_err(), - OcrError::FileRead { path: failed, source } if failed == path && source.kind() == std::io::ErrorKind::NotFound - )); - assert!(seen.lock().unwrap().is_empty()); -} - -#[tokio::test] -async fn cancellation_at_before_send_prevents_execution_and_further_resumption() { - let host = LocalOcrHost::new(wire_request( - "mistral/model", - "http://127.0.0.1:1", - json!({}), - )); - let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, |_| { - Err(HostFailure::Cancelled(OcrError::InvalidRequest( - "cancelled".into(), - ))) - }) - .await; - assert!(matches!(outcome, Err(OcrError::InvalidRequest(message)) if message == "cancelled")); - assert_eq!(ops, ["ProjectRequest", "BeforeSend"]); - assert!(machine.resume(Some(HostResult::Emitted)).await.is_err()); -} - -#[tokio::test] -async fn missing_host_result_preserves_pending_operation() { - let request = wire_request("mistral/model", "http://127.0.0.1:1", json!({})); - let mut machine = ocr_machine(ocr_client()); - assert!(matches!( - machine.resume(None).await.unwrap(), - MachineStep::Host(HostOp::Route(OcrOp::ProjectRequest)) - )); - assert!(machine.resume(None).await.is_err()); - assert!(matches!( - machine - .resume(Some(HostResult::Route(OcrOpResult::Request { - request: Box::new(request), - caller_token: false, - }))) - .await - .unwrap(), - MachineStep::Host(HostOp::BeforeSend { .. }) - )); -} - -async fn read_bounded_response(response: Vec, limit: usize) -> Result { - use tokio::io::{AsyncReadExt, AsyncWriteExt}; - - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut request = [0; 4096]; - assert!(socket.read(&mut request).await.unwrap() > 0); - socket.write_all(&response).await.unwrap(); - std::future::pending::<()>().await; - }); - let response = reqwest::Client::new() - .get(format!("http://{address}")) - .send() - .await - .unwrap(); - let result = tokio::time::timeout( - std::time::Duration::from_secs(2), - litellm_llms::base_llm::ocr::handler::read_response_bytes(response, limit), - ) - .await; - server.abort(); - let _ = server.await; - result.expect("bounded reads must finish without waiting for the rest of an oversized body") -} - -#[tokio::test] -async fn response_limit_accepts_exact_size_and_rejects_declared_and_chunked_overflow() { - use litellm_llms::base_llm::ocr::error::Error; - - for response in [ - "HTTP/1.1 200 OK\r\nContent-Length: 8\r\n\r\nabcdefgh", - "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n4\r\nefgh\r\n0\r\n\r\n", - ] { - assert_eq!( - read_bounded_response(response.as_bytes().to_vec(), 8) - .await - .unwrap(), - "abcdefgh" - ); - } - for response in [ - "HTTP/1.1 200 OK\r\nContent-Length: 9\r\n\r\n", - "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n5\r\nefghi\r\n", - ] { - assert!(matches!( - read_bounded_response(response.as_bytes().to_vec(), 8).await, - Err(Error::TooLarge { limit: 8 }) - )); - } -} - -#[rstest] -#[case::declared("Content-Length: 1000000")] -#[case::chunked("Transfer-Encoding: chunked")] -#[tokio::test] -async fn oversized_error_retains_http_status_and_bounded_diagnostics_without_draining( - #[case] headers: &str, -) { - let prefix = "x".repeat(4096); - let body = if headers.starts_with("Transfer") { - format!("{:x}\r\n{prefix}\r\n", prefix.len()) - } else { - prefix.clone() - }; - let response = format!("HTTP/1.1 429 Too Many Requests\r\n{headers}\r\n\r\n{body}"); - let error = read_bounded_response(response.into_bytes(), prefix.len()) - .await - .unwrap_err(); - match error { - OcrError::Transport(litellm_http::transport::Error::Http { status, body }) => { - assert_eq!(status, 429); - assert_eq!(body, prefix); - } - error => panic!("unexpected error: {error}"), - } -} - -#[test] -fn response_limit_is_validated_and_not_forwarded_to_the_provider() { - let request = wire_request( - "mistral/model", - "http://localhost", - json!({"max_response_bytes": 123}), - ); - assert_eq!(request.transport.max_response_bytes, 123); - assert!(!request.optional_params.contains_key("max_response_bytes")); - for value in [ - json!(0), - json!(-1), - json!(true), - json!("123"), - json!(1.5), - json!(OCR_RESPONSE_MAX_BYTES + 1), - Value::Null, - ] { - let wire = serde_json::from_value(json!({ - "model": "mistral/model", "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "optional_params": {"max_response_bytes": value} - })).unwrap(); - let Err(error) = decode_request(wire) else { - panic!("invalid response limit accepted") - }; - assert!(error.to_string().contains("max_response_bytes")); - } -} - -#[derive(Debug)] -struct PendingToken { - entered: Arc, - dropped: Arc, -} - -struct TokenFutureDrop(Arc); - -impl Drop for TokenFutureDrop { - fn drop(&mut self) { - self.0.store(true, std::sync::atomic::Ordering::SeqCst); - } -} - -impl litellm_auth::TokenProvider for PendingToken { - fn acquire(&self) -> litellm_auth::TokenFuture<'_> { - Box::pin(async move { - let _guard = TokenFutureDrop(self.dropped.clone()); - self.entered.notify_one(); - std::future::pending().await - }) - } -} - -#[tokio::test] -async fn interrupt_drops_provider_captures_before_returning() { - use std::sync::atomic::{AtomicBool, Ordering}; - - let entered = Arc::new(tokio::sync::Notify::new()); - let dropped = Arc::new(AtomicBool::new(false)); - let request = wire_request("azure_ai/mistral-ocr", "https://example.invalid", json!({})); - let request = crate::ocr::types::LiteLLMOcrRequest { - transport: OcrTransportConfig { - extra_headers: vec![("authorization".into(), "Bearer test-key".into())], - ..request.transport - }, - azure_ad_token_provider: Some(litellm_auth::TokenProviderHandle::new(Arc::new( - PendingToken { - entered: entered.clone(), - dropped: dropped.clone(), - }, - ))), - ..request - }; - let host = LocalOcrHost::new(request); - let mut machine = ocr_machine(ocr_client()); - let mut result = None; - tokio::time::timeout(std::time::Duration::from_secs(2), async { - loop { - tokio::select! { - _ = entered.notified() => break, - step = machine.resume(result.take()) => { - result = Some(match step.unwrap() { - MachineStep::Host(HostOp::Route(op)) => HostResult::Route(host.route(op).await.unwrap()), - MachineStep::Host(HostOp::BeforeSend { wire, .. }) => { - HostResult::BeforeSend(wire) - } - MachineStep::Host(HostOp::Emit(_)) => HostResult::Emitted, - MachineStep::Complete(_) => panic!("pending provider completed"), - }); - } - } - } - }) - .await - .unwrap(); - assert!(!dropped.load(Ordering::SeqCst)); - let selected = OcrError::InvalidRequest("cancelled".into()); - let acknowledgement = machine.interrupt(HostFailure::Cancelled(selected.clone())); - assert!( - dropped.load(Ordering::SeqCst), - "interrupt returned while provider captures were still alive" - ); - assert!( - matches!(acknowledgement.await, Err(OcrError::InvalidRequest(message)) if message == "cancelled") - ); -} - -struct CallerTokenHost { - request: Mutex>, - trace: Mutex>, -} - -impl Host for CallerTokenHost { - async fn route(&self, op: OcrOp) -> Result { - match op { - OcrOp::ProjectRequest => { - self.trace.lock().unwrap().push("project".into()); - Ok(OcrOpResult::Request { - request: Box::new(self.request.lock().unwrap().take().unwrap()), - caller_token: true, - }) - } - OcrOp::AcquireAzureAdToken => { - self.trace.lock().unwrap().push("token".into()); - Ok(OcrOpResult::AzureAdToken( - litellm_auth::ResolvedCredential::Static(litellm_auth::SecretValue::new( - "caller-token", - )), - )) - } - OcrOp::ReadDocument => Err(OcrError::InvalidRequest("no reader".into())), - } - } - - async fn before_send( - &self, - wire: WireRequest, - _: &litellm_host::event::RequestContext, - ) -> Result { - let is_authorization = |name: &str| name.eq_ignore_ascii_case("authorization"); - let authorization = wire - .headers - .iter() - .find(|(name, _)| is_authorization(name)) - .map(|(_, value)| value.clone()) - .unwrap_or_default(); - self.trace - .lock() - .unwrap() - .push(format!("before_send:{authorization}")); - let headers = wire - .headers - .into_iter() - .map(|(name, value)| match is_authorization(&name) { - true => (name, "Bearer edited".to_string()), - false => (name, value), - }) - .collect(); - Ok(WireRequest { headers, ..wire }) - } -} - -#[tokio::test] -async fn the_callers_azure_token_is_acquired_before_before_send_which_can_still_replace_it() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let mut request = wire_request("azure_ai/model", &base, json!({})); - request.credentials.api_key = None; - let host = CallerTokenHost { - request: Mutex::new(Some(request)), - trace: Mutex::new(Vec::new()), - }; - - litellm_host::run::run(ocr_machine(ocr_client()), &host) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!( - *host.trace.lock().unwrap(), - ["project", "token", "before_send:Bearer caller-token"] - ); - assert!( - seen.lock().unwrap()[0] - .to_ascii_lowercase() - .contains("authorization: bearer edited\r\n") - ); -} - -#[tokio::test] -async fn interrupting_an_in_flight_provider_request_closes_its_connection() { - use tokio::io::AsyncReadExt; - - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://{}", listener.local_addr().unwrap()); - let received = Arc::new(tokio::sync::Notify::new()); - let server_received = received.clone(); - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut request = Vec::new(); - let mut buffer = [0u8; 4096]; - while !request.windows(4).any(|window| window == b"\r\n\r\n") { - let read = socket.read(&mut buffer).await.unwrap(); - request.extend_from_slice(&buffer[..read]); - } - server_received.notify_one(); - loop { - if socket.read(&mut buffer).await.unwrap() == 0 { - break; - } - } - }); - let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))); - let mut machine = ocr_machine(ocr_client()); - let mut result = None; - tokio::time::timeout(std::time::Duration::from_secs(2), async { - loop { - tokio::select! { - _ = received.notified() => break, - step = machine.resume(result.take()) => { - result = Some(match step.unwrap() { - MachineStep::Host(HostOp::Route(op)) => HostResult::Route(host.route(op).await.unwrap()), - MachineStep::Host(HostOp::BeforeSend { wire, .. }) => HostResult::BeforeSend(wire), - MachineStep::Host(HostOp::Emit(_)) => HostResult::Emitted, - MachineStep::Complete(_) => panic!("the stalled provider completed"), - }); - } - } - } - }) - .await - .unwrap(); - - let cancelled = OcrError::InvalidRequest("cancelled".into()); - assert!( - machine - .interrupt(HostFailure::Cancelled(cancelled)) - .await - .is_err() - ); - tokio::time::timeout(std::time::Duration::from_secs(1), server) - .await - .expect("the provider connection stayed open after the interrupt") - .unwrap(); -} diff --git a/litellm-rust/crates/core/tests/ocr/document.rs b/litellm-rust/crates/core/tests/ocr/document.rs deleted file mode 100644 index 855548dc6bf..00000000000 --- a/litellm-rust/crates/core/tests/ocr/document.rs +++ /dev/null @@ -1,152 +0,0 @@ -use litellm_host::event::WireRequest; -use litellm_llms::base_llm::ocr::error::Error; -use rstest::rstest; -use serde_json::{Value, json}; - -use super::test_support::{ - MockResponse, SERVED_DOCUMENT, document_server, mock_server, perform_ocr_with, request_body, - wire_request_with_document, -}; -use crate::ocr::route::LocalOcrHost; - -#[derive(Clone, Copy, Debug)] -enum Route { - Mistral, - AzureAi, - VertexMistral, - AzureCohereParse, - Cohere, -} - -impl Route { - fn model(self) -> &'static str { - match self { - Self::Mistral => "mistral/model", - Self::AzureAi => "azure_ai/model", - Self::VertexMistral => "vertex_ai/mistral-ocr-maas", - Self::AzureCohereParse => "azure_ai/cohere-parse", - Self::Cohere => "cohere/model", - } - } - - fn document_type(self) -> &'static str { - match self { - Self::Mistral | Self::AzureAi | Self::VertexMistral => "document_url", - Self::AzureCohereParse | Self::Cohere => "image_url", - } - } - - fn options(self) -> Value { - match self { - Self::Mistral | Self::AzureAi => json!({"pages": [0]}), - Self::VertexMistral => json!({"pages": [0], "vertex_project": "project-1"}), - Self::AzureCohereParse | Self::Cohere => json!({"output_format": "markdown"}), - } - } -} - -/// What the host does to the wire request in `before_send`. -#[derive(Clone, Copy, Debug)] -enum Host { - Detached, - ReplacesDocument, -} - -const REPLACED_DOCUMENT: &str = "data:image/png;base64,cmVwbGFjZWQ="; - -impl Host { - fn before_send(self, wire: WireRequest) -> WireRequest { - let Value::Object(fields) = wire.body else { - return wire; - }; - let body = fields - .into_iter() - .map(|(name, value)| match self { - Self::Detached => (name, value), - Self::ReplacesDocument if name == "document" => { - let document_type = value["type"].clone(); - let key = document_type.as_str().unwrap_or_default().to_string(); - (name, json!({"type": document_type, key: REPLACED_DOCUMENT})) - } - Self::ReplacesDocument => (name, value), - }) - .collect(); - WireRequest { - body: Value::Object(body), - ..wire - } - } -} - -struct Sent { - result: Result<(), Error>, - provider_body: Option, -} - -async fn send(route: Route, host: Host, document_base: &str) -> Sent { - let (base, seen, provider) = mock_server(vec![MockResponse::json(json!({"pages": []}))]).await; - let document_type = route.document_type(); - let document = - json!({"type": document_type, document_type: format!("{document_base}/scan.png")}); - let request = wire_request_with_document(route.model(), &base, document, route.options()); - let local = - LocalOcrHost::new(request).with_before_send(move |wire, _| Ok(host.before_send(wire))); - let result = perform_ocr_with(local).await.map(|_| ()); - match result { - Ok(()) => provider.await.unwrap(), - Err(_) => provider.abort(), - } - let provider_body = seen - .lock() - .unwrap() - .first() - .map(|request| request_body(request)); - Sent { - result, - provider_body, - } -} - -fn served_document_uri() -> String { - use base64::Engine; - format!( - "data:image/png;base64,{}", - base64::engine::general_purpose::STANDARD.encode(SERVED_DOCUMENT) - ) -} - -#[rstest] -#[case::azure_ai(Route::AzureAi)] -#[case::vertex_mistral(Route::VertexMistral)] -#[case::azure_cohere_parse(Route::AzureCohereParse)] -#[tokio::test] -async fn inlining_routes_send_the_downloaded_document(#[case] route: Route) { - let (document_base, _documents) = document_server().await; - let sent = send(route, Host::Detached, &document_base).await; - sent.result.unwrap(); - assert_eq!( - sent.provider_body.unwrap()["document"][route.document_type()], - json!(served_document_uri()) - ); -} - -#[rstest] -#[tokio::test] -async fn document_replaced_by_the_host_reaches_the_provider( - #[values( - Route::Mistral, - Route::AzureAi, - Route::VertexMistral, - Route::AzureCohereParse, - Route::Cohere - )] - route: Route, -) { - let (document_base, _documents) = document_server().await; - let sent = send(route, Host::ReplacesDocument, &document_base).await; - sent.result.unwrap(); - assert_eq!( - sent.provider_body.unwrap()["document"][route.document_type()], - json!(REPLACED_DOCUMENT) - ); -} diff --git a/litellm-rust/crates/core/tests/ocr/support.rs b/litellm-rust/crates/core/tests/ocr/support.rs deleted file mode 100644 index 974fa3d6655..00000000000 --- a/litellm-rust/crates/core/tests/ocr/support.rs +++ /dev/null @@ -1,203 +0,0 @@ -use std::sync::{Arc, Mutex}; - -use futures_util::future::BoxFuture; -use litellm_host::event::WireRequest; -use litellm_llms::base_llm::ocr::{ - error::Error, - handler::{CallHooks, OcrClient}, - transformation::LiteLLMOcrResponse, -}; -use serde_json::{Value, json}; -use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::TcpListener, -}; - -use crate::ocr::{ - route::{LocalOcrHost, ocr_machine}, - types::LiteLLMOcrRequest, - wire::{OcrWireRequest, decode_request}, -}; - -/// Stands in for a host with no hooks registered: the wire request goes out unchanged -/// and response events go nowhere. -pub(crate) struct NoHooks; - -impl CallHooks for NoHooks { - fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result> { - Box::pin(async move { Ok(wire) }) - } - - fn response_received<'a>(&'a self, _body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> { - Box::pin(async { Ok(()) }) - } -} - -pub(crate) fn ocr_client() -> OcrClient { - let document_http = reqwest::Client::builder() - .redirect(reqwest::redirect::Policy::none()) - .build() - .expect("test document client builds"); - OcrClient::for_test(reqwest::Client::new(), document_http) -} - -pub(crate) async fn perform_ocr(request: LiteLLMOcrRequest) -> Result { - crate::ocr::client::perform(&ocr_client(), request).await -} - -pub(crate) async fn perform_ocr_with(host: LocalOcrHost) -> Result { - litellm_host::run::run(ocr_machine(ocr_client()), &host).await -} - -pub(crate) fn wire_request(model: &str, base: &str, options: Value) -> LiteLLMOcrRequest { - wire_request_with_document( - model, - base, - json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), - options, - ) -} - -pub(crate) fn wire_request_with_document( - model: &str, - base: &str, - document: Value, - options: Value, -) -> LiteLLMOcrRequest { - decode_request(OcrWireRequest { - model: model.into(), - document, - api_key: Some(litellm_auth::SecretValue::new("test-key")), - api_base: Some(base.into()), - custom_llm_provider: None, - extra_headers: None, - optional_params: options.as_object().unwrap().clone(), - input_sources: Default::default(), - timeout_seconds: Some(2.0), - }) - .unwrap() -} - -pub(crate) fn resolved_request( - request: LiteLLMOcrRequest, -) -> crate::ocr::types::ResolvedOcrRequest { - request - .map_document(crate::ocr::document::prepare_document) - .unwrap() -} - -pub(crate) fn with_source(request: LiteLLMOcrRequest, source: &str) -> LiteLLMOcrRequest { - let request = resolved_request(request); - let document = request.document.clone().with_source(source.into()); - request.with_document(document.into()) -} - -pub(crate) fn request_body(request: &str) -> Value { - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() -} - -pub(crate) const SERVED_DOCUMENT: &[u8] = b"\x89PNG served document"; - -/// Serves [`SERVED_DOCUMENT`] as `image/png` to every connection until aborted. -pub(crate) async fn document_server() -> (String, tokio::task::JoinHandle<()>) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://{}", listener.local_addr().unwrap()); - let task = tokio::spawn(async move { - loop { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut buffer = [0u8; 4096]; - let _ = socket.read(&mut buffer).await.unwrap(); - let head = format!( - "HTTP/1.1 200 OK\r\nContent-Type: image/png\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", - SERVED_DOCUMENT.len() - ); - socket.write_all(head.as_bytes()).await.unwrap(); - socket.write_all(SERVED_DOCUMENT).await.unwrap(); - } - }); - (base, task) -} - -pub(crate) struct MockResponse { - pub status: u16, - pub headers: Vec<(&'static str, String)>, - pub body: Value, -} - -impl MockResponse { - pub fn json(body: Value) -> Self { - Self { - status: 200, - headers: vec![], - body, - } - } -} - -pub(crate) async fn mock_server( - responses: Vec, -) -> (String, Arc>>, tokio::task::JoinHandle<()>) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://{}", listener.local_addr().unwrap()); - let requests = Arc::new(Mutex::new(Vec::new())); - let seen = requests.clone(); - let server_base = base.clone(); - let task = tokio::spawn(async move { - for response in responses { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut bytes = Vec::new(); - let mut buffer = [0u8; 4096]; - let header_end = loop { - let n = socket.read(&mut buffer).await.unwrap(); - assert!(n > 0); - bytes.extend_from_slice(&buffer[..n]); - if let Some(index) = bytes.windows(4).position(|s| s == b"\r\n\r\n") { - break index + 4; - } - }; - let length = String::from_utf8_lossy(&bytes[..header_end]) - .lines() - .find_map(|line| { - let (name, value) = line.split_once(':')?; - name.eq_ignore_ascii_case("content-length") - .then(|| value.trim().parse::().unwrap()) - }) - .unwrap_or(0); - while bytes.len() < header_end + length { - let n = socket.read(&mut buffer).await.unwrap(); - assert!(n > 0); - bytes.extend_from_slice(&buffer[..n]); - } - seen.lock() - .unwrap() - .push(String::from_utf8_lossy(&bytes).into_owned()); - let body = serde_json::to_vec(&response.body).unwrap(); - let headers = response - .headers - .into_iter() - .map(|(name, value)| { - format!("{name}: {}\r\n", value.replace("{base}", &server_base)) - }) - .collect::(); - let head = format!( - "HTTP/1.1 {} OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n{}\r\n", - response.status, - body.len(), - headers - ); - socket.write_all(head.as_bytes()).await.unwrap(); - socket.write_all(&body).await.unwrap(); - } - }); - (base, requests, task) -} - -pub(crate) fn header<'a>(request: &'a str, name: &str) -> Option<&'a str> { - request - .lines() - .take_while(|line| !line.is_empty()) - .find_map(|line| { - let (key, value) = line.split_once(':')?; - key.eq_ignore_ascii_case(name).then(|| value.trim()) - }) -} diff --git a/litellm-rust/crates/core/tests/reducto_ocr.rs b/litellm-rust/crates/core/tests/reducto_ocr.rs deleted file mode 100644 index 83e7754122b..00000000000 --- a/litellm-rust/crates/core/tests/reducto_ocr.rs +++ /dev/null @@ -1,584 +0,0 @@ -use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; -use litellm_llms::base_llm::ocr::{error::Error, transformation::OcrDocument}; -use rstest::rstest; -use serde_json::{Value, json}; - -use super::test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request}; -use crate::ocr::route::LocalOcrHost; - -fn request_body(request: &str) -> Value { - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() -} - -#[rstest] -#[case( - "reducto/parse-v3", - json!({ - "formatting":{"table_output_format":"html"}, - "retrieval":{"chunk_mode":"section"}, - "settings":{"ocr_system":"standard"}, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - "reducto://already.pdf", - json!({ - "input":"reducto://already.pdf", - "formatting":{"table_output_format":"html"}, - "retrieval":{"chunk_mode":"section"}, - "settings":{"ocr_system":"standard"}, - "future_ocr_option":true, - "provider_option":"value" - }) -)] -#[case( - "reducto/parse-legacy", - json!({ - "enhance":{"agentic":[{"type":"table"}]}, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - "reducto://legacy.pdf", - json!({ - "document_url":"reducto://legacy.pdf", - "options":{"enhance":{"agentic":[{"type":"table"}]}}, - "future_ocr_option":true, - "provider_option":"value" - }) -)] -#[tokio::test] -async fn request_mapping_matches_python( - #[case] model: &str, - #[case] options: Value, - #[case] source: &str, - #[case] expected: Value, -) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "result":{"chunks":[]} - }))]) - .await; - let request = super::test_support::with_source(wire_request(model, &base, options), source); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /parse ")); - assert_eq!(request_body(&requests[0]), expected); -} - -#[rstest] -#[case("parse-v3")] -#[case("parse-legacy")] -#[tokio::test] -async fn data_uri_upload_preserves_multipart_headers( - #[case] model: &str, - #[values("application/pdf", "image/png")] mime_type: &str, -) { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})), - ]) - .await; - let document = if mime_type.starts_with("image/") { - json!({"type":"image_url","image_url":format!("data:{mime_type};base64,YWJj")}) - } else { - json!({"type":"document_url","document_url":format!("data:{mime_type};base64,YWJj")}) - }; - let mut request = crate::ocr::types::LiteLLMOcrRequest { - document: serde_json::from_value::(document) - .unwrap() - .into(), - ..wire_request(&format!("reducto/{model}"), &base, json!({})) - }; - request.transport.extra_headers = vec![ - ("Content-Type".into(), "application/json".into()), - ("X-Trace".into(), "upload-test".into()), - ]; - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "hello"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 2); - assert!(requests[0].starts_with("POST /upload ")); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("content-type: multipart/form-data; boundary=") - ); - assert!(requests[0].contains("x-trace: upload-test")); - let multipart = requests[0].split_once("\r\n\r\n").unwrap().1; - assert!(multipart.contains(&format!("Content-Type: {mime_type}\r\n"))); - assert!(multipart.contains("\r\n\r\nabc\r\n--")); - assert!(requests[1].starts_with("POST /parse ")); - let source_field = if model == "parse-legacy" { - "document_url" - } else { - "input" - }; - assert_eq!( - request_body(&requests[1]), - json!({source_field:"reducto://uploaded.pdf"}) - ); - for request in requests.iter() { - assert!( - request - .to_ascii_lowercase() - .contains("authorization: bearer test-key\r\n") - ); - } -} - -#[tokio::test] -async fn response_received_stays_after_reducto_upload_and_parse() { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[]}})), - ]) - .await; - let request_count = seen.clone(); - let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))).with_observer( - move |event| { - if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { - assert_eq!(request_count.lock().unwrap().len(), 2); - assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); - } - }, - ); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 2); -} - -#[rstest] -#[case(json!({"file_id":""}))] -#[case(json!({}))] -#[case(json!({"file_id":null}))] -#[tokio::test] -async fn invalid_upload_ids_stop_before_parse(#[case] response: Value) { - let (base, seen, server) = mock_server(vec![MockResponse::json(response)]).await; - let error = perform_ocr(wire_request("reducto/parse-v3", &base, json!({}))) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("file_id")); - assert_eq!(seen.lock().unwrap().len(), 1); -} - -#[tokio::test] -async fn upload_failure_stops_before_parse() { - let (base, seen, server) = mock_server(vec![MockResponse { - status: 503, - headers: vec![], - body: json!({"error":"unavailable"}), - }]) - .await; - assert!( - perform_ocr(wire_request("reducto/parse-v3", &base, json!({}))) - .await - .is_err() - ); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 1); -} - -#[rstest] -#[case("https://example.com/a.pdf", Error::ReductoSource)] -#[case("reducto://", Error::RequestField { path: "document file id".into() })] -#[case("data:application/pdf;base64", Error::InvalidDataUri)] -#[case("data:application/pdf;base64,INVALID!", Error::InvalidDataUri)] -#[tokio::test] -async fn rejects_invalid_document_sources_before_network( - #[case] source: &str, - #[case] expected: Error, -) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({}))]).await; - let request = super::test_support::with_source( - wire_request("reducto/parse-v3", &base, json!({})), - source, - ); - let result = perform_ocr(request).await; - server.abort(); - let _ = server.await; - assert!( - seen.lock().unwrap().is_empty(), - "sent invalid source: {source}" - ); - let error = result.unwrap_err(); - assert_eq!( - std::mem::discriminant(&error), - std::mem::discriminant(&expected) - ); - assert_eq!(error.http_status_code(), Some(400)); - assert_eq!(error.to_string(), expected.to_string()); -} - -#[test] -fn response_normalization_groups_blocks_and_distinguishes_null_result() { - use litellm_llms::reducto::ocr::transformation::{ - ReductoResponse, normalize_response as transform_ocr_response, - }; - - let raw = json!({"usage":{"num_pages":"2","credits":"3"},"result":{"type":"full","chunks":[ - {"blocks":[{ - "type":"Table", - "content":"B", - "bbox":{"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}, - "confidence":"high", - "granular_confidence":{"parse_confidence":0.95,"extract_confidence":null}, - "image_url":null - }]}, - {"blocks":[{"content":"A","bbox":{"page":1},"type":"Text"},{"content":"C","bbox":{"page":1}}]} - ]}}); - let response: ReductoResponse = serde_json::from_value(raw).unwrap(); - let normalized = transform_ocr_response("parse-v3", response) - .unwrap() - .into_json(); - assert_eq!(normalized["pages"][0]["markdown"], "A\n\nC"); - assert_eq!(normalized["pages"][1]["markdown"], "B"); - assert_eq!(normalized["pages"][1]["blocks"][0]["type"], "Table"); - assert_eq!( - normalized["pages"][1]["blocks"][0]["bbox"], - json!({"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}) - ); - assert_eq!(normalized["pages"][1]["blocks"][0]["confidence"], "high"); - assert_eq!( - normalized["pages"][1]["blocks"][0]["granular_confidence"]["parse_confidence"], - 0.95 - ); - assert!(normalized["pages"][1]["blocks"][0]["image_url"].is_null()); - assert_eq!(normalized["usage_info"]["pages_processed"], 2); - assert_eq!(normalized["usage_info"]["credits"], 3.0); - - let missing: ReductoResponse = - serde_json::from_value(json!({"chunks":[{"content":"text"}]})).unwrap(); - let missing = transform_ocr_response("parse-v3", missing).unwrap(); - assert_eq!(missing.pages[0].markdown, "text"); - let null: ReductoResponse = serde_json::from_value( - json!({"result":null,"chunks":[{"content":"ignored"}],"usage":null}), - ) - .unwrap(); - let null = transform_ocr_response("parse-v3", null).unwrap(); - assert!(null.pages.is_empty()); -} - -#[tokio::test] -async fn facade_omits_native_response_by_default_and_preserves_auth_priority() { - let raw = json!({"job_id":"job-1","result":{"chunks":[]}}); - let (base, seen, server) = mock_server(vec![MockResponse::json(raw)]).await; - let mut request = super::test_support::with_source( - wire_request("reducto/parse-v3", &base, json!({})), - "reducto://ready.pdf", - ); - request.transport.extra_headers = vec![("authorization".into(), "Bearer existing".into())]; - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.provider_native_response, None); - assert!( - seen.lock().unwrap()[0] - .to_ascii_lowercase() - .contains("authorization: bearer existing") - ); -} - -#[tokio::test] -async fn native_format_retains_the_provider_response() { - let raw = json!({ - "result":{"chunks":[{"content":"native OCR response"}]}, - "usage":{"num_pages":1} - }); - let (base, _, server) = mock_server(vec![MockResponse::json(raw.clone())]).await; - let request = super::test_support::with_source( - wire_request("reducto/parse-v3", &base, json!({"req_format":"native"})), - "reducto://ready.pdf", - ); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - - assert_eq!(response.pages[0].markdown, "native OCR response"); - assert_eq!(response.provider_native_response.as_ref(), raw.as_object()); -} - -#[tokio::test] -async fn unknown_model_reaches_parse_and_keeps_its_name() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "result":{"chunks":[{"content":"future model response"}]} - }))]) - .await; - let request = super::test_support::with_source( - wire_request("reducto/future-parse-model", &base, json!({})), - "reducto://ready.pdf", - ); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - - assert_eq!(response.model, "future-parse-model"); - assert_eq!(response.pages[0].markdown, "future model response"); - let requests = seen.lock().unwrap(); - assert!(requests[0].starts_with("POST /parse ")); - assert_eq!( - request_body(&requests[0]), - json!({"input":"reducto://ready.pdf"}) - ); -} - -#[tokio::test] -async fn guardrail_rewrites_document_before_upload() { - let (base, seen, server) = - mock_server(vec![MockResponse::json(json!({"result":{"chunks":[]}}))]).await; - let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) - .with_before_send(|wire, _| { - assert_eq!( - wire.body["document_url"], - "data:application/pdf;base64,YWJj" - ); - Ok(WireRequest { - body: json!({"type":"document_url","document_url":"reducto://guarded.pdf"}), - ..wire - }) - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /parse ")); - assert!(requests[0].contains("reducto://guarded.pdf")); -} - -mod transformation { - use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; - use litellm_llms::{ - base_llm::ocr::transformation::{BaseOcrConfig, OcrConnection, OcrRequestContext}, - reducto::ocr::transformation::*, - }; - use rstest::rstest; - - use super::*; - use crate::ocr::{ - route::LocalOcrHost, - test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request}, - }; - - #[tokio::test] - async fn v3_options_preserve_explicit_null() { - let overrides = - serde_json::from_value(json!({"formatting":null,"settings":{},"unknown":true})) - .unwrap(); - let params = ReductoParseV3Config - .map_ocr_params(&overrides, "parse-v3") - .unwrap(); - let client = crate::ocr::test_support::ocr_client(); - let connection = OcrConnection::default(); - let document = serde_json::from_value( - json!({"type":"document_url","document_url":"reducto://ready.pdf"}), - ) - .unwrap(); - let body = ReductoParseV3Config - .async_transform_ocr_request( - "parse-v3", - document, - ¶ms, - &[], - OcrRequestContext { - client: &client, - connection: &connection, - }, - ) - .await - .unwrap(); - assert_eq!( - serde_json::to_value(body).unwrap(), - json!({ - "input":"reducto://ready.pdf", "formatting":null, "settings":{} - }) - ); - let absent = ReductoParseV3Config - .map_ocr_params( - &litellm_core_utils::call_arguments::CallArguments::default(), - "parse-v3", - ) - .unwrap(); - assert_eq!(serde_json::to_value(absent).unwrap(), json!({})); - } - - #[rstest] - #[case( - "reducto/parse-v3", - json!({ - "formatting":{"table_output_format":"html"}, - "retrieval":{"chunk_mode":"section"}, - "settings":{"ocr_system":"standard"}, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - "reducto://already.pdf", - json!({ - "input":"reducto://already.pdf", - "formatting":{"table_output_format":"html"}, - "retrieval":{"chunk_mode":"section"}, - "settings":{"ocr_system":"standard"}, - "future_ocr_option":true, - "provider_option":"value" - }) - )] - #[case( - "reducto/parse-legacy", - json!({ - "enhance":{"agentic":[{"type":"table"}]}, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - "reducto://legacy.pdf", - json!({ - "document_url":"reducto://legacy.pdf", - "options":{"enhance":{"agentic":[{"type":"table"}]}}, - "future_ocr_option":true, - "provider_option":"value" - }) - )] - #[tokio::test] - async fn request_mapping_matches_python( - #[case] model: &str, - #[case] options: Value, - #[case] source: &str, - #[case] expected: Value, - ) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "result":{"chunks":[]} - }))]) - .await; - let request = - crate::ocr::test_support::with_source(wire_request(model, &base, options), source); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /parse ")); - assert_eq!(request_body(&requests[0]), expected); - } - - #[rstest] - #[case("parse-v3")] - #[case("parse-legacy")] - #[tokio::test] - async fn data_uri_upload_preserves_multipart_headers(#[case] model: &str) { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})), - ]) - .await; - let mut request = wire_request(&format!("reducto/{model}"), &base, json!({})); - request.transport.extra_headers = vec![ - ("Content-Type".into(), "application/json".into()), - ("X-Trace".into(), "upload-test".into()), - ]; - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "hello"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 2); - assert!(requests[0].starts_with("POST /upload ")); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("content-type: multipart/form-data; boundary=") - ); - assert!(requests[0].contains("x-trace: upload-test")); - assert!(requests[0].contains("application/pdf")); - assert!(requests[0].contains("abc")); - assert!(requests[1].starts_with("POST /parse ")); - } - - #[tokio::test] - async fn response_received_stays_after_reducto_upload_and_parse() { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[]}})), - ]) - .await; - let request_count = seen.clone(); - let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) - .with_observer(move |event| { - if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { - assert_eq!(request_count.lock().unwrap().len(), 2); - assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); - } - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 2); - } - - #[rstest] - #[case("https://example.com/a.pdf")] - #[case("reducto://")] - #[case("data:application/pdf;base64")] - #[case("data:application/pdf;base64,INVALID!")] - #[tokio::test] - async fn rejects_invalid_document_sources_before_network(#[case] source: &str) { - let request = crate::ocr::test_support::with_source( - wire_request("reducto/parse-v3", "http://127.0.0.1:1", json!({})), - source, - ); - assert!(perform_ocr(request).await.is_err()); - } - - #[tokio::test] - async fn facade_omits_native_response_by_default_and_preserves_auth_priority() { - let raw = json!({"job_id":"job-1","result":{"chunks":[]}}); - let (base, seen, server) = mock_server(vec![MockResponse::json(raw)]).await; - let mut request = crate::ocr::test_support::with_source( - wire_request("reducto/parse-v3", &base, json!({})), - "reducto://ready.pdf", - ); - request.transport.extra_headers = vec![("authorization".into(), "Bearer existing".into())]; - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.provider_native_response, None); - assert!( - seen.lock().unwrap()[0] - .to_ascii_lowercase() - .contains("authorization: bearer existing") - ); - } - - #[rstest] - #[case("reducto/parse-v3")] - #[case("reducto/parse-legacy")] - #[tokio::test] - async fn guardrail_headers_reach_upload_and_parse(#[case] model: &str) { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[]}})), - ]) - .await; - let mut request = wire_request(model, &base, json!({})); - request.transport.extra_headers = vec![("authorization".into(), "Bearer original".into())]; - let host = LocalOcrHost::new(request).with_before_send(|wire, _| { - Ok(WireRequest { - headers: vec![("authorization".into(), "Bearer guarded".into())], - ..wire - }) - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 2); - assert!(requests[0].starts_with("POST /upload ")); - assert!(requests[1].starts_with("POST /parse ")); - for request in requests.iter() { - assert!(request.contains("authorization: Bearer guarded")); - assert!(!request.contains("Bearer original")); - } - } -} diff --git a/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs deleted file mode 100644 index 2e8d69f5f64..00000000000 --- a/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs +++ /dev/null @@ -1,143 +0,0 @@ -use litellm_auth::InputSource; -use serde_json::{Value, json}; - -use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; - -fn request_body(request: &str) -> Value { - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() -} - -#[tokio::test] -async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "choices":[{"message":{"content":"recognized"}}], - "usage":{"prompt_tokens":1} - }))]) - .await; - let request = wire_request( - "vertex_ai/deepseek-ocr-maas", - &base, - json!({ - "vertex_project":"project-1", - "vertex_location":"europe-west4", - "temperature":0.1, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - ); - let request = super::test_support::with_source(request, "gs://bucket/document.pdf"); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "recognized"); - assert_eq!( - response.usage_info.unwrap().extra_fields["prompt_tokens"], - 1 - ); - let requests = seen.lock().unwrap(); - assert!(requests[0].starts_with( - "POST /v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions " - )); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer test-key") - ); - let body = request_body(&requests[0]); - assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); - assert_eq!(body["temperature"], 0.1); - assert_eq!(body["future_ocr_option"], true); - assert!(body.get("extra_body").is_none()); - assert_eq!( - body["messages"][0]["content"][0], - json!({"type":"image_url","image_url":"gs://bucket/document.pdf"}) - ); -} - -#[test] -fn host_registration_selects_deepseek_without_affecting_mistral() { - assert!(crate::ocr::arguments::is_supported_request( - "deepseek-ocr-maas", - Some("vertex_ai") - )); - assert!(crate::ocr::arguments::is_supported_request( - "mistral-ocr-maas", - Some("vertex_ai") - )); -} - -#[tokio::test] -async fn request_controlled_api_base_is_rejected_before_vertex_auth() { - let mut request = wire_request( - "vertex_ai/deepseek-ocr-maas", - "https://caller.example", - json!({"vertex_project":"project-1"}), - ); - request.credentials.api_base = Some(litellm_auth::Sourced::new( - "https://caller.example".into(), - InputSource::Request, - )); - - let error = perform_ocr(request).await.unwrap_err(); - assert!( - error - .to_string() - .contains("request-controlled Vertex AI endpoint") - ); -} - -mod deepseek_transformation { - use serde_json::json; - - use super::*; - use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; - - #[tokio::test] - async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "choices":[{"message":{"content":"recognized"}}], - "usage":{"prompt_tokens":1} - }))]) - .await; - let request = wire_request( - "vertex_ai/deepseek-ocr-maas", - &base, - json!({ - "vertex_project":"project-1", - "vertex_location":"europe-west4", - "temperature":0.1, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - ); - let request = crate::ocr::test_support::with_source(request, "gs://bucket/document.pdf"); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "recognized"); - assert_eq!( - response.usage_info.unwrap().extra_fields["prompt_tokens"], - 1 - ); - let requests = seen.lock().unwrap(); - assert!(requests[0].starts_with( - "POST /v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions " - )); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer test-key") - ); - let body = request_body(&requests[0]); - assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); - assert_eq!(body["temperature"], 0.1); - assert_eq!(body["future_ocr_option"], true); - assert_eq!(body["provider_option"], "value"); - assert!(body.get("vertex_project").is_none()); - assert!(body.get("extra_body").is_none()); - assert_eq!( - body["messages"][0]["content"][0], - json!({"type":"image_url","image_url":"gs://bucket/document.pdf"}) - ); - } -} diff --git a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs deleted file mode 100644 index 035f3fe944d..00000000000 --- a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs +++ /dev/null @@ -1,293 +0,0 @@ -use litellm_auth::InputSource; -use litellm_llms::base_llm::ocr::{settings::OcrSettings, transformation::OcrResponseFormat}; -use serde_json::{Value, json}; - -use super::test_support::{MockResponse, mock_server, ocr_client, perform_ocr, wire_request}; - -fn request_body(request: &str) -> Value { - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() -} - -#[tokio::test] -async fn facade_executes_vertex_mistral_with_resolved_project_and_location() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"hello"}], - "usage_info":{"pages_processed":1} - }))]) - .await; - let request = wire_request( - "vertex_ai/mistral-ocr-maas", - &base, - json!({ - "vertex_project":"project-1", - "vertex_location":"europe-west4", - "extract_footer":true - }), - ); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "hello"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with( - "POST /v1/projects/project-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict " - )); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer test-key") - ); - assert_eq!( - request_body(&requests[0]), - json!({ - "model":"mistral-ocr-maas", - "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, - "extract_footer":true - }) - ); -} - -#[tokio::test] -async fn configured_project_and_location_apply_when_the_call_sets_neither() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let client = ocr_client().with_settings(OcrSettings { - vertex_project: Some("configured-project".into()), - vertex_location: Some("europe-west4".into()), - ..OcrSettings::default() - }); - - crate::ocr::client::perform( - &client, - wire_request("vertex_ai/mistral-ocr-maas", &base, json!({})), - ) - .await - .unwrap(); - server.await.unwrap(); - assert!(seen.lock().unwrap()[0].starts_with( - "POST /v1/projects/configured-project/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict " - )); -} - -#[tokio::test] -async fn supplied_authorization_is_forwarded_without_a_static_token() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let mut request = wire_request( - "vertex_ai/model", - &base, - json!({"vertex_project":"project-1"}), - ); - request.credentials.api_key = None; - request.transport.extra_headers = vec![("authorization".into(), "Bearer supplied".into())]; - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert!( - seen.lock().unwrap()[0] - .to_ascii_lowercase() - .contains("authorization: bearer supplied") - ); -} - -#[tokio::test] -async fn invalid_credentials_fail_before_provider_http() { - let request = wire_request( - "vertex_ai/model", - "http://127.0.0.1:1", - json!({"vertex_credentials": true}), - ); - let error = perform_ocr(request).await.unwrap_err(); - assert!(error.to_string().contains("vertex_credentials")); -} - -#[tokio::test] -async fn request_controlled_api_base_is_rejected_before_vertex_auth() { - let mut request = wire_request( - "vertex_ai/mistral-ocr-maas", - "https://caller.example", - json!({"vertex_project":"project-1"}), - ); - request.credentials.api_base = Some(litellm_auth::Sourced::new( - "https://caller.example".into(), - InputSource::Request, - )); - - let error = perform_ocr(request).await.unwrap_err(); - assert!( - error - .to_string() - .contains("request-controlled Vertex AI endpoint") - ); -} - -#[tokio::test] -async fn adapters_build_complete_requests_and_share_mistral_normalization() { - use std::time::Duration; - - use litellm_llms::{ - base_llm::ocr::transformation::BaseOcrConfig, - mistral::ocr::transformation::MistralOcrConfig, - vertex_ai::ocr::transformation::VertexAiOcrConfig, - }; - - use crate::ocr::test_support::ocr_client; - - let client = ocr_client(); - let options = json!({ - "pages": [0, 2], - "include_image_base64": true, - "vertex_project": "project-1", - "vertex_location": "us-central1", - "unknown": "ignored" - }); - let direct = wire_request( - "mistral/mistral-ocr-maas", - "https://mistral.test", - options.clone(), - ); - let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options); - let direct = crate::ocr::prepare::prepare_request_for_test( - super::test_support::resolved_request(direct), - ); - let vertex = crate::ocr::prepare::prepare_request_for_test( - super::test_support::resolved_request(vertex), - ); - let direct_http = MistralOcrConfig - .prepare_request(&direct, &client, &crate::ocr::test_support::NoHooks) - .await - .unwrap(); - let vertex_http = VertexAiOcrConfig - .prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks) - .await - .unwrap(); - assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr"); - assert_eq!( - vertex_http.url(), - "https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict" - ); - for http in [&direct_http, &vertex_http] { - assert_eq!(http.header("authorization").unwrap(), "Bearer test-key"); - assert_eq!(http.header("content-type").unwrap(), "application/json"); - assert_eq!(http.timeout(), Some(Duration::from_secs(2))); - let body: Value = serde_json::from_slice(http.body()).unwrap(); - assert_eq!( - body, - json!({ - "model": "mistral-ocr-maas", - "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "pages": [0, 2], - "include_image_base64": true, - "unknown": "ignored" - }) - ); - } - let payload = json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}); - let raw = serde_json::to_vec(&payload).unwrap(); - let direct_response = MistralOcrConfig - .transform_ocr_response(&direct.model, &raw, OcrResponseFormat::Litellm) - .unwrap() - .into_json(); - let vertex_response = VertexAiOcrConfig - .transform_ocr_response(&vertex.model, &raw, OcrResponseFormat::Litellm) - .unwrap() - .into_json(); - assert_eq!(direct_response, vertex_response); - assert_eq!(direct_response["model"], "mistral-ocr-maas"); - assert_eq!(direct_response["object"], "ocr"); - assert_eq!(direct_response["extra"], "preserved"); -} - -mod transformation { - - use rstest::rstest; - use serde_json::{Value, json}; - - use crate::ocr::test_support::wire_request; - - #[rstest] - #[case::mistral(false)] - #[case::vertex(true)] - #[tokio::test] - async fn configs_build_complete_requests_and_share_mistral_normalization( - #[case] use_vertex: bool, - ) { - use std::time::Duration; - - use litellm_llms::{ - base_llm::ocr::transformation::BaseOcrConfig, - mistral::ocr::transformation::MistralOcrConfig, - vertex_ai::ocr::transformation::VertexAiOcrConfig, - }; - - use crate::ocr::test_support::ocr_client; - - let client = ocr_client(); - let options = json!({ - "pages": [0, 2], - "include_image_base64": true, - "vertex_project": "project-1", - "vertex_location": "us-central1", - "unknown": "preserved" - }); - let direct = wire_request( - "mistral/mistral-ocr-maas", - "https://mistral.test", - options.clone(), - ); - let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options); - let direct = crate::ocr::prepare::prepare_request_for_test( - crate::ocr::test_support::resolved_request(direct), - ); - let vertex = crate::ocr::prepare::prepare_request_for_test( - crate::ocr::test_support::resolved_request(vertex), - ); - let direct_http = MistralOcrConfig - .prepare_request(&direct, &client, &crate::ocr::test_support::NoHooks) - .await - .unwrap(); - let vertex_http = VertexAiOcrConfig - .prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks) - .await - .unwrap(); - assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr"); - assert_eq!( - vertex_http.url(), - "https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict" - ); - let http = if use_vertex { - &vertex_http - } else { - &direct_http - }; - assert_eq!(http.header("authorization").unwrap(), "Bearer test-key"); - assert_eq!(http.header("content-type").unwrap(), "application/json"); - assert_eq!(http.timeout(), Some(Duration::from_secs(2))); - let body: Value = serde_json::from_slice(http.body()).unwrap(); - assert_eq!( - body, - json!({ - "model": "mistral-ocr-maas", - "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "pages": [0, 2], - "include_image_base64": true, - "unknown": "preserved" - }) - ); - let payload = serde_json::to_vec( - &json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}), - ) - .unwrap(); - let direct_response = MistralOcrConfig - .transform_ocr_response(&direct.model, &payload, Default::default()) - .unwrap() - .into_json(); - let vertex_response = VertexAiOcrConfig - .transform_ocr_response(&vertex.model, &payload, Default::default()) - .unwrap() - .into_json(); - assert_eq!(direct_response, vertex_response); - assert_eq!(direct_response["model"], "mistral-ocr-maas"); - assert_eq!(direct_response["object"], "ocr"); - assert_eq!(direct_response["extra"], "preserved"); - } -} diff --git a/litellm-rust/crates/http/src/request.rs b/litellm-rust/crates/http/src/request.rs index 874a0f3abf9..fcf296793a5 100644 --- a/litellm-rust/crates/http/src/request.rs +++ b/litellm-rust/crates/http/src/request.rs @@ -217,13 +217,25 @@ mod tests { "Authorization".to_string(), "Bearer abc".to_string() )])); + assert!(has_bearer_auth(&[( + "authorization".to_string(), + "bearer abc".to_string() + )])); assert!(!has_bearer_auth(&[( "Authorization".to_string(), "Bearer ".to_string() )])); + assert!(!has_bearer_auth(&[( + "authorization".to_string(), + String::new() + )])); assert!(!has_bearer_auth(&[( "Authorization".to_string(), "Basic abc".to_string() )])); + assert!(!has_bearer_auth(&[( + "x-api-key".to_string(), + "abc".to_string() + )])); } } diff --git a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs index 6fc4f00b981..fd86c5ca25a 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs @@ -218,7 +218,3 @@ fn anthropic_body( ); Value::Object(body) } - -#[cfg(test)] -#[path = "tests.rs"] -mod tests; diff --git a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs index 09c456f1d0a..b5db88d7dc4 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs @@ -302,7 +302,3 @@ fn has_blank_text(message: &ChatMessage) -> bool { }), } } - -#[cfg(test)] -#[path = "tests.rs"] -mod tests; diff --git a/litellm-rust/crates/llms/src/anthropic/chat/tests.rs b/litellm-rust/crates/llms/tests/anthropic_chat_transformation.rs similarity index 96% rename from litellm-rust/crates/llms/src/anthropic/chat/tests.rs rename to litellm-rust/crates/llms/tests/anthropic_chat_transformation.rs index 3777347d240..ed22a1d141d 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/tests.rs +++ b/litellm-rust/crates/llms/tests/anthropic_chat_transformation.rs @@ -1,7 +1,11 @@ -use serde_json::json; - -use super::*; -use crate::base_llm::chat::transformation::Error; +use litellm_llms::{ + anthropic::chat::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG, + base_llm::chat::transformation::{ + BaseConfig, Error, ProviderChatResponseData, RequestAuth, Unsupported, + }, +}; +use litellm_types::{llms::openai::ChatMessage, utils::ChatCompletionsResponse}; +use serde_json::{Map, Value, json}; fn messages(value: Value) -> Vec { serde_json::from_value(value).expect("valid messages") @@ -205,7 +209,8 @@ fn declines_tool_calls_tool_results_and_multimodal_content() { ); assert_eq!( reason( - json!([{"role": "user", "content": [ + json!([ + {"role": "user", "content": [ {"type": "image_url", "image_url": {"url": "https://x/y.png"}} ]}]), json!({}) @@ -214,7 +219,8 @@ fn declines_tool_calls_tool_results_and_multimodal_content() { ); assert_eq!( reason( - json!([{"role": "user", "content": [ + json!([ + {"role": "user", "content": [ {"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}} ]}]), json!({}) diff --git a/litellm-rust/crates/llms/src/bedrock/chat/tests.rs b/litellm-rust/crates/llms/tests/bedrock_converse_transformation.rs similarity index 98% rename from litellm-rust/crates/llms/src/bedrock/chat/tests.rs rename to litellm-rust/crates/llms/tests/bedrock_converse_transformation.rs index d7ecde47c6b..4127bcfa19d 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/tests.rs +++ b/litellm-rust/crates/llms/tests/bedrock_converse_transformation.rs @@ -1,7 +1,11 @@ -use serde_json::json; - -use super::*; -use crate::base_llm::chat::transformation::Error; +use litellm_llms::{ + base_llm::chat::transformation::{ + BaseConfig, Error, ProviderChatResponseData, RequestAuth, Unsupported, + }, + bedrock::chat::converse_transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG, +}; +use litellm_types::{llms::openai::ChatMessage, utils::ChatCompletionsResponse}; +use serde_json::{Map, Value, json}; fn messages(value: Value) -> Vec { serde_json::from_value(value).expect("valid messages") diff --git a/litellm-rust/crates/types/src/utils.rs b/litellm-rust/crates/types/src/utils.rs index 5ca56ec9e49..af0ba2c01c9 100644 --- a/litellm-rust/crates/types/src/utils.rs +++ b/litellm-rust/crates/types/src/utils.rs @@ -56,7 +56,7 @@ pub struct ChatCompletionsChoice { /// /// There is deliberately no `id`: Python mints the `chatcmpl-…` id on the /// `ModelResponse` it already created, and echoing the provider's own id here -/// would change it. Pinned by `response_carries_no_id` in `tests.rs`. +/// would change it. Pinned by `response_carries_no_id` in the Anthropic chat transformation tests. #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct ChatCompletionsResponse { pub created: u64,