diff --git a/litellm/llms/vertex_ai/files/handler.py b/litellm/llms/vertex_ai/files/handler.py index 816fb82812c..bd4b2ac8bbc 100644 --- a/litellm/llms/vertex_ai/files/handler.py +++ b/litellm/llms/vertex_ai/files/handler.py @@ -199,7 +199,7 @@ class VertexAIFilesHandler(GCSBucketBase): # Apply transformation to convert Vertex AI batch outputs to OpenAI format config = VertexAIFilesConfig() - + # Create a logging object for transformation logging_obj = Logging( model="", @@ -210,11 +210,9 @@ class VertexAIFilesHandler(GCSBucketBase): litellm_call_id="", function_id="", ) - + return config.transform_file_content_response( - raw_response=mock_response, - logging_obj=logging_obj, - litellm_params={} + raw_response=mock_response, logging_obj=logging_obj, litellm_params={} ) def file_content( diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 3a140eb723f..75d2c69bf1c 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -46,15 +46,15 @@ _GCP_LABEL_VALUE_MAX_LEN = 63 def _sanitize_gcp_label_value(value: str) -> str: """ Sanitize a string to meet GCP label value constraints. - + GCP label values must: - Be lowercase - Contain only letters, numbers, underscores, and hyphens - Be max 63 characters - + Args: value: The string to sanitize - + Returns: A sanitized string that meets GCP label constraints """ @@ -293,14 +293,16 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): litellm_params={}, cached_content=None, ) - + # Add custom_id as a label for correlation in batch outputs custom_id = _openai_jsonl_content.get("custom_id") if custom_id: if "labels" not in vertex_request_body: vertex_request_body["labels"] = {} - _set_litellm_batch_custom_id_labels(vertex_request_body["labels"], custom_id) - + _set_litellm_batch_custom_id_labels( + vertex_request_body["labels"], custom_id + ) + vertex_jsonl_content.append({"request": vertex_request_body}) return vertex_jsonl_content @@ -509,10 +511,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): ) -> HttpxBinaryResponseContent: """ Transform file content response, converting Vertex AI batch output to OpenAI format if applicable. - + This method automatically detects and transforms Vertex AI batch prediction outputs (predictions.jsonl files) into OpenAI-compatible batch response format. - + If the file is not a batch output or transformation fails, the original content is returned as-is to maintain backward compatibility. """ @@ -526,11 +528,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): if transformed_content != content: # Create a new response with transformed content and updated Content-Length import httpx - + # Update headers with correct Content-Length new_headers = dict(raw_response.headers) new_headers["content-length"] = str(len(transformed_content)) - + mock_response = httpx.Response( status_code=raw_response.status_code, content=transformed_content, @@ -541,16 +543,14 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): except Exception: # If transformation fails, return as-is pass - + return HttpxBinaryResponseContent(response=raw_response) - - def _try_transform_vertex_batch_output_to_openai( - self, content: bytes - ) -> bytes: + + def _try_transform_vertex_batch_output_to_openai(self, content: bytes) -> bytes: """ Try to transform Vertex AI batch output to OpenAI format. If conversion fails at any point, return the original content as-is. - + Vertex AI batch output format (predictions.jsonl): { "request": {"contents": [...], "labels": {"litellm_custom_id": "request-1", "litellm_custom_id_raw": "..."}}, @@ -558,7 +558,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): "response": {"candidates": [...], "modelVersion": "gemini-2.5-flash", ...}, "processed_time": "2026-04-13T10:18:18.102004+00:00" } - + OpenAI batch output format: { "id": "batch_req_...", @@ -574,15 +574,15 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): try: # Decode content content_str = content.decode("utf-8") - + # Check if it's JSONL (multiple lines) lines = content_str.strip().split("\n") if not lines: return content - + # Try to parse the first line to see if it's Vertex AI batch output first_line = json.loads(lines[0]) - + # Check if it has Vertex AI batch output structure with discriminating fields # Must have request, response, and processed_time # Plus either candidates (success) or status (error) @@ -592,37 +592,38 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): and "processed_time" in first_line ) has_success_or_error = ( - "candidates" in first_line.get("response", {}) - or "status" in first_line + "candidates" in first_line.get("response", {}) or "status" in first_line ) - + if not (has_base_structure and has_success_or_error): # Not a Vertex AI batch output, return as-is return content - + # Transform all lines transformed_lines = [] for line in lines: if not line.strip(): continue - + try: vertex_output = json.loads(line) - openai_output = self._transform_single_vertex_batch_output_to_openai( - vertex_output + openai_output = ( + self._transform_single_vertex_batch_output_to_openai( + vertex_output + ) ) transformed_lines.append(json.dumps(openai_output)) except Exception: # If any line fails, return original content return content - + # Return transformed content return "\n".join(transformed_lines).encode("utf-8") - + except Exception: # If anything fails, return original content return content - + def _transform_single_vertex_batch_output_to_openai( self, vertex_output: Dict[str, Any] ) -> Dict[str, Any]: @@ -634,11 +635,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): request_data = vertex_output.get("request", {}) labels = request_data.get("labels", {}) or {} custom_id = _get_litellm_batch_custom_id_from_labels(labels) - + # Check if there's an error status = vertex_output.get("status", "") has_error = bool(status) - + if has_error: # Return error response in OpenAI format return { @@ -651,25 +652,25 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): "error": { "message": status, "type": "vertex_ai_error", - "code": "vertex_ai_error" + "code": "vertex_ai_error", } - } + }, }, "error": { "message": status, "type": "vertex_ai_error", - "code": "vertex_ai_error" - } + "code": "vertex_ai_error", + }, } - + # Transform successful response using existing transformation vertex_response = vertex_output.get("response", {}) - + # Extract model from response model = vertex_response.get("modelVersion", "gemini-1.5-flash-001") if "@" in model: model = model.split("@")[0] - + # Create logging object for transformation logging_obj = Logging( model=model, @@ -681,7 +682,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): function_id="", ) logging_obj.optional_params = {} - + # Create mock httpx response for transformation mock_httpx_response = httpx.Response( status_code=200, @@ -689,12 +690,12 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): headers={"content-type": "application/json"}, request=httpx.Request(method="POST", url="https://example.com"), ) - + try: # Use existing VertexGeminiConfig transformation vertex_gemini_config = VertexGeminiConfig() model_response = ModelResponse() - + transformed_response = vertex_gemini_config._transform_google_generate_content_to_openai_model_response( completion_response=vertex_response, model_response=model_response, @@ -702,10 +703,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): logging_obj=logging_obj, raw_response=mock_httpx_response, ) - + # Convert ModelResponse to dict response_dict = transformed_response.model_dump() - + # Return in OpenAI batch format return { "id": f"batch_req_{uuid.uuid4()}", @@ -713,11 +714,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): "response": { "status_code": 200, "request_id": response_dict.get("id", ""), - "body": response_dict + "body": response_dict, }, - "error": None + "error": None, } - + except Exception as e: # If transformation fails, return error return { @@ -730,15 +731,15 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): "error": { "message": f"Failed to transform response: {str(e)}", "type": "transformation_error", - "code": "transformation_error" + "code": "transformation_error", } - } + }, }, "error": { "message": f"Failed to transform response: {str(e)}", "type": "transformation_error", - "code": "transformation_error" - } + "code": "transformation_error", + }, } @@ -799,14 +800,16 @@ class VertexAIJsonlFilesTransformation(VertexGeminiConfig): litellm_params={}, cached_content=None, ) - + # Add custom_id as a label for correlation in batch outputs custom_id = _openai_jsonl_content.get("custom_id") if custom_id: if "labels" not in vertex_request_body: vertex_request_body["labels"] = {} - _set_litellm_batch_custom_id_labels(vertex_request_body["labels"], custom_id) - + _set_litellm_batch_custom_id_labels( + vertex_request_body["labels"], custom_id + ) + vertex_jsonl_content.append({"request": vertex_request_body}) return vertex_jsonl_content