[Bug]: `/v1/messages/count_tokens` for Vertex AI Claude models fails on missing `vertexai` (google-cloud-aiplatform) import

[mateo-berri]

What happened

POST /v1/messages/count_tokens against a Vertex AI–routed Claude model returns HTTP 500:

{"detail":{"error":"Internal server error: vertexai import failed please run `pip install -U \"google-cloud-aiplatform>=1.38\"`. Got error: No module named 'vertexai'"}}

The proxy is up, all other Vertex AI routes (/v1/messages, tool calls, streaming, vision, web_search, etc.) work for the same model alias — only count_tokens requires the vertexai SDK module. The same count_tokens call against claude-sonnet-4-6 (anthropic-native) returns {"input_tokens": 9} as expected from the same proxy.

google-cloud-aiplatform is not a LiteLLM core dependency, nor is it in the proxy-dev/proxy extras used by uv sync --frozen --group proxy-dev --extra proxy. So a fresh stable-tag install can never count tokens on Vertex AI Claude unless the operator separately installs the Gemini SDK.

Reproducer

# proxy launched via `uv run litellm --config config.yaml --port 4101`
curl -sS -X POST http://127.0.0.1:4101/v1/messages/count_tokens \
  -H "Authorization: Bearer sk-1234" \
  -H "Content-Type: application/json" \
  -H "anthropic-version: 2023-06-01" \
  -d '{"model":"claude-sonnet-4-6-vertex","messages":[{"role":"user","content":"hello world"}]}'

Returns the 500 above. The claude-sonnet-4-6-vertex alias routes to Vertex AI's Claude API and works for every other endpoint without google-cloud-aiplatform.

Expected

Counting tokens for Claude-on-Vertex should not require the Gemini-specific vertexai Python module. Vertex AI's Claude models are addressed via the Anthropic Messages API protocol; the token counter should use the same Anthropic-protocol path the rest of the vertex_ai Claude routing does (or call Vertex's :rawPredict count-tokens equivalent directly via HTTPS, like LiteLLM's other Vertex Claude transformations do).

Either:

1. Make count_tokens for Claude-on-Vertex avoid import vertexai entirely (preferred — keeps deps minimal); or
2. Add google-cloud-aiplatform to the default proxy extras with a clear error message pointing the operator at the install when it's absent (worse — pulls in a heavy Gemini SDK for a feature that should be SDK-free).

Why this matters

Claude Code's headless mode (and any tooling that uses count_tokens for budget display) is silently broken against any LiteLLM proxy that routes Claude to Vertex AI unless the operator knows to install the Gemini SDK.

Relevant LiteLLM version

v1.83.14-stable (also reproduced against current main).

Surfaced by

Daily Claude Code × LiteLLM compatibility-matrix cron (PR BerriAI/litellm-docs#142 — see count_tokens × vertex_ai cell).

[VANDRANKI]

Verified against the local install (v1.84.0).

Root cause is a spurious import vertexai guard in litellm/llms/vertex_ai/vertex_ai_partner_models/main.py, inside VertexAIPartnerModels.count_tokens at lines 295-301:

async def count_tokens(self, model, messages, litellm_params, ...):
    try:
        import vertexai         # <-- guard that raises if SDK missing
    except Exception as e:
        raise VertexAIError(
            status_code=400,
            message=f"vertexai import failed please run `pip install -U "
                    f'"google-cloud-aiplatform>=1.38"`. Got error: {e}',
        )
    # ... checks vertexai.preview attributes ...

    # actual work (lines 311-338):
    from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import (
        VertexAIPartnerModelsTokenCounter,
    )
    token_counter = VertexAIPartnerModelsTokenCounter()
    result = await token_counter.handle_count_tokens_request(...)  # pure HTTPS call

The VertexAIPartnerModelsTokenCounter in vertex_ai_partner_models/count_tokens/handler.py uses get_async_httpx_client (plain HTTPS, no vertexai SDK) — it calls Vertex AI's publisher-specific count-tokens endpoint directly. The import vertexai guard is copy-pasted from other methods in the same class that do need the Gemini SDK (e.g., acompletion), but count_tokens never calls any vertexai.* API after the guard.

Impact: Every call to /v1/messages/count_tokens for a Claude-on-Vertex model fails with HTTP 500 on any proxy deployment that hasn't installed google-cloud-aiplatform — which is the majority, since it's not in LiteLLM's core or proxy extras.

Fix — remove the import vertexai guard and the vertexai.preview attribute check from count_tokens. The method doesn't use the SDK at all:

async def count_tokens(self, model, messages, litellm_params, ...):
    # No import vertexai needed — VertexAIPartnerModelsTokenCounter uses HTTPS directly
    try:
        from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import (
            VertexAIPartnerModelsTokenCounter,
        )
        request_data = {"model": model, "messages": messages}
        _litellm_params = litellm_params.copy()
        ...
        result = await VertexAIPartnerModelsTokenCounter().handle_count_tokens_request(...)
        return result
    except Exception as e:
        if hasattr(e, "status_code"):
            raise e
        raise VertexAIError(status_code=500, message=str(e))

The VertexAIPartnerModelsTokenCounter already handles credential resolution and endpoint construction without the vertexai module, so this change has no behavioral impact on environments that do have google-cloud-aiplatform installed.