From 5a474fd7996e96f90e18c539108381d811cd5e63 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 23:31:27 +0000 Subject: [PATCH] refactor(rust): inject VertexAuth into OcrClient so the bridge keeps one token cache Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 2 ++ litellm-rust/crates/core/Cargo.toml | 1 + litellm-rust/crates/core/src/ocr/client.rs | 4 +++- litellm-rust/crates/core/tests/ocr.rs | 3 +++ .../crates/llms/src/custom_httpx/llm_http_handler.rs | 8 ++++++-- litellm-rust/crates/python-bridge/Cargo.toml | 1 + litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs | 7 ++++++- 7 files changed, 22 insertions(+), 4 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ffbdc80efd3..ed0fee3411f 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2056,6 +2056,7 @@ dependencies = [ "futures-util", "litellm-auth", "litellm-auth-aws", + "litellm-auth-gcp", "litellm-callbacks", "litellm-core-utils", "litellm-http", @@ -2175,6 +2176,7 @@ dependencies = [ "criterion", "futures-util", "litellm-auth", + "litellm-auth-gcp", "litellm-callbacks-legacy", "litellm-core", "litellm-host-python", diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 047188c74f9..ce8463affc1 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -15,6 +15,7 @@ futures-util.workspace = true base64.workspace = true litellm-auth.workspace = true litellm-auth-aws.workspace = true +litellm-auth-gcp.workspace = true litellm-http.workspace = true litellm-llms.workspace = true moka.workspace = true diff --git a/litellm-rust/crates/core/src/ocr/client.rs b/litellm-rust/crates/core/src/ocr/client.rs index 21c81505cff..a380037487e 100644 --- a/litellm-rust/crates/core/src/ocr/client.rs +++ b/litellm-rust/crates/core/src/ocr/client.rs @@ -1,3 +1,4 @@ +use litellm_auth_gcp::VertexAuth; use litellm_http::{HttpClientConfig, HttpClientPool}; use litellm_llms::{ base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse}, @@ -19,7 +20,8 @@ pub async fn perform( pub async fn ocr( pool: &HttpClientPool, config: &HttpClientConfig, + vertex_auth: VertexAuth, request: LiteLLMOcrRequest, ) -> Result { - perform(&OcrClient::new(pool, config)?, request).await + perform(&OcrClient::new(pool, config, vertex_auth)?, request).await } diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index d59c6179aa5..4ce6d4ee8f6 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -1,5 +1,6 @@ use std::sync::{Arc, Mutex}; +use litellm_auth_gcp::VertexAuth; use litellm_callbacks::{ event::{CallEvent, WireRequest}, host::{Host, HostOp, HostResult}, @@ -182,6 +183,7 @@ async fn facade_uses_the_injected_http_pool_configuration() { crate::ocr::client::ocr( &HttpClientPool::new(), &config, + VertexAuth::default(), wire_request("mistral/model", &base, json!({})), ) .await @@ -200,6 +202,7 @@ async fn unbuildable_http_configuration_fails_before_dispatch() { let error = crate::ocr::client::ocr( &HttpClientPool::new(), &config, + VertexAuth::default(), wire_request("mistral/model", &base, json!({})), ) .await diff --git a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs index 38ddec08da9..9826d73a061 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs +++ b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs @@ -42,12 +42,16 @@ pub struct OcrClient { } impl OcrClient { - pub fn new(pool: &HttpClientPool, config: &HttpClientConfig) -> Result { + pub fn new( + pool: &HttpClientPool, + config: &HttpClientConfig, + vertex_auth: VertexAuth, + ) -> Result { Ok(Self { provider_http: pool.client(config, ClientVariant::Provider)?, polling_http: pool.client(config, ClientVariant::NoRedirect)?, document_fetcher: MediaFetcher::new(pool, config)?, - vertex_auth: VertexAuth::default(), + vertex_auth, }) } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 762e22e433b..c66701548d1 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -20,6 +20,7 @@ bytes.workspace = true litellm-auth.workspace = true litellm-callbacks-legacy.workspace = true litellm-core.workspace = true +litellm-auth-gcp.workspace = true litellm-http.workspace = true litellm-llms.workspace = true litellm-types.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index f252e1dc47b..be188434275 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -3,7 +3,10 @@ mod errors; mod host; mod project; +use std::sync::LazyLock; + use host::OcrRouteHost; +use litellm_auth_gcp::VertexAuth; use litellm_callbacks_legacy::{LegacySurface, PublicCall, run_legacy_call}; use litellm_core::ocr::route::ocr_machine; use litellm_llms::custom_httpx::llm_http_handler::OcrClient; @@ -24,6 +27,8 @@ const ASYNC_SURFACE: LegacySurface = LegacySurface { ..SURFACE }; +static VERTEX_AUTH: LazyLock = LazyLock::new(VertexAuth::default); + fn run_ocr( py: Python<'_>, request: Bound<'_, PyAny>, @@ -32,7 +37,7 @@ fn run_ocr( asynchronous: bool, ) -> PyResult> { let config = http::call_config(py, &kwargs)?; - let client = OcrClient::new(http::pool(), &config) + let client = OcrClient::new(http::pool(), &config, VERTEX_AUTH.clone()) .map_err(|error| RustBridgeDeclined::new_err(error.to_string()))?; run_legacy_call( py,