diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 4580ad17a19..686bbc89467 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -108,6 +108,7 @@ jobs: tests/test_litellm/completion_extras tests/test_litellm/containers tests/test_litellm/endpoints + tests/test_litellm/files tests/test_litellm/images tests/test_litellm/interactions tests/test_litellm/messages diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 7209ac6a1e7..819a279a43c 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -11,7 +11,10 @@ from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_ from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo from litellm.llms.bedrock.batches.transformation import titan_embedding_usage_from_batch_output -from litellm.llms.vertex_ai.batches.transformation import vertex_prompt_tokens_details +from litellm.llms.vertex_ai.batches.transformation import ( + is_native_vertex_batch_output_row, + native_vertex_batch_row_stats, +) from litellm.types.llms.openai import Batch from litellm.types.utils import ModelInfo, Usage from litellm.utils import token_counter @@ -31,6 +34,20 @@ class BatchCostUsageResult: _COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"}) + + +def _uses_native_vertex_output( + custom_llm_provider: str, + model_name: str | None, + first_row: Mapping[str, object] | None, +) -> bool: + if custom_llm_provider != "vertex_ai": + return False + if model_name and getattr(litellm, "disable_vertex_batch_output_transformation", False): + return True + return first_row is not None and is_native_vertex_batch_output_row(first_row) + + _TERMINAL_BATCH_STATUSES: Final = _COMPLETED_BATCH_STATUSES | frozenset({"failed", "cancelled", "expired"}) @@ -66,12 +83,9 @@ async def calculate_batch_cost_and_usage( deployment-specific pricing (e.g. input_cost_per_token_batches) is used instead of the global cost map. """ - if ( - custom_llm_provider == "vertex_ai" - and model_name - and getattr(litellm, "disable_vertex_batch_output_transformation", False) - ): - return calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name) + first_row: Final = file_content_dictionary[0] if file_content_dictionary else None + if _uses_native_vertex_output(custom_llm_provider, model_name, first_row): + return calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name, model_info=model_info) return _aggregate_batch_cost_usage_models( entries=file_content_dictionary, @@ -126,11 +140,11 @@ async def _handle_completed_batch( ) output_file_result: Final = ( - calculate_vertex_ai_batch_cost_and_usage(_get_file_content_as_dictionary(file_content), model_name) - if ( - custom_llm_provider == "vertex_ai" - and model_name - and getattr(litellm, "disable_vertex_batch_output_transformation", False) + calculate_vertex_ai_batch_cost_and_usage( + _iter_batch_output_entries(file_content), model_name, model_info=model_info + ) + if _uses_native_vertex_output( + custom_llm_provider, model_name, next(_iter_batch_output_entries(file_content), None) ) else _aggregate_batch_cost_usage_models( entries=_iter_batch_output_entries(file_content), @@ -332,69 +346,36 @@ def _aggregate_batch_cost_usage_models( def calculate_vertex_ai_batch_cost_and_usage( - vertex_ai_batch_responses: list[dict], + vertex_ai_batch_responses: Iterable[dict], model_name: str | None = None, + model_info: ModelInfo | None = None, ) -> BatchCostUsageResult: """ - Calculate both cost and usage from raw Vertex AI batch responses. - - Used only when ``litellm.disable_vertex_batch_output_transformation = True``. - In that case the GCS predictions.jsonl is returned as-is, with each line in - the native Vertex format: - - {"request": ..., "response": {"candidates": [...], "usageMetadata": {...}}} - - usageMetadata contains promptTokenCount, candidatesTokenCount, totalTokenCount. - - A row with no ``response`` is counted as failed - the same signal already - used to skip it from cost/usage aggregation, since Vertex batch prediction - output doesn't establish a distinct error shape in this (non-default) path. + Cost and usage of a native Vertex predictions.jsonl, one + `{"request": ..., "response": {"candidates": [...], "usageMetadata": {...}, "modelVersion": ...}}` + generateContent row or `{"request": ..., "response": {"embedding": {...}, "usageMetadata": {...}}}` + embedding row per line. `model_name` (the deployment model) prices every row, else each row's own + `modelVersion` does; a row without a usable response counts as failed. """ from litellm.cost_calculator import batch_cost_calculator + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig - total_prompt_cost = 0.0 # rebind-ok: loop accumulator, matches total_tokens below - total_completion_cost = 0.0 # rebind-ok: loop accumulator, matches total_tokens below - total_tokens = 0 - prompt_tokens = 0 - completion_tokens = 0 - successful_requests = 0 # rebind-ok: loop accumulator, matches total_cost/total_tokens above - failed_requests = 0 # rebind-ok: loop accumulator, matches total_cost/total_tokens above - actual_model_name: Final = model_name or "gemini-2.0-flash-001" - - for response in vertex_ai_batch_responses: - response_body = response.get("response") - if response_body is None: - failed_requests += 1 - continue - successful_requests += 1 - - 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, - prompt_tokens_details=vertex_prompt_tokens_details(usage_metadata), + row_stats: Final = tuple( + native_vertex_batch_row_stats( + row, + model_name, + model_info=model_info, + calculate_usage=VertexGeminiConfig._calculate_usage, + cost_calculator=batch_cost_calculator, ) - - try: - p_cost, c_cost = batch_cost_calculator( - usage=line_usage, - model=actual_model_name, - custom_llm_provider="vertex_ai", - ) - total_prompt_cost += p_cost - total_completion_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 - + for row in vertex_ai_batch_responses + ) + priced: Final = tuple(stats for stats in row_stats if stats is not None) + total_prompt_cost: Final = sum(stats.prompt_cost for stats in priced) + total_completion_cost: Final = sum(stats.completion_cost for stats in priced) + prompt_tokens: Final = sum(stats.usage.prompt_tokens for stats in priced) + completion_tokens: Final = sum(stats.usage.completion_tokens for stats in priced) + total_tokens: Final = sum(stats.total_tokens for stats in priced) total_cost: Final = total_prompt_cost + total_completion_cost verbose_logger.info( "vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d, successful=%d, failed=%d", @@ -402,8 +383,8 @@ def calculate_vertex_ai_batch_cost_and_usage( prompt_tokens, completion_tokens, total_tokens, - successful_requests, - failed_requests, + len(priced), + len(row_stats) - len(priced), ) return BatchCostUsageResult( @@ -413,9 +394,13 @@ def calculate_vertex_ai_batch_cost_and_usage( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, ), - models=[actual_model_name], - successful_requests=successful_requests, - failed_requests=failed_requests, + models=( + [model_name] + if model_name + else list(dict.fromkeys(stats.model for stats in priced if stats.model is not None)) + ), + successful_requests=len(priced), + failed_requests=len(row_stats) - len(priced), prompt_cost=total_prompt_cost, completion_cost=total_completion_cost, ) diff --git a/litellm/files/main.py b/litellm/files/main.py index e0804244ff7..72832aeccc9 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -176,6 +176,22 @@ def create_file( if logging_obj is None: raise ValueError("logging_obj is required") client: Final = kwargs.get("client") + if litellm_params_dict.get("passthrough") is True and ( + custom_llm_provider != "vertex_ai" or purpose != "batch" + ): + raise litellm.exceptions.BadRequestError( + message=( + "`passthrough=True` uploads the file bytes unchanged for a native Vertex AI batch, so it needs " + f"custom_llm_provider='vertex_ai' and purpose='batch', got '{custom_llm_provider}' and '{purpose}'." + ), + model="n/a", + llm_provider=custom_llm_provider or "n/a", + response=httpx.Response( + status_code=400, + content="passthrough needs a vertex_ai batch", + request=httpx.Request(method="create_file", url="https://github.com/BerriAI/litellm"), + ), + ) ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index f5f1ab2068a..a7dbb058465 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -1,7 +1,11 @@ -from collections.abc import Mapping -from typing import Any, Final +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Any, Final, Protocol from urllib.parse import unquote +from pydantic import TypeAdapter, ValidationError + +from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.llms.vertex_ai.common_utils import ( VertexAIError, @@ -9,35 +13,128 @@ from litellm.llms.vertex_ai.common_utils import ( ) from litellm.types.llms.openai import BatchJobStatus, CreateBatchRequest from litellm.types.llms.vertex_ai import * -from litellm.types.utils import LiteLLMBatch, PromptTokensDetailsWrapper +from litellm.types.llms.vertex_ai import GenerateContentResponseBody +from litellm.types.utils import LiteLLMBatch, ModelInfo, Usage + +_NATIVE_VERTEX_RESPONSE: Final = TypeAdapter(GenerateContentResponseBody) -def vertex_prompt_tokens_details( - usage_metadata: Mapping[str, object], -) -> PromptTokensDetailsWrapper | None: - raw_details: Final = usage_metadata.get("promptTokensDetails") - if not isinstance(raw_details, list): - return None +def _int_field(mapping: Mapping[str, object], key: str) -> int: + value: Final = mapping.get(key) + if isinstance(value, int): + return value + return int(value) if isinstance(value, str) and value.isdigit() else 0 - def _normalize(detail: object) -> tuple[str, int] | None: - if not isinstance(detail, Mapping): + +def vertex_embedding_prompt_token_count(vertex_response: Mapping[str, object]) -> 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: Final = vertex_response.get("usageMetadata") + if isinstance(usage_metadata, Mapping): + return _int_field(usage_metadata, "promptTokenCount") + return _int_field(vertex_response, "tokenCount") + + +def is_vertex_embedding_batch_output_response(response_body: Mapping[str, object]) -> bool: + return isinstance(response_body.get("embedding"), dict) + + +def is_native_vertex_batch_output_row(row: Mapping[str, object]) -> bool: + return isinstance(row.get("request"), dict) + + +class NativeVertexBatchCostCalculator(Protocol): + def __call__( + self, + usage: Usage, + model: str, + custom_llm_provider: str | None = None, + model_info: ModelInfo | None = None, + ) -> tuple[float, float]: ... + + +@dataclass(frozen=True, slots=True) +class NativeVertexBatchRowStats: + usage: Usage + total_tokens: int + model: str | None + prompt_cost: float + completion_cost: float + + +def _native_vertex_row_usage( + response_body: Mapping[str, object], + calculate_usage: Callable[[GenerateContentResponseBody], Usage], +) -> Usage | None: + if "usageMetadata" not in response_body: + if not is_vertex_embedding_batch_output_response(response_body): return None - modality: Final = detail.get("modality") - token_count: Final = detail.get("tokenCount") - if not isinstance(modality, str) or not isinstance(token_count, int): - return None - return modality.upper(), token_count - - parsed_details: Final = tuple(_normalize(detail) for detail in raw_details) - normalized: Final = tuple(detail for detail in parsed_details if detail is not None) - if len(normalized) != len(parsed_details): + prompt_tokens: Final = vertex_embedding_prompt_token_count(response_body) + return Usage(prompt_tokens=prompt_tokens, completion_tokens=0, total_tokens=prompt_tokens) + try: + completion_response: Final = _NATIVE_VERTEX_RESPONSE.validate_python(response_body) + except ValidationError as e: + verbose_logger.debug("vertex_ai batch row response is not a GenerateContentResponse: %s", str(e)) return None + return calculate_usage(completion_response) - return PromptTokensDetailsWrapper( - text_tokens=sum(token_count for modality, token_count in normalized if modality in ("TEXT", "DOCUMENT")), - audio_tokens=sum(token_count for modality, token_count in normalized if modality == "AUDIO"), - image_tokens=sum(token_count for modality, token_count in normalized if modality == "IMAGE"), - video_tokens=sum(token_count for modality, token_count in normalized if modality == "VIDEO"), + +def native_vertex_batch_row_stats( + row: Mapping[str, object], + model_name: str | None, + *, + model_info: ModelInfo | None, + calculate_usage: Callable[[GenerateContentResponseBody], Usage], + cost_calculator: NativeVertexBatchCostCalculator, +) -> NativeVertexBatchRowStats | None: + """ + Usage and cost of one native Vertex predictions.jsonl row, a + `{"request": ..., "response": {"candidates": [...], "usageMetadata": {...}, "modelVersion": ...}}` + generateContent object or a `{"request": ..., "response": {"embedding": {...}, "usageMetadata": {...}}}` + embedding object (an embedding row without `usageMetadata` is billed from its documented `tokenCount`). + `model_name` (the deployment model) prices the row unless it is a wildcard, else its own `modelVersion` + does, else the wildcard name so explicit deployment prices still apply; a row without a response, a + generateContent row without `response.usageMetadata`, and a row whose response fails validation are + None (failed). + """ + response_body: Final = row.get("response") + if not isinstance(response_body, dict): + return None + usage: Final = _native_vertex_row_usage(response_body, calculate_usage) + if usage is None: + return None + total_tokens: Final = usage.total_tokens or (usage.prompt_tokens + usage.completion_tokens) + model_version: Final = response_body.get("modelVersion") + deployment_model: Final = model_name if model_name and "*" not in model_name else None + model: Final = deployment_model or (model_version if isinstance(model_version, str) else model_name) + if model is None: + verbose_logger.warning( + "vertex_ai batch output row could not be costed, so it is billed at $0 and the rest of the batch " + "is still billed: the row has no modelVersion and the batch has no deployment model" + ) + return NativeVertexBatchRowStats( + usage=usage, total_tokens=total_tokens, model=None, prompt_cost=0.0, completion_cost=0.0 + ) + try: + prompt_cost, completion_cost = cost_calculator( + usage=usage, model=model, custom_llm_provider="vertex_ai", model_info=model_info + ) + except Exception as e: # noqa: BLE001 # one unpriceable row must not abort the batch's cost accounting + verbose_logger.warning( + "vertex_ai batch output row could not be costed, so it is billed at $0 and the rest of the batch " + "is still billed. model=%s error=%s", + model, + str(e), + ) + return NativeVertexBatchRowStats( + usage=usage, total_tokens=total_tokens, model=model, prompt_cost=0.0, completion_cost=0.0 + ) + return NativeVertexBatchRowStats( + usage=usage, total_tokens=total_tokens, model=model, prompt_cost=prompt_cost, completion_cost=completion_cost ) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 80d32289c94..789b36ef3d0 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -9,7 +9,7 @@ from collections.abc import AsyncGenerator, Callable, Iterable, Iterator, Mappin from contextlib import aclosing from dataclasses import dataclass from types import MappingProxyType -from typing import Any, Final, TypedDict +from typing import IO, Any, Final, TypedDict from urllib.parse import quote, unquote import httpx @@ -41,6 +41,7 @@ from litellm.llms.base_llm.files.transformation import ( BaseFileUploadStream, LiteLLMLoggingObj, ) +from litellm.llms.vertex_ai.batches.transformation import vertex_embedding_prompt_token_count from litellm.llms.vertex_ai.common_utils import ( _convert_vertex_datetime_to_openai_datetime, get_vertex_ai_fine_tuned_endpoint_id, @@ -56,6 +57,7 @@ from litellm.types.files import StreamingMediaUploadConfig from litellm.types.llms.openai import ( AllMessageValues, CreateFileRequest, + FileContent, FileTypes, HttpxBinaryResponseContent, OpenAICreateFileRequestOptionalParams, @@ -87,6 +89,8 @@ _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM: Final = ( _VERTEX_BATCH_FANNED_OUT_KEY_PATTERN: Final = re.compile(r"(?P[^#]*)#(?P\d+)/(?P\d+)") _JSONL_NEWLINE: Final = b"\n" _BATCH_OUTPUT_FIRST_ROW_PEEK_LIMIT_BYTES: Final = 32 * 1024 * 1024 +_PASSTHROUGH_MANAGED_GCS_PREFIX: Final = f"{VERTEX_AI_MANAGED_GCS_PREFIX}passthrough/" +_RAW_UPLOAD_CHUNK_BYTES: Final = 1024 * 1024 class _GcsObjectMetadataJson(TypedDict, total=False): @@ -418,19 +422,6 @@ def _split_vertex_batch_key(vertex_output_row: Mapping[str, object]) -> tuple[st return unquote(match["custom_id"]), int(match["index"]), int(match["total"]) -def _embedding_prompt_token_count(vertex_response: _VertexEmbeddingResponse) -> 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[_VertexEmbeddingBatchRow, ...], @@ -471,7 +462,7 @@ def _vertex_embeddings_rows_to_openai_batch_output_row( ) responses = tuple(row["response"] for row in vertex_output_rows) - token_count = sum(_embedding_prompt_token_count(response) for response in responses) + token_count = sum(vertex_embedding_prompt_token_count(response) for response in responses) body = EmbeddingResponse( model=model or "", data=[ @@ -528,6 +519,16 @@ def _model_from_managed_gcs_url(url: str) -> str | None: return match.group(1) if match else None +def is_passthrough_managed_gcs_url(url: str) -> bool: + decoded_url: Final = unquote(url) + managed_prefix_start: Final = decoded_url.find(VERTEX_AI_MANAGED_GCS_PREFIX) + return managed_prefix_start >= 0 and decoded_url.startswith(_PASSTHROUGH_MANAGED_GCS_PREFIX, managed_prefix_start) + + +def is_passthrough_batch_upload(create_file_data: Mapping[str, object], litellm_params: Mapping[str, object]) -> bool: + return create_file_data.get("purpose") == "batch" and litellm_params.get("passthrough") is True + + def _is_embeddings_batch_entry(openai_entry: Mapping[str, object]) -> bool: """ Whether an OpenAI batch JSONL line targets the embeddings endpoint. @@ -791,6 +792,58 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): return self._iter_vertex_jsonl_chunks() +def _read_chunk_as_bytes(handle: IO[bytes]) -> bytes: + chunk: Final[bytes | str] = handle.read(_RAW_UPLOAD_CHUNK_BYTES) + return chunk.encode("utf-8") if isinstance(chunk, str) else bytes(chunk) + + +def _iter_raw_file_chunks(file_content: FileTypes) -> Iterator[bytes]: + content: Final[FileContent | str] = file_content[1] if isinstance(file_content, tuple) else file_content + if isinstance(content, (bytes, bytearray)): + yield from ( + bytes(content[offset : offset + _RAW_UPLOAD_CHUNK_BYTES]) + for offset in range(0, len(content), _RAW_UPLOAD_CHUNK_BYTES) + ) + return + if isinstance(content, str): + yield content.encode("utf-8") + return + if isinstance(content, PathLike): + with open(str(content), "rb") as handle: + yield from iter(lambda: handle.read(_RAW_UPLOAD_CHUNK_BYTES), b"") + return + if not hasattr(content, "read"): + raise ValueError("Unsupported file content type") + seek: Final = getattr(content, "seek", None) + if seek is None: + raise ValueError( + "Batch upload file handle must be seekable; got a non-seekable " + "stream. Pass bytes, a path, or a seekable handle." + ) + seek(0) + yield from iter(lambda: _read_chunk_as_bytes(content), b"") + + +class _RawFileUploadStream(BaseFileUploadStream): + def __init__(self, file_content: FileTypes) -> None: + self._file_content = file_content + + def iter_bytes(self) -> Iterator[bytes]: + return _iter_raw_file_chunks(self._file_content) + + +def _managed_batch_object_name(raw_model: str, *, passthrough: bool) -> str: + endpoint_id: Final = get_vertex_ai_fine_tuned_endpoint_id(raw_model) + model_path: Final = ( + f"endpoints/{endpoint_id}" + if endpoint_id is not None + else (raw_model if "publishers/google/models" in raw_model else f"publishers/google/models/{raw_model}") + ) + safe_model_path: Final = sanitize_cloud_object_path(model_path, fallback="model") + prefix: Final = _PASSTHROUGH_MANAGED_GCS_PREFIX if passthrough else VERTEX_AI_MANAGED_GCS_PREFIX + return f"{prefix}{safe_model_path}/{uuid.uuid4()}" + + class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ Config for VertexAI Files @@ -848,23 +901,34 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): if deployment_model else openai_jsonl_content[0].get("body", {}).get("model", "") ) - endpoint_id: Final = get_vertex_ai_fine_tuned_endpoint_id(raw_model) - model_path: Final = ( - f"endpoints/{endpoint_id}" - if endpoint_id is not None - else (raw_model if "publishers/google/models" in raw_model else f"publishers/google/models/{raw_model}") - ) - safe_model_path: Final = sanitize_cloud_object_path(model_path, fallback="model") - object_name: Final = f"{VERTEX_AI_MANAGED_GCS_PREFIX}{safe_model_path}/{uuid.uuid4()}" - return object_name + return _managed_batch_object_name(raw_model, passthrough=False) - def get_object_name(self, file_data: FileTypes, purpose: str, deployment_model: str | None = None) -> str: + def _get_passthrough_gcs_object_name(self, deployment_model: str | None) -> str: + if not deployment_model: + raise VertexAIError( + status_code=400, + message=( + "Native Vertex batch passthrough uploads need the deployment model to name the GCS object, " + "since native rows carry no model: pass `target_model_names` (proxy) or `model` (SDK)." + ), + ) + return _managed_batch_object_name(deployment_model.removeprefix("vertex_ai/"), passthrough=True) + + def get_object_name( + self, + file_data: FileTypes, + purpose: str, + deployment_model: str | None = None, + passthrough: bool = False, + ) -> str: """ Get the object name for the request. Reads only the first JSONL entry (streamed) for batch files, so a large upload is never materialized just to derive the GCS object name. """ + if purpose == "batch" and passthrough: + return self._get_passthrough_gcs_object_name(deployment_model) if purpose == "batch": ## 1. If jsonl, derive the object name from the deployment model (or the first entry's) first_entry: Final = next(_iter_openai_jsonl_entries(file_data), None) @@ -922,6 +986,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): file_data, purpose, deployment_model=configured_model if isinstance(configured_model, str) else None, + passthrough=is_passthrough_batch_upload(data, litellm_params), ) if object_prefix: object_name = f"{object_prefix}/{object_name}" @@ -984,6 +1049,14 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): if file_data is None: raise ValueError("file is required") + if is_passthrough_batch_upload(create_file_data, litellm_params): + return { + "streaming_media_upload": StreamingMediaUploadConfig( + body_stream=_RawFileUploadStream(file_data), + content_type="application/json", + ) + } + _, content_type = extract_file_metadata(file_data) if FilesAPIUtils.is_batch_jsonl_request( create_file_data=create_file_data, @@ -1164,6 +1237,8 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): # transformation, e.g. if they consume raw `predictions.jsonl` directly. if getattr(litellm, "disable_vertex_batch_output_transformation", False): return HttpxBinaryResponseContent(response=raw_response) + if is_passthrough_managed_gcs_url(str(raw_response.request.url)): + return HttpxBinaryResponseContent(response=raw_response) # Try to transform batch output if it's a JSONL file content: Final = raw_response.content @@ -1209,7 +1284,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): Everything else is passed through unchanged, including a row that fails to transform mid-stream. """ - if litellm.disable_vertex_batch_output_transformation: + if litellm.disable_vertex_batch_output_transformation or is_passthrough_managed_gcs_url(request_url): return FileContentStreamingResult(stream_iterator=stream_iterator, headers=headers) first_line, buffered = await _peek_first_jsonl_line( diff --git a/litellm/proxy/openai_files_endpoints/batch_file_validation.py b/litellm/proxy/openai_files_endpoints/batch_file_validation.py index a41bd36d510..fcd3be56ae7 100644 --- a/litellm/proxy/openai_files_endpoints/batch_file_validation.py +++ b/litellm/proxy/openai_files_endpoints/batch_file_validation.py @@ -8,10 +8,25 @@ from typing_extensions import assert_never from litellm.proxy._types import ProxyException -BATCH_LINE_REQUIRED_KEYS: Final = ("custom_id", "method", "url", "body") _MB: Final = 1024 * 1024 +@dataclass(frozen=True, slots=True) +class BatchLineShape: + required_keys: tuple[str, ...] + hint: str + + +BATCH_LINE_SHAPE: Final = BatchLineShape( + required_keys=("custom_id", "method", "url", "body"), + hint="Each line must be a JSON object with keys custom_id, method, url, body", +) +PASSTHROUGH_BATCH_LINE_SHAPE: Final = BatchLineShape( + required_keys=("request",), + hint="A passthrough upload takes native Vertex batch rows, so each line must be a JSON object with a request key", +) + + @dataclass(frozen=True, slots=True) class BatchFileTooLarge: size_bytes: int @@ -42,6 +57,7 @@ class BatchFileLineNotObject: class BatchFileMissingLineKey: line_number: int key: str + line_shape: BatchLineShape = BATCH_LINE_SHAPE BatchFileValidationFailure = ( @@ -70,20 +86,20 @@ def _iter_lines(file_source: bytes | BinaryIO) -> Iterator[bytes]: return iter(file_source) -def _check_line(line_number: int, raw_line: bytes) -> BatchFileValidationFailure | None: +def _check_line(line_number: int, raw_line: bytes, line_shape: BatchLineShape) -> BatchFileValidationFailure | None: try: parsed: Final = json.loads(raw_line) except (json.JSONDecodeError, UnicodeDecodeError): return BatchFileInvalidJsonLine(line_number=line_number) if not isinstance(parsed, dict): return BatchFileLineNotObject(line_number=line_number) - missing: Final = next((key for key in BATCH_LINE_REQUIRED_KEYS if key not in parsed), None) + missing: Final = next((key for key in line_shape.required_keys if key not in parsed), None) if missing is None: return None - return BatchFileMissingLineKey(line_number=line_number, key=missing) + return BatchFileMissingLineKey(line_number=line_number, key=missing, line_shape=line_shape) -def _scan_lines(file_source: bytes | BinaryIO) -> BatchFileValidationFailure | None: +def _scan_lines(file_source: bytes | BinaryIO, line_shape: BatchLineShape) -> BatchFileValidationFailure | None: content_lines: Final = ( (line_number, raw_line) for line_number, raw_line in enumerate(_iter_lines(file_source), start=1) @@ -96,7 +112,7 @@ def _scan_lines(file_source: bytes | BinaryIO) -> BatchFileValidationFailure | N ( failure for line_number, raw_line in chain((first_line,), content_lines) - for failure in (_check_line(line_number, raw_line),) + for failure in (_check_line(line_number, raw_line, line_shape),) if failure is not None ), None, @@ -107,6 +123,7 @@ def check_batch_file_upload( filename: str | None, file_source: bytes | BinaryIO, max_batch_file_size_mb: int | None, + line_shape: BatchLineShape = BATCH_LINE_SHAPE, ) -> BatchFileValidationFailure | None: if filename is None or not filename.lower().endswith(".jsonl"): return BatchFileWrongExtension(filename=filename or "") @@ -114,7 +131,7 @@ def check_batch_file_upload( size_bytes: Final = _file_size_bytes(file_source) if size_bytes > max_batch_file_size_mb * _MB: return BatchFileTooLarge(size_bytes=size_bytes, limit_mb=max_batch_file_size_mb) - scan_failure: Final = _scan_lines(file_source) + scan_failure: Final = _scan_lines(file_source, line_shape) if not isinstance(file_source, bytes): file_source.seek(0) return scan_failure @@ -169,11 +186,11 @@ def raise_batch_file_validation_failure(failure: BatchFileValidationFailure) -> param="file", code=400, ) - case BatchFileMissingLineKey(line_number=line_number, key=key): + case BatchFileMissingLineKey(line_number=line_number, key=key, line_shape=line_shape): raise ProxyException( message=( f"Missing required parameter: '{key}' (batch input file line {line_number}). " - f"Each line must be a JSON object with keys {', '.join(BATCH_LINE_REQUIRED_KEYS)}. " + f"{line_shape.hint}. " "The file was not forwarded to the provider." ), type="invalid_request_error", diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index ea2cad558c2..f5e62da5962 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -57,6 +57,8 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_type, ) from litellm.proxy.openai_files_endpoints.batch_file_validation import ( + BATCH_LINE_SHAPE, + PASSTHROUGH_BATCH_LINE_SHAPE, check_batch_file_upload, raise_batch_file_validation_failure, ) @@ -207,10 +209,91 @@ def get_files_provider_config( return None +def _deployment_provider(llm_router: Router, model_id: str, team_id: str | None) -> str | None: + credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=model_id, team_id=team_id) + return None if credentials is None else credentials.get("custom_llm_provider") + + +def _resolves_to_vertex_deployments_only(llm_router: Router | None, model_name: str, team_id: str | None) -> bool: + if llm_router is None or _deployment_provider(llm_router, model_name, team_id) != "vertex_ai": + return False + return all( + _deployment_provider(llm_router, str(deployment["model_info"]["id"]), team_id) == "vertex_ai" + for deployment in llm_router.get_model_list(model_name=model_name, team_id=team_id) or () + if "id" in deployment.get("model_info", {}) + ) + + +def _validate_passthrough_upload( + *, + purpose: str, + target_model_names: Sequence[str], + model: str | None, + target_storage: str | None, + llm_router: Router | None, + team_id: str | None, +) -> None: + if purpose != "batch": + raise ProxyException( + message=( + "`passthrough` uploads the file bytes unchanged for a native Vertex batch, " + f"so purpose must be 'batch', got '{purpose}'." + ), + type="invalid_request_error", + param="passthrough", + code=400, + ) + if target_storage and target_storage != "default": + raise ProxyException( + message=( + "`passthrough` writes the native batch file to the Vertex AI deployment's GCS bucket, " + f"so it cannot be combined with target_storage='{target_storage}'." + ), + type="invalid_request_error", + param="target_storage", + code=400, + ) + named_deployments: Final = ( + *(("target_model_names", name) for name in target_model_names), + *((("model", model),) if model else ()), + ) + if not named_deployments: + raise ProxyException( + message=( + "`passthrough` needs the Vertex AI deployment that will run the batch, " + "since native rows carry no model: pass `target_model_names` or `model`." + ), + type="invalid_request_error", + param="target_model_names", + code=400, + ) + offending: Final = next( + ( + (param, name) + for param, name in named_deployments + if not _resolves_to_vertex_deployments_only(llm_router, name, team_id) + ), + None, + ) + if offending is None: + return + param, name = offending + raise ProxyException( + message=( + f"`passthrough` is only supported for Vertex AI deployments; '{name}' does not resolve " + "to vertex_ai deployments only." + ), + type="invalid_request_error", + param=param, + code=400, + ) + + async def _scan_batch_upload( *, file_source: bytes | BinaryIO, purpose: str, + passthrough: bool, request_metadata: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging, @@ -222,6 +305,17 @@ async def _scan_batch_upload( or not proxy_logging_obj.has_pre_call_guardrails(request_metadata) ): return None + if passthrough: + raise ProxyException( + message=( + "Batch guardrails cannot scan native Vertex batch rows, so a `passthrough` upload is refused " + "when the key, team, or request has pre-call guardrails configured. " + "The file was not forwarded to the provider." + ), + type="invalid_request_error", + param="passthrough", + code=400, + ) outcome: Final = await scan_batch_input_file( file_source=file_source, request_metadata=request_metadata, @@ -458,6 +552,7 @@ async def create_file( custom_llm_provider: str = Form(default="openai"), file: UploadFile = File(...), litellm_metadata: str | None = Form(default=None), + passthrough: bool = Form(default=False), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -560,17 +655,28 @@ async def create_file( if blocked_extension_failure is not None: raise_upload_validation_failure(blocked_extension_failure) + if passthrough: + _validate_passthrough_upload( + purpose=purpose, + target_model_names=target_model_names_list, + model=model_param, + target_storage=target_storage, + llm_router=llm_router, + team_id=user_api_key_dict.team_id, + ) + if purpose == "batch": batch_file_failure: Final = await asyncio.to_thread( check_batch_file_upload, file.filename, file_source, _MAX_BATCH_FILE_SIZE_MB_ADAPTER.validate_python(general_settings.get("max_batch_file_size_mb")), + PASSTHROUGH_BATCH_LINE_SHAPE if passthrough else BATCH_LINE_SHAPE, ) if batch_file_failure is not None: raise_batch_file_validation_failure(batch_file_failure) - data = {} + data = {"passthrough": True} if passthrough else {} # Parse expires_after if provided expires_after: FileExpiresAfter | None = None @@ -673,6 +779,7 @@ async def create_file( scan_result: Final = await _scan_batch_upload( file_source=file_source, purpose=purpose, + passthrough=passthrough, request_metadata=request_metadata, user_api_key_dict=user_api_key_dict, proxy_logging_obj=proxy_logging_obj, diff --git a/litellm/router.py b/litellm/router.py index d328fbbb12f..8960cd92cd8 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -5980,6 +5980,7 @@ class Router: replace_model_in_jsonl_bool: Final = should_replace_model_in_jsonl( purpose=purpose, + passthrough=kwargs.get("passthrough") is True, ) if replace_model_in_jsonl_bool: file = replace_model_in_jsonl( diff --git a/litellm/router_utils/batch_utils.py b/litellm/router_utils/batch_utils.py index be20c358202..386a2135239 100644 --- a/litellm/router_utils/batch_utils.py +++ b/litellm/router_utils/batch_utils.py @@ -62,15 +62,15 @@ def parse_jsonl_with_embedded_newlines(content: str) -> list[dict]: def should_replace_model_in_jsonl( purpose: OpenAIFilesPurpose, + passthrough: bool = False, ) -> bool: """ Check if the model name should be replaced in the JSONL file for the deployment model name. Azure raises an error on create batch if the model name for deployment is not in the .jsonl. + A passthrough upload keeps the caller's bytes untouched, so its rows are never rewritten. """ - if purpose == "batch": - return True - return False + return purpose == "batch" and not passthrough def replace_model_in_jsonl(file_content: FileTypes, new_model_name: str) -> FileTypes: diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 9bb6d05bec8..6e9cf45e787 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -61,6 +61,7 @@ from e2e_http import ( StreamingResponse, Success, UnknownApiError, + proxy_error, require_successful_call, unwrap, ) @@ -1752,3 +1753,114 @@ class TestBatchTerminalState: assert (cost_row.total_tokens or 0) > 0, ( f"batch cost row has no token usage: {cost_row.total_tokens!r}" ) + + +NATIVE_VERTEX_BATCH_ROWS: Final = b"".join( + json.dumps( + { + "request": { + "contents": [{"role": "user", "parts": [{"text": text}]}], + "tools": [{"googleSearch": {"excludeDomains": ["example.com"]}}], + } + } + ).encode() + + b"\n" + for text in ("What is the tallest building in the world?", "Who won the last FIFA World Cup?") +) +VERTEX_BATCH_PROVIDER: Final = next(p for p in PROVIDERS if p.name == "vertex_ai") + + +class TestVertexNativePassthrough: + """`passthrough=true` on POST /v1/files uploads native Vertex batch JSONL byte for + byte (no OpenAI-to-Vertex translation, so `googleSearch` tools and the grounding + metadata they produce survive), and a batch created from that file is accepted. + + Terminal-state assertions (native output rows with groundingMetadata, the spend + row) are deliberately not here: retrieving a non-terminal batch books a $0 spend + row that blocks the real-cost row, the same reason TestBatchTerminalState polls + the list endpoint only. Those are proven by the PR's live curl proof instead. + """ + + @pytest.mark.covers( + "llm.files.vertex.native_passthrough.nonstream.works", + "llm.batches.vertex.native_passthrough.nonstream.works", + exercised_on=["files", "batches"], + ) + def test_native_jsonl_round_trips_untouched_and_starts_a_batch( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = resources.key() + file = unwrap( + client.upload_file( + content=NATIVE_VERTEX_BATCH_ROWS, + form=FileUploadForm( + purpose="batch", target_model_names=VERTEX_BATCH_PROVIDER.model, passthrough=True + ), + key=key, + ) + ) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) + assert_file_object(file, provider="vertex_ai") + assert is_managed_id(file.id), f"passthrough upload must return a managed file id, got {file.id!r}" + assert file.bytes == len(NATIVE_VERTEX_BATCH_ROWS), ( + f"passthrough upload must report the caller's byte count, got {file.bytes}" + ) + + downloaded = client.proxy.transport.download( + f"/v1/files/{file.id}/content", headers=client.proxy.transport.bearer(key) + ) + assert downloaded.status_code == 200, ( + f"file content must be 200, got {downloaded.status_code}: {downloaded.body[:300]}" + ) + assert downloaded.body.encode() == NATIVE_VERTEX_BATCH_ROWS, ( + "passthrough file content must be the uploaded native rows byte for byte" + ) + + created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) + require_successful_call(created) + batch = BatchObject.model_validate_json(created.body) + resources.defer(lambda: cleanup_batch(client, batch.id, key=key, delete_output_files=True)) + assert is_managed_id(batch.id), f"passthrough batch must be LiteLLM-managed, got {batch.id!r}" + assert batch.status in CREATED_BATCH_STATUSES, f"passthrough batch has non-transitional status {batch.status!r}" + assert batch.input_file_id == file.id + + @pytest.mark.covers("llm.files.vertex.native_passthrough_validation.nonstream.works", exercised_on=["files"]) + @pytest.mark.parametrize( + "content, form, expected_param", + [ + pytest.param( + NATIVE_VERTEX_BATCH_ROWS, + FileUploadForm(purpose="batch", passthrough=True), + "target_model_names", + id="no-target-model", + ), + pytest.param( + NATIVE_VERTEX_BATCH_ROWS, + FileUploadForm(purpose="batch", target_model_names=OPENAI_BATCH_MODEL, passthrough=True), + "target_model_names", + id="non-vertex-target-model", + ), + pytest.param( + render_jsonl(VERTEX_BATCH_PROVIDER.raw_model), + FileUploadForm(purpose="batch", target_model_names=VERTEX_BATCH_PROVIDER.model, passthrough=True), + "request", + id="openai-shaped-rows", + ), + ], + ) + def test_passthrough_upload_is_rejected_outside_a_native_vertex_batch( + self, + content: bytes, + form: FileUploadForm, + expected_param: str, + client: BatchClient, + resources: ResourceManager, + batch_deployments: None, + ) -> None: + key = resources.key() + result = client.upload_file(content=content, form=form, key=key) + assert isinstance(result, UnknownApiError), f"expected a 400, got {result!r}" + assert result.status_code == 400, f"expected 400, got {result.status_code}: {result.body[:300]}" + error = proxy_error(result.body) + assert error.param == expected_param, f"unexpected error param in {error!r}" + assert "passthrough" in error.message diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index 3e389acc2a9..7d334ed41ff 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -21,6 +21,7 @@ - {id: llm.batches.openai_provider_fallback.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py", rationale: "Provider-fallback raw-id scenario"} - {id: llm.batches.azure_openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:98", rationale: "Azure batches all scenarios"} - {id: llm.batches.vertex.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:98", rationale: "Vertex batches"} +- {id: llm.batches.vertex.native_passthrough.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: vertex, capability: native_passthrough, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-4790", rationale: "A batch created from a passthrough-uploaded native Vertex JSONL file is accepted and starts on the deployment named at upload"} - {id: llm.batches.bedrock.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:98", rationale: "Bedrock batches (encoded/unified only)"} - {id: llm.batches.bedrock.assume_role.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: bedrock_converse, capability: assume_role, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Bedrock batch create under STS assume-role credentials"} - {id: llm.batches.bedrock.govcloud_partition.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: bedrock_converse, capability: govcloud_partition, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Bedrock batch create in the us-gov-west-1 partition"} @@ -46,6 +47,8 @@ - {id: llm.files.openai.passthrough.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_passthrough_e2e.py", rationale: "POST/DELETE /openai_passthrough/v1/files relay OpenAI's own file object; the dedicated prefix must not bind as a provider name on the /{provider}/v1/files route (GitHub issue #36086)"} - {id: llm.files.azure_openai.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:45", rationale: "Azure file upload managed backend"} - {id: llm.files.vertex.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:52", rationale: "Vertex file upload to GCS"} +- {id: llm.files.vertex.native_passthrough.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: vertex, capability: native_passthrough, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-4790", rationale: "POST /v1/files with passthrough=true ships native Vertex batch JSONL (googleSearch tools and all) to GCS untouched and GET /v1/files/{id}/content returns the same bytes"} +- {id: llm.files.vertex.native_passthrough_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: vertex, capability: input_validation, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-4790", rationale: "passthrough=true without a Vertex target_model_names, or with OpenAI-shaped rows, is a 400 naming the offending field and nothing is uploaded"} - {id: llm.files.bedrock.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:59", rationale: "Bedrock file upload to S3"} - {id: llm.files.bedrock.govcloud_partition.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: govcloud_partition, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Bedrock file upload to an S3 bucket in the us-gov-west-1 partition"} - {id: llm.files.bedrock.split_s3_credentials.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: split_s3_credentials, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-8297", rationale: "Bedrock file upload, content and delete sign S3 with s3_access_key_id / s3_secret_access_key when they differ from the aws_* identity"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index e009b02b69c..e5626144fad 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -73,6 +73,7 @@ LlmCapability = Literal[ "long_context_1m", "mid_conversation_system", "multi_turn", + "native_passthrough", "pdf_input", "prompt_cache_1h", "prompt_cache_5m", diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 022caddd42a..e5d50d05c87 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -69,6 +69,7 @@ class FileUploadForm(BaseModel): purpose: str = "batch" target_model_names: str | None = None custom_llm_provider: str | None = None + passthrough: bool | None = None # ---------- Result types ---------- @@ -376,12 +377,18 @@ class ProxyErrorDetail(BaseModel): message: str type: str code: str + param: str | None = None class _ProxyErrorBody(BaseModel): error: ProxyErrorDetail +def proxy_error(body: str) -> ProxyErrorDetail: + """The proxy's own error envelope (`{"error": {message, type, param, code}}`) parsed off a rejected call.""" + return _ProxyErrorBody.model_validate_json(body).error + + def relayed_provider_rate_limit(outcome: RateLimitedError) -> ProxyErrorDetail | None: """The provider's own 429 as the proxy relayed it, or None when the 429 is the proxy's own.""" if PROVIDER_RATE_LIMIT_MARKER not in outcome.body: diff --git a/tests/router_unit_tests/test_router_batch_utils.py b/tests/router_unit_tests/test_router_batch_utils.py index e274ac61a01..4336185a07f 100644 --- a/tests/router_unit_tests/test_router_batch_utils.py +++ b/tests/router_unit_tests/test_router_batch_utils.py @@ -141,6 +141,7 @@ def test_should_replace_model_in_jsonl(): from litellm.router_utils.batch_utils import should_replace_model_in_jsonl assert should_replace_model_in_jsonl(purpose="batch") is True + assert should_replace_model_in_jsonl(purpose="batch", passthrough=True) is False assert should_replace_model_in_jsonl(purpose="test") is False assert should_replace_model_in_jsonl(purpose="user_data") is False diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py new file mode 100644 index 00000000000..0b2bfe9d266 --- /dev/null +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -0,0 +1,387 @@ +import json + +import pytest + +import litellm +import litellm.batches.batch_utils as bu +from litellm.types.llms.openai import Batch + +GROUNDED_USAGE_METADATA = { + "promptTokenCount": 19, + "candidatesTokenCount": 59, + "thoughtsTokenCount": 406, + "toolUsePromptTokenCount": 73, + "totalTokenCount": 557, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 19}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 59}], + "toolUsePromptTokensDetails": [{"modality": "TEXT", "tokenCount": 73}], + "trafficType": "ON_DEMAND", +} +PASSTHROUGH_OUTPUT_URI = ( + "gs://litellm-bucket/litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/u/" + "predictions.jsonl" +) +UNGROUNDED_USAGE_METADATA = { + "promptTokenCount": 20, + "candidatesTokenCount": 48, + "thoughtsTokenCount": 195, + "toolUsePromptTokenCount": 73, + "totalTokenCount": 336, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 20}], + "trafficType": "ON_DEMAND", +} + + +def _batch(output_file_id: str) -> Batch: + return Batch( + id="b", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="f", + object="batch", + status="completed", + output_file_id=output_file_id, + ) + + +def _vertex_jsonl(rows: list[dict]) -> bytes: + return "\n".join(json.dumps(row) for row in rows).encode() + + +def _vertex_openai_row(custom_id: str, model: str, prompt_tokens: int, completion_tokens: int) -> dict: + return { + "id": f"batch_req_{custom_id}", + "custom_id": custom_id, + "response": { + "status_code": 200, + "request_id": custom_id, + "body": { + "id": f"chatcmpl-{custom_id}", + "object": "chat.completion", + "model": model, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + }, + }, + }, + "error": None, + } + + +def _native_vertex_row(usage_metadata: dict, *, grounded: bool, model_version: str | None = "gemini-2.5-flash"): + candidate = {"content": {"role": "model", "parts": [{"text": "ok"}]}, "finishReason": "STOP"} + grounding = {"groundingMetadata": {"webSearchQueries": ["q"]}} if grounded else {} + response = {"candidates": [{**candidate, **grounding}], "usageMetadata": usage_metadata} + return { + "request": {"contents": [{"role": "user", "parts": [{"text": "q"}]}], "tools": [{"googleSearch": {}}]}, + "status": "", + "response": {**response, **({"modelVersion": model_version} if model_version else {})}, + "processed_time": "2026-09-23T19:02:00.000+00:00", + } + + +def _capture_cost_calls(monkeypatch, prompt_cost=0.5, completion_cost=0.25) -> list: + import litellm.cost_calculator as cc + + calls: list = [] + + def _calc(**kw): + calls.append(kw) + return (prompt_cost, completion_cost) + + monkeypatch.setattr(cc, "batch_cost_calculator", _calc) + return calls + + +def test_vertex_native_cost_bills_embedding_rows(monkeypatch): + monkeypatch.setitem(litellm.model_cost, "vertex_ai/gemini-embedding-2", {"input_cost_per_token_batches": 1e-7}) + rows = [ + { + "key": "id_1", + "status": "", + "request": {"content": {"parts": [{"text": "hello world"}]}}, + "response": {"embedding": {"values": [0.1, 0.2]}, "usageMetadata": {"promptTokenCount": 2}}, + }, + { + "key": "id_2", + "status": "", + "request": {"content": {"parts": [{"text": "hello"}]}}, + "response": {"embedding": {"values": [0.3]}, "tokenCount": "3"}, + }, + {"key": "id_3", "status": "INVALID_ARGUMENT", "request": {"content": {"parts": [{"text": ""}]}}}, + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-embedding-2") + + assert (result.successful_requests, result.failed_requests) == (2, 1) + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (5, 0, 5) + assert result.cost == pytest.approx(5 * 1e-7) + assert result.models == ["gemini-embedding-2"] + + +@pytest.mark.asyncio +async def test_native_vertex_rows_route_to_vertex_cost_path_without_flag(monkeypatch): + monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False) + monkeypatch.setattr( + bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run") + ) + calls = _capture_cost_calls(monkeypatch) + rows = [ + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False), + ] + + result = await bu.calculate_batch_cost_and_usage( + file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash" + ) + + assert result.cost == pytest.approx(1.5) + assert (result.successful_requests, result.failed_requests) == (2, 0) + assert result.models == ["gemini-2.5-flash"] + assert {(call["model"], call["custom_llm_provider"]) for call in calls} == {("gemini-2.5-flash", "vertex_ai")} + + +@pytest.mark.asyncio +async def test_openai_shaped_vertex_rows_keep_the_generic_path_without_flag(monkeypatch): + monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False) + monkeypatch.setattr( + bu, "calculate_vertex_ai_batch_cost_and_usage", lambda *a, **kw: pytest.fail("native path should not run") + ) + _capture_cost_calls(monkeypatch) + rows = [_vertex_openai_row("request-1", "gemini-2.5-flash", 10, 5)] + + result = await bu.calculate_batch_cost_and_usage( + file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash" + ) + + assert result.successful_requests == 1 + + +@pytest.mark.asyncio +async def test_native_vertex_rows_on_another_provider_keep_the_generic_path(monkeypatch): + monkeypatch.setattr( + bu, "calculate_vertex_ai_batch_cost_and_usage", lambda *a, **kw: pytest.fail("native path should not run") + ) + _capture_cost_calls(monkeypatch) + + result = await bu.calculate_batch_cost_and_usage( + file_content_dictionary=[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)], + custom_llm_provider="openai", + ) + + assert result.successful_requests == 0 + + +@pytest.mark.asyncio +async def test_handle_completed_batch_routes_native_rows_without_flag(monkeypatch): + monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False) + raw_rows = [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)] + + async def fake_fetch(batch, custom_llm_provider, litellm_params=None): + return _vertex_jsonl(raw_rows) + + monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) + monkeypatch.setattr( + bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run") + ) + calls = _capture_cost_calls(monkeypatch, prompt_cost=0.7, completion_cost=0.3) + deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6} + + result = await bu._handle_completed_batch( + _batch(PASSTHROUGH_OUTPUT_URI), + custom_llm_provider="vertex_ai", + model_name="gemini-2.5-flash", + model_info=deployment_model_info, + ) + + assert result.cost == pytest.approx(1.0) + assert result.usage.total_tokens == 557 + assert [call["model_info"] for call in calls] == [deployment_model_info] + + +def test_native_vertex_usage_is_billed_like_the_online_path(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + grounded = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True) + ungrounded = _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False) + + result = bu.calculate_vertex_ai_batch_cost_and_usage([grounded, ungrounded], "gemini-2.5-flash") + + grounded_usage, ungrounded_usage = (call["usage"] for call in calls) + assert grounded_usage.prompt_tokens == 19 + assert grounded_usage.completion_tokens == 59 + 406 + assert grounded_usage.completion_tokens_details.reasoning_tokens == 406 + assert ungrounded_usage.prompt_tokens == 20 + 73 + assert ungrounded_usage.completion_tokens == 48 + 195 + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == ( + 19 + 93, + 465 + 243, + 557 + 336, + ) + + +def test_native_vertex_rows_are_priced_by_model_version_without_a_model_name(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + rows = [ + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash"), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version="gemini-2.5-pro"), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version=None), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows) + + assert [call["model"] for call in calls] == ["gemini-2.5-flash", "gemini-2.5-pro"] + assert result.models == ["gemini-2.5-flash", "gemini-2.5-pro"] + assert result.cost == pytest.approx(1.5) + assert result.successful_requests == 3 + assert result.usage.total_tokens == 557 + 336 + 336 + + +def test_native_vertex_rows_without_usage_metadata_count_as_failed(monkeypatch): + _capture_cost_calls(monkeypatch) + rows = [ + {"request": {"contents": []}, "status": "Error: bad request", "processed_time": "t"}, + {"request": {"contents": []}, "response": {"candidates": []}}, + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") + + assert (result.successful_requests, result.failed_requests) == (1, 2) + assert result.usage.total_tokens == 557 + + +def test_native_vertex_batch_whose_rows_all_failed_still_names_the_deployment_model(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + rows = [{"request": {"contents": []}, "status": "Error: quota exceeded", "processed_time": "t"}] * 2 + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") + + assert result.models == ["gemini-2.5-flash"] + assert (result.successful_requests, result.failed_requests, result.cost) == (0, 2, 0.0) + assert calls == [] + + +def test_native_vertex_rows_are_priced_with_the_deployment_model_info(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6} + + bu.calculate_vertex_ai_batch_cost_and_usage( + [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)], + "gemini-2.5-flash", + model_info=deployment_model_info, + ) + + assert [call["model_info"] for call in calls] == [deployment_model_info] + + +@pytest.mark.asyncio +async def test_native_vertex_rows_keep_the_deployment_model_info_through_the_batch_entrypoint(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + deployment_model_info = {"input_cost_per_token_batches": 1e-6} + + await bu.calculate_batch_cost_and_usage( + file_content_dictionary=[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)], + custom_llm_provider="vertex_ai", + model_name="gemini-2.5-flash", + model_info=deployment_model_info, + ) + + assert [call["model_info"] for call in calls] == [deployment_model_info] + + +def test_native_vertex_rows_are_priced_by_the_deployment_model_over_model_version(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + rows = [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-pro")] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") + + assert [call["model"] for call in calls] == ["gemini-2.5-flash"] + assert result.models == ["gemini-2.5-flash"] + + +def test_native_vertex_rows_that_fail_response_validation_count_as_failed(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + rows = [ + {"request": {"contents": []}, "response": {"candidates": "nope", "usageMetadata": GROUNDED_USAGE_METADATA}}, + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") + + assert (result.successful_requests, result.failed_requests) == (1, 1) + assert result.usage.total_tokens == 557 + assert len(calls) == 1 + + +@pytest.mark.parametrize("wildcard_model", ["*", "vertex_ai/*"]) +def test_native_vertex_rows_under_a_wildcard_deployment_are_priced_by_model_version(monkeypatch, wildcard_model): + calls = _capture_cost_calls(monkeypatch) + rows = [ + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash"), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version=None), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, wildcard_model) + + assert [call["model"] for call in calls] == ["gemini-2.5-flash", wildcard_model] + assert result.cost == pytest.approx(1.5) + assert (result.successful_requests, result.failed_requests) == (2, 0) + assert result.usage.total_tokens == 557 + 336 + + +def test_native_vertex_row_without_model_version_under_a_wildcard_deployment_bills_its_explicit_prices(): + deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6} + with_version = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash") + without_version = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version=None) + + twin = bu.calculate_vertex_ai_batch_cost_and_usage([with_version], "vertex_ai/*", model_info=deployment_model_info) + both = bu.calculate_vertex_ai_batch_cost_and_usage( + [with_version, without_version], "vertex_ai/*", model_info=deployment_model_info + ) + + assert twin.cost > 0 + assert both.cost == pytest.approx(2 * twin.cost) + assert (both.successful_requests, both.failed_requests) == (2, 0) + + +def test_native_vertex_row_the_cost_map_cannot_price_is_billed_at_zero_and_the_rest_still_bills(monkeypatch): + import litellm.cost_calculator as cc + + def _calc(**kw): + if kw["model"] == "gemini-unpriced": + raise ValueError("no pricing") + return (0.5, 0.25) + + monkeypatch.setattr(cc, "batch_cost_calculator", _calc) + rows = [ + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-unpriced"), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version="gemini-2.5-flash"), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows) + + assert result.cost == pytest.approx(0.75) + assert (result.successful_requests, result.failed_requests) == (2, 0) + assert result.usage.total_tokens == 557 + 336 + assert result.models == ["gemini-unpriced", "gemini-2.5-flash"] + + +@pytest.mark.asyncio +async def test_flag_sends_every_vertex_row_down_the_native_path_when_a_model_is_known(monkeypatch): + monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False) + monkeypatch.setattr( + bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run") + ) + calls = _capture_cost_calls(monkeypatch) + rows = [_vertex_openai_row("request-1", "gemini-2.5-flash", 10, 5)] + + result = await bu.calculate_batch_cost_and_usage( + file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash" + ) + + assert calls == [] + assert (result.successful_requests, result.failed_requests) == (0, 1) diff --git a/tests/test_litellm/files/__init__.py b/tests/test_litellm/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/files/test_main.py b/tests/test_litellm/files/test_main.py new file mode 100644 index 00000000000..2704bdfb6ff --- /dev/null +++ b/tests/test_litellm/files/test_main.py @@ -0,0 +1,71 @@ +from typing import Final +from urllib.parse import parse_qs, urlparse + +import httpx +import pytest + +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler + +NATIVE_VERTEX_ROWS: Final = ( + b'{"request": {"contents": [{"role": "user", "parts": [{"text": "Who won the 2024 Tour de France?"}]}],' + b' "tools": [{"googleSearch": {"excludeDomains": ["example.com"]}}]}}\n' + b'{"request": {"contents": [{"role": "user", "parts": [{"text": "What is the tallest building in Tokyo?"}]}],' + b' "tools": [{"googleSearch": {}}]}}\n' +) + + +@pytest.mark.parametrize( + "custom_llm_provider, purpose", + [("openai", "batch"), ("vertex_ai", "assistants")], + ids=["non-vertex-provider", "non-batch-purpose"], +) +def test_create_file_passthrough_is_rejected_outside_a_vertex_batch(custom_llm_provider, purpose): + with pytest.raises(litellm.BadRequestError) as exc_info: + litellm.create_file( + file=("batch.jsonl", b'{"request": {"contents": []}}\n', "application/jsonl"), + purpose=purpose, + custom_llm_provider=custom_llm_provider, + passthrough=True, + api_key="sk-test", + api_base="http://127.0.0.1:9", + ) + + assert "vertex_ai" in str(exc_info.value) + assert "batch" in str(exc_info.value) + + +def _gcs_upload_transport(uploads: list[httpx.Request]) -> httpx.MockTransport: + def respond(request: httpx.Request) -> httpx.Response: + uploads.append(request) + object_name: Final = parse_qs(urlparse(str(request.url)).query)["name"][0] + return httpx.Response( + 200, + json={ + "id": f"my-bucket/{object_name}/1758585600000000", + "name": object_name, + "size": str(len(request.read())), + "timeCreated": "2026-09-23T00:00:00.000Z", + }, + ) + + return httpx.MockTransport(respond) + + +def test_create_file_passthrough_kwarg_ships_native_rows_byte_for_byte_under_the_passthrough_prefix(): + uploads: Final[list[httpx.Request]] = [] + file_object = litellm.create_file( + file=("batch.jsonl", NATIVE_VERTEX_ROWS, "application/jsonl"), + purpose="batch", + custom_llm_provider="vertex_ai", + passthrough=True, + model="vertex_ai/gemini-2.5-flash", + gcs_bucket_name="my-bucket", + api_key="test-token", + client=HTTPHandler(client=httpx.Client(transport=_gcs_upload_transport(uploads))), + ) + (upload,) = uploads + object_name: Final = parse_qs(urlparse(str(upload.url)).query)["name"][0] + assert upload.read() == NATIVE_VERTEX_ROWS + assert object_name.startswith("litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/") + assert file_object.id == f"gs://my-bucket/{object_name}" diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py index e6126b02790..ae045edbec5 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py @@ -19,7 +19,6 @@ import pytest from litellm.llms.vertex_ai.batches.transformation import ( # noqa: E402 VertexAIBatchTransformation, - vertex_prompt_tokens_details, ) from litellm.llms.vertex_ai.common_utils import ( # noqa: E402 VertexAIError, @@ -41,27 +40,6 @@ ENDPOINT_INPUT_FILE = ( ) -def test_vertex_prompt_tokens_details_rejects_malformed_details(): - assert vertex_prompt_tokens_details({"promptTokensDetails": [1]}) is None - assert vertex_prompt_tokens_details({"promptTokensDetails": [{"modality": "AUDIO"}]}) is None - assert ( - vertex_prompt_tokens_details( - { - "promptTokensDetails": [ - {"modality": "AUDIO", "tokenCount": 1}, - "malformed", - ] - } - ) - is None - ) - - -# =========================================================================== # -# transform_openai_batch_request_to_vertex_ai_batch_request -# =========================================================================== # - - def test_transform_openai_request_builds_full_vertex_job(): with patch( "litellm.llms.vertex_ai.batches.transformation.uuid.uuid4", @@ -477,3 +455,19 @@ def test_list_response_none_jobs_treated_as_empty(): out = T.transform_vertex_ai_batch_list_response_to_openai_list_response({"batchPredictionJobs": None}) assert out["data"] == [] assert out["first_id"] is None + + +PASSTHROUGH_INPUT_FILE = ( + "gs://litellm-testing-bucket/litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/uuid-1" +) + + +def test_get_model_from_passthrough_gcs_file(): + assert T._get_model_from_gcs_file(PASSTHROUGH_INPUT_FILE) == "publishers/google/models/gemini-2.5-flash" + + +def test_get_gcs_uri_prefix_keeps_passthrough_segment_so_output_lands_beside_input(): + assert ( + T._get_gcs_uri_prefix_from_file(PASSTHROUGH_INPUT_FILE) + == "gs://litellm-testing-bucket/litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash" + ) diff --git a/tests/test_litellm/llms/vertex_ai/files/__init__.py b/tests/test_litellm/llms/vertex_ai/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/files/test_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_transformation.py new file mode 100644 index 00000000000..6958576b0c8 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/files/test_transformation.py @@ -0,0 +1,310 @@ +import io +import json +import urllib.parse +from pathlib import Path +from unittest.mock import MagicMock +from urllib.parse import parse_qs, urlparse + +import httpx +import pytest + +from litellm.llms.vertex_ai.common_utils import VertexAIError +from litellm.llms.vertex_ai.files.transformation import VertexAIFilesConfig, is_passthrough_managed_gcs_url + +NATIVE_VERTEX_ROW = json.dumps( + { + "request": { + "contents": [{"role": "user", "parts": [{"text": "What is the tallest building in the world?"}]}], + "tools": [{"googleSearch": {"excludeDomains": ["example.com"]}}], + } + } +).encode() +NATIVE_VERTEX_JSONL = NATIVE_VERTEX_ROW + b"\n" + NATIVE_VERTEX_ROW + b"\n" +OPENAI_BATCH_JSONL = ( + b'{"custom_id": "r1", "method": "POST", "url": "/v1/chat/completions",' + b' "body": {"model": "gemini-2.5-flash", "messages": [{"role": "user", "content": "hi"}]}}\n' +) +PASSTHROUGH_OBJECT = ( + "litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/uuid-1/predictions.jsonl" +) +TRANSFORMED_OBJECT = "litellm-vertex-files/publishers/google/models/gemini-2.5-flash/uuid-1/predictions.jsonl" +UPLOAD_CHUNK_BYTES = 1024 * 1024 + + +@pytest.fixture +def config() -> VertexAIFilesConfig: + return VertexAIFilesConfig() + + +def _gcs_media_url(object_name: str) -> str: + return ( + f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{urllib.parse.quote(object_name, safe='')}?alt=media" + ) + + +def _native_output_jsonl() -> bytes: + return ( + json.dumps( + { + "request": json.loads(NATIVE_VERTEX_ROW)["request"], + "status": "", + "response": { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "The Burj Khalifa."}]}, + "finishReason": "STOP", + "groundingMetadata": {"webSearchQueries": ["tallest building in the world"]}, + } + ], + "modelVersion": "gemini-2.5-flash", + "usageMetadata": {"promptTokenCount": 20, "candidatesTokenCount": 48, "totalTokenCount": 68}, + }, + "processed_time": "2026-09-23T19:02:00.000+00:00", + } + ).encode() + + b"\n" + ) + + +def _upload_chunks(config: VertexAIFilesConfig, file: object, litellm_params: dict) -> list[bytes]: + body = config.transform_create_file_request( + model="", + create_file_data={"file": file, "purpose": "batch"}, + optional_params={}, + litellm_params=litellm_params, + ) + return list(body["streaming_media_upload"]["body_stream"].iter_bytes()) + + +def _upload_body_bytes(config: VertexAIFilesConfig, file: object, litellm_params: dict) -> bytes: + return b"".join(_upload_chunks(config, file, litellm_params)) + + +class TestPassthroughBatchUpload: + """`passthrough=True` on a batch upload ships the caller's native Vertex JSONL + to GCS byte for byte, filed under a `passthrough/` object path so the batch + output that lands beside it is recognized and returned untouched as well.""" + + def _upload_url(self, config, litellm_params, file, purpose="batch") -> str: + return config.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params=litellm_params, + data={"file": file, "purpose": purpose}, + ) + + def test_passthrough_object_is_filed_under_passthrough_prefix_named_by_deployment_model(self, config): + url = self._upload_url( + config, + {"gcs_bucket_name": "my-bucket", "model": "vertex_ai/gemini-2.5-flash", "passthrough": True}, + ("batch.jsonl", NATIVE_VERTEX_JSONL, "application/jsonl"), + ) + object_name = parse_qs(urlparse(url).query)["name"][0] + assert object_name.startswith("litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/") + + def test_passthrough_upload_without_deployment_model_is_rejected(self, config): + with pytest.raises(VertexAIError) as exc_info: + self._upload_url( + config, + {"gcs_bucket_name": "my-bucket", "passthrough": True}, + ("batch.jsonl", NATIVE_VERTEX_JSONL, "application/jsonl"), + ) + assert exc_info.value.status_code == 400 + assert "target_model_names" in exc_info.value.message + + def test_passthrough_flag_does_not_ship_a_non_batch_upload_raw(self, config): + result = config.transform_create_file_request( + model="", + create_file_data={"file": ("notes.txt", b"plain text", "text/plain"), "purpose": "user_data"}, + optional_params={}, + litellm_params={"gcs_bucket_name": "my-bucket", "passthrough": True}, + ) + assert result == b"plain text" + + def test_passthrough_flag_is_ignored_for_non_batch_purposes(self, config): + url = self._upload_url( + config, + {"gcs_bucket_name": "my-bucket", "model": "vertex_ai/gemini-2.5-flash", "passthrough": True}, + ("notes.txt", b"plain text", "text/plain"), + purpose="user_data", + ) + object_name = parse_qs(urlparse(url).query)["name"][0] + assert object_name.startswith("litellm-vertex-files/uploads/") + assert "passthrough" not in object_name + + @pytest.mark.parametrize( + "file", + [ + ("batch.jsonl", NATIVE_VERTEX_JSONL, "application/jsonl"), + NATIVE_VERTEX_JSONL, + ("batch.jsonl", io.BytesIO(NATIVE_VERTEX_JSONL), "application/jsonl"), + ("batch.jsonl", NATIVE_VERTEX_JSONL.decode(), "application/jsonl"), + ], + ids=["bytes-tuple", "bare-bytes", "handle-tuple", "text-tuple"], + ) + def test_passthrough_upload_body_is_the_callers_bytes(self, config, file): + body = config.transform_create_file_request( + model="", + create_file_data={"file": file, "purpose": "batch"}, + optional_params={}, + litellm_params={"passthrough": True}, + ) + stream = body["streaming_media_upload"]["body_stream"] + assert b"".join(stream.iter_bytes()) == NATIVE_VERTEX_JSONL + assert b"".join(stream.iter_bytes()) == NATIVE_VERTEX_JSONL + assert body["streaming_media_upload"]["content_type"] == "application/json" + + def test_passthrough_upload_streams_a_large_handle_in_bounded_chunks(self, config): + content = NATIVE_VERTEX_ROW * (3 * UPLOAD_CHUNK_BYTES // len(NATIVE_VERTEX_ROW) + 1) + chunks = _upload_chunks( + config, ("batch.jsonl", io.BytesIO(content), "application/jsonl"), {"passthrough": True} + ) + assert len(chunks) >= 3 + assert max(len(chunk) for chunk in chunks) <= UPLOAD_CHUNK_BYTES + assert b"".join(chunks) == content + + def test_passthrough_upload_streams_a_path_in_bounded_chunks(self, config, tmp_path: Path): + content = NATIVE_VERTEX_ROW * (2 * UPLOAD_CHUNK_BYTES // len(NATIVE_VERTEX_ROW) + 1) + batch_path = tmp_path / "batch.jsonl" + batch_path.write_bytes(content) + chunks = _upload_chunks(config, ("batch.jsonl", batch_path, "application/jsonl"), {"passthrough": True}) + assert len(chunks) >= 2 + assert max(len(chunk) for chunk in chunks) <= UPLOAD_CHUNK_BYTES + assert b"".join(chunks) == content + + def test_passthrough_upload_rejects_a_non_seekable_handle(self, config): + class _Pipe: + def read(self, size=-1): + return b"" + + with pytest.raises(ValueError, match="seekable"): + _upload_body_bytes(config, ("batch.jsonl", _Pipe(), "application/jsonl"), {"passthrough": True}) + + def test_passthrough_upload_rejects_content_that_is_neither_bytes_path_nor_handle(self, config): + with pytest.raises(ValueError, match="Unsupported file content type"): + _upload_body_bytes(config, ("batch.jsonl", 42, "application/jsonl"), {"passthrough": True}) + + def test_openai_rows_are_translated_unless_passthrough_is_set(self, config): + file = ("batch.jsonl", OPENAI_BATCH_JSONL, "application/jsonl") + translated = _upload_body_bytes(config, file, {}) + untouched = _upload_body_bytes(config, file, {"passthrough": True}) + assert untouched == OPENAI_BATCH_JSONL + assert translated != OPENAI_BATCH_JSONL + assert b'"contents"' in translated + + def test_passthrough_output_content_is_returned_untouched(self, config): + raw_jsonl = _native_output_jsonl() + + def _download(object_name: str) -> bytes: + raw_response = httpx.Response( + status_code=200, + content=raw_jsonl, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request("GET", _gcs_media_url(object_name)), + ) + result = config.transform_file_content_response( + raw_response=raw_response, logging_obj=MagicMock(), litellm_params={} + ) + return result.response.content + + assert _download(PASSTHROUGH_OBJECT) == raw_jsonl + assert _download(f"team-a/{PASSTHROUGH_OBJECT}") == raw_jsonl + transformed = _download(TRANSFORMED_OBJECT) + assert transformed != raw_jsonl + assert json.loads(transformed.splitlines()[0])["response"]["body"]["choices"] + nested = _download(f"litellm-vertex-files/{PASSTHROUGH_OBJECT}") + assert nested != raw_jsonl + assert json.loads(nested.splitlines()[0])["response"]["body"]["choices"] + + def test_output_of_an_upload_whose_model_smuggles_the_passthrough_segment_is_still_transformed(self, config): + smuggled_model = b"litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash" + upload_url = self._upload_url( + config, + {"gcs_bucket_name": "my-bucket"}, + ("batch.jsonl", OPENAI_BATCH_JSONL.replace(b"gemini-2.5-flash", smuggled_model), "application/jsonl"), + ) + object_name = parse_qs(urlparse(upload_url).query)["name"][0] + raw_jsonl = _native_output_jsonl() + raw_response = httpx.Response( + status_code=200, + content=raw_jsonl, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request("GET", _gcs_media_url(f"{object_name}/predictions.jsonl")), + ) + + result = config.transform_file_content_response( + raw_response=raw_response, logging_obj=MagicMock(), litellm_params={} + ) + + assert object_name.startswith("litellm-vertex-files/litellm-vertex-files/passthrough/") + assert json.loads(result.response.content.splitlines()[0])["response"]["body"]["choices"] + + @pytest.mark.parametrize( + "url, expected", + [ + (f"gs://my-bucket/{PASSTHROUGH_OBJECT}", True), + (f"gs://my-bucket/team-a/{PASSTHROUGH_OBJECT}", True), + (f"gs://my-bucket/litellm-vertex-files/{PASSTHROUGH_OBJECT}", False), + (_gcs_media_url(f"team-a/sub/{PASSTHROUGH_OBJECT}"), True), + (_gcs_media_url(f"litellm-vertex-files/publishers/google/models/x/{PASSTHROUGH_OBJECT}"), False), + (_gcs_media_url(TRANSFORMED_OBJECT), False), + ], + ids=["gs", "gs-prefixed", "gs-smuggled", "https-prefixed", "https-model-path-smuggled", "https-transformed"], + ) + def test_passthrough_detection_anchors_on_the_first_managed_segment(self, url, expected): + assert is_passthrough_managed_gcs_url(url) is expected + + @pytest.mark.asyncio + async def test_passthrough_output_stream_is_returned_untouched(self, config): + stream_iterator = object() + headers = {"content-type": "application/octet-stream"} + result = await config.transform_file_content_stream( + stream_iterator=stream_iterator, + headers=headers, + request_url=_gcs_media_url(f"team-a/{PASSTHROUGH_OBJECT}"), + logging_obj=MagicMock(), + litellm_params={}, + ) + assert result.stream_iterator is stream_iterator + assert result.headers == headers + + +class TestEmbeddingOutputTranslation: + EMBEDDING_OBJECT = ( + "litellm-vertex-files/publishers/google/models/gemini-embedding-2/prediction-model-1/predictions.jsonl" + ) + + def _transform(self, config: VertexAIFilesConfig, rows: list[dict]) -> list[dict]: + raw_response = httpx.Response( + status_code=200, + content="\n".join(json.dumps(row) for row in rows).encode(), + headers={"content-type": "application/octet-stream"}, + request=httpx.Request("GET", _gcs_media_url(self.EMBEDDING_OBJECT)), + ) + result = config.transform_file_content_response( + raw_response=raw_response, logging_obj=MagicMock(), litellm_params={} + ) + return [json.loads(line) for line in result.response.content.decode().splitlines()] + + def test_embedding_rows_become_openai_batch_rows_billed_by_their_prompt_tokens(self, config): + live_row = { + "key": "request-1", + "request": {"content": {"parts": [{"text": "hello world"}]}}, + "response": {"embedding": {"values": [-0.015, 0.024]}, "usageMetadata": {"promptTokenCount": 2}}, + } + documented_row = { + "key": "request-2", + "request": {"content": {"parts": [{"text": "hello"}]}}, + "response": {"embedding": {"values": [0.5]}, "tokenCount": "3"}, + } + + live, documented = self._transform(config, [live_row, documented_row]) + + assert (live["custom_id"], live["error"], live["response"]["status_code"]) == ("request-1", None, 200) + assert live["response"]["body"]["model"] == "gemini-embedding-2" + assert live["response"]["body"]["data"] == [{"embedding": [-0.015, 0.024], "index": 0, "object": "embedding"}] + live_usage, documented_usage = (row["response"]["body"]["usage"] for row in (live, documented)) + assert (live_usage["prompt_tokens"], live_usage["total_tokens"]) == (2, 2) + assert (documented_usage["prompt_tokens"], documented_usage["total_tokens"]) == (3, 3) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_batch_file_validation.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_batch_file_validation.py index f5542fc0446..3a73c39e177 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_batch_file_validation.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_batch_file_validation.py @@ -4,7 +4,8 @@ import pytest from litellm.proxy._types import ProxyException from litellm.proxy.openai_files_endpoints.batch_file_validation import ( - BATCH_LINE_REQUIRED_KEYS, + BATCH_LINE_SHAPE, + PASSTHROUGH_BATCH_LINE_SHAPE, BatchFileEmpty, BatchFileInvalidJsonLine, BatchFileLineNotObject, @@ -96,7 +97,7 @@ def test_non_object_line_rejected(): assert check_batch_file_upload("batch.jsonl", content, None) == BatchFileLineNotObject(line_number=2) -@pytest.mark.parametrize("missing_key", BATCH_LINE_REQUIRED_KEYS) +@pytest.mark.parametrize("missing_key", BATCH_LINE_SHAPE.required_keys) def test_missing_required_key_rejected(missing_key): import json @@ -174,3 +175,36 @@ def test_failures_map_to_openai_shaped_proxy_exceptions(failure, expected_code, assert exc_info.value.param == expected_param for fragment in expected_fragments: assert fragment in exc_info.value.message + + +NATIVE_VERTEX_LINE = b'{"request": {"contents": [{"role": "user", "parts": [{"text": "hi"}]}]}}' + + +def test_passthrough_keys_accept_native_vertex_rows(): + content = NATIVE_VERTEX_LINE + b"\n" + NATIVE_VERTEX_LINE + b"\n" + assert check_batch_file_upload("batch.jsonl", content, None, PASSTHROUGH_BATCH_LINE_SHAPE) is None + + +def test_passthrough_keys_reject_openai_rows(): + content = NATIVE_VERTEX_LINE + b"\n" + VALID_LINE + b"\n" + assert check_batch_file_upload( + "batch.jsonl", content, None, PASSTHROUGH_BATCH_LINE_SHAPE + ) == BatchFileMissingLineKey(line_number=2, key="request", line_shape=PASSTHROUGH_BATCH_LINE_SHAPE) + + +def test_default_keys_still_reject_native_vertex_rows(): + assert check_batch_file_upload("batch.jsonl", NATIVE_VERTEX_LINE, None) == BatchFileMissingLineKey( + line_number=1, key="custom_id" + ) + + +def test_passthrough_missing_key_message_says_what_a_passthrough_upload_takes(): + with pytest.raises(ProxyException) as exc_info: + raise_batch_file_validation_failure( + BatchFileMissingLineKey(line_number=3, key="request", line_shape=PASSTHROUGH_BATCH_LINE_SHAPE) + ) + assert exc_info.value.param == "request" + assert "line 3" in exc_info.value.message + assert "passthrough upload takes native Vertex batch rows" in exc_info.value.message + assert "with a request key." in exc_info.value.message + assert "custom_id" not in exc_info.value.message diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 48699b47e7f..84cf4ea7c32 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -5669,3 +5669,241 @@ def test_model_routed_file_retrieve_allows_key_with_model_grant(mocker: MockerFi assert response.status_code == 200, response.text assert captured_kwargs["api_key"] == "mistral-key" assert captured_kwargs["custom_llm_provider"] == "mistral" + + +NATIVE_VERTEX_BATCH_LINE = ( + b'{"request": {"contents": [{"role": "user", "parts": [{"text": "What is the tallest building?"}]}],' + b' "tools": [{"googleSearch": {"excludeDomains": ["example.com"]}}]}}\n' +) + + +def _passthrough_router() -> Router: + return Router( + model_list=[ + { + "model_name": "vertex-batch", + "litellm_params": { + "model": "vertex_ai/gemini-2.5-flash", + "vertex_project": "proj", + "vertex_location": "us-central1", + }, + "model_info": {"id": "vertex-batch-id"}, + }, + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "openai/gpt-3.5-turbo", "api_key": "openai_api_key"}, + "model_info": {"id": "gpt-3.5-turbo-id"}, + }, + ] + ) + + +def _setup_passthrough_upload_endpoint(monkeypatch, llm_router: Router) -> list: + """Like _setup_batch_upload_endpoint, but reads the forwarded file bytes while the spool is open.""" + from litellm.proxy.openai_files_endpoints import files_endpoints as fe + + forwarded_calls = _setup_batch_upload_endpoint(monkeypatch, llm_router) + + async def fake_route_create_file(**kwargs): + upload_source = kwargs["_create_file_request"]["file"][1] + upload_source.seek(0) + forwarded_calls.append({**kwargs, "file_bytes": upload_source.read()}) + return OpenAIFileObject( + id="dummy-id", + object="file", + bytes=0, + created_at=1234567890, + filename="batch.jsonl", + purpose="batch", + status="uploaded", + ) + + monkeypatch.setattr(fe, "route_create_file", fake_route_create_file) + return forwarded_calls + + +def _upload(content: bytes, form: dict): + return client.post( + "/v1/files", + files={"file": ("batch.jsonl", content, "application/jsonl")}, + data=form, + headers={"Authorization": "Bearer test-key"}, + ) + + +def test_create_file_passthrough_forwards_native_vertex_rows_untouched(monkeypatch): + forwarded_calls = _setup_passthrough_upload_endpoint(monkeypatch, _passthrough_router()) + content = NATIVE_VERTEX_BATCH_LINE * 2 + + try: + response = _upload(content, {"purpose": "batch", "target_model_names": "vertex-batch", "passthrough": "true"}) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 200, response.text + (call,) = forwarded_calls + assert call["_create_file_request"]["passthrough"] is True + assert call["file_bytes"] == content + assert call["target_model_names_list"] == ["vertex-batch"] + + +def test_create_file_passthrough_rejects_rows_without_a_request(monkeypatch): + forwarded_calls = _setup_passthrough_upload_endpoint(monkeypatch, _passthrough_router()) + + try: + response = _upload( + NATIVE_VERTEX_BATCH_LINE + VALID_BATCH_LINE, + {"purpose": "batch", "target_model_names": "vertex-batch", "passthrough": "true"}, + ) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 400, response.text + error = response.json()["error"] + assert error["param"] == "request" + assert "line 2" in error["message"] + assert forwarded_calls == [] + + +def test_create_file_without_passthrough_still_rejects_native_vertex_rows(monkeypatch): + forwarded_calls = _setup_passthrough_upload_endpoint(monkeypatch, _passthrough_router()) + + try: + response = _upload(NATIVE_VERTEX_BATCH_LINE, {"purpose": "batch", "target_model_names": "vertex-batch"}) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 400, response.text + assert response.json()["error"]["param"] == "custom_id" + assert forwarded_calls == [] + + +@pytest.mark.parametrize( + "form, expected_param, expected_fragment", + [ + ({"purpose": "batch", "passthrough": "true"}, "target_model_names", "target_model_names"), + ( + {"purpose": "batch", "target_model_names": "gpt-3.5-turbo", "passthrough": "true"}, + "target_model_names", + "'gpt-3.5-turbo'", + ), + ( + {"purpose": "batch", "target_model_names": "vertex-batch,gpt-3.5-turbo", "passthrough": "true"}, + "target_model_names", + "'gpt-3.5-turbo'", + ), + ({"purpose": "user_data", "target_model_names": "vertex-batch", "passthrough": "true"}, "passthrough", "batch"), + ( + {"purpose": "batch", "target_model_names": "vertex-batch", "passthrough": "true", "target_storage": "s3"}, + "target_storage", + "'s3'", + ), + ( + {"purpose": "batch", "model": "gpt-3.5-turbo", "passthrough": "true"}, + "model", + "'gpt-3.5-turbo'", + ), + ], + ids=[ + "no-model", + "non-vertex-model", + "mixed-models", + "non-batch-purpose", + "target-storage", + "non-vertex-model-param", + ], +) +def test_create_file_passthrough_rejected_outside_a_vertex_batch(monkeypatch, form, expected_param, expected_fragment): + forwarded_calls = _setup_passthrough_upload_endpoint(monkeypatch, _passthrough_router()) + + try: + response = _upload(NATIVE_VERTEX_BATCH_LINE, form) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 400, response.text + error = response.json()["error"] + assert error["type"] == "invalid_request_error" + assert error["param"] == expected_param + assert expected_fragment in error["message"] + assert forwarded_calls == [] + + +def test_create_file_passthrough_accepts_the_model_param_as_the_deployment(monkeypatch): + forwarded_calls = _setup_passthrough_upload_endpoint(monkeypatch, _passthrough_router()) + + try: + response = _upload( + NATIVE_VERTEX_BATCH_LINE, {"purpose": "batch", "model": "vertex-batch", "passthrough": "true"} + ) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 200, response.text + (call,) = forwarded_calls + assert call["model"] == "vertex-batch" + assert call["_create_file_request"]["passthrough"] is True + + +def test_create_file_passthrough_rejects_a_model_group_with_a_non_vertex_deployment(monkeypatch): + mixed_router = Router( + model_list=[ + { + "model_name": "vertex-batch", + "litellm_params": { + "model": "vertex_ai/gemini-2.5-flash", + "vertex_project": "proj", + "vertex_location": "us-central1", + }, + "model_info": {"id": "vertex-batch-id"}, + }, + { + "model_name": "vertex-batch", + "litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": "openai_api_key"}, + "model_info": {"id": "vertex-batch-openai-id"}, + }, + ] + ) + forwarded_calls = _setup_passthrough_upload_endpoint(monkeypatch, mixed_router) + + try: + response = _upload( + NATIVE_VERTEX_BATCH_LINE, {"purpose": "batch", "target_model_names": "vertex-batch", "passthrough": "true"} + ) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 400, response.text + error = response.json()["error"] + assert error["param"] == "target_model_names" + assert "'vertex-batch'" in error["message"] + assert forwarded_calls == [] + + +def test_create_file_passthrough_fails_closed_when_guardrails_would_scan_the_batch(monkeypatch): + """Batch guardrails read OpenAI-shaped rows, so a passthrough upload on a guardrailed + key is refused rather than forwarded unscanned.""" + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy.utils import ProxyLogging + + class _Redactor(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + return data + + forwarded_calls = _setup_passthrough_upload_endpoint(monkeypatch, _passthrough_router()) + monkeypatch.setattr(litellm, "callbacks", [_Redactor(guardrail_name="g", default_on=True)]) + ProxyLogging._callback_capabilities_cache.clear() + + try: + response = _upload( + NATIVE_VERTEX_BATCH_LINE, {"purpose": "batch", "target_model_names": "vertex-batch", "passthrough": "true"} + ) + finally: + _teardown_batch_upload_endpoint() + ProxyLogging._callback_capabilities_cache.clear() + + assert response.status_code == 400, response.text + error = response.json()["error"] + assert error["param"] == "passthrough" + assert "guardrails" in error["message"] + assert forwarded_calls == [] diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 411377f29cf..f985b212b01 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -528,6 +528,39 @@ async def test_async_router_acreate_file_with_jsonl(): assert first_call_content == non_jsonl_content +@pytest.mark.asyncio +async def test_async_router_acreate_file_passthrough_keeps_the_file_and_forwards_the_flag(): + """A passthrough batch upload must reach the provider byte for byte: the router + neither rewrites body.model to the deployment model nor drops the flag.""" + from io import BytesIO + from unittest.mock import MagicMock, patch + + jsonl_content = b'{"custom_id": "r1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "vertex-batch"}}\n' + router = litellm.Router( + model_list=[ + { + "model_name": "vertex-batch", + "litellm_params": {"model": "vertex_ai/gemini-2.5-flash", "vertex_project": "p"}, + } + ], + ) + + with patch("litellm.acreate_file", return_value=MagicMock()) as mock_acreate_file: + await router.acreate_file( + model="vertex-batch", purpose="batch", file=BytesIO(jsonl_content), passthrough=True + ) + forwarded = mock_acreate_file.call_args.kwargs + assert forwarded["passthrough"] is True + forwarded["file"].seek(0) + assert forwarded["file"].read() == jsonl_content + + mock_acreate_file.reset_mock() + await router.acreate_file(model="vertex-batch", purpose="batch", file=BytesIO(jsonl_content)) + rewritten = mock_acreate_file.call_args.kwargs["file"] + rewritten.seek(0) + assert b'"gemini-2.5-flash"' in rewritten.read() + + @pytest.mark.asyncio async def test_async_router_acreate_file_does_not_fall_back_across_model_groups(): """A file created for batches only exists under the credentials of the model group diff --git a/tests/unit/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index a4de8eee23c..d1572f4a7c9 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -640,7 +640,7 @@ async def test_calculate_vertex_disable_transform_path(monkeypatch): monkeypatch.setattr( bu, "calculate_vertex_ai_batch_cost_and_usage", - lambda content, model: bu.BatchCostUsageResult( + lambda content, model, model_info=None: bu.BatchCostUsageResult( cost=9.9, usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3), models=["gemini-2.0-flash-001"], @@ -671,7 +671,7 @@ async def test_calculate_vertex_disable_transform_needs_model_name(monkeypatch): monkeypatch.setattr( bu, "calculate_vertex_ai_batch_cost_and_usage", - lambda content, model: pytest.fail("raw vertex path should not run"), + lambda content, model, model_info=None: pytest.fail("raw vertex path should not run"), ) result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=[], custom_llm_provider="vertex_ai") @@ -735,7 +735,11 @@ def test_vertex_batch_usage_preserves_modality_token_details(monkeypatch): ) responses = [ { + "key": "id_1", + "status": "", + "request": {"content": {"parts": [{"text": "hello"}, {"fileData": {"mimeType": "audio/wav"}}]}}, "response": { + "embedding": {"values": [0.1, 0.2]}, "usageMetadata": { "promptTokenCount": 84, "candidatesTokenCount": 0, @@ -744,13 +748,14 @@ def test_vertex_batch_usage_preserves_modality_token_details(monkeypatch): {"modality": "AUDIO", "tokenCount": 64}, {"modality": "TEXT", "tokenCount": 20}, ], - } - } + }, + }, } ] result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-embedding-2") + assert (result.successful_requests, result.usage.prompt_tokens) == (1, 84) assert result.prompt_cost == pytest.approx(64 * 3.25e-6 + 20 * 1e-7) @@ -1336,7 +1341,7 @@ async def test_handle_completed_batch_vertex_disable_transform_path(monkeypatch) monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False) seen: dict = {} - def fake_vertex_calc(content, model): + def fake_vertex_calc(content, model, model_info=None): seen["content"] = content seen["model"] = model return bu.BatchCostUsageResult( @@ -1358,7 +1363,7 @@ async def test_handle_completed_batch_vertex_disable_transform_path(monkeypatch) assert result.cost == 7.7 assert result.usage.total_tokens == 3 assert result.models == ["gemini-x"] - assert seen["content"] == raw_rows + assert list(seen["content"]) == raw_rows assert seen["model"] == "gemini-x" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 2938e65cfde..052cc693f8b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25526,6 +25526,11 @@ export interface components { file: string; /** Litellm Metadata */ litellm_metadata?: string | null; + /** + * Passthrough + * @default false + */ + passthrough: boolean; /** Purpose */ purpose: string; /** @@ -25550,6 +25555,11 @@ export interface components { file: string; /** Litellm Metadata */ litellm_metadata?: string | null; + /** + * Passthrough + * @default false + */ + passthrough: boolean; /** Purpose */ purpose: string; /** @@ -25574,6 +25584,11 @@ export interface components { file: string; /** Litellm Metadata */ litellm_metadata?: string | null; + /** + * Passthrough + * @default false + */ + passthrough: boolean; /** Purpose */ purpose: string; /**