tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-07 04:24:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

test_vertex_ai_partner_models_token_counting_endpoint

Browse Source
This commit is contained in:
Ishaan Jaffer
2025-11-26 11:37:55 -08:00
parent 19a6693dfc
commit 85d4000af6
1 changed files with 8 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/proxy_unit_tests/test_proxy_token_counter.py
+8 -4
View File
@@ -821,18 +821,22 @@ async def test_vertex_ai_anthropic_token_counting():
assert response.original_response["input_tokens"] == 15
@pytest.mark.parametrize("vertex_location", ["global", "us-central1"])
def test_vertex_ai_gemini_token_counting_endpoint(vertex_location):
def test_vertex_ai_partner_models_token_counting_endpoint(vertex_location):
"""
Test that the VertexAIPartnerModelsTokenCounter builds the correct endpoint URL
for different vertex locations, including the special 'global' location.
"""
from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import (
VertexAIPartnerModelsTokenCounter,
)
endpoint = VertexAIPartnerModelsTokenCounter()._build_count_tokens_endpoint(
model="gemini-2.5-pro",
model="claude-3-5-sonnet-20241022",
project_id="test-project",
vertex_location=vertex_location,
api_base=None,
)
if vertex_location == "global":
assert endpoint == "https://aiplatform.googleapis.com"
assert endpoint.startswith("https://aiplatform.googleapis.com")
else:
assert endpoint == f"https://{vertex_location}-aiplatform.googleapis.com"
assert endpoint.startswith(f"https://{vertex_location}-aiplatform.googleapis.com")