fix(ocr): forward the supplied client on the Python path and build pooled clients outside the lock

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm-rust/crates/http/src/pool.rs

@@ -1,6 +1,6 @@
 use std::{
     collections::HashMap,
-    sync::{Arc, Mutex, PoisonError},
+    sync::{Arc, Mutex, MutexGuard, PoisonError},
 };
 
 use reqwest::dns::Resolve;
@@ -38,13 +38,15 @@ impl HttpClientPool {
         variant: ClientVariant,
     ) -> Result<reqwest::Client, Error> {
         let key = (config.clone(), variant);
-        let mut clients = self.clients.lock().unwrap_or_else(PoisonError::into_inner);
-        if let Some(client) = clients.get(&key) {
+        if let Some(client) = self.lock().get(&key) {
             return Ok(client.clone());
         }
         let client = self.apply(variant, config.client_builder()?).build()?;
-        clients.insert(key, client.clone());
-        Ok(client)
+        Ok(self.lock().entry(key).or_insert(client).clone())
+    }
+
+    fn lock(&self) -> MutexGuard<'_, HashMap<(HttpClientConfig, ClientVariant), reqwest::Client>> {
+        self.clients.lock().unwrap_or_else(PoisonError::into_inner)
     }
 
     fn apply(

litellm/ocr/main.py

@@ -25,6 +25,7 @@ from litellm.llms.base_llm.ocr.transformation import (
     OCRResponse,
     parse_ocr_request_format,
 )
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
 from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.utils import CustomPricingLiteLLMParams
@@ -52,6 +53,11 @@ class _PreparedOCRRequest:
     litellm_logging_obj: LiteLLMLoggingObj
 
 
+def _supplied_client(kwargs: Mapping[str, object]) -> HTTPHandler | AsyncHTTPHandler | None:
+    candidate: Final = kwargs.get("client")
+    return candidate if isinstance(candidate, (HTTPHandler, AsyncHTTPHandler)) else None
+
+
 def _prepare_ocr_request(
     model: str,
     document: Mapping[str, object],
@@ -238,6 +244,7 @@ async def aocr(
             api_key=prepared.api_key,
             api_base=prepared.api_base,
             custom_llm_provider=prepared.custom_llm_provider,
+            client=_supplied_client(kwargs),
             aocr=True,
             headers=prepared.extra_headers,
             provider_config=prepared.provider_config,
@@ -404,6 +411,7 @@ def ocr(
             api_key=prepared.api_key,
             api_base=prepared.api_base,
             custom_llm_provider=prepared.custom_llm_provider,
+            client=_supplied_client(kwargs),
             aocr=_is_async,
             headers=prepared.extra_headers,
             provider_config=prepared.provider_config,

tests/test_litellm/ocr/test_main.py

@@ -113,6 +113,35 @@ async def test_python_request_response_and_callbacks(
     assert logger.log_pre_api_call.call_count == 1
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True])
+async def test_python_uses_the_supplied_client(provider: Mock, asynchronous: bool) -> None:
+    supplied: Final = Mock(return_value=provider.return_value)
+    transport: Final = httpx.MockTransport(supplied)
+    arguments: Final = {
+        "model": "mistral/mistral-ocr-latest",
+        "document": dict(PRICING_DOCUMENT),
+        "api_key": "test-key",
+        "api_base": "https://ocr.test/v1",
+    }
+
+    async def call() -> OCRResponse:
+        if not asynchronous:
+            with httpx.Client(transport=transport) as sync_client:
+                return litellm.ocr(**arguments, client=HTTPHandler(client=sync_client))
+        async with httpx.AsyncClient(transport=transport) as async_client:
+            handler: Final = AsyncHTTPHandler()
+            await handler.client.aclose()
+            handler.client = async_client
+            return await litellm.aocr(**arguments, client=handler)
+
+    response: Final = await call()
+    assert response.pages[0].markdown == "parsed document"
+    assert supplied.call_count == 1
+    assert str(supplied.call_args.args[0].url) == "https://ocr.test/v1/ocr"
+    assert provider.call_count == 0
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("asynchronous", [False, True])
 async def test_python_provider_errors_keep_public_exception(provider: Mock, asynchronous: bool) -> None: