diff --git a/tests/pass_through_tests/test_vertex_ai.py b/tests/pass_through_tests/test_vertex_ai.py index e8223f2219c..35cb5f49c56 100644 --- a/tests/pass_through_tests/test_vertex_ai.py +++ b/tests/pass_through_tests/test_vertex_ai.py @@ -11,6 +11,7 @@ import os import pytest import asyncio +import requests # Path to your service account JSON file SERVICE_ACCOUNT_FILE = "path/to/your/service-account.json" @@ -57,98 +58,114 @@ def load_vertex_ai_credentials(): os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.path.abspath(temp_file.name) -async def call_spend_logs_endpoint(): - """ - Call this - curl -X GET "http://0.0.0.0:4000/spend/logs" -H "Authorization: Bearer sk-1234" - """ - import datetime - import requests +LITE_LLM_ENDPOINT = "http://localhost:4000" - todays_date = datetime.datetime.now().strftime("%Y-%m-%d") - url = f"http://0.0.0.0:4000/global/spend/logs?api_key=best-api-key-ever" - headers = {"Authorization": f"Bearer sk-1234"} - response = requests.get(url, headers=headers) - print("response from call_spend_logs_endpoint", response) +SPEND_LOG_API_KEY = "best-api-key-ever" - if response.status_code != 200: - print(f"spend logs endpoint returned {response.status_code}: {response.text}") - return None - - json_response = response.json() - # get spend for today +def get_tracked_spend() -> float: """ - json response looks like this + Total spend recorded under the pass-through key in the global spend view. - [{'date': '2024-08-30', 'spend': 0.00016600000000000002, 'api_key': 'best-api-key-ever'}] + Sums every day the endpoint returns instead of matching the runner's local + "today" so a UTC date rollover mid-test can't hide a freshly billed call, and + treats an unreachable endpoint as "nothing recorded yet" (0.0). """ - print("json_response", json_response) + url = f"{LITE_LLM_ENDPOINT}/global/spend/logs?api_key={SPEND_LOG_API_KEY}" + response = requests.get(url, headers={"Authorization": "Bearer sk-1234"}) + if response.status_code != 200: + print(f"global spend logs endpoint returned {response.status_code}: {response.text}") + return 0.0 - todays_date = datetime.datetime.now().strftime("%Y-%m-%d") - for spend_log in json_response: - if spend_log["date"] == todays_date: - return spend_log["spend"] + rows = response.json() + print("global spend logs rows", rows) + return sum(float(row.get("spend") or 0.0) for row in rows) -LITE_LLM_ENDPOINT = "http://localhost:4000" - +VERTEX_PROJECT = "litellm-ci-cd" +VERTEX_MODEL = "gemini-3.1-flash-lite" +VERTEX_GENERATE_CONTENT_URL = ( + f"{LITE_LLM_ENDPOINT}/vertex_ai/v1/projects/{VERTEX_PROJECT}" + f"/locations/global/publishers/google/models/{VERTEX_MODEL}:generateContent" +) -def _is_vertex_quota_error(exc: Exception) -> bool: - message = str(exc) - return ( - "429" in message - or "Too Many Requests" in message - or "RESOURCE_EXHAUSTED" in message - ) +def _vertex_access_token() -> str: + import google.auth + import google.auth.transport.requests -@pytest.mark.asyncio() -async def test_basic_vertex_ai_pass_through_with_spendlog(): + credentials, _ = google.auth.default( + scopes=["https://www.googleapis.com/auth/cloud-platform"] + ) + credentials.refresh(google.auth.transport.requests.Request()) + return credentials.token - spend_before = await call_spend_logs_endpoint() or 0.0 - load_vertex_ai_credentials() - vertexai.init( - project="litellm-ci-cd", - location="global", - api_endpoint=f"{LITE_LLM_ENDPOINT}/vertex_ai", - api_transport="rest", +def _spend_log_for_request(call_id: str) -> dict | None: + response = requests.get( + f"{LITE_LLM_ENDPOINT}/spend/logs?request_id={call_id}", + headers={"Authorization": "Bearer sk-1234"}, + timeout=30, ) + if response.status_code != 200: + return None + rows = response.json() + return rows[0] if rows else None - model = GenerativeModel(model_name="gemini-3.1-flash-lite") - try: - response = model.generate_content("hi") - except Exception as exc: - if _is_vertex_quota_error(exc): - pytest.skip("Vertex AI quota exhausted") - raise - print("response", response) +def _is_vertex_quota_error(response: requests.Response) -> bool: + return response.status_code == 429 or "RESOURCE_EXHAUSTED" in response.text - # Spend logging is async/batched and can lag under CI load, so poll instead of - # sleeping a fixed amount. A transient empty read is skipped, not counted as 0.0 - # spend, which would spuriously fail the assertion on an otherwise-billed call. - max_wait = 240 # total seconds to wait - poll_interval = 10 # seconds between checks - elapsed = 0 - spend_after = spend_before - while elapsed < max_wait: - await asyncio.sleep(poll_interval) - elapsed += poll_interval - latest_spend = await call_spend_logs_endpoint() - if latest_spend is None: - print(f"spend logs unavailable (elapsed={elapsed}s), retrying") - continue - spend_after = latest_spend - print(f"spend_after (elapsed={elapsed}s)", spend_after) - if spend_after > spend_before: - break - assert ( - spend_after > spend_before - ), "Spend should be greater than before after {}s. spend_before: {}, spend_after: {}".format( - elapsed, spend_before, spend_after +@pytest.mark.asyncio() +async def test_basic_vertex_ai_pass_through_with_spendlog(): + load_vertex_ai_credentials() + access_token = _vertex_access_token() + + # Drive the pass-through over HTTP instead of the vertexai SDK: the SDK intermittently + # routes generateContent to the public Vertex endpoint rather than the proxy override, + # so the call never reaches LiteLLM and no spend is logged. A direct request always + # hits the proxy. Spend logging then runs on a best-effort background worker that can + # drop a single event, so retry a few billed calls and assert that one specific call's + # spend log lands. Failing every attempt still fails hard, which is the signal we want + # if cost tracking is broken. + max_attempts = 3 + poll_seconds = 60 + poll_interval = 5 + + for attempt in range(1, max_attempts + 1): + response = requests.post( + VERTEX_GENERATE_CONTENT_URL, + headers={ + "Authorization": f"Bearer {access_token}", + "Content-Type": "application/json", + }, + json={"contents": [{"role": "user", "parts": [{"text": "hi"}]}]}, + timeout=60, + ) + if _is_vertex_quota_error(response): + pytest.skip("Vertex AI quota exhausted") + assert ( + response.status_code == 200 + ), f"vertex pass-through call failed: {response.status_code} {response.text}" + + call_id = response.headers.get("x-litellm-call-id") + assert call_id, "proxy response missing x-litellm-call-id header" + + for _ in range(poll_seconds // poll_interval): + await asyncio.sleep(poll_interval) + row = _spend_log_for_request(call_id) + if row is not None and float(row.get("spend") or 0) > 0: + assert "gemini" in row["model"], f"unexpected model in spend log: {row}" + assert ( + row["custom_llm_provider"] == "vertex_ai" + ), f"unexpected provider in spend log: {row}" + return + + print(f"attempt {attempt}: spend log for call {call_id} not found yet, re-billing") + + pytest.fail( + f"Vertex pass-through spend never recorded after {max_attempts} billed calls" ) @@ -156,7 +173,7 @@ async def test_basic_vertex_ai_pass_through_with_spendlog(): @pytest.mark.skip(reason="skip flaky test - vertex pass through streaming is flaky") async def test_basic_vertex_ai_pass_through_streaming_with_spendlog(): - spend_before = await call_spend_logs_endpoint() or 0.0 + spend_before = get_tracked_spend() print("spend_before", spend_before) load_vertex_ai_credentials() @@ -176,7 +193,7 @@ async def test_basic_vertex_ai_pass_through_streaming_with_spendlog(): print("response", response) await asyncio.sleep(20) - spend_after = await call_spend_logs_endpoint() + spend_after = get_tracked_spend() print("spend_after", spend_after) assert ( spend_after > spend_before