diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 3db94211032..b7f91bfba0d 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -7,6 +7,7 @@ import re import time from collections.abc import Callable, Iterable, Iterator, Mapping from typing import Any, Final, TypedDict +from urllib.parse import quote, unquote import httpx from httpx import Headers, Response @@ -43,6 +44,9 @@ from litellm.llms.vertex_ai.gemini.transformation import _transform_request_body from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) +from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import ( + transform_openai_input_gemini_embed_content, +) from litellm.types.files import StreamingMediaUploadConfig from litellm.types.llms.openai import ( AllMessageValues, @@ -54,14 +58,28 @@ from litellm.types.llms.openai import ( OpenAIFilesPurpose, PathLike, ) -from litellm.types.llms.vertex_ai import GcsBucketResponse -from litellm.types.utils import LlmProviders, ModelResponse +from litellm.types.llms.vertex_ai import GcsBucketResponse, GeminiEmbeddingInput +from litellm.types.utils import ( + Embedding, + EmbeddingResponse, + LlmProviders, + ModelResponse, + Usage, +) from ..common_utils import VertexAIError from ..vertex_llm_base import VertexBase _GCP_LABEL_VALUE_MAX_LEN: Final = 63 _CUSTOM_ID_RAW_LABEL_PREFIX: Final = "b32_" +_VERTEX_BATCH_KEY_FIELD: Final = "key" +_MANAGED_GCS_MODEL_PATH_PATTERN: Final = re.compile(r"publishers/[^/]+/models/([^/?]+)") +_EMBED_REQUEST_FIELD_BY_GEMINI_PARAM: Final = ( + ("outputDimensionality", "output_dimensionality"), + ("taskType", "task_type"), + ("title", "title"), +) +_VERTEX_BATCH_FANNED_OUT_KEY_PATTERN: Final = re.compile(r"(?P[^#]*)#(?P\d+)/(?P\d+)") class _GcsObjectMetadataJson(TypedDict, total=False): @@ -164,8 +182,26 @@ def _set_litellm_batch_custom_id_labels(labels: dict[str, str], custom_id: objec labels[f"litellm_custom_id_raw_{index}"] = raw_label_chunk -def _get_litellm_batch_custom_id_from_labels(labels: Mapping[str, object]) -> str: +def _get_litellm_batch_custom_id(vertex_output_row: Mapping[str, object]) -> str: + """ + Resolve the OpenAI `custom_id` for a Vertex batch output row. + + Embedding rows carry it in the top-level `key` field that Vertex echoes back; + `generateContent` rows have no such field, so it is smuggled through request + labels instead (see `_set_litellm_batch_custom_id_labels`). + """ + key = vertex_output_row.get(_VERTEX_BATCH_KEY_FIELD) + if key is not None: + return unquote(str(key)) + request_data = vertex_output_row.get("request") + labels = request_data.get("labels") if isinstance(request_data, Mapping) else None + return _get_litellm_batch_custom_id_from_labels(labels) + + +def _get_litellm_batch_custom_id_from_labels(labels: Mapping[str, object] | None) -> str: """Prefer encoded custom_id when present (see _set_litellm_batch_custom_id_labels).""" + if not labels: + return "unknown" raw: Final = labels.get("litellm_custom_id_raw") if raw: raw_chunks: Final = [str(raw)] @@ -182,17 +218,311 @@ def _get_litellm_batch_custom_id_from_labels(labels: Mapping[str, object]) -> st return str(labels.get("litellm_custom_id", "unknown")) -def _openai_batch_jsonl_entry_to_vertex_wrapped_request( +def _is_vertex_embeddings_batch_output_row(vertex_output_row: Mapping[str, Any]) -> bool: + """ + Whether a Vertex batch output row came from an `EmbedContentRequest`. + + Successful rows hold the vector under `response.embedding.values`; failed rows only + carry `status`, so they are recognized from the singular `content` that the + embeddings request shape echoes back. + """ + if "request" not in vertex_output_row: + return False + response = vertex_output_row.get("response") + if isinstance(response, dict) and isinstance(response.get("embedding"), dict): + return True + request_data = vertex_output_row.get("request") + return bool(vertex_output_row.get("status")) and isinstance(request_data, dict) and "content" in request_data + + +def _openai_batch_output_row( + custom_id: str, + body: Mapping[str, Any] | None = None, + error_code: str | None = None, + error_message: str = "", +) -> _OpenAIBatchOutputRow: + """ + One row of an OpenAI batch output file. Per the OpenAI Batch spec, failed rows set + `response` to null and populate `error` instead. + """ + return { + "id": f"batch_req_{uuid.uuid4()}", + "custom_id": custom_id, + "response": None + if body is None + else { + "status_code": 200, + "request_id": body.get("id", ""), + "body": body, + }, + "error": None if error_code is None else {"code": error_code, "message": error_message}, + } + + +def _split_vertex_batch_key(vertex_output_row: Mapping[str, Any]) -> tuple[str, int, int]: + """ + Resolve `(custom_id, index within that custom_id, group size)` for a Vertex batch + output row. + + A `/v1/embeddings` entry whose `input` is an array fans out into one Vertex row per + element, tagged `#/` (see + `_vertex_batch_embeddings_key`), so the rows can be reassembled into a single OpenAI + response. + """ + key = vertex_output_row.get(_VERTEX_BATCH_KEY_FIELD) + if key is None: + return _get_litellm_batch_custom_id(vertex_output_row), 0, 1 + match = _VERTEX_BATCH_FANNED_OUT_KEY_PATTERN.fullmatch(str(key)) + if match is None: + return unquote(str(key)), 0, 1 + return unquote(match["custom_id"]), int(match["index"]), int(match["total"]) + + +def _embedding_prompt_token_count(vertex_response: Mapping[str, Any]) -> int: + """ + Prompt tokens billed for one Vertex Gemini Embedding batch row. + + Live rows report usage under `usageMetadata`; the documented `tokenCount` is kept as + a fallback. + """ + usage_metadata = vertex_response.get("usageMetadata") + if isinstance(usage_metadata, Mapping): + return int(usage_metadata.get("promptTokenCount") or 0) + return int(vertex_response.get("tokenCount") or 0) + + +def _vertex_embeddings_rows_to_openai_batch_output_row( + custom_id: str, + vertex_output_rows: tuple[Mapping[str, Any], ...], + element_indices: tuple[int, ...], + element_count: int, + model: str | None, +) -> _OpenAIBatchOutputRow: + """ + Transforms the Vertex Gemini Embedding batch output rows belonging to one OpenAI + batch entry into an OpenAI batch output row holding an `/v1/embeddings` response. + + Example Vertex jsonl + {"key": "id_1", "request": {...}, "response": {"embedding": {"values": [-0.015, 0.024]}, "usageMetadata": {"promptTokenCount": 2}}} + + An entry that asked for several embeddings at once maps to several rows here, which + become the indexed elements of a single `data` array. One failed or missing element + fails the whole entry, since an OpenAI batch row is either a response or an error and + a partial `data` array would silently shift the remaining embeddings onto the wrong + input positions. Rows carry no `modelVersion`, so the model comes from the batch they + belong to. + """ + status = next((row["status"] for row in vertex_output_rows if row.get("status")), "") + if status: + return _openai_batch_output_row( + custom_id=custom_id, + error_code="vertex_ai_error", + error_message=status, + ) + + if element_indices != tuple(range(element_count)): + return _openai_batch_output_row( + custom_id=custom_id, + error_code="vertex_ai_error", + error_message=( + f"Vertex returned embeddings for input positions {list(element_indices)} " + f"of the {element_count} requested" + ), + ) + + responses = tuple(row["response"] for row in vertex_output_rows) + token_count = sum(_embedding_prompt_token_count(response) for response in responses) + body = EmbeddingResponse( + model=model or "", + data=[ + Embedding( + embedding=response["embedding"]["values"], + index=index, + object="embedding", + ) + for index, response in enumerate(responses) + ], + usage=Usage(prompt_tokens=token_count, total_tokens=token_count), + ).model_dump() + return _openai_batch_output_row(custom_id=custom_id, body=body) + + +def _transform_vertex_embeddings_batch_output_to_openai( + vertex_output_rows: Iterable[Mapping[str, Any]], + model: str | None, +) -> tuple[_OpenAIBatchOutputRow, ...]: + """ + Transforms a whole Vertex Gemini Embedding batch output into OpenAI batch output + rows, one per OpenAI batch entry, in the order the entries first appear. + + Rows are grouped rather than mapped one to one because a single entry can fan out + into several Vertex rows, and Vertex returns them in arbitrary order. + """ + keyed_rows = tuple((_split_vertex_batch_key(row), row) for row in vertex_output_rows) + grouped_rows = { + custom_id: tuple(group) + for custom_id, group in itertools.groupby(sorted(keyed_rows, key=lambda kr: kr[0]), key=lambda kr: kr[0][0]) + } + return tuple( + _vertex_embeddings_rows_to_openai_batch_output_row( + custom_id=custom_id, + vertex_output_rows=tuple(row for _, row in grouped_rows[custom_id]), + element_indices=tuple(index for (_, index, _), _ in grouped_rows[custom_id]), + element_count=max(total for (_, _, total), _ in grouped_rows[custom_id]), + model=model, + ) + for custom_id in dict.fromkeys(custom_id for (custom_id, _, _), _ in keyed_rows) + ) + + +def _model_from_managed_gcs_url(url: str) -> str | None: + """ + Extracts the model from a LiteLLM-managed Vertex batch GCS url. + + Batch inputs and their sibling outputs are stored under + `.../publishers/google/models//...`, which is the only place the model of an + embeddings batch output row can be recovered from; unlike `generateContent` + responses, embedding rows carry no `modelVersion`. + """ + match = _MANAGED_GCS_MODEL_PATH_PATTERN.search(unquote(url)) + return match.group(1) if match else None + + +def _is_embeddings_batch_entry(openai_entry: Mapping[str, Any]) -> bool: + """ + Whether an OpenAI batch JSONL line targets the embeddings endpoint. + + OpenAI puts the target route on each line's `url` (e.g. `/v1/embeddings`); Vertex + has no equivalent per-line field, so the route decides which Vertex request shape + the line has to be translated into. + """ + url = openai_entry.get("url") + if not isinstance(url, str): + return False + path = url.split("?")[0].rstrip("/") + return path == "embeddings" or path.endswith("/embeddings") + + +def _openai_embedding_input_elements( + embedding_input: GeminiEmbeddingInput, +) -> tuple[str | list[str], ...]: + """ + Split an OpenAI `input` into the elements that each get their own embedding. + + A string is one embedding, a flat array is one embedding per element, and a nested + array is one combined embedding per inner array, matching the online + `batchEmbedContents` path. + """ + if isinstance(embedding_input, list): + return tuple(embedding_input) + return (embedding_input,) + + +def _vertex_batch_embeddings_key(custom_id: str, index: int, total: int) -> str: + """ + The top-level `key` Vertex echoes back on an embeddings row. + + An entry asking for several embeddings needs several Vertex rows, so its key also + carries the element index and the group size; `_split_vertex_batch_key` reads them + back out. The `custom_id` is percent-encoded so that a customer one ending in + `#/` cannot be mistaken for that tag, which would merge two entries. + """ + encoded_custom_id = quote(custom_id, safe="") + return encoded_custom_id if total < 2 else f"{encoded_custom_id}#{index}/{total}" + + +def _vertex_embeddings_row(key: str | None, embed_content_request: Mapping[str, Any]) -> Mapping[str, Any]: + """ + One Vertex Gemini Embedding batch input row. + + The config fields live inside the `EmbedContentRequest` under their snake_case batch + names, and the OpenAI `custom_id` rides along in the top-level `key` that Vertex + echoes back. + """ + request = { + "content": embed_content_request["content"], + **{ + request_field: embed_content_request[gemini_param] + for gemini_param, request_field in _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM + if gemini_param in embed_content_request + }, + } + if key is None: + return {"request": request} + return {_VERTEX_BATCH_KEY_FIELD: key, "request": request} + + +def _openai_batch_jsonl_entry_to_vertex_embeddings_rows( + openai_entry: Mapping[str, Any], +) -> tuple[Mapping[str, Any], ...]: + """ + Transforms a single OpenAI `/v1/embeddings` batch entry into Vertex Gemini Embedding + batch rows, one per requested embedding. + + Example Vertex jsonl + {"key": "id_1", "request": {"content": {"parts": [{"text": "Hello World"}]}, "output_dimensionality": 768, "task_type": "RETRIEVAL_DOCUMENT"}} + + Note that `content` is singular (an `EmbedContentRequest`, not a + `GenerateContentRequest`) and that the `custom_id` round-trips through the top-level + `key`. An `EmbedContentRequest` returns exactly one vector, so an entry whose `input` + is an array fans out into one row per element and is reassembled on the way back. + The docs put the per-row config in an `embed_content_config` sibling of `request`, + but the API rejects that key outright and fails the whole batch job, so the config + fields go inside the `EmbedContentRequest` itself. + + API Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/batch-prediction-genai-embeddings + """ + openai_request_body = openai_entry.get("body") + if not isinstance(openai_request_body, dict): + raise TypeError( + "`body` on /v1/embeddings batch requests must be a JSON object, but was missing or not an object" + ) + embedding_input = openai_request_body.get("input") + if embedding_input is None: + raise ValueError("`input` is required on /v1/embeddings batch requests, but was not provided") + + elements = _openai_embedding_input_elements(embedding_input) + if not elements: + raise ValueError("`input` on /v1/embeddings batch requests must not be empty") + + embed_content_requests = tuple( + transform_openai_input_gemini_embed_content( + input=element, + model=openai_request_body.get("model", ""), + optional_params=openai_request_body, + ) + for element in elements + ) + custom_id = openai_entry.get("custom_id") + return tuple( + _vertex_embeddings_row( + key=None + if custom_id is None + else _vertex_batch_embeddings_key( + custom_id=str(custom_id), + index=index, + total=len(embed_content_requests), + ), + embed_content_request=embed_content_request, + ) + for index, embed_content_request in enumerate(embed_content_requests) + ) + + +def _openai_batch_jsonl_entry_to_vertex_rows( openai_entry: dict[str, Any], map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, Any]], -) -> dict[str, Any]: +) -> tuple[Mapping[str, Any], ...]: """ - Transforms a single OpenAI JSONL batch entry into its Vertex wrapped request. + Transforms a single OpenAI JSONL batch entry into the Vertex rows it maps to. jsonl body for vertex is {"request": } Example Vertex jsonl {"request":{"contents": [{"role": "user", "parts": [{"text": "What is the relation between the following video and image samples?"}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/animals.mp4", "mimeType": "video/mp4"}}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/image/cricket.jpeg", "mimeType": "image/jpeg"}}]}]}} """ + if _is_embeddings_batch_entry(openai_entry): + return _openai_batch_jsonl_entry_to_vertex_embeddings_rows(openai_entry) + openai_request_body: Final = openai_entry.get("body") or {} vertex_request_body: Final = _transform_request_body( messages=openai_request_body.get("messages", []), @@ -209,7 +539,7 @@ def _openai_batch_jsonl_entry_to_vertex_wrapped_request( vertex_request_body["labels"] = {} _set_litellm_batch_custom_id_labels(vertex_request_body["labels"], custom_id) - return {"request": vertex_request_body} + return ({"request": vertex_request_body},) def _iter_stripped_lines(raw_lines: Iterable[str | bytes]) -> Iterator[str]: @@ -312,10 +642,10 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): def _iter_vertex_jsonl_chunks(self) -> Iterator[bytes]: first = True for entry in _iter_openai_jsonl_entries(self._openai_file_content): - wrapped = _openai_batch_jsonl_entry_to_vertex_wrapped_request(entry, self._map_openai_to_vertex_params) - prefix = b"" if first else b"\n" - first = False - yield prefix + json.dumps(wrapped).encode("utf-8") + for wrapped in _openai_batch_jsonl_entry_to_vertex_rows(entry, self._map_openai_to_vertex_params): + prefix = b"" if first else b"\n" + first = False + yield prefix + json.dumps(wrapped).encode("utf-8") def iter_bytes(self) -> Iterator[bytes]: return self._iter_vertex_jsonl_chunks() @@ -667,6 +997,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): transformed_content: Final = self._try_transform_vertex_batch_output_to_openai( content=content, logging_obj=logging_obj, + model=_model_from_managed_gcs_url(str(raw_response.request.url)), ) if transformed_content != content: # Create a new response with transformed content and updated Content-Length @@ -688,7 +1019,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): return HttpxBinaryResponseContent(response=raw_response) def _try_transform_vertex_batch_output_to_openai( - self, content: bytes, logging_obj: LiteLLMLoggingObj | None = None + self, + content: bytes, + logging_obj: LiteLLMLoggingObj | None = None, + model: str | None = None, ) -> bytes: """ Try to transform Vertex AI batch output to OpenAI format. @@ -730,7 +1064,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): # first line is not valid UTF-8/JSON) raises and falls through to the # passthrough below, leaving the content untouched. first_row: Final = _parse_vertex_batch_output_row(first_line) - is_vertex_batch_output: Final = ( + is_vertex_batch_output: Final = _is_vertex_embeddings_batch_output_row(first_row) or ( "request" in first_row and "response" in first_row and "processed_time" in first_row @@ -763,11 +1097,23 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): request=httpx.Request(method="POST", url="https://example.com"), ) + all_lines = itertools.chain((first_line,), lines) + + # Embedding rows are grouped by `custom_id` rather than transformed one at a + # time, since an entry that asked for several embeddings comes back as + # several rows, in arbitrary order. + if _is_vertex_embeddings_batch_output_row(first_row): + openai_outputs = _transform_vertex_embeddings_batch_output_to_openai( + vertex_output_rows=(json.loads(line) for line in all_lines), + model=model, + ) + return b"\n".join(json.dumps(openai_output).encode("utf-8") for openai_output in openai_outputs) + # Transform each row straight into the output buffer, so peak memory # stays at ~one row plus the output. If any row fails, return the # original content unchanged. output = bytearray() - for line in itertools.chain([first_line], lines): + for line in all_lines: try: openai_output = self._transform_single_vertex_batch_output_to_openai( vertex_output=_parse_vertex_batch_output_row(line), @@ -798,25 +1144,18 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): Transform a single Vertex AI batch output line to OpenAI format. Uses the existing VertexGeminiConfig transformation for the response. """ - # Extract custom_id from request labels (prefer raw for OpenAI round-trip) - request_data: Final = vertex_output.get("request", {}) - labels: Final[Mapping[str, object]] = request_data.get("labels", {}) or {} - custom_id: Final = _get_litellm_batch_custom_id_from_labels(labels) + custom_id: Final = _get_litellm_batch_custom_id(vertex_output) # Check if there's an error status: Final = vertex_output.get("status", "") has_error: Final = bool(status) if has_error: - return { - "id": f"batch_req_{uuid.uuid4()}", - "custom_id": custom_id, - "response": None, - "error": { - "code": "vertex_ai_error", - "message": status, - }, - } + return _openai_batch_output_row( + custom_id=custom_id, + error_code="vertex_ai_error", + error_message=status, + ) # Transform successful response using existing transformation vertex_response: Final = vertex_output.get("response", {}) @@ -842,24 +1181,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): response_dict: Final = transformed_response.model_dump() # Return in OpenAI batch format - return { - "id": f"batch_req_{uuid.uuid4()}", - "custom_id": custom_id, - "response": { - "status_code": 200, - "request_id": response_dict.get("id", ""), - "body": response_dict, - }, - "error": None, - } + return _openai_batch_output_row(custom_id=custom_id, body=response_dict) except Exception as e: - return { - "id": f"batch_req_{uuid.uuid4()}", - "custom_id": custom_id, - "response": None, - "error": { - "code": "transformation_error", - "message": f"Failed to transform response: {e}", - }, - } + return _openai_batch_output_row( + custom_id=custom_id, + error_code="transformation_error", + error_message=f"Failed to transform response: {e}", + ) diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py index 2e3280c0ed1..957fc7dbcf4 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py @@ -37,7 +37,7 @@ from litellm.llms.vertex_ai.files.transformation import ( _get_litellm_batch_custom_id_from_labels, _iter_openai_jsonl_entries, _iter_openai_jsonl_lines, - _openai_batch_jsonl_entry_to_vertex_wrapped_request, + _openai_batch_jsonl_entry_to_vertex_rows, ) from litellm.types.llms.openai import CreateFileRequest @@ -84,8 +84,9 @@ def _reference_vertex_jsonl_string(cfg: VertexAIFilesConfig, content: str) -> st transform, so the streaming path can be checked against it for parity.""" entries = [json.loads(line) for line in content.splitlines() if line.strip()] return "\n".join( - json.dumps(_openai_batch_jsonl_entry_to_vertex_wrapped_request(entry, cfg._map_openai_to_vertex_params)) + json.dumps(row) for entry in entries + for row in _openai_batch_jsonl_entry_to_vertex_rows(entry, cfg._map_openai_to_vertex_params) ) diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 8c5305ee67b..3c2d56997b7 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -15,7 +15,7 @@ from unittest.mock import MagicMock from litellm.llms.vertex_ai.files.transformation import ( VertexAIFilesConfig, _get_litellm_batch_custom_id_from_labels, - _openai_batch_jsonl_entry_to_vertex_wrapped_request, + _openai_batch_jsonl_entry_to_vertex_rows, _sanitize_gcp_label_value, ) from litellm.types.llms.openai import OpenAIFileObject, HttpxBinaryResponseContent @@ -32,40 +32,26 @@ class TestParseGcsUri: def test_should_parse_standard_gs_uri(self, config): file_id = "gs://my-bucket/litellm-vertex-files/path/to/object.jsonl" - bucket, encoded = config._parse_gcs_uri( - file_id, litellm_params={"gcs_bucket_name": "my-bucket"} - ) + bucket, encoded = config._parse_gcs_uri(file_id, litellm_params={"gcs_bucket_name": "my-bucket"}) assert bucket == "my-bucket" - assert encoded == urllib.parse.quote( - "litellm-vertex-files/path/to/object.jsonl", safe="" - ) + assert encoded == urllib.parse.quote("litellm-vertex-files/path/to/object.jsonl", safe="") def test_should_parse_uri_with_nested_publisher_path(self, config): uri = "gs://litellm-local/litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc-123" - bucket, encoded = config._parse_gcs_uri( - uri, litellm_params={"gcs_bucket_name": "litellm-local"} - ) + bucket, encoded = config._parse_gcs_uri(uri, litellm_params={"gcs_bucket_name": "litellm-local"}) assert bucket == "litellm-local" - expected_path = ( - "litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc-123" - ) + expected_path = "litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc-123" assert encoded == urllib.parse.quote(expected_path, safe="") def test_should_handle_url_encoded_input(self, config): - encoded_uri = urllib.parse.quote( - "gs://my-bucket/litellm-vertex-files/some/path", safe="" - ) - bucket, encoded = config._parse_gcs_uri( - encoded_uri, litellm_params={"gcs_bucket_name": "my-bucket"} - ) + encoded_uri = urllib.parse.quote("gs://my-bucket/litellm-vertex-files/some/path", safe="") + bucket, encoded = config._parse_gcs_uri(encoded_uri, litellm_params={"gcs_bucket_name": "my-bucket"}) assert bucket == "my-bucket" assert encoded == urllib.parse.quote("litellm-vertex-files/some/path", safe="") def test_should_reject_bucket_only(self, config): with pytest.raises(ValueError, match="object name"): - config._parse_gcs_uri( - "gs://my-bucket", litellm_params={"gcs_bucket_name": "my-bucket"} - ) + config._parse_gcs_uri("gs://my-bucket", litellm_params={"gcs_bucket_name": "my-bucket"}) def test_should_reject_no_gs_prefix(self, config): with pytest.raises(ValueError, match="gs://"): @@ -110,9 +96,7 @@ class TestParseGcsUri: "gs://my-bucket/private/object.txt", litellm_params={ "gcs_bucket_name": "my-bucket", - "_litellm_internal_model_credentials": { - "allow_legacy_cloud_file_ids": True - }, + "_litellm_internal_model_credentials": {"allow_legacy_cloud_file_ids": True}, }, ) @@ -176,7 +160,6 @@ class TestCreateFileUrl: class TestTransformRetrieveFile: - def test_should_build_correct_gcs_metadata_url(self, config): file_id = "gs://my-bucket/litellm-vertex-files/path/to/file.jsonl" url, params = config.transform_retrieve_file_request( @@ -184,13 +167,8 @@ class TestTransformRetrieveFile: optional_params={}, litellm_params={"gcs_bucket_name": "my-bucket"}, ) - expected_encoded = urllib.parse.quote( - "litellm-vertex-files/path/to/file.jsonl", safe="" - ) - assert ( - url - == f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{expected_encoded}" - ) + expected_encoded = urllib.parse.quote("litellm-vertex-files/path/to/file.jsonl", safe="") + assert url == f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{expected_encoded}" assert params == {} def test_should_return_openai_file_object_from_gcs_response(self, config): @@ -237,7 +215,6 @@ class TestTransformRetrieveFile: class TestTransformFileContent: - def test_should_build_gcs_media_download_url(self, config): file_id = "gs://my-bucket/litellm-vertex-files/path/to/file.jsonl" url, params = config.transform_file_content_request( @@ -246,10 +223,7 @@ class TestTransformFileContent: litellm_params={"gcs_bucket_name": "my-bucket"}, ) encoded = urllib.parse.quote("litellm-vertex-files/path/to/file.jsonl", safe="") - assert ( - url - == f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{encoded}?alt=media" - ) + assert url == f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{encoded}?alt=media" assert params == {} def test_should_return_binary_response_content(self, config): @@ -269,9 +243,7 @@ class TestTransformFileContent: assert isinstance(result, HttpxBinaryResponseContent) assert result.response.content == b'{"line": 1}\n{"line": 2}\n' - def test_should_not_mutate_caller_logging_obj_for_batch_output_transform( - self, config, monkeypatch - ): + def test_should_not_mutate_caller_logging_obj_for_batch_output_transform(self, config, monkeypatch): original_model = "vertex_ai/original-model" original_start_time = 123.456 original_optional_params = {"temperature": 0.1} @@ -283,9 +255,7 @@ class TestTransformFileContent: "processed_time": "2024-11-01T18:13:16.826+00:00", "request": {"labels": {"litellm_custom_id": "request-1"}}, "response": { - "candidates": [ - {"content": {"parts": [{"text": "ok"}], "role": "model"}} - ], + "candidates": [{"content": {"parts": [{"text": "ok"}], "role": "model"}}], "modelVersion": "gemini-2.0-flash-001@default", }, } @@ -308,9 +278,7 @@ class TestTransformFileContent: captured["logging_obj"] = logging_obj logging_obj.model = "gemini-2.0-flash-001" logging_obj.start_time = 789.0 - return { - "custom_id": vertex_output["request"]["labels"]["litellm_custom_id"] - } + return {"custom_id": vertex_output["request"]["labels"]["litellm_custom_id"]} monkeypatch.setattr( config, @@ -330,9 +298,7 @@ class TestTransformFileContent: assert logging_obj.optional_params == original_optional_params assert result.response is not raw_response - def test_should_skip_batch_output_transformation_when_opt_out_flag_set( - self, config, monkeypatch - ): + def test_should_skip_batch_output_transformation_when_opt_out_flag_set(self, config, monkeypatch): """When `litellm.disable_vertex_batch_output_transformation` is True the Vertex predictions.jsonl content must be returned untouched, so callers that parse raw `candidates`/`modelVersion` keep working.""" @@ -344,9 +310,7 @@ class TestTransformFileContent: "processed_time": "2024-11-01T18:13:16.826+00:00", "request": {"labels": {"litellm_custom_id": "request-1"}}, "response": { - "candidates": [ - {"content": {"parts": [{"text": "ok"}], "role": "model"}} - ], + "candidates": [{"content": {"parts": [{"text": "ok"}], "role": "model"}}], "modelVersion": "gemini-2.0-flash-001@default", }, } @@ -358,9 +322,7 @@ class TestTransformFileContent: request=httpx.Request("GET", "https://example.com"), ) - monkeypatch.setattr( - litellm, "disable_vertex_batch_output_transformation", True, raising=False - ) + monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False) result = config.transform_file_content_response( raw_response=raw_response, @@ -381,9 +343,7 @@ class TestTransformDeleteFile: litellm_params={"gcs_bucket_name": "my-bucket"}, ) encoded = urllib.parse.quote("litellm-vertex-files/path/to/file.jsonl", safe="") - assert ( - url == f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{encoded}" - ) + assert url == f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{encoded}" assert params == {} def test_should_return_file_deleted_with_reconstructed_id(self, config): @@ -393,9 +353,7 @@ class TestTransformDeleteFile: "litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc", safe="", ) - mock_request.url = ( - f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{encoded_name}" - ) + mock_request.url = f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{encoded_name}" raw_response.request = mock_request result = config.transform_delete_file_response( @@ -407,10 +365,7 @@ class TestTransformDeleteFile: assert isinstance(result, FileDeleted) assert result.deleted is True assert result.object == "file" - assert ( - result.id - == "gs://my-bucket/litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc" - ) + assert result.id == "gs://my-bucket/litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc" def test_should_fallback_to_deleted_id_when_no_request(self, config): raw_response = MagicMock(spec=httpx.Response) @@ -435,9 +390,7 @@ class TestTransformDeleteFile: raw_response = MagicMock(spec=httpx.Response) mock_request = MagicMock() encoded_object = urllib.parse.quote("path/to/file.jsonl", safe="") - mock_request.url = ( - f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{encoded_object}" - ) + mock_request.url = f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{encoded_object}" raw_response.request = mock_request result = config.transform_delete_file_response( @@ -466,8 +419,7 @@ class TestTransformDeleteFile: ) assert result.id == ( - "gs://prod-bucket/litellm-vertex-files/publishers/google/" - "models/gemini-2.0-flash-001/abc-123" + "gs://prod-bucket/litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc-123" ) @@ -504,9 +456,7 @@ class TestVertexBatchOutputTransformation: } content = json.dumps(vertex_output).encode("utf-8") - transformed_content = config._try_transform_vertex_batch_output_to_openai( - content - ) + transformed_content = config._try_transform_vertex_batch_output_to_openai(content) result = json.loads(transformed_content.decode("utf-8")) # Verify OpenAI format @@ -548,9 +498,7 @@ class TestVertexBatchOutputTransformation: } content = json.dumps(vertex_output).encode("utf-8") - transformed_content = config._try_transform_vertex_batch_output_to_openai( - content - ) + transformed_content = config._try_transform_vertex_batch_output_to_openai(content) result = json.loads(transformed_content.decode("utf-8")) # Per OpenAI Batch output spec, error entries set response to null @@ -584,9 +532,7 @@ class TestVertexBatchOutputTransformation: } class _RaisingGeminiConfig(VertexGeminiConfig): - def _transform_google_generate_content_to_openai_model_response( - self, *args, **kwargs - ): + def _transform_google_generate_content_to_openai_model_response(self, *args, **kwargs): raise ValueError("simulated transform failure") mock_response = httpx.Response( @@ -637,9 +583,7 @@ class TestVertexBatchOutputTransformation: } content = json.dumps(vertex_output).encode("utf-8") - transformed_content = config._try_transform_vertex_batch_output_to_openai( - content - ) + transformed_content = config._try_transform_vertex_batch_output_to_openai(content) result = json.loads(transformed_content.decode("utf-8")) assert result["custom_id"] == "myrequest-1" @@ -651,9 +595,7 @@ class TestVertexBatchOutputTransformation: "status": "", "processed_time": "2024-11-01T18:13:16.826+00:00", "request": { - "contents": [ - {"role": "user", "parts": [{"text": "First request"}]} - ], + "contents": [{"role": "user", "parts": [{"text": "First request"}]}], "labels": {"litellm_custom_id": "request-1"}, }, "response": { @@ -678,9 +620,7 @@ class TestVertexBatchOutputTransformation: "status": "", "processed_time": "2024-11-01T18:13:17.826+00:00", "request": { - "contents": [ - {"role": "user", "parts": [{"text": "Second request"}]} - ], + "contents": [{"role": "user", "parts": [{"text": "Second request"}]}], "labels": {"litellm_custom_id": "request-2"}, }, "response": { @@ -703,12 +643,8 @@ class TestVertexBatchOutputTransformation: }, ] - content = "\n".join(json.dumps(output) for output in vertex_outputs).encode( - "utf-8" - ) - transformed_content = config._try_transform_vertex_batch_output_to_openai( - content - ) + content = "\n".join(json.dumps(output) for output in vertex_outputs).encode("utf-8") + transformed_content = config._try_transform_vertex_batch_output_to_openai(content) lines = transformed_content.decode("utf-8").strip().split("\n") assert len(lines) == 2 @@ -718,14 +654,12 @@ class TestVertexBatchOutputTransformation: assert "id" in result assert "response" in result assert result["response"]["status_code"] == 200 - assert result["custom_id"] == f"request-{i+1}" + assert result["custom_id"] == f"request-{i + 1}" body = result["response"]["body"] assert "choices" in body assert len(body["choices"]) > 0 - def test_transform_vertex_batch_output_with_first_line_prompt_feedback( - self, config, monkeypatch - ): + def test_transform_vertex_batch_output_with_first_line_prompt_feedback(self, config, monkeypatch): """Test that promptFeedback-only first lines are detected as Vertex batch output.""" vertex_outputs = [ { @@ -751,9 +685,7 @@ class TestVertexBatchOutputTransformation: logging_obj, mock_httpx_response, ): - return { - "custom_id": vertex_output["request"]["labels"]["litellm_custom_id"] - } + return {"custom_id": vertex_output["request"]["labels"]["litellm_custom_id"]} monkeypatch.setattr( config, @@ -761,15 +693,9 @@ class TestVertexBatchOutputTransformation: mock_transform_single, ) - content = "\n".join(json.dumps(output) for output in vertex_outputs).encode( - "utf-8" - ) - transformed_content = config._try_transform_vertex_batch_output_to_openai( - content - ) - results = [ - json.loads(line) for line in transformed_content.decode("utf-8").split("\n") - ] + content = "\n".join(json.dumps(output) for output in vertex_outputs).encode("utf-8") + transformed_content = config._try_transform_vertex_batch_output_to_openai(content) + results = [json.loads(line) for line in transformed_content.decode("utf-8").split("\n")] assert [result["custom_id"] for result in results] == [ "blocked-request", @@ -786,9 +712,7 @@ class TestVertexBatchOutputTransformation: } content = json.dumps(non_batch_output).encode("utf-8") - transformed_content = config._try_transform_vertex_batch_output_to_openai( - content - ) + transformed_content = config._try_transform_vertex_batch_output_to_openai(content) assert transformed_content == content @@ -818,9 +742,7 @@ class TestVertexBatchOutputTransformation: id(mock_httpx_response), ) ) - return { - "custom_id": vertex_output["request"]["labels"]["litellm_custom_id"] - } + return {"custom_id": vertex_output["request"]["labels"]["litellm_custom_id"]} monkeypatch.setattr( config, @@ -828,12 +750,8 @@ class TestVertexBatchOutputTransformation: mock_transform_single, ) - content = "\n".join(json.dumps(output) for output in vertex_outputs).encode( - "utf-8" - ) - transformed_content = config._try_transform_vertex_batch_output_to_openai( - content - ) + content = "\n".join(json.dumps(output) for output in vertex_outputs).encode("utf-8") + transformed_content = config._try_transform_vertex_batch_output_to_openai(content) assert len(transformed_content.decode("utf-8").strip().split("\n")) == 2 assert len(set(helper_ids)) == 1 @@ -841,17 +759,13 @@ class TestVertexBatchOutputTransformation: def test_non_batch_output_passthrough(self, config): """Test that non-batch output is returned as-is""" regular_content = b"This is just a regular file content" - transformed_content = config._try_transform_vertex_batch_output_to_openai( - regular_content - ) + transformed_content = config._try_transform_vertex_batch_output_to_openai(regular_content) assert transformed_content == regular_content def test_invalid_json_passthrough(self, config): """Test that invalid JSON is returned as-is""" invalid_content = b'{"invalid": json content}' - transformed_content = config._try_transform_vertex_batch_output_to_openai( - invalid_content - ) + transformed_content = config._try_transform_vertex_batch_output_to_openai(invalid_content) assert transformed_content == invalid_content def test_binary_content_passthrough(self, config): @@ -903,9 +817,7 @@ class TestVertexBatchOutputTransformation: }, } - content = ("\n".join(json.dumps(vertex_row(i)) for i in range(4000))).encode( - "utf-8" - ) + content = ("\n".join(json.dumps(vertex_row(i)) for i in range(4000))).encode("utf-8") def list_pipeline() -> bytes: gemini_config = VertexGeminiConfig() @@ -944,9 +856,7 @@ class TestVertexBatchOutputTransformation: finally: tracemalloc.stop() - streaming_peak = peak_of( - lambda: config._try_transform_vertex_batch_output_to_openai(content) - ) + streaming_peak = peak_of(lambda: config._try_transform_vertex_batch_output_to_openai(content)) list_peak = peak_of(list_pipeline) assert streaming_peak < list_peak * 0.75, ( @@ -999,9 +909,9 @@ class TestTryTransformDoesNotMutateCallerLoggingObj: logging_obj=logging_obj, ) - assert ( - logging_obj.model == sentinel_model - ), "logging_obj.model was mutated by _try_transform_vertex_batch_output_to_openai" + assert logging_obj.model == sentinel_model, ( + "logging_obj.model was mutated by _try_transform_vertex_batch_output_to_openai" + ) def test_should_not_overwrite_start_time_on_caller_logging_obj(self, config): sentinel_start = 1234567890.0 @@ -1014,9 +924,9 @@ class TestTryTransformDoesNotMutateCallerLoggingObj: logging_obj=logging_obj, ) - assert ( - logging_obj.start_time == sentinel_start - ), "logging_obj.start_time was mutated by _try_transform_vertex_batch_output_to_openai" + assert logging_obj.start_time == sentinel_start, ( + "logging_obj.start_time was mutated by _try_transform_vertex_batch_output_to_openai" + ) def test_should_not_overwrite_optional_params_on_caller_logging_obj(self, config): sentinel_params = {"temperature": 0.5, "top_p": 0.9} @@ -1028,9 +938,9 @@ class TestTryTransformDoesNotMutateCallerLoggingObj: logging_obj=logging_obj, ) - assert ( - logging_obj.optional_params is sentinel_params - ), "logging_obj.optional_params was replaced by _try_transform_vertex_batch_output_to_openai" + assert logging_obj.optional_params is sentinel_params, ( + "logging_obj.optional_params was replaced by _try_transform_vertex_batch_output_to_openai" + ) assert logging_obj.optional_params == { "temperature": 0.5, "top_p": 0.9, @@ -1054,14 +964,13 @@ class TestTryTransformDoesNotMutateCallerLoggingObj: def _wrap_entries(openai_jsonl_content): - """Vertex-wrapped requests for a list of OpenAI batch entries, built via the - live single-entry transform that the streaming upload path uses.""" + """Vertex rows for a list of OpenAI batch entries, built via the live + single-entry transform that the streaming upload path uses.""" cfg = VertexAIFilesConfig() return [ - _openai_batch_jsonl_entry_to_vertex_wrapped_request( - entry, cfg._map_openai_to_vertex_params - ) + row for entry in openai_jsonl_content + for row in _openai_batch_jsonl_entry_to_vertex_rows(entry, cfg._map_openai_to_vertex_params) ] @@ -1122,9 +1031,7 @@ class TestVertexBatchCustomIdLabels: assert "litellm_custom_id_raw_1" in labels_a assert "litellm_custom_id_raw_1" in labels_b assert labels_a["litellm_custom_id_raw"] == labels_b["litellm_custom_id_raw"] - assert ( - labels_a["litellm_custom_id_raw_1"] != labels_b["litellm_custom_id_raw_1"] - ) + assert labels_a["litellm_custom_id_raw_1"] != labels_b["litellm_custom_id_raw_1"] assert _get_litellm_batch_custom_id_from_labels(labels_a) == custom_id_a assert _get_litellm_batch_custom_id_from_labels(labels_b) == custom_id_b @@ -1133,12 +1040,12 @@ class TestVertexBatchCustomIdLabels: openai_jsonl_content = [ { - "custom_id": f"request-{i+1}", + "custom_id": f"request-{i + 1}", "method": "POST", "url": "/v1/chat/completions", "body": { "model": "gemini-1.5-flash-001", - "messages": [{"role": "user", "content": f"Question {i+1}"}], + "messages": [{"role": "user", "content": f"Question {i + 1}"}], }, } for i in range(3) @@ -1149,11 +1056,8 @@ class TestVertexBatchCustomIdLabels: assert len(vertex_jsonl_content) == 3 for i, vertex_request in enumerate(vertex_jsonl_content): - expected_custom_id = f"request-{i+1}" - assert ( - vertex_request["request"]["labels"]["litellm_custom_id"] - == expected_custom_id - ) + expected_custom_id = f"request-{i + 1}" + assert vertex_request["request"]["labels"]["litellm_custom_id"] == expected_custom_id raw_label = vertex_request["request"]["labels"]["litellm_custom_id_raw"] assert raw_label != expected_custom_id assert _sanitize_gcp_label_value(raw_label) == raw_label @@ -1200,9 +1104,7 @@ class TestVertexBatchCustomIdLabels: vertex_input = _wrap_entries(openai_input) # Verify both labels are GCP-safe and encoded raw preserves round-trip. - assert ( - vertex_input[0]["request"]["labels"]["litellm_custom_id"] == "myrequest-1" - ) + assert vertex_input[0]["request"]["labels"]["litellm_custom_id"] == "myrequest-1" raw_label = vertex_input[0]["request"]["labels"]["litellm_custom_id_raw"] assert raw_label != "MyRequest-1" assert _sanitize_gcp_label_value(raw_label) == raw_label @@ -1230,9 +1132,7 @@ class TestVertexBatchCustomIdLabels: # Step 3: Transform Vertex AI output back to OpenAI format content = json.dumps(vertex_output).encode("utf-8") - transformed_content = config._try_transform_vertex_batch_output_to_openai( - content - ) + transformed_content = config._try_transform_vertex_batch_output_to_openai(content) openai_output = json.loads(transformed_content.decode("utf-8")) # Step 4: Verify custom_id was preserved (original casing, not sanitized label) @@ -1268,9 +1168,7 @@ class TestVertexBatchCustomIdLabels: vertex_input = _wrap_entries(openai_input) # Verify both labels are safe for GCP labels. - assert ( - vertex_input[0]["request"]["labels"]["litellm_custom_id"] == "myrequest-1" - ) + assert vertex_input[0]["request"]["labels"]["litellm_custom_id"] == "myrequest-1" raw_label = vertex_input[0]["request"]["labels"]["litellm_custom_id_raw"] assert raw_label != "MyRequest-1" assert _sanitize_gcp_label_value(raw_label) == raw_label @@ -1279,26 +1177,15 @@ class TestVertexBatchCustomIdLabels: class TestConfiguredBucketNameResolution: def test_should_resolve_new_gcs_bucket_name_key(self, config, monkeypatch): monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) - assert ( - config._get_configured_bucket_name({"gcs_bucket_name": "my-new-bucket"}) - == "my-new-bucket" - ) + assert config._get_configured_bucket_name({"gcs_bucket_name": "my-new-bucket"}) == "my-new-bucket" def test_should_resolve_legacy_bucket_name_key(self, config, monkeypatch): monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) - assert ( - config._get_configured_bucket_name({"bucket_name": "my-legacy-bucket"}) - == "my-legacy-bucket" - ) + assert config._get_configured_bucket_name({"bucket_name": "my-legacy-bucket"}) == "my-legacy-bucket" def test_should_prefer_new_key_over_legacy(self, config, monkeypatch): monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) - assert ( - config._get_configured_bucket_name( - {"gcs_bucket_name": "new", "bucket_name": "legacy"} - ) - == "new" - ) + assert config._get_configured_bucket_name({"gcs_bucket_name": "new", "bucket_name": "legacy"}) == "new" def test_should_fall_back_to_env(self, config, monkeypatch): monkeypatch.setenv("GCS_BUCKET_NAME", "env-bucket") @@ -1318,3 +1205,558 @@ class TestConfiguredBucketNameResolution: assert "bucket_name" in OPTIONAL_KWARGS_KEYS params = get_litellm_params(bucket_name="my-legacy-bucket") assert params.get("bucket_name") == "my-legacy-bucket" + + +def _embeddings_entry(**overrides): + entry = { + "custom_id": "request-1", + "method": "POST", + "url": "/v1/embeddings", + "body": {"model": "gemini-embedding-2", "input": "hello world"}, + } + entry.update(overrides) + return entry + + +class TestVertexEmbeddingsBatchInputTranslation: + """ + /v1/embeddings batch lines must be translated to Vertex's Gemini Embedding batch + shape, not the generateContent shape. + + Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/batch-prediction-genai-embeddings + """ + + def test_should_emit_embed_content_request_shape(self): + (row,) = _wrap_entries([_embeddings_entry()]) + + assert row["request"] == {"content": {"parts": [{"text": "hello world"}]}} + assert "contents" not in row["request"] + assert "labels" not in row["request"] + + def test_should_round_trip_custom_id_through_top_level_key(self): + (row,) = _wrap_entries([_embeddings_entry(custom_id="MyRequest-1")]) + + assert row["key"] == "MyRequest-1" + + def test_should_omit_key_when_no_custom_id(self): + entry = _embeddings_entry() + del entry["custom_id"] + + (row,) = _wrap_entries([entry]) + + assert "key" not in row + + def test_should_map_openai_params_into_the_embed_content_request(self): + """ + The docs put these in an `embed_content_config` sibling of `request`, but Vertex + rejects that key and fails the whole job, so they belong inside the request. + """ + (row,) = _wrap_entries( + [ + _embeddings_entry( + body={ + "model": "gemini-embedding-001", + "input": "hello world", + "dimensions": 768, + "task_type": "RETRIEVAL_DOCUMENT", + "title": "some_title", + } + ) + ] + ) + + assert row == { + "key": "request-1", + "request": { + "content": {"parts": [{"text": "hello world"}]}, + "output_dimensionality": 768, + "task_type": "RETRIEVAL_DOCUMENT", + "title": "some_title", + }, + } + + def test_should_omit_config_fields_when_no_params_given(self): + (row,) = _wrap_entries([_embeddings_entry()]) + + assert set(row["request"]) == {"content"} + + def test_should_translate_multimodal_gcs_input(self): + (row,) = _wrap_entries( + [ + _embeddings_entry( + body={ + "model": "gemini-embedding-2", + "input": "gs://cloud-samples-data/generative-ai/image/benchmark.jpeg", + } + ) + ] + ) + + assert row["request"]["content"]["parts"] == [ + { + "file_data": { + "mime_type": "image/jpeg", + "file_uri": "gs://cloud-samples-data/generative-ai/image/benchmark.jpeg", + } + } + ] + + @pytest.mark.parametrize("url", ["/v1/embeddings", "v1/embeddings", "/v1/embeddings/"]) + def test_should_detect_embeddings_route_variants(self, url): + (row,) = _wrap_entries([_embeddings_entry(url=url)]) + + assert "content" in row["request"] + + def test_should_raise_when_input_missing(self): + with pytest.raises(ValueError, match="`input` is required"): + _wrap_entries([_embeddings_entry(body={"model": "gemini-embedding-2"})]) + + def test_should_raise_when_input_empty(self): + with pytest.raises(ValueError, match="must not be empty"): + _wrap_entries([_embeddings_entry(body={"model": "gemini-embedding-2", "input": []})]) + + def test_should_fan_an_input_array_out_into_one_row_per_element(self): + """ + An `EmbedContentRequest` returns exactly one vector, so an OpenAI entry asking + for several embeddings needs several Vertex rows. + """ + rows = _wrap_entries( + [ + _embeddings_entry( + body={ + "model": "gemini-embedding-001", + "input": ["first", "second"], + "dimensions": 768, + } + ) + ] + ) + + assert rows == [ + { + "key": "request-1#0/2", + "request": { + "content": {"parts": [{"text": "first"}]}, + "output_dimensionality": 768, + }, + }, + { + "key": "request-1#1/2", + "request": { + "content": {"parts": [{"text": "second"}]}, + "output_dimensionality": 768, + }, + }, + ] + + def test_should_keep_the_bare_custom_id_for_single_element_arrays(self): + (row,) = _wrap_entries([_embeddings_entry(body={"model": "gemini-embedding-2", "input": ["only one"]})]) + + assert row["key"] == "request-1" + + def test_should_encode_a_custom_id_that_looks_like_a_fan_out_tag(self): + """A customer custom_id ending in `#/` must not read back as fan-out metadata.""" + (row,) = _wrap_entries( + [ + _embeddings_entry( + custom_id="request-1#0/2", + body={"model": "gemini-embedding-2", "input": "hello world"}, + ) + ] + ) + + assert row["key"] == "request-1%230%2F2" + + def test_should_combine_a_nested_input_into_one_multipart_row(self): + """Nested arrays are the combined-embedding shape, as on the online path.""" + (row,) = _wrap_entries( + [ + _embeddings_entry( + body={ + "model": "gemini-embedding-2", + "input": [ + [ + "a caption", + "gs://cloud-samples-data/generative-ai/image/benchmark.jpeg", + ] + ], + } + ) + ] + ) + + assert row["key"] == "request-1" + assert row["request"]["content"]["parts"] == [ + {"text": "a caption"}, + { + "file_data": { + "mime_type": "image/jpeg", + "file_uri": "gs://cloud-samples-data/generative-ai/image/benchmark.jpeg", + } + }, + ] + + def test_should_keep_chat_completions_lines_on_generate_content_path(self): + (row,) = _wrap_entries( + [ + { + "custom_id": "request-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gemini-2.0-flash-001", + "messages": [{"role": "user", "content": "Hello"}], + }, + } + ] + ) + + assert row["request"]["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}] + assert row["request"]["labels"]["litellm_custom_id"] == "request-1" + assert "key" not in row + + def test_should_keep_lines_without_a_url_on_generate_content_path(self): + """`url` is optional on a batch line, and chat is the shape LiteLLM has always assumed.""" + (row,) = _wrap_entries( + [ + { + "custom_id": "request-1", + "body": { + "model": "gemini-2.0-flash-001", + "messages": [{"role": "user", "content": "Hello"}], + }, + } + ] + ) + + assert row["request"]["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}] + + def test_should_translate_each_line_by_its_own_url(self): + chat_row, embeddings_row = _wrap_entries( + [ + { + "custom_id": "chat-1", + "url": "/v1/chat/completions", + "body": { + "model": "gemini-2.0-flash-001", + "messages": [{"role": "user", "content": "Hello"}], + }, + }, + _embeddings_entry(custom_id="embed-1"), + ] + ) + + assert "contents" in chat_row["request"] + assert "content" in embeddings_row["request"] + + +class TestVertexEmbeddingsBatchOutputTranslation: + """Vertex Gemini Embedding batch output rows must come back as OpenAI batch rows.""" + + def _vertex_embeddings_output_row(self, **overrides): + row = { + "key": "request-1", + "request": {"content": {"parts": [{"text": "hello world"}]}}, + "response": { + "embedding": {"values": [-0.015, 0.024]}, + "usageMetadata": {"promptTokenCount": 2}, + }, + } + row.update(overrides) + return row + + def _transform(self, config, rows, url="https://example.com"): + content = "\n".join(json.dumps(row) for row in rows).encode("utf-8") + result = config.transform_file_content_response( + raw_response=httpx.Response( + status_code=200, + content=content, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request("GET", url), + ), + logging_obj=MagicMock(), + litellm_params={}, + ) + return [json.loads(line) for line in result.response.content.decode("utf-8").split("\n")] + + def test_should_transform_embeddings_output_to_openai_batch_row(self, config): + (result,) = self._transform(config, [self._vertex_embeddings_output_row()]) + + assert result["custom_id"] == "request-1" + assert result["error"] is None + assert result["response"]["status_code"] == 200 + body = result["response"]["body"] + assert body["object"] == "list" + assert body["data"] == [{"embedding": [-0.015, 0.024], "index": 0, "object": "embedding"}] + assert body["usage"]["prompt_tokens"] == 2 + assert body["usage"]["total_tokens"] == 2 + + def test_should_fall_back_to_documented_token_count_field(self, config): + (result,) = self._transform( + config, + [ + self._vertex_embeddings_output_row( + response={ + "embedding": {"values": [-0.015, 0.024]}, + "tokenCount": "2", + } + ) + ], + ) + + assert result["response"]["body"]["usage"]["prompt_tokens"] == 2 + + def test_should_resolve_model_from_managed_gcs_object_path(self, config): + object_path = urllib.parse.quote( + "litellm-vertex-files/publishers/google/models/gemini-embedding-2/" + "prediction-model-2026-07-29T05:55:52Z/predictions.jsonl", + safe="", + ) + url = f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{object_path}?alt=media" + + (result,) = self._transform(config, [self._vertex_embeddings_output_row()], url=url) + + assert result["response"]["body"]["model"] == "gemini-embedding-2" + + def test_should_surface_failed_embeddings_row_as_error(self, config): + (result,) = self._transform( + config, + [self._vertex_embeddings_output_row(status="Failed to parse JSON into proto", response={})], + ) + + assert result["custom_id"] == "request-1" + assert result["response"] is None + assert result["error"]["code"] == "vertex_ai_error" + assert "Failed to parse JSON into proto" in result["error"]["message"] + + def test_should_transform_every_row_of_a_multi_row_file(self, config): + results = self._transform( + config, + [self._vertex_embeddings_output_row(key=f"request-{index}") for index in range(3)], + ) + + assert [result["custom_id"] for result in results] == [ + "request-0", + "request-1", + "request-2", + ] + + def test_should_reassemble_a_fanned_out_input_array_into_one_row(self, config): + """Vertex returns the rows of one entry in arbitrary order.""" + (result,) = self._transform( + config, + [ + self._vertex_embeddings_output_row( + key="request-1#1/2", + response={ + "embedding": {"values": [0.3, 0.4]}, + "usageMetadata": {"promptTokenCount": 5}, + }, + ), + self._vertex_embeddings_output_row( + key="request-1#0/2", + response={ + "embedding": {"values": [0.1, 0.2]}, + "usageMetadata": {"promptTokenCount": 3}, + }, + ), + ], + ) + + assert result["custom_id"] == "request-1" + assert result["response"]["body"]["data"] == [ + {"embedding": [0.1, 0.2], "index": 0, "object": "embedding"}, + {"embedding": [0.3, 0.4], "index": 1, "object": "embedding"}, + ] + assert result["response"]["body"]["usage"]["prompt_tokens"] == 8 + + def test_should_keep_fanned_out_entries_apart_and_in_file_order(self, config): + results = self._transform( + config, + [ + self._vertex_embeddings_output_row(key="request-2#0/2"), + self._vertex_embeddings_output_row(key="request-1"), + self._vertex_embeddings_output_row(key="request-2#1/2"), + ], + ) + + assert [result["custom_id"] for result in results] == ["request-2", "request-1"] + assert len(results[0]["response"]["body"]["data"]) == 2 + assert len(results[1]["response"]["body"]["data"]) == 1 + + def test_should_not_merge_an_entry_whose_custom_id_looks_like_a_fan_out_tag(self, config): + """`request-1#0/2` is a legal custom_id, and a distinct entry from `request-1`.""" + lookalike_row, plain_row = _wrap_entries( + [ + _embeddings_entry( + custom_id="request-1#0/2", + body={"model": "gemini-embedding-2", "input": "lookalike"}, + ), + _embeddings_entry( + custom_id="request-1", + body={"model": "gemini-embedding-2", "input": "plain"}, + ), + ] + ) + + results = self._transform( + config, + [ + {**row, "status": "", "response": {"embedding": {"values": values}}} + for row, values in ((lookalike_row, [0.1]), (plain_row, [0.2])) + ], + ) + + assert [result["custom_id"] for result in results] == [ + "request-1#0/2", + "request-1", + ] + assert [result["response"]["body"]["data"][0]["embedding"] for result in results] == [[0.1], [0.2]] + + def test_should_round_trip_a_fan_out_of_a_custom_id_holding_the_separator(self, config): + rows = _wrap_entries( + [ + _embeddings_entry( + custom_id="request#1/1", + body={ + "model": "gemini-embedding-2", + "input": ["first", "second"], + }, + ) + ] + ) + + assert [row["key"] for row in rows] == [ + "request%231%2F1#0/2", + "request%231%2F1#1/2", + ] + + (result,) = self._transform( + config, + [ + {**row, "status": "", "response": {"embedding": {"values": values}}} + for row, values in zip(reversed(rows), ([0.3], [0.1])) + ], + ) + + assert result["custom_id"] == "request#1/1" + assert [embedding["embedding"] for embedding in result["response"]["body"]["data"]] == [[0.1], [0.3]] + + def test_should_fail_the_whole_entry_when_one_of_its_rows_failed(self, config): + """An OpenAI batch row is either a response or an error, never both.""" + (result,) = self._transform( + config, + [ + self._vertex_embeddings_output_row(key="request-1#0/2"), + self._vertex_embeddings_output_row(key="request-1#1/2", status="Quota exceeded", response={}), + ], + ) + + assert result["custom_id"] == "request-1" + assert result["response"] is None + assert result["error"]["message"] == "Quota exceeded" + + def test_should_fail_the_whole_entry_when_a_fanned_out_row_is_missing(self, config): + """A partial `data` array would shift embeddings onto the wrong input positions.""" + (result,) = self._transform( + config, + [self._vertex_embeddings_output_row(key="request-1#1/2")], + ) + + assert result["custom_id"] == "request-1" + assert result["response"] is None + assert result["error"]["code"] == "vertex_ai_error" + assert result["error"]["message"] == ("Vertex returned embeddings for input positions [1] of the 2 requested") + + def test_should_fail_the_whole_entry_when_a_fanned_out_row_is_duplicated(self, config): + (result,) = self._transform( + config, + [ + self._vertex_embeddings_output_row(key="request-1#0/2"), + self._vertex_embeddings_output_row(key="request-1#0/2"), + ], + ) + + assert result["custom_id"] == "request-1" + assert result["response"] is None + assert result["error"]["code"] == "vertex_ai_error" + assert result["error"]["message"] == ( + "Vertex returned embeddings for input positions [0, 0] of the 2 requested" + ) + + def test_should_end_to_end_round_trip_a_fanned_out_embeddings_batch(self, config): + first_row, second_row = _wrap_entries( + [ + _embeddings_entry( + custom_id="MyRequest-1", + body={ + "model": "gemini-embedding-2", + "input": ["hello world", "goodbye world"], + }, + ) + ] + ) + + (result,) = self._transform( + config, + [ + { + **row, + "status": "", + "response": {"embedding": {"values": values}}, + } + for row, values in ((second_row, [0.3]), (first_row, [0.1])) + ], + ) + + assert result["custom_id"] == "MyRequest-1" + assert [embedding["embedding"] for embedding in result["response"]["body"]["data"]] == [[0.1], [0.3]] + + def test_should_end_to_end_round_trip_openai_embeddings_batch(self, config): + (vertex_row,) = _wrap_entries( + [ + _embeddings_entry( + custom_id="MyRequest-1", + body={ + "model": "gemini-embedding-2", + "input": "hello world", + "dimensions": 2, + }, + ) + ] + ) + + (result,) = self._transform( + config, + [ + { + **vertex_row, + "status": "", + "processed_time": "2026-07-29T05:55:52.379528Z", + "response": { + "embedding": {"values": [-0.015, 0.024]}, + "usageMetadata": {"promptTokenCount": 2}, + }, + } + ], + ) + + assert result["custom_id"] == "MyRequest-1" + assert result["response"]["body"]["data"][0]["embedding"] == [-0.015, 0.024] + + def test_should_leave_legacy_predict_embeddings_output_untouched(self, config): + legacy_row = { + "instance": {"content": "hello world"}, + "predictions": [ + { + "embeddings": { + "statistics": {"token_count": 2, "truncated": False}, + "values": [0.2], + } + } + ], + "status": "", + } + content = json.dumps(legacy_row).encode("utf-8") + + assert config._try_transform_vertex_batch_output_to_openai(content) == content