From 35f6fd4223065eb2c1710964ba43c60b25f9edff Mon Sep 17 00:00:00 2001 From: Ephrim Stanley Date: Sat, 28 Feb 2026 14:42:24 -0500 Subject: [PATCH] Managed batches fixes for Gemini/Vertex --- .../proxy/hooks/managed_files.py | 7 ++ litellm/batches/batch_utils.py | 101 ++++++++---------- litellm/files/main.py | 2 +- litellm/llms/vertex_ai/batches/handler.py | 17 ++- .../llms/vertex_ai/batches/transformation.py | 2 +- .../llms/vertex_ai/files/transformation.py | 62 +++++++++-- litellm/types/llms/vertex_ai.py | 2 +- 7 files changed, 121 insertions(+), 72 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 4fa050a84aa..37ca341fdf2 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -589,7 +589,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_file_id_mapping = cast( Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping") ) + # model_info may be at top-level or nested under litellm_metadata + # (batch/file operations use litellm_metadata) model_id = cast(Optional[str], kwargs.get("model_info", {}).get("id", None)) + if model_id is None: + model_id = cast( + Optional[str], + kwargs.get("litellm_metadata", {}).get("model_info", {}).get("id", None), + ) mapped_file_id: Optional[str] = None if input_file_id and model_file_id_mapping and model_id: mapped_file_id = model_file_id_mapping.get(input_file_id, {}).get( diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 29bd99c2a60..80351664dfe 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -128,73 +128,58 @@ def calculate_vertex_ai_batch_cost_and_usage( model_name: Optional[str] = None, ) -> Tuple[float, Usage]: """ - Calculate both cost and usage from Vertex AI batch responses + Calculate both cost and usage from Vertex AI batch responses. + + Vertex AI batch output lines have format: + {"request": ..., "status": "", "response": {"candidates": [...], "usageMetadata": {...}}} + + usageMetadata contains promptTokenCount, candidatesTokenCount, totalTokenCount. """ - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexGeminiConfig, - ) + from litellm.cost_calculator import batch_cost_calculator + total_cost = 0.0 total_tokens = 0 prompt_tokens = 0 completion_tokens = 0 - - for response in vertex_ai_batch_responses: - if response.get("status") == "JOB_STATE_SUCCEEDED": # Check if response was successful - # Transform Vertex AI response to OpenAI format if needed + actual_model_name = model_name or "gemini-2.0-flash-001" - # Create required arguments for the transformation method - model_response = ModelResponse() - - # Ensure model_name is not None - actual_model_name = model_name or "gemini-2.5-flash" - - # Create a real LiteLLM logging object - logging_obj = Logging( + for response in vertex_ai_batch_responses: + response_body = response.get("response") + if response_body is None: + continue + + usage_metadata = response_body.get("usageMetadata", {}) + _prompt = usage_metadata.get("promptTokenCount", 0) or 0 + _completion = usage_metadata.get("candidatesTokenCount", 0) or 0 + _total = usage_metadata.get("totalTokenCount", 0) or (_prompt + _completion) + + line_usage = Usage( + prompt_tokens=_prompt, + completion_tokens=_completion, + total_tokens=_total, + ) + + try: + p_cost, c_cost = batch_cost_calculator( + usage=line_usage, model=actual_model_name, - messages=[{"role": "user", "content": "batch_request"}], - stream=False, - call_type=CallTypes.aretrieve_batch, - start_time=time.time(), - litellm_call_id="batch_" + str(uuid.uuid4()), - function_id="batch_processing", - litellm_trace_id=str(uuid.uuid4()), - kwargs={"optional_params": {}} - ) - - # Add the optional_params attribute that the Vertex AI transformation expects - logging_obj.optional_params = {} - raw_response = httpx.Response(200) # Mock response object - - openai_format_response = VertexGeminiConfig()._transform_google_generate_content_to_openai_model_response( - completion_response=response["response"], - model_response=model_response, - model=actual_model_name, - logging_obj=logging_obj, - raw_response=raw_response, - ) - - # Calculate cost using existing function - cost = litellm.completion_cost( - completion_response=openai_format_response, custom_llm_provider="vertex_ai", - call_type=CallTypes.aretrieve_batch.value, ) - total_cost += cost - - # Extract usage from the transformed response - usage_obj = getattr(openai_format_response, 'usage', None) - if usage_obj: - usage = usage_obj - else: - # Fallback: create usage from response dict - response_dict = openai_format_response.dict() if hasattr(openai_format_response, 'dict') else {} - usage = _get_batch_job_usage_from_response_body(response_dict) - - total_tokens += usage.total_tokens - prompt_tokens += usage.prompt_tokens - completion_tokens += usage.completion_tokens - + total_cost += p_cost + c_cost + except Exception as e: + verbose_logger.debug( + "vertex_ai batch cost calculation error for line: %s", str(e) + ) + + prompt_tokens += _prompt + completion_tokens += _completion + total_tokens += _total + + verbose_logger.info( + "vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d", + total_cost, prompt_tokens, completion_tokens, total_tokens, + ) + return total_cost, Usage( total_tokens=total_tokens, prompt_tokens=prompt_tokens, diff --git a/litellm/files/main.py b/litellm/files/main.py index 78e41bb5a68..66d3a97468d 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -295,7 +295,7 @@ def create_file( @client async def afile_retrieve( file_id: str, - custom_llm_provider: Literal["openai", "azure", "gemini", "hosted_vllm", "manus"] = "openai", + custom_llm_provider: Literal["openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "manus"] = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 36f5e65e7a2..ba3b5fb7a2c 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -108,11 +108,18 @@ 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(vertex_batch_request), - ) + try: + response = await client.post( + url=api_base, + headers=headers, + data=json.dumps(vertex_batch_request), + ) + except httpx.HTTPStatusError as e: + error_body = e.response.text if hasattr(e, 'response') else "N/A" + litellm.verbose_logger.error( + f"Vertex AI batch create failed: status={e.response.status_code}, body={error_body[:1000]}" + ) + raise if response.status_code != 200: raise Exception(f"Error: {response.status_code} {response.text}") diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index a0adb3e55a8..7cb06fea9e2 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -29,7 +29,7 @@ class VertexAIBatchTransformation: if input_file_id is None: raise ValueError("input_file_id is required, but not provided") input_config: InputConfig = InputConfig( - gcsSource=GcsSource(uris=input_file_id), instancesFormat="jsonl" + gcsSource=GcsSource(uris=[input_file_id]), instancesFormat="jsonl" ) model: str = cls._get_model_from_gcs_file(input_file_id) output_config: OutputConfig = OutputConfig( diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 2470c59bbac..f0493cd6be9 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -335,13 +335,37 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): status_code=status_code, message=error_message, headers=headers ) + def _parse_gcs_uri(self, file_id: str) -> Tuple[str, str]: + """ + Parse a GCS URI (gs://bucket/path/to/object) into (bucket, url-encoded-object-path). + Handles both raw and URL-encoded input. + """ + import urllib.parse + + decoded = urllib.parse.unquote(file_id) + if decoded.startswith("gs://"): + full_path = decoded[5:] + else: + full_path = decoded + + if "/" in full_path: + bucket_name, object_path = full_path.split("/", 1) + else: + bucket_name = full_path + object_path = "" + + encoded_object = urllib.parse.quote(object_path, safe="") + return bucket_name, encoded_object + def transform_retrieve_file_request( self, file_id: str, optional_params: dict, litellm_params: dict, ) -> tuple[str, dict]: - raise NotImplementedError("VertexAIFilesConfig does not support file retrieval") + bucket, encoded_object = self._parse_gcs_uri(file_id) + url = f"https://storage.googleapis.com/storage/v1/b/{bucket}/o/{encoded_object}" + return url, {} def transform_retrieve_file_response( self, @@ -349,7 +373,21 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): logging_obj: LiteLLMLoggingObj, litellm_params: dict, ) -> OpenAIFileObject: - raise NotImplementedError("VertexAIFilesConfig does not support file retrieval") + response_json = raw_response.json() + gcs_id = response_json.get("id", "") + gcs_id = "/".join(gcs_id.split("/")[:-1]) if gcs_id else "" + return OpenAIFileObject( + id=f"gs://{gcs_id}", + bytes=int(response_json.get("size", 0)), + created_at=_convert_vertex_datetime_to_openai_datetime( + vertex_datetime=response_json.get("timeCreated", "") + ), + filename=response_json.get("name", ""), + object="file", + purpose=response_json.get("metadata", {}).get("purpose", "batch"), + status="processed", + status_details=None, + ) def transform_delete_file_request( self, @@ -357,7 +395,9 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): optional_params: dict, litellm_params: dict, ) -> tuple[str, dict]: - raise NotImplementedError("VertexAIFilesConfig does not support file deletion") + bucket, encoded_object = self._parse_gcs_uri(file_id) + url = f"https://storage.googleapis.com/storage/v1/b/{bucket}/o/{encoded_object}" + return url, {} def transform_delete_file_response( self, @@ -365,7 +405,14 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): logging_obj: LiteLLMLoggingObj, litellm_params: dict, ) -> FileDeleted: - raise NotImplementedError("VertexAIFilesConfig does not support file deletion") + file_id = "deleted" + if hasattr(raw_response, "request") and raw_response.request: + url = str(raw_response.request.url) + if "/o/" in url: + import urllib.parse + encoded_name = url.split("/o/")[-1].split("?")[0] + file_id = f"gs://{urllib.parse.unquote(encoded_name)}" + return FileDeleted(id=file_id, deleted=True, object="file") def transform_list_files_request( self, @@ -389,7 +436,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): optional_params: dict, litellm_params: dict, ) -> tuple[str, dict]: - raise NotImplementedError("VertexAIFilesConfig does not support file content retrieval") + file_id = file_content_request.get("file_id", "") + bucket, encoded_object = self._parse_gcs_uri(file_id) + url = f"https://storage.googleapis.com/storage/v1/b/{bucket}/o/{encoded_object}?alt=media" + return url, {} def transform_file_content_response( self, @@ -397,7 +447,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): logging_obj: LiteLLMLoggingObj, litellm_params: dict, ) -> HttpxBinaryResponseContent: - raise NotImplementedError("VertexAIFilesConfig does not support file content retrieval") + return HttpxBinaryResponseContent(response=raw_response) class VertexAIJsonlFilesTransformation(VertexGeminiConfig): diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 049a5010c79..190e680b7b9 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -560,7 +560,7 @@ class VertexAIBatchEmbeddingsResponseObject(TypedDict): class GcsSource(TypedDict): - uris: str + uris: List[str] class InputConfig(TypedDict):