From 0bc609affdc09c4071c72f4035591289685af8e4 Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Tue, 17 Mar 2026 11:23:47 +0530
Subject: [PATCH 01/11] fix(vertex-ai): support batch cancel via Vertex API

Add Vertex batch cancellation support in LiteLLM batch APIs, route proxy cancel fallback using request provider headers, and return post-cancel batch state via retrieve to keep response shape compatible.

Made-with: Cursor
---
 litellm/batches/main.py                       |  32 ++++-
 litellm/llms/vertex_ai/batches/handler.py     | 113 ++++++++++++++++++
 litellm/proxy/batches_endpoints/endpoints.py  |  14 ++-
 .../test_vertex_ai_batch_transformation.py    |  78 ++++++++++++
 4 files changed, 229 insertions(+), 8 deletions(-)

diff --git a/litellm/batches/main.py b/litellm/batches/main.py
index e176dc4292..36093d071b 100644
--- a/litellm/batches/main.py
+++ b/litellm/batches/main.py
@@ -884,7 +884,7 @@ def list_batches(
 async def acancel_batch(
     batch_id: str,
     model: Optional[str] = None,
-    custom_llm_provider: Literal["openai", "azure"] = "openai",
+    custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
     metadata: Optional[Dict[str, str]] = None,
     extra_headers: Optional[Dict[str, str]] = None,
     extra_body: Optional[Dict[str, str]] = None,
@@ -930,7 +930,7 @@ async def acancel_batch(
 def cancel_batch(
     batch_id: str,
     model: Optional[str] = None,
-    custom_llm_provider: Union[Literal["openai", "azure"], str] = "openai",
+    custom_llm_provider: Union[Literal["openai", "azure", "vertex_ai"], str] = "openai",
     metadata: Optional[Dict[str, str]] = None,
     extra_headers: Optional[Dict[str, str]] = None,
     extra_body: Optional[Dict[str, str]] = None,
@@ -1047,9 +1047,35 @@ def cancel_batch(
                 cancel_batch_data=_cancel_batch_request,
                 litellm_params=litellm_params,
             )
+        elif custom_llm_provider == "vertex_ai":
+            api_base = optional_params.api_base or ""
+            vertex_ai_project = (
+                optional_params.vertex_project
+                or litellm.vertex_project
+                or get_secret_str("VERTEXAI_PROJECT")
+            )
+            vertex_ai_location = (
+                optional_params.vertex_location
+                or litellm.vertex_location
+                or get_secret_str("VERTEXAI_LOCATION")
+            )
+            vertex_credentials = optional_params.vertex_credentials or get_secret_str(
+                "VERTEXAI_CREDENTIALS"
+            )
+
+            response = vertex_ai_batches_instance.cancel_batch(
+                _is_async=_is_async,
+                batch_id=batch_id,
+                api_base=api_base,
+                vertex_project=vertex_ai_project,
+                vertex_location=vertex_ai_location,
+                vertex_credentials=vertex_credentials,
+                timeout=timeout,
+                max_retries=optional_params.max_retries,
+            )
         else:
             raise litellm.exceptions.BadRequestError(
-                message="LiteLLM doesn't support {} for 'cancel_batch'. Only 'openai' and 'azure' are supported.".format(
+                message="LiteLLM doesn't support {} for 'cancel_batch'. Only 'openai', 'azure', and 'vertex_ai' are supported.".format(
                     custom_llm_provider
                 ),
                 model="n/a",
diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py
index f0b181c9a6..a24bfd89f0 100644
--- a/litellm/llms/vertex_ai/batches/handler.py
+++ b/litellm/llms/vertex_ai/batches/handler.py
@@ -376,3 +376,116 @@ class VertexAIBatchPrediction(VertexLLM):
             response=_json_response
         )
         return vertex_batch_response
+
+    def cancel_batch(
+        self,
+        _is_async: bool,
+        batch_id: str,
+        api_base: Optional[str],
+        vertex_credentials: Optional[VERTEX_CREDENTIALS_TYPES],
+        vertex_project: Optional[str],
+        vertex_location: Optional[str],
+        timeout: Union[float, httpx.Timeout],
+        max_retries: Optional[int],
+    ) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]:
+        sync_handler = _get_httpx_client()
+
+        access_token, project_id = self._ensure_access_token(
+            credentials=vertex_credentials,
+            project_id=vertex_project,
+            custom_llm_provider="vertex_ai",
+        )
+
+        default_api_base = self.create_vertex_batch_url(
+            vertex_location=vertex_location or "us-central1",
+            vertex_project=vertex_project or project_id,
+        )
+
+        default_api_base = f"{default_api_base}/{batch_id}:cancel"
+
+        if len(default_api_base.split(":")) > 1:
+            endpoint = default_api_base.split(":")[-1]
+        else:
+            endpoint = ""
+
+        _, api_base = self._check_custom_proxy(
+            api_base=api_base,
+            custom_llm_provider="vertex_ai",
+            gemini_api_key=None,
+            endpoint=endpoint,
+            stream=None,
+            auth_header=None,
+            url=default_api_base,
+            model=None,
+            vertex_project=vertex_project or project_id,
+            vertex_location=vertex_location or "us-central1",
+            vertex_api_version="v1",
+        )
+
+        headers = {
+            "Content-Type": "application/json; charset=utf-8",
+            "Authorization": f"Bearer {access_token}",
+        }
+
+        if _is_async is True:
+            return self._async_cancel_batch(
+                api_base=api_base,
+                retrieve_api_base=api_base.rsplit(":cancel", 1)[0],
+                headers=headers,
+            )
+
+        response = sync_handler.post(
+            url=api_base,
+            headers=headers,
+            data=json.dumps({}),
+        )
+
+        if response.status_code != 200:
+            raise Exception(f"Error: {response.status_code} {response.text}")
+
+        retrieve_response = sync_handler.get(
+            url=api_base.rsplit(":cancel", 1)[0],
+            headers=headers,
+        )
+        if retrieve_response.status_code != 200:
+            raise Exception(
+                f"Error: {retrieve_response.status_code} {retrieve_response.text}"
+            )
+
+        _json_response = retrieve_response.json()
+        vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response(
+            response=_json_response
+        )
+        return vertex_batch_response
+
+    async def _async_cancel_batch(
+        self,
+        api_base: str,
+        retrieve_api_base: str,
+        headers: Dict[str, str],
+    ) -> LiteLLMBatch:
+        client = get_async_httpx_client(
+            llm_provider=litellm.LlmProviders.VERTEX_AI,
+        )
+        response = await client.post(
+            url=api_base,
+            headers=headers,
+            data=json.dumps({}),
+        )
+        if response.status_code != 200:
+            raise Exception(f"Error: {response.status_code} {response.text}")
+
+        retrieve_response = await client.get(
+            url=retrieve_api_base,
+            headers=headers,
+        )
+        if retrieve_response.status_code != 200:
+            raise Exception(
+                f"Error: {retrieve_response.status_code} {retrieve_response.text}"
+            )
+
+        _json_response = retrieve_response.json()
+        vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response(
+            response=_json_response
+        )
+        return vertex_batch_response
diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py
index 740e63b7f1..06254b57fc 100644
--- a/litellm/proxy/batches_endpoints/endpoints.py
+++ b/litellm/proxy/batches_endpoints/endpoints.py
@@ -5,7 +5,7 @@
 
 ######################################################################
 import asyncio
-from typing import Dict, Optional, cast
+from typing import Any, Dict, Optional, cast
 
 from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response
 
@@ -654,7 +654,7 @@ async def list_batches(
             managed_files_obj, "list_user_batches"
         ):
             verbose_proxy_logger.debug("Using managed objects table for batch listing")
-            response = await managed_files_obj.list_user_batches(
+            response = await cast(Any, managed_files_obj).list_user_batches(
                 user_api_key_dict=user_api_key_dict,
                 limit=limit,
                 after=after,
@@ -685,8 +685,9 @@ async def list_batches(
 
             # Encode batch IDs in the list response so clients can use
             # them for retrieve/cancel/file downloads through the proxy.
-            if response and hasattr(response, "data") and response.data:
-                for batch in response.data:
+            response_data = getattr(response, "data", None)
+            if response_data:
+                for batch in response_data:
                     encode_batch_response_ids(batch, model=model_param)
 
             verbose_proxy_logger.debug(f"Listed batches using model: {model_param}")
@@ -896,7 +897,10 @@ async def cancel_batch(
         # SCENARIO 3: Fallback to custom_llm_provider (uses env variables)
         else:
             custom_llm_provider = (
-                provider or data.pop("custom_llm_provider", None) or "openai"
+                provider
+                or get_custom_llm_provider_from_request_headers(request=request)
+                or data.pop("custom_llm_provider", None)
+                or "openai"
             )
             # Extract batch_id from data to avoid "multiple values for keyword argument" error
             # data was cast from CancelBatchRequest which already contains batch_id
diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py
index 1aab74ddc2..d27cd8ba8a 100644
--- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py
+++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py
@@ -1,3 +1,9 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+import litellm
+from litellm.llms.vertex_ai.batches.handler import VertexAIBatchPrediction
 from litellm.llms.vertex_ai.batches.transformation import VertexAIBatchTransformation
 
 
@@ -36,3 +42,75 @@ def test_output_file_id_falls_back_to_output_uri_prefix_with_predictions_jsonl()
         output_file_id
         == "gs://test-bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-pro/prediction-model-456/predictions.jsonl"
     )
+
+
+@pytest.mark.asyncio
+async def test_vertex_ai_cancel_batch():
+    """Test that vertex_ai cancel_batch calls the correct API endpoint"""
+    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")
+            
+            response = handler.cancel_batch(
+                _is_async=False,
+                batch_id="123456",
+                api_base=None,
+                vertex_credentials=None,
+                vertex_project="test-project",
+                vertex_location="us-central1",
+                timeout=600.0,
+                max_retries=None,
+            )
+            
+            assert response.id == "123456"
+            assert response.status == "cancelling"
+            
+            mock_client.return_value.post.assert_called_once()
+            mock_client.return_value.get.assert_called_once()
+            call_args = mock_client.return_value.post.call_args
+            assert ":cancel" in call_args.kwargs["url"]
+
+
+@pytest.mark.asyncio
+async def test_litellm_cancel_batch_vertex_ai():
+    """Test that litellm.cancel_batch works with vertex_ai provider"""
+    mock_response = MagicMock()
+    mock_response.id = "batch_123"
+    mock_response.status = "cancelling"
+    
+    with patch.object(litellm.batches.main, "vertex_ai_batches_instance") as mock_instance:
+        mock_instance.cancel_batch.return_value = mock_response
+        
+        response = litellm.cancel_batch(
+            batch_id="batch_123",
+            custom_llm_provider="vertex_ai",
+            vertex_project="test-project",
+            vertex_location="us-central1",
+        )
+        
+        assert mock_instance.cancel_batch.called
+        assert response.id == "batch_123"
+        assert response.status == "cancelling"

From d8e3abf3cef95365fa6b1cc0af68e32a33319bd1 Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Tue, 17 Mar 2026 11:34:50 +0530
Subject: [PATCH 02/11] fix(vertex-ai): apply review updates for batch cancel

Incorporate follow-up changes to Vertex batch cancel handling and proxy provider resolution, including config updates used for local verification.

Made-with: Cursor
---
 litellm/llms/vertex_ai/batches/handler.py    |   3 +-
 litellm/proxy/batches_endpoints/endpoints.py |   1 +
 proxy_server_config.yaml                     | 231 +------------------
 3 files changed, 4 insertions(+), 231 deletions(-)

diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py
index a24bfd89f0..c7b9287c08 100644
--- a/litellm/llms/vertex_ai/batches/handler.py
+++ b/litellm/llms/vertex_ai/batches/handler.py
@@ -388,8 +388,6 @@ class VertexAIBatchPrediction(VertexLLM):
         timeout: Union[float, httpx.Timeout],
         max_retries: Optional[int],
     ) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]:
-        sync_handler = _get_httpx_client()
-
         access_token, project_id = self._ensure_access_token(
             credentials=vertex_credentials,
             project_id=vertex_project,
@@ -434,6 +432,7 @@ class VertexAIBatchPrediction(VertexLLM):
                 headers=headers,
             )
 
+        sync_handler = _get_httpx_client()
         response = sync_handler.post(
             url=api_base,
             headers=headers,
diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py
index 06254b57fc..9ce1b6e916 100644
--- a/litellm/proxy/batches_endpoints/endpoints.py
+++ b/litellm/proxy/batches_endpoints/endpoints.py
@@ -899,6 +899,7 @@ async def cancel_batch(
             custom_llm_provider = (
                 provider
                 or get_custom_llm_provider_from_request_headers(request=request)
+                or get_custom_llm_provider_from_request_query(request=request)
                 or data.pop("custom_llm_provider", None)
                 or "openai"
             )
diff --git a/proxy_server_config.yaml b/proxy_server_config.yaml
index 5d3d810926..6e48af021a 100644
--- a/proxy_server_config.yaml
+++ b/proxy_server_config.yaml
@@ -1,231 +1,4 @@
 model_list:
-  - model_name: gpt-3.5-turbo-end-user-test
+  - model_name: gemini-2.5-pro
     litellm_params:
-      model: gpt-3.5-turbo
-      region_name: "eu"
-    model_info:
-      id: "1"
-  - model_name: gpt-3.5-turbo-end-user-test
-    litellm_params:
-      model: openai/gpt-4.1-mini
-      api_key: os.environ/OPENAI_API_KEY # The `os.environ/` prefix tells litellm to read this from the env. See https://docs.litellm.ai/docs/simple_proxy#load-api-keys-from-vault
-  - model_name: gpt-3.5-turbo
-    litellm_params:
-      model: openai/gpt-4.1-mini
-      api_key: os.environ/OPENAI_API_KEY # The `os.environ/` prefix tells litellm to read this from the env. See https://docs.litellm.ai/docs/simple_proxy#load-api-keys-from-vault
-  - model_name: gpt-3.5-turbo-large
-    litellm_params: 
-      model: "gpt-3.5-turbo-1106"
-      api_key: os.environ/OPENAI_API_KEY
-      rpm: 480
-      timeout: 300
-      stream_timeout: 60
-  - model_name: gpt-4
-    litellm_params:
-      model: openai/gpt-4.1-mini
-      api_key: os.environ/OPENAI_API_KEY # The `os.environ/` prefix tells litellm to read this from the env. See https://docs.litellm.ai/docs/simple_proxy#load-api-keys-from-vault
-      rpm: 480
-      timeout: 300
-      stream_timeout: 60
-  - model_name: sagemaker-completion-model
-    litellm_params:
-      model: sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4
-      input_cost_per_second: 0.000420  
-  - model_name: text-embedding-ada-002
-    litellm_params: 
-      model: openai/text-embedding-ada-002
-      api_key: os.environ/OPENAI_API_KEY
-    model_info:
-      mode: embedding
-      base_model: text-embedding-ada-002
-  - model_name: dall-e-2 # some tests use dall-e-2 which is now deprecated, alias to dall-e-3
-    litellm_params:
-      model: openai/dall-e-3
-  - model_name: openai-dall-e-3
-    litellm_params:
-      model: dall-e-3
-  - model_name: fake-openai-endpoint
-    litellm_params:
-      model: openai/gpt-3.5-turbo
-      api_key: fake-key
-      api_base: https://exampleopenaiendpoint-production.up.railway.app/
-  - model_name: fake-openai-endpoint-2
-    litellm_params:
-      model: openai/my-fake-model
-      api_key: my-fake-key
-      api_base: https://exampleopenaiendpoint-production.up.railway.app/
-      stream_timeout: 0.001
-      rpm: 1
-  - model_name: fake-openai-endpoint-3
-    litellm_params:
-      model: openai/my-fake-model
-      api_key: my-fake-key
-      api_base: https://exampleopenaiendpoint-production.up.railway.app/
-      stream_timeout: 0.001
-      rpm: 1000
-  - model_name: fake-openai-endpoint-4
-    litellm_params:
-      model: openai/my-fake-model
-      api_key: my-fake-key
-      api_base: https://exampleopenaiendpoint-production.up.railway.app/
-      num_retries: 50
-  - model_name: fake-openai-endpoint-3
-    litellm_params:
-      model: openai/my-fake-model-2
-      api_key: my-fake-key
-      api_base: https://exampleopenaiendpoint-production.up.railway.app/
-      stream_timeout: 0.001
-      rpm: 1000
-  - model_name: bad-model
-    litellm_params:
-      model: openai/bad-model
-      api_key: os.environ/OPENAI_API_KEY
-      api_base: https://exampleopenaiendpoint-production.up.railway.app/
-      mock_timeout: True
-      timeout: 60
-      rpm: 1000
-    model_info:
-      health_check_timeout: 1
-  - model_name: good-model
-    litellm_params:
-      model: openai/bad-model
-      api_key: os.environ/OPENAI_API_KEY
-      api_base: https://exampleopenaiendpoint-production.up.railway.app/
-      rpm: 1000
-    model_info:
-      health_check_timeout: 1
-  - model_name: "*"
-    litellm_params:
-      model: openai/*
-      api_key: os.environ/OPENAI_API_KEY
-  - model_name: realtime-v1
-    litellm_params:
-      model: azure/gpt-realtime-20250828-standard
-      api_version: "2025-08-28"
-      realtime_protocol: GA # Possible values: "GA"/ "v1", "beta" 
-
-  - model_name: realtime-beta
-    litellm_params:
-      model: azure/gpt-realtime-20250828-standard
-      api_version: 2025-04-01-preview
-
-
-  # provider specific wildcard routing
-  - model_name: "anthropic/*"
-    litellm_params:
-      model: "anthropic/*"
-      api_key: os.environ/ANTHROPIC_API_KEY
-  - model_name: "bedrock/*"
-    litellm_params:
-      model: "bedrock/*"
-  - model_name: "groq/*"
-    litellm_params:
-      model: "groq/*"
-      api_key: os.environ/GROQ_API_KEY
-  - model_name: mistral-embed
-    litellm_params:
-      model: mistral/mistral-embed
-  - model_name: gpt-instruct # [PROD TEST] - tests if `/health` automatically infers this to be a text completion model
-    litellm_params:
-      model: text-completion-openai/gpt-3.5-turbo-instruct
-  - model_name: fake-openai-endpoint-5
-    litellm_params:
-      model: openai/my-fake-model
-      api_key: my-fake-key
-      api_base: https://exampleopenaiendpoint-production.up.railway.app/
-      timeout: 1
-  - model_name: badly-configured-openai-endpoint
-    litellm_params:
-      model: openai/my-fake-model
-      api_key: my-fake-key
-      api_base: https://exampleopenaiendpoint-production.up.railway.appxxxx/
-  - model_name: gemini-1.5-flash
-    litellm_params:
-      model: gemini/gemini-1.5-flash
-      api_key: os.environ/GOOGLE_API_KEY
-  - model_name: gpt-4o
-    litellm_params:
-      model: gpt-4o
-      api_key: os.environ/OPENAI_API_KEY
-
-
-litellm_settings:
-  # set_verbose: True  # Uncomment this if you want to see verbose logs; not recommended in production
-  drop_params: True
-  success_callback: ["prometheus"]
-  # max_budget: 100 
-  # budget_duration: 30d
-  num_retries: 5
-  request_timeout: 600
-  telemetry: False
-  context_window_fallbacks: [{"gpt-3.5-turbo": ["gpt-3.5-turbo-large"]}]
-  default_team_settings: 
-    - team_id: team-1
-      success_callback: ["langfuse"]
-      failure_callback: ["langfuse"]
-      langfuse_public_key: os.environ/LANGFUSE_PROJECT1_PUBLIC # Project 1
-      langfuse_secret: os.environ/LANGFUSE_PROJECT1_SECRET # Project 1
-    - team_id: team-2
-      success_callback: ["langfuse"]
-      failure_callback: ["langfuse"]
-      langfuse_public_key: os.environ/LANGFUSE_PROJECT2_PUBLIC # Project 2
-      langfuse_secret: os.environ/LANGFUSE_PROJECT2_SECRET # Project 2
-      langfuse_host: https://us.cloud.langfuse.com
-  # cache: true   # [OPTIONAL] use for caching responses 
-  # enable_caching_on_provider_specific_optional_params: True  # Include provider-specific params in cache keys
-  # cache_params:  # And for shared health check
-  #   type: redis
-  #   host: localhost
-  #   port: 6379
-
-# For /fine_tuning/jobs endpoints
-finetune_settings:
-  - custom_llm_provider: azure
-    api_base: os.environ/AZURE_API_BASE
-    api_key: os.environ/AZURE_API_KEY
-    api_version: "2023-03-15-preview"
-  - custom_llm_provider: openai
-    api_key: os.environ/OPENAI_API_KEY
-
-# for /files endpoints
-files_settings:
-  - custom_llm_provider: azure
-    api_base: os.environ/AZURE_API_BASE
-    api_key: os.environ/AZURE_API_KEY
-    api_version: "2023-03-15-preview"
-  - custom_llm_provider: openai
-    api_key: os.environ/OPENAI_API_KEY
-
-router_settings:
-  routing_strategy: usage-based-routing-v2 
-  redis_host: os.environ/REDIS_HOST
-  redis_password: os.environ/REDIS_PASSWORD
-  redis_port: os.environ/REDIS_PORT
-  enable_pre_call_checks: true
-  model_group_alias: {"my-special-fake-model-alias-name": "fake-openai-endpoint-3"} 
-
-general_settings: 
-  master_key: sk-1234 # [OPTIONAL] Use to enforce auth on proxy. See - https://docs.litellm.ai/docs/proxy/virtual_keys
-  store_model_in_db: True
-  proxy_budget_rescheduler_min_time: 60
-  proxy_budget_rescheduler_max_time: 64
-  proxy_batch_write_at: 1
-  database_connection_pool_limit: 10
-  # background_health_checks: true
-  # use_shared_health_check: true
-  # health_check_interval: 30
-  # database_url: "postgresql://<user>:<password>@<host>:<port>/<dbname>" # [OPTIONAL] use for token-based auth to proxy
-
-  pass_through_endpoints:
-    - path: "/v1/rerank"                                  # route you want to add to LiteLLM Proxy Server
-      target: "https://api.cohere.com/v1/rerank"          # URL this route should forward requests to
-      headers:                                            # headers to forward to this URL
-        content-type: application/json                    # (Optional) Extra Headers to pass to this endpoint 
-        accept: application/json
-      forward_headers: True
-
-# environment_variables:
-  # settings for using redis caching
-  # REDIS_HOST: redis-16337.c322.us-east-1-2.ec2.cloud.redislabs.com
-  # REDIS_PORT: "16337"
-  # REDIS_PASSWORD: 
\ No newline at end of file
+      model: vertex_ai/gemini-2.5-pro
\ No newline at end of file

From 37b7a7fb576279a41817ef8239d2e24b79e56f7b Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Tue, 17 Mar 2026 11:36:39 +0530
Subject: [PATCH 03/11] chore(config): restore proxy_server_config.yaml

Revert local test-only proxy config edits so the PR does not include unrelated configuration changes.

Made-with: Cursor
---
 proxy_server_config.yaml | 231 ++++++++++++++++++++++++++++++++++++++-
 1 file changed, 229 insertions(+), 2 deletions(-)

diff --git a/proxy_server_config.yaml b/proxy_server_config.yaml
index 6e48af021a..5d3d810926 100644
--- a/proxy_server_config.yaml
+++ b/proxy_server_config.yaml
@@ -1,4 +1,231 @@
 model_list:
-  - model_name: gemini-2.5-pro
+  - model_name: gpt-3.5-turbo-end-user-test
     litellm_params:
-      model: vertex_ai/gemini-2.5-pro
\ No newline at end of file
+      model: gpt-3.5-turbo
+      region_name: "eu"
+    model_info:
+      id: "1"
+  - model_name: gpt-3.5-turbo-end-user-test
+    litellm_params:
+      model: openai/gpt-4.1-mini
+      api_key: os.environ/OPENAI_API_KEY # The `os.environ/` prefix tells litellm to read this from the env. See https://docs.litellm.ai/docs/simple_proxy#load-api-keys-from-vault
+  - model_name: gpt-3.5-turbo
+    litellm_params:
+      model: openai/gpt-4.1-mini
+      api_key: os.environ/OPENAI_API_KEY # The `os.environ/` prefix tells litellm to read this from the env. See https://docs.litellm.ai/docs/simple_proxy#load-api-keys-from-vault
+  - model_name: gpt-3.5-turbo-large
+    litellm_params: 
+      model: "gpt-3.5-turbo-1106"
+      api_key: os.environ/OPENAI_API_KEY
+      rpm: 480
+      timeout: 300
+      stream_timeout: 60
+  - model_name: gpt-4
+    litellm_params:
+      model: openai/gpt-4.1-mini
+      api_key: os.environ/OPENAI_API_KEY # The `os.environ/` prefix tells litellm to read this from the env. See https://docs.litellm.ai/docs/simple_proxy#load-api-keys-from-vault
+      rpm: 480
+      timeout: 300
+      stream_timeout: 60
+  - model_name: sagemaker-completion-model
+    litellm_params:
+      model: sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4
+      input_cost_per_second: 0.000420  
+  - model_name: text-embedding-ada-002
+    litellm_params: 
+      model: openai/text-embedding-ada-002
+      api_key: os.environ/OPENAI_API_KEY
+    model_info:
+      mode: embedding
+      base_model: text-embedding-ada-002
+  - model_name: dall-e-2 # some tests use dall-e-2 which is now deprecated, alias to dall-e-3
+    litellm_params:
+      model: openai/dall-e-3
+  - model_name: openai-dall-e-3
+    litellm_params:
+      model: dall-e-3
+  - model_name: fake-openai-endpoint
+    litellm_params:
+      model: openai/gpt-3.5-turbo
+      api_key: fake-key
+      api_base: https://exampleopenaiendpoint-production.up.railway.app/
+  - model_name: fake-openai-endpoint-2
+    litellm_params:
+      model: openai/my-fake-model
+      api_key: my-fake-key
+      api_base: https://exampleopenaiendpoint-production.up.railway.app/
+      stream_timeout: 0.001
+      rpm: 1
+  - model_name: fake-openai-endpoint-3
+    litellm_params:
+      model: openai/my-fake-model
+      api_key: my-fake-key
+      api_base: https://exampleopenaiendpoint-production.up.railway.app/
+      stream_timeout: 0.001
+      rpm: 1000
+  - model_name: fake-openai-endpoint-4
+    litellm_params:
+      model: openai/my-fake-model
+      api_key: my-fake-key
+      api_base: https://exampleopenaiendpoint-production.up.railway.app/
+      num_retries: 50
+  - model_name: fake-openai-endpoint-3
+    litellm_params:
+      model: openai/my-fake-model-2
+      api_key: my-fake-key
+      api_base: https://exampleopenaiendpoint-production.up.railway.app/
+      stream_timeout: 0.001
+      rpm: 1000
+  - model_name: bad-model
+    litellm_params:
+      model: openai/bad-model
+      api_key: os.environ/OPENAI_API_KEY
+      api_base: https://exampleopenaiendpoint-production.up.railway.app/
+      mock_timeout: True
+      timeout: 60
+      rpm: 1000
+    model_info:
+      health_check_timeout: 1
+  - model_name: good-model
+    litellm_params:
+      model: openai/bad-model
+      api_key: os.environ/OPENAI_API_KEY
+      api_base: https://exampleopenaiendpoint-production.up.railway.app/
+      rpm: 1000
+    model_info:
+      health_check_timeout: 1
+  - model_name: "*"
+    litellm_params:
+      model: openai/*
+      api_key: os.environ/OPENAI_API_KEY
+  - model_name: realtime-v1
+    litellm_params:
+      model: azure/gpt-realtime-20250828-standard
+      api_version: "2025-08-28"
+      realtime_protocol: GA # Possible values: "GA"/ "v1", "beta" 
+
+  - model_name: realtime-beta
+    litellm_params:
+      model: azure/gpt-realtime-20250828-standard
+      api_version: 2025-04-01-preview
+
+
+  # provider specific wildcard routing
+  - model_name: "anthropic/*"
+    litellm_params:
+      model: "anthropic/*"
+      api_key: os.environ/ANTHROPIC_API_KEY
+  - model_name: "bedrock/*"
+    litellm_params:
+      model: "bedrock/*"
+  - model_name: "groq/*"
+    litellm_params:
+      model: "groq/*"
+      api_key: os.environ/GROQ_API_KEY
+  - model_name: mistral-embed
+    litellm_params:
+      model: mistral/mistral-embed
+  - model_name: gpt-instruct # [PROD TEST] - tests if `/health` automatically infers this to be a text completion model
+    litellm_params:
+      model: text-completion-openai/gpt-3.5-turbo-instruct
+  - model_name: fake-openai-endpoint-5
+    litellm_params:
+      model: openai/my-fake-model
+      api_key: my-fake-key
+      api_base: https://exampleopenaiendpoint-production.up.railway.app/
+      timeout: 1
+  - model_name: badly-configured-openai-endpoint
+    litellm_params:
+      model: openai/my-fake-model
+      api_key: my-fake-key
+      api_base: https://exampleopenaiendpoint-production.up.railway.appxxxx/
+  - model_name: gemini-1.5-flash
+    litellm_params:
+      model: gemini/gemini-1.5-flash
+      api_key: os.environ/GOOGLE_API_KEY
+  - model_name: gpt-4o
+    litellm_params:
+      model: gpt-4o
+      api_key: os.environ/OPENAI_API_KEY
+
+
+litellm_settings:
+  # set_verbose: True  # Uncomment this if you want to see verbose logs; not recommended in production
+  drop_params: True
+  success_callback: ["prometheus"]
+  # max_budget: 100 
+  # budget_duration: 30d
+  num_retries: 5
+  request_timeout: 600
+  telemetry: False
+  context_window_fallbacks: [{"gpt-3.5-turbo": ["gpt-3.5-turbo-large"]}]
+  default_team_settings: 
+    - team_id: team-1
+      success_callback: ["langfuse"]
+      failure_callback: ["langfuse"]
+      langfuse_public_key: os.environ/LANGFUSE_PROJECT1_PUBLIC # Project 1
+      langfuse_secret: os.environ/LANGFUSE_PROJECT1_SECRET # Project 1
+    - team_id: team-2
+      success_callback: ["langfuse"]
+      failure_callback: ["langfuse"]
+      langfuse_public_key: os.environ/LANGFUSE_PROJECT2_PUBLIC # Project 2
+      langfuse_secret: os.environ/LANGFUSE_PROJECT2_SECRET # Project 2
+      langfuse_host: https://us.cloud.langfuse.com
+  # cache: true   # [OPTIONAL] use for caching responses 
+  # enable_caching_on_provider_specific_optional_params: True  # Include provider-specific params in cache keys
+  # cache_params:  # And for shared health check
+  #   type: redis
+  #   host: localhost
+  #   port: 6379
+
+# For /fine_tuning/jobs endpoints
+finetune_settings:
+  - custom_llm_provider: azure
+    api_base: os.environ/AZURE_API_BASE
+    api_key: os.environ/AZURE_API_KEY
+    api_version: "2023-03-15-preview"
+  - custom_llm_provider: openai
+    api_key: os.environ/OPENAI_API_KEY
+
+# for /files endpoints
+files_settings:
+  - custom_llm_provider: azure
+    api_base: os.environ/AZURE_API_BASE
+    api_key: os.environ/AZURE_API_KEY
+    api_version: "2023-03-15-preview"
+  - custom_llm_provider: openai
+    api_key: os.environ/OPENAI_API_KEY
+
+router_settings:
+  routing_strategy: usage-based-routing-v2 
+  redis_host: os.environ/REDIS_HOST
+  redis_password: os.environ/REDIS_PASSWORD
+  redis_port: os.environ/REDIS_PORT
+  enable_pre_call_checks: true
+  model_group_alias: {"my-special-fake-model-alias-name": "fake-openai-endpoint-3"} 
+
+general_settings: 
+  master_key: sk-1234 # [OPTIONAL] Use to enforce auth on proxy. See - https://docs.litellm.ai/docs/proxy/virtual_keys
+  store_model_in_db: True
+  proxy_budget_rescheduler_min_time: 60
+  proxy_budget_rescheduler_max_time: 64
+  proxy_batch_write_at: 1
+  database_connection_pool_limit: 10
+  # background_health_checks: true
+  # use_shared_health_check: true
+  # health_check_interval: 30
+  # database_url: "postgresql://<user>:<password>@<host>:<port>/<dbname>" # [OPTIONAL] use for token-based auth to proxy
+
+  pass_through_endpoints:
+    - path: "/v1/rerank"                                  # route you want to add to LiteLLM Proxy Server
+      target: "https://api.cohere.com/v1/rerank"          # URL this route should forward requests to
+      headers:                                            # headers to forward to this URL
+        content-type: application/json                    # (Optional) Extra Headers to pass to this endpoint 
+        accept: application/json
+      forward_headers: True
+
+# environment_variables:
+  # settings for using redis caching
+  # REDIS_HOST: redis-16337.c322.us-east-1-2.ec2.cloud.redislabs.com
+  # REDIS_PORT: "16337"
+  # REDIS_PASSWORD: 
\ No newline at end of file

From 018ccff23f1df9e4d223741e0f084447f08bcade Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Wed, 18 Mar 2026 09:11:20 +0530
Subject: [PATCH 04/11] fix(vertex-ai): address greptile review feedback on
 batch cancel

- Add try/except httpx.HTTPStatusError blocks in _async_cancel_batch for
  both POST cancel and GET retrieve calls, with verbose_logger error logging
- Fix endpoint extraction inconsistency: compute endpoint from URL without
  :cancel suffix so it matches behaviour of create_batch/retrieve_batch
- Add explicit validation that api_base ends with ':cancel' before
  stripping it, raising a descriptive error for unsupported custom proxy
  URL rewriting scenarios
- Use string-based patch() in test instead of patch.object() for robustness
  against import order changes

Made-with: Cursor
---
 litellm/llms/vertex_ai/batches/handler.py     | 54 ++++++++++++++-----
 .../test_vertex_ai_batch_transformation.py    |  2 +-
 2 files changed, 41 insertions(+), 15 deletions(-)

diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py
index c7b9287c08..36728499b9 100644
--- a/litellm/llms/vertex_ai/batches/handler.py
+++ b/litellm/llms/vertex_ai/batches/handler.py
@@ -399,10 +399,12 @@ class VertexAIBatchPrediction(VertexLLM):
             vertex_project=vertex_project or project_id,
         )
 
-        default_api_base = f"{default_api_base}/{batch_id}:cancel"
+        # Compute endpoint from the URL without :cancel for consistency with other methods
+        base_without_cancel = f"{default_api_base}/{batch_id}"
+        default_api_base = f"{base_without_cancel}:cancel"
 
-        if len(default_api_base.split(":")) > 1:
-            endpoint = default_api_base.split(":")[-1]
+        if len(base_without_cancel.split(":")) > 1:
+            endpoint = base_without_cancel.split(":")[-1]
         else:
             endpoint = ""
 
@@ -420,6 +422,14 @@ class VertexAIBatchPrediction(VertexLLM):
             vertex_api_version="v1",
         )
 
+        if not api_base.endswith(":cancel"):
+            raise ValueError(
+                f"cancel_batch: expected api_base to end with ':cancel', got: {api_base!r}. "
+                "Custom proxy URL rewriting is not supported for this operation."
+            )
+
+        retrieve_api_base = api_base.rsplit(":cancel", 1)[0]
+
         headers = {
             "Content-Type": "application/json; charset=utf-8",
             "Authorization": f"Bearer {access_token}",
@@ -428,7 +438,7 @@ class VertexAIBatchPrediction(VertexLLM):
         if _is_async is True:
             return self._async_cancel_batch(
                 api_base=api_base,
-                retrieve_api_base=api_base.rsplit(":cancel", 1)[0],
+                retrieve_api_base=retrieve_api_base,
                 headers=headers,
             )
 
@@ -443,7 +453,7 @@ class VertexAIBatchPrediction(VertexLLM):
             raise Exception(f"Error: {response.status_code} {response.text}")
 
         retrieve_response = sync_handler.get(
-            url=api_base.rsplit(":cancel", 1)[0],
+            url=retrieve_api_base,
             headers=headers,
         )
         if retrieve_response.status_code != 200:
@@ -466,18 +476,34 @@ class VertexAIBatchPrediction(VertexLLM):
         client = get_async_httpx_client(
             llm_provider=litellm.LlmProviders.VERTEX_AI,
         )
-        response = await client.post(
-            url=api_base,
-            headers=headers,
-            data=json.dumps({}),
-        )
+        try:
+            response = await client.post(
+                url=api_base,
+                headers=headers,
+                data=json.dumps({}),
+            )
+        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 = await client.get(
-            url=retrieve_api_base,
-            headers=headers,
-        )
+        try:
+            retrieve_response = await client.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}"
diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py
index d27cd8ba8a..e03555a9c0 100644
--- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py
+++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py
@@ -101,7 +101,7 @@ async def test_litellm_cancel_batch_vertex_ai():
     mock_response.id = "batch_123"
     mock_response.status = "cancelling"
     
-    with patch.object(litellm.batches.main, "vertex_ai_batches_instance") as mock_instance:
+    with patch("litellm.batches.main.vertex_ai_batches_instance") as mock_instance:
         mock_instance.cancel_batch.return_value = mock_response
         
         response = litellm.cancel_batch(

From 74382f1c89cf91b3ddc58dd82d6298407465c51c Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Wed, 18 Mar 2026 09:28:19 +0530
Subject: [PATCH 05/11] fix(vertex-ai): address greptile review feedback on
 batch cancel
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Replace misleading endpoint extraction with explicit endpoint = "cancel"
- Compute retrieve_api_base from URL components directly instead of
  stripping ":cancel" from the post-proxy URL, removing the hard ValueError
  that broke any custom Vertex AI proxy configuration
- Align cancel_batch provider priority in proxy endpoints to match
  create_batch order: body field → request headers → query params → default

Made-with: Cursor
---
 litellm/llms/vertex_ai/batches/handler.py    | 21 +++++++-------------
 litellm/proxy/batches_endpoints/endpoints.py |  2 +-
 2 files changed, 8 insertions(+), 15 deletions(-)

diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py
index 36728499b9..b30130688d 100644
--- a/litellm/llms/vertex_ai/batches/handler.py
+++ b/litellm/llms/vertex_ai/batches/handler.py
@@ -399,14 +399,11 @@ class VertexAIBatchPrediction(VertexLLM):
             vertex_project=vertex_project or project_id,
         )
 
-        # Compute endpoint from the URL without :cancel for consistency with other methods
-        base_without_cancel = f"{default_api_base}/{batch_id}"
-        default_api_base = f"{base_without_cancel}:cancel"
+        retrieve_api_base_default = f"{default_api_base}/{batch_id}"
+        default_api_base = f"{retrieve_api_base_default}:cancel"
 
-        if len(base_without_cancel.split(":")) > 1:
-            endpoint = base_without_cancel.split(":")[-1]
-        else:
-            endpoint = ""
+        # The Vertex AI action suffix for this operation
+        endpoint = "cancel"
 
         _, api_base = self._check_custom_proxy(
             api_base=api_base,
@@ -422,13 +419,9 @@ class VertexAIBatchPrediction(VertexLLM):
             vertex_api_version="v1",
         )
 
-        if not api_base.endswith(":cancel"):
-            raise ValueError(
-                f"cancel_batch: expected api_base to end with ':cancel', got: {api_base!r}. "
-                "Custom proxy URL rewriting is not supported for this operation."
-            )
-
-        retrieve_api_base = api_base.rsplit(":cancel", 1)[0]
+        # Use the canonical retrieve URL built from components rather than stripping
+        # ":cancel" from api_base, so custom proxy URL rewriting does not break retrieval.
+        retrieve_api_base = retrieve_api_base_default
 
         headers = {
             "Content-Type": "application/json; charset=utf-8",
diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py
index 9ce1b6e916..3b26e6c096 100644
--- a/litellm/proxy/batches_endpoints/endpoints.py
+++ b/litellm/proxy/batches_endpoints/endpoints.py
@@ -898,9 +898,9 @@ async def cancel_batch(
         else:
             custom_llm_provider = (
                 provider
+                or data.pop("custom_llm_provider", None)
                 or get_custom_llm_provider_from_request_headers(request=request)
                 or get_custom_llm_provider_from_request_query(request=request)
-                or data.pop("custom_llm_provider", None)
                 or "openai"
             )
             # Extract batch_id from data to avoid "multiple values for keyword argument" error

From 74ae17d15305334d0614b107e866bb5f6f7ef0e9 Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Wed, 18 Mar 2026 09:41:46 +0530
Subject: [PATCH 06/11]  greptile comments

---
 litellm/batches/main.py                       |  2 +-
 litellm/llms/vertex_ai/batches/handler.py     | 31 +++++++++++++------
 .../test_vertex_ai_batch_transformation.py    |  2 +-
 3 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/litellm/batches/main.py b/litellm/batches/main.py
index 36093d071b..ae79469dd1 100644
--- a/litellm/batches/main.py
+++ b/litellm/batches/main.py
@@ -1048,7 +1048,7 @@ def cancel_batch(
                 litellm_params=litellm_params,
             )
         elif custom_llm_provider == "vertex_ai":
-            api_base = optional_params.api_base or ""
+            api_base = optional_params.api_base or None
             vertex_ai_project = (
                 optional_params.vertex_project
                 or litellm.vertex_project
diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py
index b30130688d..f4ba0533c8 100644
--- a/litellm/llms/vertex_ai/batches/handler.py
+++ b/litellm/llms/vertex_ai/batches/handler.py
@@ -400,28 +400,41 @@ class VertexAIBatchPrediction(VertexLLM):
         )
 
         retrieve_api_base_default = f"{default_api_base}/{batch_id}"
-        default_api_base = f"{retrieve_api_base_default}:cancel"
+        cancel_api_base_default = f"{retrieve_api_base_default}:cancel"
 
-        # The Vertex AI action suffix for this operation
-        endpoint = "cancel"
+        # Save the caller-supplied value before _check_custom_proxy overwrites api_base,
+        # so we can pass it unchanged to the second proxy-check for the retrieve URL.
+        caller_api_base = api_base
 
         _, api_base = self._check_custom_proxy(
-            api_base=api_base,
+            api_base=caller_api_base,
             custom_llm_provider="vertex_ai",
             gemini_api_key=None,
-            endpoint=endpoint,
+            endpoint="cancel",
             stream=None,
             auth_header=None,
-            url=default_api_base,
+            url=cancel_api_base_default,
             model=None,
             vertex_project=vertex_project or project_id,
             vertex_location=vertex_location or "us-central1",
             vertex_api_version="v1",
         )
 
-        # Use the canonical retrieve URL built from components rather than stripping
-        # ":cancel" from api_base, so custom proxy URL rewriting does not break retrieval.
-        retrieve_api_base = retrieve_api_base_default
+        # Route the retrieve GET through the same proxy as the cancel POST by running
+        # _check_custom_proxy a second time with the non-cancel default URL.
+        _, retrieve_api_base = self._check_custom_proxy(
+            api_base=caller_api_base,
+            custom_llm_provider="vertex_ai",
+            gemini_api_key=None,
+            endpoint="",
+            stream=None,
+            auth_header=None,
+            url=retrieve_api_base_default,
+            model=None,
+            vertex_project=vertex_project or project_id,
+            vertex_location=vertex_location or "us-central1",
+            vertex_api_version="v1",
+        )
 
         headers = {
             "Content-Type": "application/json; charset=utf-8",
diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py
index e03555a9c0..1cf6fa3266 100644
--- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py
+++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py
@@ -45,7 +45,7 @@ def test_output_file_id_falls_back_to_output_uri_prefix_with_predictions_jsonl()
 
 
 @pytest.mark.asyncio
-async def test_vertex_ai_cancel_batch():
+def test_vertex_ai_cancel_batch():
     """Test that vertex_ai cancel_batch calls the correct API endpoint"""
     handler = VertexAIBatchPrediction(gcs_bucket_name="test-bucket")
     

From d0d593beb8027b2b467b1ef9cbdd302981d658a4 Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Wed, 18 Mar 2026 09:48:24 +0530
Subject: [PATCH 07/11] Update
 tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
---
 .../llms/vertex_ai/test_vertex_ai_batch_transformation.py        | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py
index 1cf6fa3266..1fa8dd2e16 100644
--- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py
+++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py
@@ -44,7 +44,6 @@ def test_output_file_id_falls_back_to_output_uri_prefix_with_predictions_jsonl()
     )
 
 
-@pytest.mark.asyncio
 def test_vertex_ai_cancel_batch():
     """Test that vertex_ai cancel_batch calls the correct API endpoint"""
     handler = VertexAIBatchPrediction(gcs_bucket_name="test-bucket")

From 547db8f5d1f1408e598088380fa007497efb3f23 Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Wed, 18 Mar 2026 10:01:47 +0530
Subject: [PATCH 08/11] Fix greptile comments

---
 litellm/llms/vertex_ai/batches/handler.py | 24 +++++------------------
 1 file changed, 5 insertions(+), 19 deletions(-)

diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py
index f4ba0533c8..9c760f628e 100644
--- a/litellm/llms/vertex_ai/batches/handler.py
+++ b/litellm/llms/vertex_ai/batches/handler.py
@@ -402,12 +402,8 @@ class VertexAIBatchPrediction(VertexLLM):
         retrieve_api_base_default = f"{default_api_base}/{batch_id}"
         cancel_api_base_default = f"{retrieve_api_base_default}:cancel"
 
-        # Save the caller-supplied value before _check_custom_proxy overwrites api_base,
-        # so we can pass it unchanged to the second proxy-check for the retrieve URL.
-        caller_api_base = api_base
-
         _, api_base = self._check_custom_proxy(
-            api_base=caller_api_base,
+            api_base=api_base,
             custom_llm_provider="vertex_ai",
             gemini_api_key=None,
             endpoint="cancel",
@@ -420,20 +416,10 @@ class VertexAIBatchPrediction(VertexLLM):
             vertex_api_version="v1",
         )
 
-        # Route the retrieve GET through the same proxy as the cancel POST by running
-        # _check_custom_proxy a second time with the non-cancel default URL.
-        _, retrieve_api_base = self._check_custom_proxy(
-            api_base=caller_api_base,
-            custom_llm_provider="vertex_ai",
-            gemini_api_key=None,
-            endpoint="",
-            stream=None,
-            auth_header=None,
-            url=retrieve_api_base_default,
-            model=None,
-            vertex_project=vertex_project or project_id,
-            vertex_location=vertex_location or "us-central1",
-            vertex_api_version="v1",
+        retrieve_api_base = (
+            api_base.removesuffix(":cancel")
+            if api_base.endswith(":cancel")
+            else retrieve_api_base_default
         )
 
         headers = {

From c4d27cb239d89d1cf934d27b52285d5a6b490f95 Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Wed, 18 Mar 2026 10:30:05 +0530
Subject: [PATCH 09/11] =?UTF-8?q?fix(vertex-ai):=20address=20greptile=20re?=
 =?UTF-8?q?view=20=E2=80=93=20proxy=20retrieve=20URL,=20timeout=20forwardi?=
 =?UTF-8?q?ng,=20sync=20logging?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 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
---
 litellm/llms/vertex_ai/batches/handler.py     | 49 ++++++++----
 .../test_vertex_ai_batch_transformation.py    | 78 +++++++++++++++++++
 2 files changed, 113 insertions(+), 14 deletions(-)

diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py
index 9c760f628e..416c5fc69f 100644
--- a/litellm/llms/vertex_ai/batches/handler.py
+++ b/litellm/llms/vertex_ai/batches/handler.py
@@ -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(
diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py
index 1fa8dd2e16..b934afd103 100644
--- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py
+++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py
@@ -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"""

From 1181adbaf3ac51ad33df74c9686736462c509f6d Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Wed, 18 Mar 2026 10:41:31 +0530
Subject: [PATCH 10/11] address greptile review feedback (greploop iteration 1)

- Remove dead elif branch in retrieve_api_base derivation
- Replace unreachable try/except httpx.HTTPStatusError around GET
  calls with logging inside the status_code check (HTTPHandler.get()
  does not call raise_for_status())
- Add comments noting HTTPHandler.get()/AsyncHTTPHandler.get() do not
  accept a timeout parameter

Made-with: Cursor
---
 litellm/llms/vertex_ai/batches/handler.py | 38 ++++++++++-------------
 1 file changed, 16 insertions(+), 22 deletions(-)

diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py
index 416c5fc69f..2cb0294206 100644
--- a/litellm/llms/vertex_ai/batches/handler.py
+++ b/litellm/llms/vertex_ai/batches/handler.py
@@ -418,8 +418,6 @@ class VertexAIBatchPrediction(VertexLLM):
 
         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("/")
 
@@ -455,19 +453,17 @@ class VertexAIBatchPrediction(VertexLLM):
         if response.status_code != 200:
             raise Exception(f"Error: {response.status_code} {response.text}")
 
-        try:
-            retrieve_response = sync_handler.get(
-                url=retrieve_api_base,
-                headers=headers,
-            )
-        except httpx.HTTPStatusError as e:
+        # HTTPHandler.get() does not accept a timeout parameter
+        retrieve_response = sync_handler.get(
+            url=retrieve_api_base,
+            headers=headers,
+        )
+        if retrieve_response.status_code != 200:
             litellm.verbose_logger.error(
                 "Vertex AI batch retrieve-after-cancel failed: status=%s, body=%s",
-                e.response.status_code,
-                e.response.text[:1000],
+                retrieve_response.status_code,
+                retrieve_response.text[:1000],
             )
-            raise
-        if retrieve_response.status_code != 200:
             raise Exception(
                 f"Error: {retrieve_response.status_code} {retrieve_response.text}"
             )
@@ -505,19 +501,17 @@ class VertexAIBatchPrediction(VertexLLM):
         if response.status_code != 200:
             raise Exception(f"Error: {response.status_code} {response.text}")
 
-        try:
-            retrieve_response = await client.get(
-                url=retrieve_api_base,
-                headers=headers,
-            )
-        except httpx.HTTPStatusError as e:
+        # AsyncHTTPHandler.get() does not accept a timeout parameter
+        retrieve_response = await client.get(
+            url=retrieve_api_base,
+            headers=headers,
+        )
+        if retrieve_response.status_code != 200:
             litellm.verbose_logger.error(
                 "Vertex AI batch retrieve-after-cancel failed: status=%s, body=%s",
-                e.response.status_code,
-                e.response.text[:1000],
+                retrieve_response.status_code,
+                retrieve_response.text[:1000],
             )
-            raise
-        if retrieve_response.status_code != 200:
             raise Exception(
                 f"Error: {retrieve_response.status_code} {retrieve_response.text}"
             )

From 694cf22c9e13cd6b901d3bb5c83a67f6c691aa9c Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Wed, 18 Mar 2026 11:09:20 +0530
Subject: [PATCH 11/11] Update
 tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
---
 .../test_vertex_ai_batch_transformation.py    | 36 +++----------------
 1 file changed, 4 insertions(+), 32 deletions(-)

diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py
index b934afd103..7310c68b4e 100644
--- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py
+++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py
@@ -94,39 +94,11 @@ def test_vertex_ai_cancel_batch():
 
 
 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")
+    """Test that timeout is forwarded to the POST (cancel) HTTP call.
 
-    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
+    Note: the follow-up GET (retrieve) call does not accept a timeout
+    parameter in the underlying HTTP handler, so it is intentionally omitted.
+    """
 
 
 def test_vertex_ai_cancel_batch_custom_proxy_retrieve_url():