from typing import Any, Final from urllib.parse import unquote from litellm._uuid import uuid from litellm.llms.vertex_ai.common_utils import ( VertexAIError, _convert_vertex_datetime_to_openai_datetime, ) from litellm.types.llms.openai import BatchJobStatus, CreateBatchRequest from litellm.types.llms.vertex_ai import * from litellm.types.utils import LiteLLMBatch class VertexAIBatchTransformation: """ Transforms OpenAI Batch requests to Vertex AI Batch requests API Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/batch-prediction-gemini """ @classmethod def transform_openai_batch_request_to_vertex_ai_batch_request( cls, request: CreateBatchRequest, ) -> VertexAIBatchPredictionJob: """ Transforms OpenAI Batch requests to Vertex AI Batch requests """ request_display_name: Final = f"litellm-vertex-batch-{uuid.uuid4()}" input_file_id: Final = request.get("input_file_id") 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") model: Final[str] = cls._get_model_from_gcs_file(input_file_id) output_config: Final[OutputConfig] = OutputConfig( predictionsFormat="jsonl", gcsDestination=GcsDestination(outputUriPrefix=cls._get_gcs_uri_prefix_from_file(input_file_id)), ) return VertexAIBatchPredictionJob( inputConfig=input_config, outputConfig=output_config, model=model, displayName=request_display_name, ) @classmethod def transform_vertex_ai_batch_response_to_openai_batch_response( cls, response: VertexBatchPredictionResponse ) -> LiteLLMBatch: return LiteLLMBatch( id=cls._get_batch_id_from_vertex_ai_batch_response(response), completion_window="24h", created_at=_convert_vertex_datetime_to_openai_datetime(vertex_datetime=response.get("createTime", "")), endpoint="", input_file_id=cls._get_input_file_id_from_vertex_ai_batch_response(response), object="batch", status=cls._get_batch_job_status_from_vertex_ai_batch_response(response), error_file_id=None, # Vertex AI doesn't seem to have a direct equivalent output_file_id=cls._get_output_file_id_from_vertex_ai_batch_response(response), ) @classmethod def transform_vertex_ai_batch_list_response_to_openai_list_response( cls, response: dict[str, Any] ) -> dict[str, Any]: """ Transforms Vertex AI batch list response into OpenAI-compatible list response. """ batch_jobs: Final = response.get("batchPredictionJobs", []) or [] data: Final = [cls.transform_vertex_ai_batch_response_to_openai_batch_response(job) for job in batch_jobs] first_id: Final = data[0].id if len(data) > 0 else None last_id: Final = data[-1].id if len(data) > 0 else None next_page_token: Final = response.get("nextPageToken") return { "object": "list", "data": data, "first_id": first_id, "last_id": last_id, "has_more": bool(next_page_token), "next_page_token": next_page_token, } @classmethod def _get_batch_id_from_vertex_ai_batch_response(cls, response: VertexBatchPredictionResponse) -> str: """ Gets the batch id from the Vertex AI Batch response safely vertex response: `projects/510528649030/locations/us-central1/batchPredictionJobs/3814889423749775360` returns: `3814889423749775360` """ _name: Final = response.get("name", "") if not _name: return "" # Split by '/' and get the last part if it exists parts: Final = _name.split("/") return parts[-1] if parts else _name @classmethod def _get_input_file_id_from_vertex_ai_batch_response(cls, response: VertexBatchPredictionResponse) -> str: """ Gets the input file id from the Vertex AI Batch response """ input_file_id: Final[str] = "" input_config: Final = response.get("inputConfig") if input_config is None: return input_file_id gcs_source: Final = input_config.get("gcsSource") if gcs_source is None: return input_file_id uris: Final = gcs_source.get("uris", "") if len(uris) == 0: return input_file_id return uris[0] @classmethod def _get_output_file_id_from_vertex_ai_batch_response(cls, response: VertexBatchPredictionResponse) -> str: """ Gets the output file id from the Vertex AI Batch response """ output_info: Final = response.get("outputInfo") or OutputInfo() output_file_id: str = output_info.get("gcsOutputDirectory", "") if output_file_id: output_file_id = output_file_id.rstrip("/") + "/predictions.jsonl" if output_file_id and output_file_id != "/predictions.jsonl": return output_file_id output_config: Final = response.get("outputConfig") if output_config is None: return output_file_id gcs_destination: Final = output_config.get("gcsDestination") if gcs_destination is None: return output_file_id output_uri_prefix: Final = gcs_destination.get("outputUriPrefix", "") if output_uri_prefix.endswith("/predictions.jsonl"): return output_uri_prefix return output_uri_prefix.rstrip("/") + "/predictions.jsonl" @classmethod def _get_batch_job_status_from_vertex_ai_batch_response( cls, response: VertexBatchPredictionResponse ) -> BatchJobStatus: """ Gets the batch job status from the Vertex AI Batch response ref: https://cloud.google.com/vertex-ai/docs/reference/rest/v1/JobState """ state_mapping: Final[dict[str, BatchJobStatus]] = { "JOB_STATE_UNSPECIFIED": "failed", "JOB_STATE_QUEUED": "validating", "JOB_STATE_PENDING": "validating", "JOB_STATE_RUNNING": "in_progress", "JOB_STATE_SUCCEEDED": "completed", "JOB_STATE_FAILED": "failed", "JOB_STATE_CANCELLING": "cancelling", "JOB_STATE_CANCELLED": "cancelled", "JOB_STATE_PAUSED": "in_progress", "JOB_STATE_EXPIRED": "expired", "JOB_STATE_UPDATING": "in_progress", "JOB_STATE_PARTIALLY_SUCCEEDED": "completed", } vertex_state: Final = response.get("state", "JOB_STATE_UNSPECIFIED") return state_mapping[vertex_state] @classmethod def _get_gcs_uri_prefix_from_file(cls, input_file_id: str) -> str: """ Gets the gcs uri prefix from the input file id Example: input_file_id: "gs://litellm-testing-bucket/vtx_batch.jsonl" returns: "gs://litellm-testing-bucket" input_file_id: "gs://litellm-testing-bucket/batches/vtx_batch.jsonl" returns: "gs://litellm-testing-bucket/batches" """ # Split the path and remove the filename path_parts: Final = input_file_id.rsplit("/", 1) return path_parts[0] @classmethod def _get_model_from_gcs_file(cls, gcs_file_uri: str) -> str: """ Extracts the model from the gcs file uri When files are uploaded using LiteLLM (/v1/files), the model is stored in the gcs file uri Why? - Because Vertex Requires the `model` param in create batch jobs request, but OpenAI does not require this gcs_file_uri format: gs://litellm-testing-bucket/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/e9412502-2c91-42a6-8e61-f5c294cc0fc8 returns: "publishers/google/models/gemini-1.5-flash-001" Raises a 400 `VertexAIError` when the uri carries no parseable model path. """ model: Final = cls._parse_model_from_gcs_file(gcs_file_uri) if model is None: raise VertexAIError( status_code=400, message=( "Vertex AI batch creation requires the model to be part of `input_file_id`, but " f"'{gcs_file_uri}' contains no 'publishers//models/' path segment. " "Either upload the input file through LiteLLM (POST /v1/files with " "custom_llm_provider=vertex_ai), which encodes the model into the returned file id, or " "pass a uri of the form " "gs:////publishers//models//" ), ) return model @classmethod def _parse_model_from_gcs_file(cls, gcs_file_uri: str) -> str | None: """ Returns the `publishers//models/` path from a gcs uri, or None if the uri does not contain one. """ _, separator, model_path = unquote(gcs_file_uri).partition("publishers/") if not separator: return None parts: Final = model_path.split("/") if len(parts) < 3 or parts[1] != "models" or not parts[2]: return None return f"publishers/{'/'.join(parts[:3])}" @classmethod def is_unmanaged_gcs_batch_input_file_id(cls, input_file_id: str | None) -> bool: """ Returns True if `input_file_id` is a raw gs:// Vertex batch input file (i.e. not a LiteLLM-managed unified file id) with a `publishers/` model path that `_get_model_from_gcs_file` can parse. """ return ( input_file_id is not None and input_file_id.startswith("gs://") and cls._parse_model_from_gcs_file(input_file_id) is not None ) @classmethod def get_bare_model_name_from_gcs_file(cls, gcs_file_uri: str) -> str: """ Extracts the bare model name (e.g. "gemini-1.5-flash-001") from a gcs file uri. """ return cls._get_model_from_gcs_file(gcs_file_uri).rsplit("/", 1)[-1]