fix(vertex_ai): reuse one Vertex search OAuth token across searches

Every vector store search built a new search config, so the token cache it carries was thrown away and each search fetched a new Google OAuth token with a blocking call on the event loop. Under a burst of searches that blocking fetch held back every HTTP/2 response. The provider config is now shared, so searches with the same credentials reuse one token until it expires.

litellm/llms/vertex_ai/vector_stores/search_api/transformation.py

@@ -461,3 +461,6 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
 
         input_cost_per_query: Final = model_info.get("input_cost_per_query") or 0.0
         return input_cost_per_query, 0.0
+
+
+vertex_search_api_vector_store_config: Final = VertexSearchAPIVectorStoreConfig()

litellm/utils.py

@@ -10120,10 +10120,10 @@ class ProviderConfigManager:
                 return VertexVectorStoreConfig()
             elif api_type == "search_api":
                 from litellm.llms.vertex_ai.vector_stores.search_api.transformation import (
-                    VertexSearchAPIVectorStoreConfig,
+                    vertex_search_api_vector_store_config,
                 )
 
-                return VertexSearchAPIVectorStoreConfig()
+                return vertex_search_api_vector_store_config
         elif litellm.LlmProviders.BEDROCK == provider:
             from litellm.llms.bedrock.vector_stores.transformation import (
                 BedrockVectorStoreConfig,

tests/integration/sdk/test_http2_wire.py

@@ -40,6 +40,7 @@ STREAM_CHUNKS: Final = 3
 HELD_SEARCHES: Final = 150
 HOLD_SECONDS: Final = 10.0
 _HOLD_ROOMS: Final = MappingProxyType({"hold-async": queue.SimpleQueue[int](), "hold-sync": queue.SimpleQueue[int]()})
+_TOKEN_REQUESTS: Final = queue.SimpleQueue[str]()
 VERTEX_SERVING_CONFIG_PATH: Final = (
     "/v1/projects/p/locations/global/collections/default_collection/engines/e/servingConfigs/default_search"
 )
@@ -108,7 +109,8 @@ async def _peer(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallab
             send, {"results": [{"id": await _hold_until_all_arrive(hold_room, client[1] if client else 0)}]}
         )
         return
-    if scope["path"] == "/_oauth/token":
+    if scope["path"].endswith("/_oauth/token"):
+        _TOKEN_REQUESTS.put(scope["path"])
         await _send_json(send, {"access_token": "scripted", "expires_in": 3600, "token_type": "Bearer"})
         return
     if scope["path"].endswith(":search"):
@@ -327,6 +329,35 @@ def test_default_async_client_stays_http1_after_vertex_search_used_http2(
     assert asyncio.run(_run()) == (["http/2"], "1.1")
 
 
+def test_vertex_search_fetches_one_token_for_many_searches(
+    monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer
+) -> None:
+    import litellm
+
+    _trust_peer_without_global_http2(monkeypatch, http2_tls_peer)
+    credentials: Final = service_account_json("p", f"{http2_tls_peer.url}/token-reuse")
+    search: Final = MappingProxyType(
+        {
+            "vector_store_id": "default_search",
+            "query": "ping",
+            "custom_llm_provider": "vertex_ai/search_api",
+            "api_base": f"{http2_tls_peer.url}{VERTEX_SERVING_CONFIG_PATH}",
+            "vertex_credentials": credentials,
+            "vertex_project": "p",
+        }
+    )
+
+    async def _searches() -> None:
+        await litellm.vector_stores.asearch(**search)
+        await asyncio.gather(*(litellm.vector_stores.asearch(**search) for _ in range(5)))
+
+    asyncio.run(_searches())
+    litellm.vector_stores.search(**search)
+
+    token_paths: Final = tuple(iter(lambda: None if _TOKEN_REQUESTS.empty() else _TOKEN_REQUESTS.get(), None))
+    assert token_paths.count("/token-reuse/_oauth/token") == 1
+
+
 def test_vertex_search_over_forced_ipv4_trusts_the_configured_ca(
     monkeypatch: pytest.MonkeyPatch, http2_tls_peer: TlsPeer
 ) -> None: