scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

fix(vertex-ai): address greptile review – proxy retrieve URL, timeout forwarding, sync logging

Browse source 
- Fix retrieve_api_base derivation to handle custom proxies with
  path-based routing (not just :cancel suffix)
- Forward timeout to POST calls in cancel_batch (sync + async)
- Add try/except error logging to sync cancel path (parity with async)
- Add tests for timeout forwarding and custom proxy retrieve URL

Made-with: Cursor
This commit is contained in:
Sameer Kankute 2026-03-18 10:30:05 +05:30
parent 547db8f5d1
commit c4d27cb239
2 changed files with 113 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

49
litellm/llms/vertex_ai/batches/handler.py
View file

@ -416,11 +416,12 @@ class VertexAIBatchPrediction(VertexLLM):
vertex_api_version="v1",
)
retrieve_api_base = (
api_base.removesuffix(":cancel")
if api_base.endswith(":cancel")
else retrieve_api_base_default
)
if api_base.endswith(":cancel"):
retrieve_api_base = api_base.removesuffix(":cancel")
elif api_base == cancel_api_base_default:
retrieve_api_base = retrieve_api_base_default
else:
retrieve_api_base = api_base.rsplit(":cancel", 1)[0].rstrip("/")
headers = {
"Content-Type": "application/json; charset=utf-8",
@ -432,22 +433,40 @@ class VertexAIBatchPrediction(VertexLLM):
api_base=api_base,
retrieve_api_base=retrieve_api_base,
headers=headers,
timeout=timeout,
)
sync_handler = _get_httpx_client()
response = sync_handler.post(
url=api_base,
headers=headers,
data=json.dumps({}),
)
try:
response = sync_handler.post(
url=api_base,
headers=headers,
data=json.dumps({}),
timeout=timeout,
)
except httpx.HTTPStatusError as e:
litellm.verbose_logger.error(
"Vertex AI batch cancel failed: status=%s, body=%s",
e.response.status_code,
e.response.text[:1000],
)
raise
if response.status_code != 200:
raise Exception(f"Error: {response.status_code} {response.text}")
retrieve_response = sync_handler.get(
url=retrieve_api_base,
headers=headers,
)
try:
retrieve_response = sync_handler.get(
url=retrieve_api_base,
headers=headers,
)
except httpx.HTTPStatusError as e:
litellm.verbose_logger.error(
"Vertex AI batch retrieve-after-cancel failed: status=%s, body=%s",
e.response.status_code,
e.response.text[:1000],
)
raise
if retrieve_response.status_code != 200:
raise Exception(
f"Error: {retrieve_response.status_code} {retrieve_response.text}"
@ -464,6 +483,7 @@ class VertexAIBatchPrediction(VertexLLM):
api_base: str,
retrieve_api_base: str,
headers: Dict[str, str],
timeout: Union[float, httpx.Timeout] = 600.0,
) -> LiteLLMBatch:
client = get_async_httpx_client(
llm_provider=litellm.LlmProviders.VERTEX_AI,
@ -473,6 +493,7 @@ class VertexAIBatchPrediction(VertexLLM):
url=api_base,
headers=headers,
data=json.dumps({}),
timeout=timeout,
)
except httpx.HTTPStatusError as e:
litellm.verbose_logger.error(

78
tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py
View file

@ -93,6 +93,84 @@ def test_vertex_ai_cancel_batch():
assert ":cancel" in call_args.kwargs["url"]
def test_vertex_ai_cancel_batch_forwards_timeout():
"""Test that timeout is forwarded to both POST and GET HTTP calls"""
handler = VertexAIBatchPrediction(gcs_bucket_name="test-bucket")
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {
"name": "projects/test-project/locations/us-central1/batchPredictionJobs/123456",
"state": "JOB_STATE_CANCELLING",
"createTime": "2024-03-17T10:00:00.000000Z",
"inputConfig": {"gcsSource": {"uris": ["gs://test-bucket/input.jsonl"]}},
"outputConfig": {"gcsDestination": {"outputUriPrefix": "gs://test-bucket/output"}},
}
with patch("litellm.llms.vertex_ai.batches.handler._get_httpx_client") as mock_client:
mock_client.return_value.post.return_value = mock_response
mock_client.return_value.get.return_value = mock_response
with patch.object(handler, "_ensure_access_token") as mock_auth:
mock_auth.return_value = ("fake-token", "test-project")
handler.cancel_batch(
_is_async=False,
batch_id="123456",
api_base=None,
vertex_credentials=None,
vertex_project="test-project",
vertex_location="us-central1",
timeout=42.0,
max_retries=None,
)
post_kwargs = mock_client.return_value.post.call_args.kwargs
assert post_kwargs["timeout"] == 42.0
def test_vertex_ai_cancel_batch_custom_proxy_retrieve_url():
"""Retrieve URL should go through the custom proxy, not bypass it"""
handler = VertexAIBatchPrediction(gcs_bucket_name="test-bucket")
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {
"name": "projects/test-project/locations/us-central1/batchPredictionJobs/123456",
"state": "JOB_STATE_CANCELLING",
"createTime": "2024-03-17T10:00:00.000000Z",
"inputConfig": {"gcsSource": {"uris": ["gs://test-bucket/input.jsonl"]}},
"outputConfig": {"gcsDestination": {"outputUriPrefix": "gs://test-bucket/output"}},
}
with patch("litellm.llms.vertex_ai.batches.handler._get_httpx_client") as mock_client:
mock_client.return_value.post.return_value = mock_response
mock_client.return_value.get.return_value = mock_response
with patch.object(handler, "_ensure_access_token") as mock_auth:
mock_auth.return_value = ("fake-token", "test-project")
handler.cancel_batch(
_is_async=False,
batch_id="123456",
api_base="https://my-proxy.example.com",
vertex_credentials=None,
vertex_project="test-project",
vertex_location="us-central1",
timeout=600.0,
max_retries=None,
)
post_url = mock_client.return_value.post.call_args.kwargs["url"]
get_url = mock_client.return_value.get.call_args.kwargs["url"]
assert "my-proxy.example.com" in post_url
assert ":cancel" in post_url
assert "my-proxy.example.com" in get_url
assert ":cancel" not in get_url
assert "googleapis.com" not in get_url
@pytest.mark.asyncio
async def test_litellm_cancel_batch_vertex_ai():
"""Test that litellm.cancel_batch works with vertex_ai provider"""