From 54d678f51b92a3cf7f10f4ce29aeb92cfe1c6c3d Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 5 Sep 2026 17:32:50 -0700 Subject: [PATCH] fix(ocr): complete core ownership after stack rebase --- litellm-rust/crates/ai-gateway/src/ocr/mod.rs | 234 +----------------- .../python-bridge/src/routes/definition.rs | 2 +- 2 files changed, 2 insertions(+), 234 deletions(-) diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs index db4a1cd3fdc..f75a875bc78 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs @@ -1,233 +1 @@ -use crate::integrations::types::RequestHooks; -use litellm_core::Error; -use litellm_core::call_lifecycle::CallLifecycle; -use litellm_core::provider_callbacks::{NoopProviderAttemptObserver, ProviderAttemptObserver}; -use litellm_core::request_context::LiteLlmRequestContext; -use litellm_core::request_options::RequestOptions; -use serde_json::Value; - -mod handler; -mod hooks; -mod prepare; -mod types; - -pub use types::OcrRequest; - -use handler::execute_ocr_provider_call; -use prepare::{PreparedOcrCall, prepare_ocr_call}; - -#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] -pub async fn ocr( - request: OcrRequest<'_>, - options: &RequestOptions, - context: &LiteLlmRequestContext, - hooks: RequestHooks, -) -> Result { - ocr_with_observer( - request, - options, - context, - hooks, - &mut NoopProviderAttemptObserver, - ) - .await -} - -#[tracing::instrument( - name = "ocr", - target = "litellm::function_trace", - level = "trace", - skip_all -)] -pub async fn ocr_with_observer( - request: OcrRequest<'_>, - options: &RequestOptions, - context: &LiteLlmRequestContext, - hooks: RequestHooks, - observer: &mut Observer, -) -> Result -where - Observer: ProviderAttemptObserver, - Observer::Error: std::fmt::Display, -{ - let PreparedOcrCall { - request, - context: lifecycle, - hooks, - } = prepare_ocr_call(request, options.clone(), context, hooks); - CallLifecycle::default() - .run(lifecycle.clone(), request, &hooks, |request| { - execute_ocr_provider_call(request, &lifecycle, &hooks, observer) - }) - .await -} - -pub fn ocr_admitted(model: &str, provider: &str, request_format: Option<&str>) -> bool { - common_utils::ocr_provider_config(provider, model).is_some_and(|config| { - request_format != Some("native") || config.supported_ocr_params().contains(&"req_format") - }) -} - -#[cfg(test)] -mod tests { - use crate::integrations::types::RequestHooks; - use litellm_core::request_context::LiteLlmRequestContext; - use litellm_core::request_options::RequestOptions; - use serde_json::{Map, json}; - use tokio::io::{AsyncReadExt, AsyncWriteExt}; - use tokio::net::{TcpListener, TcpStream}; - - use super::{OcrRequest, ocr}; - - 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 base_ocr_request(model: &str) -> (OcrRequest<'_>, RequestOptions) { - ( - OcrRequest { - model, - document: json!({ - "type": "document_url", - "document_url": "https://example.com/doc.pdf" - }), - optional_params: Map::new(), - }, - RequestOptions { - api_key: (Some("sk-test")).map(|value| value.to_string()), - api_base: None, - custom_llm_provider: None, - extra_headers: None, - timeout: None, - ..Default::default() - }, - ) - } - - #[tokio::test] - async fn reducto_file_upload_then_parse_maps_response() { - let listener = TcpListener::bind("127.0.0.1:0") - .await - .expect("test listener binds"); - let address = listener.local_addr().expect("listener has local address"); - let server = tokio::spawn(async move { - let (mut upload_socket, _) = listener.accept().await.expect("accepts upload request"); - let upload_request = read_http_request(&mut upload_socket).await; - let upload_body = r#"{"file_id":"reducto://uploaded.pdf"}"#; - let upload_response = format!( - "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - upload_body.len(), - upload_body - ); - upload_socket - .write_all(upload_response.as_bytes()) - .await - .expect("writes upload response"); - - let (mut parse_socket, _) = listener.accept().await.expect("accepts parse request"); - let parse_request = read_http_request(&mut parse_socket).await; - let parse_body = r#"{"job_id":"job_123","usage":{"num_pages":3,"credits":3},"result":{"chunks":[{"content":"Page 1 block A","blocks":[{"content":"Page 1 block A","bbox":{"page":1},"kind":"text"}]},{"content":"Page 2 block A","blocks":[{"content":"Page 2 block A","bbox":{"page":2},"kind":"table"}]},{"content":"Page 1 block B","blocks":[{"content":"Page 1 block B","bbox":{"page":1},"kind":"text"}]},{"content":"Page 3 block A","blocks":[{"content":"Page 3 block A","bbox":{"page":3},"kind":"figure"}]}]}}"#; - let parse_response = format!( - "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - parse_body.len(), - parse_body - ); - parse_socket - .write_all(parse_response.as_bytes()) - .await - .expect("writes parse response"); - (upload_request, parse_request) - }); - let api_base = format!("http://{address}"); - let (mut request, mut options) = base_ocr_request("reducto/parse-v3"); - options.api_base = Some(&api_base).map(|value| value.to_string()); - options.api_key = None; - options.extra_headers = Some(Map::from_iter([ - ("Authorization".to_string(), json!("Bearer test-key")), - ("x-trace-id".to_string(), json!("trace-1")), - ])); - request.document = json!({ - "type": "document_url", - "document_url": "data:application/pdf;base64,JVBERi0xLjQ=" - }); - request.optional_params = Map::from_iter([ - ( - "formatting".to_string(), - json!({"table_output_format": "html"}), - ), - ("retrieval".to_string(), json!({"chunk_mode": "section"})), - ("settings".to_string(), json!({"ocr_system": "standard"})), - ]); - - let response = ocr( - request, - &options, - &LiteLlmRequestContext { - ..Default::default() - }, - RequestHooks { - ..Default::default() - }, - ) - .await - .expect("Reducto OCR succeeds"); - - assert_eq!(response["pages"].as_array().map(Vec::len), Some(3)); - assert_eq!( - response["pages"][0]["markdown"], - "Page 1 block A\n\nPage 1 block B" - ); - assert_eq!(response["pages"][1]["markdown"], "Page 2 block A"); - assert_eq!(response["pages"][2]["markdown"], "Page 3 block A"); - assert_eq!(response["usage_info"]["pages_processed"], 3); - assert_eq!(response["usage_info"]["credits"], 3); - assert_eq!(response["provider_native_response"]["job_id"], "job_123"); - let (upload_request, parse_request) = server.await.expect("server task completes"); - assert!( - upload_request - .to_ascii_lowercase() - .contains("authorization: bearer test-key") - ); - assert!(upload_request.contains("application/pdf")); - assert!(upload_request.contains("%PDF-1.4")); - assert!(upload_request.contains("x-trace-id: trace-1")); - assert!( - parse_request - .to_ascii_lowercase() - .contains("authorization: bearer test-key") - ); - assert!(parse_request.contains(r#""input":"reducto://uploaded.pdf""#)); - assert!(parse_request.contains(r#""table_output_format":"html""#)); - assert!(parse_request.contains(r#""chunk_mode":"section""#)); - assert!(parse_request.contains(r#""ocr_system":"standard""#)); - } -} +pub use litellm_core::ocr::{OcrRequest, ocr, ocr_admitted, ocr_with_observer}; diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs index b71cede1c77..f2f669f367f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/definition.rs +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -164,7 +164,7 @@ mod tests { inputs: EchoInputs, _options: crate::marshal::NativeRequestOptions, _context: crate::marshal::NativeRequestContext, - _callback_adapter: Option>, + _callback_adapter: Option>, _python_context: crate::execution::PythonCallContext<'_>, ) -> PyResult> + Send + 'static> { FUTURE_DROPPED.store(false, Ordering::SeqCst);