From 62ea38079fff208bf70071e28e9659be1d6a15e4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?d=20=F0=9F=94=B9?= Date: Sun, 17 May 2026 13:02:25 +0800 Subject: [PATCH] fix(vertex_ai/partner_models): drop unused vertexai SDK gate from count_tokens (closes #28084) (#28107) Squash-merged by litellm-agent from voidborne-d's PR. --- .../vertex_ai_partner_models/main.py | 25 ++-- .../test_count_tokens_no_vertexai_sdk.py | 121 ++++++++++++++++++ 2 files changed, 130 insertions(+), 16 deletions(-) create mode 100644 tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_no_vertexai_sdk.py diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index 123d925f7c1..eb67e3aa828 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -292,22 +292,15 @@ class VertexAIPartnerModels(VertexBase): Returns: Dict containing token count information """ - try: - import vertexai - except Exception as e: - raise VertexAIError( - status_code=400, - message=f"""vertexai import failed please run `pip install -U "google-cloud-aiplatform>=1.38"`. Got error: {e}""", - ) - - if not ( - hasattr(vertexai, "preview") or hasattr(vertexai.preview, "language_models") - ): - raise VertexAIError( - status_code=400, - message="""Upgrade vertex ai. Run `pip install "google-cloud-aiplatform>=1.38"`""", - ) - + # Note: we intentionally do not import `vertexai` (the Gemini SDK shipped + # by `google-cloud-aiplatform`) on this path. Partner models such as + # Claude on Vertex use the Anthropic Messages API protocol directly via + # `:rawPredict`, and `VertexAIPartnerModelsTokenCounter` reaches that + # endpoint with an authenticated httpx client — it never touches the + # Gemini SDK. Requiring `google-cloud-aiplatform>=1.38` here turned a + # SDK-free Anthropic-protocol call into a hard dependency on the Gemini + # SDK (see #28084), breaking `/v1/messages/count_tokens` for Claude-on- + # Vertex on any LiteLLM install without that extra. Stay SDK-free. try: from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import ( VertexAIPartnerModelsTokenCounter, diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_no_vertexai_sdk.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_no_vertexai_sdk.py new file mode 100644 index 00000000000..b483a75a939 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/test_count_tokens_no_vertexai_sdk.py @@ -0,0 +1,121 @@ +""" +Regression tests for #28084: + +`VertexAIPartnerModels.count_tokens` (for Claude / Mistral / Llama on Vertex) +used to gate on `import vertexai` even though the actual count-tokens path goes +through `VertexAIPartnerModelsTokenCounter.handle_count_tokens_request`, which +talks to the publisher's `:rawPredict` endpoint over plain httpx and never +touches the Gemini SDK. The unused gate broke `/v1/messages/count_tokens` for +any LiteLLM install that did not pull in `google-cloud-aiplatform` (which is +not in the default `proxy` / `proxy-dev` extras). + +These tests pin the absence of that gate by: + +1. simulating `vertexai` being unimportable and verifying the partner-model + path does not raise the historical "vertexai import failed" error before + reaching the network/auth layer, and +2. asserting that import of the partner-model count-tokens handler module by + itself does not pull `vertexai` into `sys.modules`. +""" + +import sys + +import pytest + +from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import ( + VertexAIPartnerModelsTokenCounter, +) +from litellm.llms.vertex_ai.vertex_ai_partner_models.main import VertexAIPartnerModels + + +@pytest.mark.asyncio +async def test_count_tokens_does_not_require_vertexai_sdk(monkeypatch): + """Even when `import vertexai` would fail, count_tokens must not raise the + historical "vertexai import failed" gate. The downstream handler talks to + `:rawPredict` over httpx with an access token — no Gemini SDK needed.""" + + # Simulate `vertexai` being unimportable, regardless of what is actually on + # the test environment's sys.path. + monkeypatch.setitem(sys.modules, "vertexai", None) + monkeypatch.setitem(sys.modules, "vertexai.preview", None) + + captured = {} + + async def fake_ensure_access_token( + self, credentials, project_id, custom_llm_provider + ): + return "fake-token", "fake-project" + + def fake_build_endpoint(self, model, project_id, vertex_location, api_base=None): + captured["model_to_endpoint"] = model + return "https://fake-endpoint" + + monkeypatch.setattr( + VertexAIPartnerModelsTokenCounter, + "_ensure_access_token_async", + fake_ensure_access_token, + ) + monkeypatch.setattr( + VertexAIPartnerModelsTokenCounter, + "_build_count_tokens_endpoint", + fake_build_endpoint, + ) + + class FakeResponse: + status_code = 200 + + def json(self): + return {"input_tokens": 9} + + class FakeClient: + async def post(self, url, headers=None, json=None, **kwargs): + captured["url"] = url + captured["headers"] = headers + captured["json"] = json + return FakeResponse() + + import litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler as handler_mod + + monkeypatch.setattr( + handler_mod, "get_async_httpx_client", lambda **kwargs: FakeClient() + ) + + result = await VertexAIPartnerModels().count_tokens( + model="claude-sonnet-4-6", + messages=[{"role": "user", "content": "hello"}], + litellm_params={"vertex_location": "us-east5"}, + vertex_project="test-project", + vertex_location="us-east5", + vertex_credentials=None, + ) + + # We should reach the publisher endpoint and parse its response, not raise + # the vertexai-import gate. + assert result == { + "input_tokens": 9, + "tokenizer_used": "vertex_ai_partner_models", + } + assert captured["headers"] == {"Authorization": "Bearer fake-token"} + assert captured["model_to_endpoint"] == "claude-sonnet-4-6" + + +def test_handler_module_does_not_import_vertexai_sdk(): + """Importing the partner-model count-tokens handler must not load the + Gemini SDK into sys.modules. Operators who only need Claude-on-Vertex + token counting should not pay for `google-cloud-aiplatform`.""" + + # Force-evict any prior load so this assertion measures what THIS module + # pulls in, not what an unrelated earlier test did. + for mod in list(sys.modules): + if mod == "vertexai" or mod.startswith("vertexai."): + sys.modules.pop(mod, None) + + # Re-import the handler module to verify it stays SDK-free. + import importlib + + import litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler as handler_mod + + importlib.reload(handler_mod) + + leaked = [m for m in sys.modules if m == "vertexai" or m.startswith("vertexai.")] + assert leaked == [], f"unexpected vertexai SDK imports: {leaked}"