diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index cdae3de47c8..62ccb07f695 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -49,6 +49,10 @@ def _uses_native_vertex_output( model_name: str | None, first_row: Mapping[str, object] | None, ) -> bool: + if custom_llm_provider == "gemini": + from litellm.llms.gemini.batches.transformation import is_native_gemini_batch_output_row + + return first_row is not None and is_native_gemini_batch_output_row(first_row) if custom_llm_provider != "vertex_ai": return False if model_name and litellm.disable_vertex_batch_output_transformation: @@ -831,6 +835,12 @@ def _get_batch_job_usage_from_response_body( _usage_dict: Final = response_body.get("usage", None) or {} if ResponseAPILoggingUtils.is_response_api_usage(_usage_dict): return ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(_usage_dict) + if custom_llm_provider == "gemini": + from litellm.llms.gemini.batches.transformation import is_gemini_batch_usage_shape, usage_from_batch_output + + gemini_usage: Final = cast(Mapping[str, object], _usage_dict) # cast-ok: usage block comes from provider JSON + if is_gemini_batch_usage_shape(gemini_usage): + return usage_from_batch_output(gemini_usage) usage: Final[Usage] = Usage(**_usage_dict) if custom_llm_provider == "xai": from litellm.llms.xai.chat.transformation import XAIChatConfig diff --git a/litellm/files/main.py b/litellm/files/main.py index d1ddf6c4c4c..3c73f6d5389 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -1091,10 +1091,10 @@ def file_content_streaming( client=client if isinstance(client, (OpenAI, AsyncOpenAI)) else None, litellm_params=optional_params.model_dump(exclude_none=True), ) - elif custom_llm_provider == LlmProviders.VERTEX_AI.value: + elif custom_llm_provider in (LlmProviders.VERTEX_AI.value, LlmProviders.GEMINI.value): if not _is_async: raise litellm.exceptions.BadRequestError( - message="Streaming 'file_content' for vertex_ai is only supported through 'afile_content'.", + message=f"Streaming 'file_content' for {custom_llm_provider} is only supported through 'afile_content'.", model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( @@ -1103,14 +1103,14 @@ def file_content_streaming( request=httpx.Request(method="file_content", url="https://github.com/BerriAI/litellm"), ), ) - vertex_files_config: Final = ProviderConfigManager.get_provider_files_config( + google_files_config: Final = ProviderConfigManager.get_provider_files_config( model="", - provider=LlmProviders.VERTEX_AI, + provider=LlmProviders(custom_llm_provider), ) - assert vertex_files_config is not None + assert google_files_config is not None response = base_llm_http_handler.async_retrieve_file_content_streaming( file_content_request=file_content_request, - provider_config=vertex_files_config, + provider_config=google_files_config, litellm_params=litellm_params, headers=extra_headers or {}, logging_obj=logging_obj, diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 011551faf0c..a5b15a24a7f 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -232,13 +232,13 @@ def _handler_may_close_client(client_refcount: int, owns_client: bool) -> bool: def _drop_streaming_anchor(_handler: object) -> None: - """Release a handler anchored to a streaming response. See ``_anchor_handler_to``. + """Release a handler anchored to a streaming response. See ``anchor_handler_to``. The work is the reference held until this point, so there is nothing to do here. """ -def _anchor_handler_to(response: httpx.Response, handler: object) -> None: +def anchor_handler_to(response: httpx.Response, handler: object) -> None: """Keep the handler alive for as long as a streaming response can still read. A body still arriving reads through the handler's connection pool, and closing @@ -927,7 +927,7 @@ class AsyncHTTPHandler: stream=True, follow_redirects=USE_CLIENT_DEFAULT if follow_redirects is None else follow_redirects, ) - _anchor_handler_to(response, self) + anchor_handler_to(response, self) response.raise_for_status() except httpx.TimeoutException as e: _raise_timeout(e, start_time) @@ -1024,7 +1024,7 @@ class AsyncHTTPHandler: ) response: Final = await self.client.send(req, stream=stream) if stream: - _anchor_handler_to(response, self) + anchor_handler_to(response, self) response.raise_for_status() return response except httpx.ConnectError: @@ -1201,7 +1201,7 @@ class AsyncHTTPHandler: ) response: Final = await self.client.send(req, stream=stream) if stream: - _anchor_handler_to(response, self) + anchor_handler_to(response, self) response.raise_for_status() return response except (httpx.RemoteProtocolError, httpx.ConnectError): @@ -1706,7 +1706,7 @@ class HTTPHandler: stream=True, follow_redirects=USE_CLIENT_DEFAULT if follow_redirects is None else follow_redirects, ) - _anchor_handler_to(response, self) + anchor_handler_to(response, self) response.raise_for_status() except httpx.TimeoutException as e: _raise_timeout(e, start_time) @@ -1771,7 +1771,7 @@ class HTTPHandler: ) response: Final = self.client.send(req, stream=stream) if stream: - _anchor_handler_to(response, self) + anchor_handler_to(response, self) response.raise_for_status() return response except httpx.TimeoutException as e: @@ -1820,7 +1820,7 @@ class HTTPHandler: ) response: Final = self.client.send(req, stream=stream) if stream: - _anchor_handler_to(response, self) + anchor_handler_to(response, self) response.raise_for_status() return response except httpx.TimeoutException as e: @@ -1869,7 +1869,7 @@ class HTTPHandler: ) response: Final = self.client.send(req, stream=stream) if stream: - _anchor_handler_to(response, self) + anchor_handler_to(response, self) return response except httpx.TimeoutException as e: _raise_timeout(e, start_time) @@ -1917,7 +1917,7 @@ class HTTPHandler: ) response: Final = self.client.send(req, stream=stream) if stream: - _anchor_handler_to(response, self) + anchor_handler_to(response, self) response.raise_for_status() return response except httpx.TimeoutException as e: diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 5b6a64d2377..efb8b302664 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -133,6 +133,7 @@ from litellm.llms.custom_httpx.container_handler import raise_for_error_status from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, + anchor_handler_to, get_async_httpx_client, get_httpx_client, header_value, @@ -5890,6 +5891,7 @@ class BaseLLMHTTPHandler: except Exception as e: # noqa: BLE001 # _handle_error maps every failure kind, like the buffered fetch raise self._handle_error(e=e, provider_config=provider_config) + anchor_handler_to(response, async_httpx_client) if response.status_code >= 400: error_body: Final = await response.aread() await response.aclose() diff --git a/litellm/llms/gemini/batches/__init__.py b/litellm/llms/gemini/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/gemini/batches/transformation.py b/litellm/llms/gemini/batches/transformation.py new file mode 100644 index 00000000000..10baec7d37d --- /dev/null +++ b/litellm/llms/gemini/batches/transformation.py @@ -0,0 +1,301 @@ +""" +Gemini Developer API batch jobs. Reference: https://ai.google.dev/gemini-api/docs/batch-api + +A file-backed job runs the uploaded JSONL unchanged. Gemini accepts OpenAI-shaped rows +(``{custom_id, method, url, body}``) and then writes the output file OpenAI-shaped too +(``{custom_id, response: {status_code, body}}``), so ``/v1/files purpose=batch`` uploads and the +shared batch cost accounting need no provider branch. The model is set on the job URL. +Ids cross the LiteLLM boundary bare: Gemini's ``batches/`` is the batch id ```` and its +``files/`` output file is the file id ````. +""" + +from collections.abc import Mapping, Sequence +from datetime import datetime +from types import MappingProxyType +from typing import Final, Literal, TypeAlias + +import httpx +from openai.types.batch import BatchRequestCounts +from openai.types.batch import Errors as BatchErrors +from openai.types.batch_error import BatchError +from pydantic import ConfigDict, Field +from typing_extensions import ReadOnly, TypedDict + +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base +from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.gemini.common_utils import GeminiError, GeminiModelInfo +from litellm.llms.gemini.files.transformation import gemini_file_name +from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.llms.openai import AllMessageValues, CreateBatchRequest +from litellm.types.utils import LiteLLMBatch, LlmProviders, Usage + +GeminiBatchState: TypeAlias = Literal[ + "BATCH_STATE_UNSPECIFIED", + "BATCH_STATE_PENDING", + "BATCH_STATE_RUNNING", + "BATCH_STATE_SUCCEEDED", + "BATCH_STATE_FAILED", + "BATCH_STATE_CANCELLED", + "BATCH_STATE_EXPIRED", +] +OpenAIBatchStatus: TypeAlias = Literal[ + "validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled" +] + +GEMINI_BATCH_RESOURCE_PREFIX: Final = "batches/" +_GEMINI_FILE_RESOURCE_PREFIX: Final = "files/" +_NO_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) +_STATUS_MAP: Final[MappingProxyType[GeminiBatchState, OpenAIBatchStatus]] = MappingProxyType( + { + "BATCH_STATE_UNSPECIFIED": "validating", + "BATCH_STATE_PENDING": "validating", + "BATCH_STATE_RUNNING": "in_progress", + "BATCH_STATE_SUCCEEDED": "completed", + "BATCH_STATE_FAILED": "failed", + "BATCH_STATE_CANCELLED": "cancelled", + "BATCH_STATE_EXPIRED": "expired", + } +) + + +class GeminiBatchInputConfig(TypedDict): + file_name: ReadOnly[str] + + +class GeminiBatchSpec(TypedDict): + display_name: ReadOnly[str] + input_config: ReadOnly[GeminiBatchInputConfig] + + +class GeminiCreateBatchRequest(TypedDict): + batch: ReadOnly[GeminiBatchSpec] + + +class GeminiPresignedRequest(TypedDict): + method: ReadOnly[Literal["GET"]] + url: ReadOnly[str] + headers: ReadOnly[Mapping[str, str]] + + +class _GeminiBatchStats(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + request_count: int = Field(default=0, alias="requestCount") + successful_request_count: int = Field(default=0, alias="successfulRequestCount") + failed_request_count: int = Field(default=0, alias="failedRequestCount") + + +class _GeminiBatchInputConfig(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + file_name: str | None = Field(default=None, alias="fileName") + + +class _GeminiBatchOutput(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + responses_file: str | None = Field(default=None, alias="responsesFile") + + +class _GeminiBatchMetadata(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + display_name: str | None = Field(default=None, alias="displayName") + input_config: _GeminiBatchInputConfig = Field(default_factory=_GeminiBatchInputConfig, alias="inputConfig") + output: _GeminiBatchOutput = Field(default_factory=_GeminiBatchOutput) + create_time: datetime = Field(alias="createTime") + end_time: datetime | None = Field(default=None, alias="endTime") + batch_stats: _GeminiBatchStats = Field(default_factory=_GeminiBatchStats, alias="batchStats") + state: GeminiBatchState = "BATCH_STATE_UNSPECIFIED" + + +class _GeminiOperationError(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + code: int | None = None + message: str = "" + + +class GeminiBatchOperation(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + name: str + metadata: _GeminiBatchMetadata + error: _GeminiOperationError | None = None + + +class _GeminiBatchOutputUsage(LiteLLMBaseModel): + """The ``usage`` block Gemini writes on an OpenAI-shaped batch output row. ``totalTokens`` + includes thinking tokens that ``completionTokens`` leaves out, and Gemini bills them as output.""" + + model_config = ConfigDict(frozen=True, extra="ignore") + + prompt_tokens: int = Field(default=0, alias="promptTokens") + completion_tokens: int = Field(default=0, alias="completionTokens") + total_tokens: int = Field(default=0, alias="totalTokens") + + +def is_native_gemini_batch_output_row(row: Mapping[str, object]) -> bool: + """Native rows carry Gemini's ``key`` (and a ``response`` or an error ``status``) where OpenAI-shaped rows + carry ``custom_id``, so the shape is known even when the first row failed.""" + return "key" in row and "custom_id" not in row + + +def is_gemini_batch_usage_shape(usage_object: Mapping[str, object]) -> bool: + return any(key in usage_object for key in ("promptTokens", "completionTokens", "totalTokens")) + + +def usage_from_batch_output(usage_object: Mapping[str, object]) -> Usage: + reported: Final = _GeminiBatchOutputUsage.model_validate(usage_object) + reasoning_tokens: Final = max(reported.total_tokens - reported.prompt_tokens - reported.completion_tokens, 0) + return Usage( + prompt_tokens=reported.prompt_tokens, + completion_tokens=reported.completion_tokens + reasoning_tokens, + total_tokens=reported.total_tokens, + reasoning_tokens=reasoning_tokens or None, + ) + + +def _bare_id(resource_name: str | None, prefix: str) -> str | None: + return None if resource_name is None else resource_name.removeprefix(prefix) + + +def _to_batch_errors(error: _GeminiOperationError | None) -> BatchErrors | None: + if error is None: + return None + return BatchErrors( + object="list", data=[BatchError(code=str(error.code) if error.code else None, message=error.message)] + ) + + +def _to_litellm_batch(operation: GeminiBatchOperation) -> LiteLLMBatch: + meta: Final = operation.metadata + status: Final = _STATUS_MAP[meta.state] + terminal_at: Final = int(meta.end_time.timestamp()) if meta.end_time is not None else None + return LiteLLMBatch( + id=_bare_id(operation.name, GEMINI_BATCH_RESOURCE_PREFIX) or operation.name, + object="batch", + endpoint="/v1/chat/completions", + input_file_id=_bare_id(meta.input_config.file_name, _GEMINI_FILE_RESOURCE_PREFIX) or "", + completion_window="24h", + status=status, + created_at=int(meta.create_time.timestamp()), + completed_at=terminal_at if status == "completed" else None, + failed_at=terminal_at if status == "failed" else None, + expired_at=terminal_at if status == "expired" else None, + cancelled_at=terminal_at if status == "cancelled" else None, + output_file_id=_bare_id(meta.output.responses_file, _GEMINI_FILE_RESOURCE_PREFIX), + errors=_to_batch_errors(operation.error), + request_counts=BatchRequestCounts( + total=meta.batch_stats.request_count, + completed=meta.batch_stats.successful_request_count, + failed=meta.batch_stats.failed_request_count, + ), + metadata={"display_name": meta.display_name} if meta.display_name else None, + ) + + +def _auth_headers(headers: Mapping[str, str], api_key: str | None) -> dict[str, str]: # mutable-ok: header contract + resolved_api_key: Final = GeminiModelInfo.get_api_key(api_key) + if not resolved_api_key: + raise ValueError("GEMINI_API_KEY is required for Gemini batch operations") + return {**headers, "x-goog-api-key": resolved_api_key} + + +class GeminiBatchesConfig(BaseBatchesConfig): + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.GEMINI + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, str]: # mutable-ok: BaseBatchesConfig signature + return _auth_headers(headers, api_key) + + def get_complete_batch_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + data: CreateBatchRequest, + ) -> str: + model_segment: Final = encode_url_path_segment(model.removeprefix("models/"), field_name="model") + return join_api_base( + GeminiModelInfo.get_api_base(api_base), f"v1beta/models/{model_segment}:batchGenerateContent" + ) + + def transform_create_batch_request( + self, + model: str, + create_batch_data: CreateBatchRequest, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> dict[str, object]: # mutable-ok: BaseBatchesConfig signature + input_file_id: Final = create_batch_data.get("input_file_id") + if input_file_id is None: + raise ValueError("input_file_id is required to create a Gemini batch job") + file_name: Final = gemini_file_name(input_file_id) + metadata: Final = create_batch_data.get("metadata") or {} + unsupported_metadata: Final = sorted(set(metadata) - {"display_name"}) + if unsupported_metadata: + raise ValueError(f"Gemini batch jobs carry no metadata besides display_name, got {unsupported_metadata}") + display_name: Final = ( + metadata.get("display_name") or f"litellm-{file_name.removeprefix(_GEMINI_FILE_RESOURCE_PREFIX)}" + ) + body: Final = GeminiCreateBatchRequest( + batch=GeminiBatchSpec(display_name=display_name, input_config=GeminiBatchInputConfig(file_name=file_name)) + ) + return dict(body) + + def transform_create_batch_response( + self, + model: str | None, + raw_response: httpx.Response, + logging_obj: object, + litellm_params: Mapping[str, object], + ) -> LiteLLMBatch: + return _to_litellm_batch(GeminiBatchOperation.model_validate(raw_response.json())) + + def transform_retrieve_batch_request( + self, + batch_id: str, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> dict[str, object]: # mutable-ok: BaseBatchesConfig signature + bare_batch_id: Final = batch_id.removeprefix(GEMINI_BATCH_RESOURCE_PREFIX) + encoded_batch_id: Final = encode_url_path_segment(bare_batch_id, field_name="batch_id") + api_base: Final = litellm_params.get("api_base") + api_key: Final = litellm_params.get("api_key") + request: Final = GeminiPresignedRequest( + method="GET", + url=join_api_base( + GeminiModelInfo.get_api_base(api_base if isinstance(api_base, str) else None), + f"v1beta/{GEMINI_BATCH_RESOURCE_PREFIX}{encoded_batch_id}", + ), + headers=_auth_headers(_NO_HEADERS, api_key if isinstance(api_key, str) else None), + ) + return dict(request) + + def transform_retrieve_batch_response( + self, + model: str | None, + raw_response: httpx.Response, + logging_obj: object, + litellm_params: Mapping[str, object], + ) -> LiteLLMBatch: + return _to_litellm_batch(GeminiBatchOperation.model_validate(raw_response.json())) + + def get_error_class( + self, error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers + ) -> BaseLLMException: + return GeminiError(status_code=status_code, message=error_message, headers=dict(headers)) diff --git a/litellm/llms/gemini/files/transformation.py b/litellm/llms/gemini/files/transformation.py index 17460c0295a..4fb5a77653d 100644 --- a/litellm/llms/gemini/files/transformation.py +++ b/litellm/llms/gemini/files/transformation.py @@ -4,8 +4,10 @@ Supports writing files to Google AI Studio Files API. For vertex ai, check out the vertex_ai/files/handler.py file. """ +import re import time from collections.abc import Mapping +from datetime import datetime from typing import Final, Literal, TypedDict from urllib.parse import urlparse @@ -25,6 +27,7 @@ from litellm.llms.base_llm.files.transformation import ( from litellm.types.llms.openai import ( AllMessageValues, CreateFileRequest, + FileContentRequest, HttpxBinaryResponseContent, OpenAICreateFileRequestOptionalParams, OpenAIFileObject, @@ -53,6 +56,34 @@ class _GeminiCreateFileResponse(TypedDict): file: ReadOnly[_GeminiFileMetadata] +_RFC3339_PARTS: Final = re.compile(r"(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2})(\.\d+)?([+-]\d{2}:\d{2})") + + +def _gemini_timestamp(rfc3339: str) -> int: + """Gemini stamps files with up to nine fractional digits and a Z suffix; Python 3.10's + fromisoformat takes neither, so trim the fraction to microseconds and spell out UTC.""" + parts: Final = _RFC3339_PARTS.fullmatch(rfc3339.replace("Z", "+00:00")) + if parts is None: + raise ValueError(f"Unrecognized Gemini timestamp: {rfc3339}") + seconds, fraction, offset = parts.groups() + return int(datetime.fromisoformat(f"{seconds}{(fraction or '')[:7]}{offset}").timestamp()) + + +def gemini_file_name(file_id: str) -> str: + """ + Normalize a file identifier into Gemini's `files/{id}` resource name. + + Supports: + - `abc123` + - `files/abc123` + - `https://generativelanguage.googleapis.com/v1beta/files/abc123` + """ + path: Final = urlparse(file_id).path.lstrip("/") if file_id.startswith(("http://", "https://")) else file_id + files_index: Final = path.find("files/") + bare_id: Final = (path[files_index:] if files_index != -1 else path).strip("/").removeprefix("files/") + return f"files/{encode_url_path_segment(bare_id, field_name='file_id')}" + + class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): def __init__(self): pass @@ -194,14 +225,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): return OpenAIFileObject( id=response_object["uri"], # Gemini uses URI as identifier bytes=int(response_object["sizeBytes"]), # Gemini doesn't return file size - created_at=int( - time.mktime( - time.strptime( - response_object["createTime"].replace("Z", "+00:00"), - "%Y-%m-%dT%H:%M:%S.%f%z", - ) - ) - ), + created_at=_gemini_timestamp(response_object["createTime"]), filename=response_object["displayName"], object="file", purpose="user_data", # Default to assistants as that's the main use case @@ -246,32 +270,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): return url, {} def _normalize_gemini_file_id(self, file_id: str) -> str: - """ - Normalize file identifier into `files/{id}` form. - - Supports: - - `abc123` - - `files/abc123` - - `https://generativelanguage.googleapis.com/v1beta/files/abc123` - """ - if file_id.startswith(("http://", "https://")): - parsed: Final = urlparse(file_id) - path: Final = parsed.path.lstrip("/") - files_index: Final = path.find("files/") - if files_index != -1: - normalized_file_id = path[files_index:] - else: - normalized_file_id = path - else: - normalized_file_id = file_id - - normalized_file_id = normalized_file_id.strip("/") - if normalized_file_id.startswith("files/"): - normalized_file_id = normalized_file_id.removeprefix("files/") - - encoded_file_id: Final = encode_url_path_segment(normalized_file_id, field_name="file_id") - - return f"files/{encoded_file_id}" + return gemini_file_name(file_id) def transform_retrieve_file_response( self, @@ -299,14 +298,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): return OpenAIFileObject( id=response_json.get("uri", ""), bytes=int(response_json.get("sizeBytes", 0)), - created_at=int( - time.mktime( - time.strptime( - response_json["createTime"].replace("Z", "+00:00"), - "%Y-%m-%dT%H:%M:%S.%f%z", - ) - ) - ), + created_at=_gemini_timestamp(response_json["createTime"]), filename=response_json.get("displayName", ""), object="file", purpose="user_data", @@ -403,11 +395,19 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): def transform_file_content_request( self, - file_content_request, - optional_params: dict, - litellm_params: dict, - ) -> tuple[str, dict]: - raise NotImplementedError("GoogleAIStudioFilesHandler does not support file content retrieval") + file_content_request: FileContentRequest, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature + file_id: Final = file_content_request.get("file_id") + if file_id is None: + raise ValueError("file_id is required to download file content") + api_base: Final = litellm_params.get("api_base") + url: Final = join_api_base( + self.get_api_base(api_base if isinstance(api_base, str) else None), + f"download/v1beta/{gemini_file_name(file_id)}:download", + ) + return url, {"alt": "media"} def transform_file_content_response( self, @@ -415,4 +415,4 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): logging_obj: LiteLLMLoggingObj, litellm_params: dict, ) -> HttpxBinaryResponseContent: - raise NotImplementedError("GoogleAIStudioFilesHandler does not support file content retrieval") + return HttpxBinaryResponseContent(response=raw_response) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index addd642eca7..81d39f848fb 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4833,7 +4833,7 @@ OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: set[str] = { } FILE_CONTENT_STREAMING_PROVIDERS: Final[frozenset[str]] = frozenset( - {*OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS, LlmProviders.VERTEX_AI.value} + {*OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS, LlmProviders.VERTEX_AI.value, LlmProviders.GEMINI.value} ) LITELLM_EXECUTED_BATCH_PROVIDERS: Final[frozenset[str]] = frozenset({LlmProviders.HOSTED_VLLM.value}) diff --git a/litellm/utils.py b/litellm/utils.py index 156247897c3..7c73d1cc4d5 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -10866,6 +10866,10 @@ class ProviderConfigManager: from litellm.llms.mistral.batches.transformation import MistralBatchesConfig return MistralBatchesConfig() + elif LlmProviders.GEMINI == provider: + from litellm.llms.gemini.batches.transformation import GeminiBatchesConfig + + return GeminiBatchesConfig() return None @staticmethod diff --git a/tests/unit/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index bef78554a36..2c8114f8736 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -14,27 +14,28 @@ maps (litellm.completion_cost, batch_cost_calculator), the tokenizer deterministic stand-ins so the arithmetic under test is the only variable. """ -import asyncio, json, time +import asyncio +import json import logging +import time from typing import Final, Literal +from unittest.mock import AsyncMock, patch import httpx import pytest import respx - import litellm import litellm.batches.batch_utils as bu -from litellm.types.utils import LiteLLMBatch, ModelInfo, Usage -from litellm.batches.batch_utils import( +from litellm.batches.batch_utils import ( _aggregate_batch_cost_usage_models, - get_file_content_as_dictionary, _get_response_from_batch_job_output_file, calculate_batch_cost_and_usage, + get_file_content_as_dictionary, ) from litellm.cost_calculator import batch_cost_calculator +from litellm.types.utils import LiteLLMBatch, ModelInfo, Usage from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome -from unittest.mock import AsyncMock, patch # --------------------------------------------------------------------------- # # Builders for batch OUTPUT file rows. @@ -1188,7 +1189,9 @@ async def test_handle_completed_vertex_batch_computes_cost_usage_and_models(monk @pytest.mark.asyncio -async def test_handle_completed_vertex_batch_regional_location_costs_uplift_over_global(monkeypatch: pytest.MonkeyPatch) -> None: +async def test_handle_completed_vertex_batch_regional_location_costs_uplift_over_global( + monkeypatch: pytest.MonkeyPatch, +) -> None: import litellm.files.main as files_main rows = [_vertex_openai_row("request-1", "gemini-3.6-flash", 10, 5)] @@ -2188,8 +2191,7 @@ GROUNDED_USAGE_METADATA = { PASSTHROUGH_OUTPUT_URI = ( - "gs://litellm-bucket/litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/u/" - "predictions.jsonl" + "gs://litellm-bucket/litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/u/predictions.jsonl" ) @@ -2555,6 +2557,7 @@ def _vcr_outcome_gate(request, vcr): yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="session") def event_loop(): try: @@ -2564,6 +2567,7 @@ def event_loop(): yield loop loop.close() + @pytest.fixture(scope="function") def setup_and_teardown(event_loop): original_state = _copy_litellm_state() @@ -2580,6 +2584,7 @@ def setup_and_teardown(event_loop): if pending: event_loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) + def _copy_litellm_state(): state = {} for attr in _CALLBACK_ATTRS: @@ -2591,6 +2596,7 @@ def _copy_litellm_state(): state[attr] = getattr(litellm, attr) return state + _CALLBACK_ATTRS = ( "callbacks", "success_callback", @@ -2613,6 +2619,7 @@ _SCALAR_ATTRS = ( "cohere_key", ) + def _clear_logging_queue(loop=None) -> None: from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -2621,6 +2628,7 @@ def _clear_logging_queue(loop=None) -> None: return asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue()) + def _reset_litellm_callbacks() -> None: for attr in _CALLBACK_ATTRS: if hasattr(litellm, attr): @@ -2630,11 +2638,13 @@ def _reset_litellm_callbacks() -> None: if callable(reset): reset() + def _restore_litellm_state(state) -> None: for attr, value in state.items(): if hasattr(litellm, attr): setattr(litellm, attr, value) + def _make_batch_output_line(prompt_tokens: int = 10, completion_tokens: int = 5): """Return a single successful batch output line (OpenAI JSONL format).""" return { @@ -2663,11 +2673,13 @@ def _make_batch_output_line(prompt_tokens: int = 10, completion_tokens: int = 5) "error": None, } + CUSTOM_MODEL_INFO = { "input_cost_per_token_batches": 0.00125, "output_cost_per_token_batches": 0.005, } + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_batch_cost_calculator_explicit_zero_pricing_not_overridden_by_global( monkeypatch, @@ -2699,6 +2711,7 @@ def test_batch_cost_calculator_explicit_zero_pricing_not_overridden_by_global( assert prompt_cost == 0.0 assert completion_cost == 0.0 + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_batch_cost_calculator_uses_custom_model_info(): """batch_cost_calculator should use model_info override when provided.""" @@ -2718,6 +2731,7 @@ def test_batch_cost_calculator_uses_custom_model_info(): f"Expected completion cost {expected_completion}, got {completion_cost}" ) + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_aggregate_batch_cost_uses_custom_model_info(): """_aggregate_batch_cost_usage_models should thread model_info to batch_cost_calculator.""" @@ -2732,6 +2746,7 @@ def test_aggregate_batch_cost_uses_custom_model_info(): expected = (10 * 0.00125) + (5 * 0.005) assert result.cost == pytest.approx(expected), f"Expected total cost {expected}, got {result.cost}" + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.parametrize("data_residency", ["eu", "us", "fedramp"]) def test_batch_cost_calculator_applies_data_residency_uplift(data_residency, monkeypatch): @@ -2779,6 +2794,7 @@ async def test_calculate_batch_cost_and_usage_uses_custom_model_info(): assert result.usage.prompt_tokens == 10 assert result.usage.completion_tokens == 5 + @pytest.fixture def sample_file_content(): return b""" @@ -2786,6 +2802,7 @@ def sample_file_content(): {"id": "batch_req_6769ca597e588190920666612634e2b4", "custom_id": "request-2", "response": {"status_code": 200, "request_id": "82e04f4c001fe2c127cbad199f5fd31b", "body": {"id": "chatcmpl-AhjSNgVB4Oa4Hq0NruTRsBaEbRWUP", "object": "chat.completion", "created": 1734986203, "model": "gpt-4o-mini-2024-07-18", "choices": [{"index": 0, "message": {"role": "assistant", "content": "Hello! What can I do for you today?", "refusal": null}, "logprobs": null, "finish_reason": "length"}], "usage": {"prompt_tokens": 22, "completion_tokens": 10, "total_tokens": 32, "prompt_tokens_details": {"cached_tokens": 0, "audio_tokens": 0}, "completion_tokens_details": {"reasoning_tokens": 0, "audio_tokens": 0, "accepted_prediction_tokens": 0, "rejected_prediction_tokens": 0}}, "system_fingerprint": "fp_0aa8d3e20b"}}, "error": null} """ + @pytest.fixture def sample_file_content_dict(): return [ @@ -2877,6 +2894,7 @@ def sample_file_content_dict(): }, ] + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_get_file_content_as_dictionary(sample_file_content): result = get_file_content_as_dictionary(sample_file_content) @@ -2886,6 +2904,7 @@ def test_get_file_content_as_dictionary(sample_file_content): assert result[0]["response"]["status_code"] == 200 assert result[0]["response"]["body"]["usage"]["total_tokens"] == 30 + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_get_batch_job_total_usage_from_file_content(sample_file_content_dict): with patch("litellm.completion_cost", return_value=0.0): @@ -2894,6 +2913,7 @@ def test_get_batch_job_total_usage_from_file_content(sample_file_content_dict): assert result.usage.prompt_tokens == 42 # 20 + 22 assert result.usage.completion_tokens == 20 # 10 + 10 + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio async def test_batch_cost_calculator(sample_file_content_dict): @@ -2913,6 +2933,7 @@ async def test_batch_cost_calculator(sample_file_content_dict): assert result.prompt_cost == pytest.approx(0.6) assert result.completion_cost == pytest.approx(0.4) + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") def test_get_response_from_batch_job_output_file(sample_file_content_dict): result = _get_response_from_batch_job_output_file(sample_file_content_dict[0]) @@ -2920,6 +2941,7 @@ def test_get_response_from_batch_job_output_file(sample_file_content_dict): assert result["object"] == "chat.completion" assert result["usage"]["total_tokens"] == 30 + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio async def test_batch_retrieve_cost_tracking_with_completed_batch_no_explicit_cost(): @@ -3016,6 +3038,7 @@ async def test_batch_retrieve_cost_tracking_with_completed_batch_no_explicit_cos assert mock_batch._hidden_params["batch_failed_requests"] == 0 assert mock_batch.usage == expected_usage + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio async def test_handle_completed_batch_computes_real_cost_from_output_file( @@ -3059,6 +3082,7 @@ async def test_handle_completed_batch_computes_real_cost_from_output_file( assert result.successful_requests == 2 assert result.failed_requests == 0 + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio async def test_batch_retrieve_cost_tracking_with_explicit_cost_data(): @@ -3146,6 +3170,7 @@ async def test_batch_retrieve_cost_tracking_with_explicit_cost_data(): assert mock_batch._hidden_params["batch_models"] == explicit_models assert mock_batch.usage == explicit_usage + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio async def test_batch_retrieve_explicit_cost_split_sets_cost_breakdown(): @@ -3196,6 +3221,7 @@ async def test_batch_retrieve_explicit_cost_split_sets_cost_breakdown(): assert logging_obj.cost_breakdown["output_cost"] == 0.04 assert logging_obj.cost_breakdown["total_cost"] == 0.10 + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio async def test_batch_retrieve_cost_tracking_with_unified_file_id_incomplete_batch(): @@ -3276,6 +3302,7 @@ async def test_batch_retrieve_cost_tracking_with_unified_file_id_incomplete_batc assert "batch_models" not in mock_batch._hidden_params assert not hasattr(mock_batch, "usage") or mock_batch.usage is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio async def test_batch_retrieve_cost_tracking_with_partial_explicit_data(): @@ -3445,7 +3472,10 @@ def test_batch_ocr_preserves_additive_page_quantities_without_changing_cost(): assert captured["pages_processed"] == 5 assert captured["doc_size_bytes"] == 300 assert result.cost == pytest.approx(5 * 0.01) - assert sorted(row["usage"]["pages_processed"] for row in result.usage.model_dump()["batch_usage_by_request"]) == [2, 3] + assert sorted(row["usage"]["pages_processed"] for row in result.usage.model_dump()["batch_usage_by_request"]) == [ + 2, + 3, + ] def test_native_vertex_capture_retains_unknown_models_without_changing_legacy_counts(): @@ -3456,3 +3486,76 @@ def test_native_vertex_capture_retains_unknown_models_without_changing_legacy_co assert captured["batch_usage_by_request"] == ({"model": "missing-usage-model", "usage": None},) assert captured["batch_unknown_usage_requests"] == 1 assert (result.cost, result.successful_requests, result.failed_requests) == (0, 0, 1) + + +def test_gemini_output_lines_bill_camelcase_usage_with_thinking_tokens_at_batch_rates(): + row = _success_row( + model="gemini-3.8-flash", + usage={"promptTokens": 9, "completionTokens": 1, "totalTokens": 75}, + ) + + result = bu._aggregate_batch_cost_usage_models( + entries=[row], + custom_llm_provider="gemini", + model_info=ModelInfo( + key="gemini/gemini-3.8-flash", + max_tokens=None, + max_input_tokens=None, + max_output_tokens=None, + input_cost_per_token=7.5e-7, + output_cost_per_token=3.75e-6, + litellm_provider="gemini", + mode="chat", + supported_openai_params=None, + input_cost_per_token_batches=3.75e-7, + output_cost_per_token_batches=1.875e-6, + ), + ) + + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (9, 66, 75) + assert result.cost == pytest.approx((9 * 3.75e-7) + (66 * 1.875e-6)) + assert (result.successful_requests, result.failed_requests, result.models) == (1, 0, ["gemini-3.8-flash"]) + + +@pytest.mark.asyncio +async def test_gemini_native_output_rows_are_priced_from_usage_metadata(monkeypatch): + native_row = { + "key": "r0", + "response": { + "candidates": [{"content": {"parts": [{"text": "0"}], "role": "model"}, "finishReason": "STOP"}], + "usageMetadata": { + "promptTokenCount": 3, + "candidatesTokenCount": 1, + "totalTokenCount": 65, + "thoughtsTokenCount": 61, + }, + "modelVersion": "gemini-3.8-flash", + }, + } + failed_first_row = {"key": "bad", "status": {"code": 3, "message": "Invalid request"}} + jsonl = "\n".join(json.dumps(row) for row in (failed_first_row, native_row)).encode() + monkeypatch.setattr(bu, "_fetch_batch_managed_file_content", AsyncMock(return_value=jsonl)) + monkeypatch.setattr(bu, "count_error_file_failed_requests", AsyncMock(return_value=0)) + + result = await bu.handle_completed_batch( + batch=_batch("batch-la6"), + custom_llm_provider="gemini", + model_name="gemini-3.8-flash", + model_info=ModelInfo( + key="gemini/gemini-3.8-flash", + max_tokens=None, + max_input_tokens=None, + max_output_tokens=None, + input_cost_per_token=7.5e-7, + output_cost_per_token=3.75e-6, + litellm_provider="gemini", + mode="chat", + supported_openai_params=None, + input_cost_per_token_batches=3.75e-7, + output_cost_per_token_batches=1.875e-6, + ), + ) + + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (3, 62, 65) + assert result.cost == pytest.approx((3 * 3.75e-7) + (62 * 1.875e-6)) + assert (result.successful_requests, result.failed_requests) == (1, 1) diff --git a/tests/unit/batches/test_main.py b/tests/unit/batches/test_main.py index 8cd74c7d83a..b12ea53bce3 100644 --- a/tests/unit/batches/test_main.py +++ b/tests/unit/batches/test_main.py @@ -1382,3 +1382,28 @@ async def test_openai_batch_list_items_come_back_as_openai_sent_them(respx_mock: assert first.usage is not None assert (first.usage.input_tokens_details.cached_tokens, first.usage.output_tokens_details.reasoning_tokens) == (4, 2) assert (second.id, second.object, second.status) == ("batch_def456", "batch", "queued") + + +# =========================================================================== # +# gemini - a provider-config provider, like mistral, so it requires `model` +# =========================================================================== # + + +def test_create__gemini_routes_to_base_http_handler_with_gemini_config(seams): + result = bm.create_batch(**CREATE_KW, custom_llm_provider="gemini", model="gemini/gemini-3.8-flash") + + assert result is seams.base_http.create_batch.return_value + _assert_only(seams.base_http.create_batch, seams, "create_batch") + forwarded = seams.base_http.create_batch.call_args.kwargs + assert type(forwarded["provider_config"]).__name__ == "GeminiBatchesConfig" + assert forwarded["model"] == "gemini-3.8-flash" + + +def test_retrieve__gemini_routes_to_base_http_handler_with_gemini_config(seams): + result = bm.retrieve_batch(batch_id="la6", custom_llm_provider="gemini", model="gemini/gemini-3.8-flash") + + assert result is seams.base_http.retrieve_batch.return_value + _assert_only(seams.base_http.retrieve_batch, seams, "retrieve_batch") + forwarded = seams.base_http.retrieve_batch.call_args.kwargs + assert type(forwarded["provider_config"]).__name__ == "GeminiBatchesConfig" + assert forwarded["batch_id"] == "la6" diff --git a/tests/unit/llms/gemini/batches/__init__.py b/tests/unit/llms/gemini/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/gemini/batches/test_gemini_batches_transformation.py b/tests/unit/llms/gemini/batches/test_gemini_batches_transformation.py new file mode 100644 index 00000000000..c8f15c9c590 --- /dev/null +++ b/tests/unit/llms/gemini/batches/test_gemini_batches_transformation.py @@ -0,0 +1,310 @@ +""" +Regression tests for ``GeminiBatchesConfig``, the BaseBatchesConfig implementation behind +``custom_llm_provider="gemini"`` on /v1/batches. + +Locks the ``models/{model}:batchGenerateContent`` create URL, the ``input_config.file_name`` +request shape, Gemini -> OpenAI state mapping, ``batchStats`` -> ``request_counts``, bare ids at +the LiteLLM boundary (``batches/`` and ``files/`` stripped) and ``x-goog-api-key`` auth. +Everything runs for real against canned httpx responses; only the API key env var is set. +""" + +import json + +import httpx +import pytest + +from litellm.llms.gemini.batches.transformation import ( + GeminiBatchesConfig, + is_gemini_batch_usage_shape, + is_native_gemini_batch_output_row, + usage_from_batch_output, +) +from litellm.llms.gemini.common_utils import GeminiError +from litellm.types.utils import LiteLLMBatch, LlmProviders + +STATUS_MAP = { + "BATCH_STATE_UNSPECIFIED": "validating", + "BATCH_STATE_PENDING": "validating", + "BATCH_STATE_RUNNING": "in_progress", + "BATCH_STATE_SUCCEEDED": "completed", + "BATCH_STATE_FAILED": "failed", + "BATCH_STATE_CANCELLED": "cancelled", + "BATCH_STATE_EXPIRED": "expired", +} + + +def _operation(**metadata_overrides): + metadata = { + "@type": "type.googleapis.com/google.ai.generativelanguage.v1main.GenerateContentBatch", + "model": "models/gemini-3.8-flash", + "displayName": "nightly-eval", + "inputConfig": {"fileName": "files/haf0y9b1kydd"}, + "output": {"responsesFile": "files/batch-la6uw1zzqrect7ugge7te7fueeslt1ms9xiw"}, + "createTime": "2026-11-10T21:18:34.714998843Z", + "endTime": "2026-11-10T21:21:01.499150334Z", + "updateTime": "2026-11-10T21:21:01.499150354Z", + "batchStats": {"requestCount": "3", "successfulRequestCount": "2", "failedRequestCount": "1"}, + "state": "BATCH_STATE_SUCCEEDED", + "name": "batches/la6uw1zzqrect7ugge7te7fueeslt1ms9xiw", + } + return { + "name": "batches/la6uw1zzqrect7ugge7te7fueeslt1ms9xiw", + "metadata": {**metadata, **metadata_overrides}, + "done": True, + } + + +def _response(payload: dict, status_code: int = 200) -> httpx.Response: + return httpx.Response( + status_code=status_code, + content=json.dumps(payload).encode(), + request=httpx.Request("GET", "https://generativelanguage.googleapis.com/v1beta/batches/x"), + ) + + +@pytest.fixture +def config() -> GeminiBatchesConfig: + return GeminiBatchesConfig() + + +@pytest.fixture +def api_key(monkeypatch) -> str: + monkeypatch.delenv("GOOGLE_API_KEY", raising=False) + monkeypatch.setenv("GEMINI_API_KEY", "AIza-test") + return "AIza-test" + + +def test_custom_llm_provider(config): + assert config.custom_llm_provider == LlmProviders.GEMINI + + +@pytest.mark.parametrize( + "input_file_id", + ["haf0y9b1kydd", "files/haf0y9b1kydd", "https://generativelanguage.googleapis.com/v1beta/files/haf0y9b1kydd"], +) +def test_create_request_names_the_uploaded_file_in_input_config(config, input_file_id): + body = config.transform_create_batch_request( + model="gemini-3.8-flash", + create_batch_data={ + "input_file_id": input_file_id, + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + optional_params={}, + litellm_params={}, + ) + assert body == { + "batch": {"display_name": "litellm-haf0y9b1kydd", "input_config": {"file_name": "files/haf0y9b1kydd"}} + } + + +def test_create_request_uses_display_name_from_metadata(config): + body = config.transform_create_batch_request( + model="gemini-3.8-flash", + create_batch_data={ + "input_file_id": "files/abc", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "metadata": {"display_name": "nightly-eval"}, + }, + optional_params={}, + litellm_params={}, + ) + assert body["batch"]["display_name"] == "nightly-eval" + + +def test_create_request_requires_input_file(config): + with pytest.raises(ValueError, match="input_file_id"): + config.transform_create_batch_request( + model="gemini-3.8-flash", + create_batch_data={"endpoint": "/v1/chat/completions", "completion_window": "24h"}, + optional_params={}, + litellm_params={}, + ) + + +@pytest.mark.parametrize( + "api_base,model,expected", + [ + ( + None, + "gemini-3.8-flash", + "https://generativelanguage.googleapis.com/v1beta/models/gemini-3.8-flash:batchGenerateContent", + ), + ( + "https://proxy.example/", + "models/gemini-3.8-flash", + "https://proxy.example/v1beta/models/gemini-3.8-flash:batchGenerateContent", + ), + ], +) +def test_create_url_targets_the_model_batch_route(config, api_base, model, expected): + data = {"input_file_id": "files/abc", "endpoint": "/v1/chat/completions", "completion_window": "24h"} + assert config.get_complete_batch_url(api_base, None, model, {}, {}, data) == expected + + +def test_validate_environment_uses_goog_api_key_header(config, api_key): + headers = config.validate_environment({"x-trace": "1"}, "gemini-3.8-flash", [], {}, {}) + assert headers == {"x-trace": "1", "x-goog-api-key": api_key} + + +def test_validate_environment_explicit_key_wins(config, api_key): + headers = config.validate_environment({}, "gemini-3.8-flash", [], {}, {}, api_key="AIza-deployment") + assert headers["x-goog-api-key"] == "AIza-deployment" + + +def test_validate_environment_without_key_raises(config, monkeypatch): + monkeypatch.delenv("GEMINI_API_KEY", raising=False) + monkeypatch.delenv("GOOGLE_API_KEY", raising=False) + with pytest.raises(ValueError, match="GEMINI_API_KEY"): + config.validate_environment({}, "gemini-3.8-flash", [], {}, {}) + + +def test_create_response_maps_operation_onto_openai_batch(config): + batch = config.transform_create_batch_response( + model="gemini-3.8-flash", raw_response=_response(_operation()), logging_obj=None, litellm_params={} + ) + assert isinstance(batch, LiteLLMBatch) + assert batch.model_dump(exclude_none=True) == { + "id": "la6uw1zzqrect7ugge7te7fueeslt1ms9xiw", + "object": "batch", + "endpoint": "/v1/chat/completions", + "input_file_id": "haf0y9b1kydd", + "output_file_id": "batch-la6uw1zzqrect7ugge7te7fueeslt1ms9xiw", + "completion_window": "24h", + "status": "completed", + "created_at": 1794345514, + "completed_at": 1794345661, + "request_counts": {"total": 3, "completed": 2, "failed": 1}, + "metadata": {"display_name": "nightly-eval"}, + } + + +def test_retrieve_request_is_presigned_get_with_bare_or_prefixed_id(config, api_key): + for batch_id in ("la6", "batches/la6"): + req = config.transform_retrieve_batch_request( + batch_id=batch_id, optional_params={}, litellm_params={"api_base": "https://proxy.example"} + ) + assert req == { + "method": "GET", + "url": "https://proxy.example/v1beta/batches/la6", + "headers": {"x-goog-api-key": api_key}, + }, batch_id + + +def test_retrieve_request_prefers_litellm_params_api_key(config, api_key): + req = config.transform_retrieve_batch_request( + batch_id="la6", optional_params={}, litellm_params={"api_key": "AIza-deployment"} + ) + assert req["headers"]["x-goog-api-key"] == "AIza-deployment" + + +@pytest.mark.parametrize("gemini_state,openai_status", sorted(STATUS_MAP.items())) +def test_retrieve_response_status_mapping(config, gemini_state, openai_status): + batch = config.transform_retrieve_batch_response( + model=None, raw_response=_response(_operation(state=gemini_state)), logging_obj=None, litellm_params={} + ) + assert batch.status == openai_status + + +@pytest.mark.parametrize( + "gemini_state,populated_field", + [ + ("BATCH_STATE_SUCCEEDED", "completed_at"), + ("BATCH_STATE_FAILED", "failed_at"), + ("BATCH_STATE_EXPIRED", "expired_at"), + ("BATCH_STATE_CANCELLED", "cancelled_at"), + ], +) +def test_end_time_lands_on_the_terminal_timestamp_for_its_state(config, gemini_state, populated_field): + batch = config.transform_retrieve_batch_response( + model=None, raw_response=_response(_operation(state=gemini_state)), logging_obj=None, litellm_params={} + ) + timestamps = {f: getattr(batch, f) for f in ("completed_at", "failed_at", "expired_at", "cancelled_at")} + assert timestamps == {f: (1794345661 if f == populated_field else None) for f in timestamps} + + +def test_running_batch_has_no_output_file_or_terminal_timestamp(config): + running = _operation(state="BATCH_STATE_RUNNING", batchStats={"requestCount": "3", "pendingRequestCount": "3"}) + del running["metadata"]["output"] + del running["metadata"]["endTime"] + batch = config.transform_retrieve_batch_response( + model=None, raw_response=_response(running), logging_obj=None, litellm_params={} + ) + assert batch.status == "in_progress" + assert batch.output_file_id is None + assert batch.completed_at is None + assert batch.request_counts.model_dump() == {"total": 3, "completed": 0, "failed": 0} + + +def test_failed_operation_error_is_reported_in_errors(config): + failed = {**_operation(state="BATCH_STATE_FAILED"), "error": {"code": 3, "message": "Input file is empty."}} + batch = config.transform_retrieve_batch_response( + model=None, raw_response=_response(failed), logging_obj=None, litellm_params={} + ) + assert batch.status == "failed" + assert [(e.code, e.message) for e in batch.errors.data] == [("3", "Input file is empty.")] + + +def test_error_class_is_gemini_error(config): + err = config.get_error_class("quota", 429, {"retry-after": "7"}) + assert isinstance(err, GeminiError) + assert (err.status_code, err.message) == (429, "quota") + + +# Gemini writes the usage of an OpenAI-shaped output row as camelCase, with thinking tokens +# counted in totalTokens but not in completionTokens (observed 2026-11-10 on gemini-3.8-flash). +@pytest.mark.parametrize( + "reported,expected", + [ + ( + {"promptTokens": 9, "completionTokens": 1, "totalTokens": 75}, + {"prompt_tokens": 9, "completion_tokens": 66, "total_tokens": 75, "reasoning_tokens": 65}, + ), + ( + {"promptTokens": 9, "completionTokens": 4, "totalTokens": 13}, + {"prompt_tokens": 9, "completion_tokens": 4, "total_tokens": 13, "reasoning_tokens": None}, + ), + ], +) +def test_batch_output_usage_folds_thinking_tokens_into_completion(reported, expected): + usage = usage_from_batch_output(reported) + details = usage.completion_tokens_details + assert { + "prompt_tokens": usage.prompt_tokens, + "completion_tokens": usage.completion_tokens, + "total_tokens": usage.total_tokens, + "reasoning_tokens": details.reasoning_tokens if details else None, + } == expected + + +@pytest.mark.parametrize( + "usage_object,expected", + [({"promptTokens": 9}, True), ({"prompt_tokens": 9, "completion_tokens": 1}, False), ({}, False)], +) +def test_batch_usage_shape_detection(usage_object, expected): + assert is_gemini_batch_usage_shape(usage_object) is expected + + +def test_create_request_rejects_metadata_gemini_cannot_carry(config): + with pytest.raises(ValueError, match=r"\['run_id'\]"): + config.transform_create_batch_request( + model="gemini-3.8-flash", + create_batch_data={ + "input_file_id": "files/abc", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "metadata": {"display_name": "nightly", "run_id": "42"}, + }, + optional_params={}, + litellm_params={}, + ) + + +def test_native_output_row_detection(): + native = {"key": "r0", "response": {"candidates": [], "usageMetadata": {"promptTokenCount": 3}}} + openai_shaped = {"custom_id": "r0", "response": {"status_code": 200, "body": {"usage": {"promptTokens": 3}}}} + assert (is_native_gemini_batch_output_row(native), is_native_gemini_batch_output_row(openai_shaped)) == ( + True, + False, + ) diff --git a/tests/unit/llms/gemini/files/test_gemini_files_transformation.py b/tests/unit/llms/gemini/files/test_gemini_files_transformation.py index a2f95724688..0baf4e7b0fb 100644 --- a/tests/unit/llms/gemini/files/test_gemini_files_transformation.py +++ b/tests/unit/llms/gemini/files/test_gemini_files_transformation.py @@ -43,9 +43,7 @@ class TestGoogleAIStudioFilesTransformation: # CRITICAL: params should be empty dict, not contain Content-Type or any other params # These would be incorrectly interpreted as query parameters assert params == {}, f"Expected empty params dict, got: {params}" - assert ( - "Content-Type" not in params - ), "Content-Type should not be in query params" + assert "Content-Type" not in params, "Content-Type should not be in query params" def test_transform_retrieve_file_request_with_file_name_only(self): """ @@ -66,9 +64,7 @@ class TestGoogleAIStudioFilesTransformation: # CRITICAL: params should be empty dict assert params == {}, f"Expected empty params dict, got: {params}" - assert ( - "Content-Type" not in params - ), "Content-Type should not be in query params" + assert "Content-Type" not in params, "Content-Type should not be in query params" def test_transform_retrieve_file_request_with_raw_id_only(self): """ @@ -86,9 +82,7 @@ class TestGoogleAIStudioFilesTransformation: litellm_params=litellm_params, ) - assert ( - url == "https://generativelanguage.googleapis.com/v1beta/files/cctqueckiggb" - ) + assert url == "https://generativelanguage.googleapis.com/v1beta/files/cctqueckiggb" assert "key=" not in url assert params == {} @@ -103,8 +97,7 @@ class TestGoogleAIStudioFilesTransformation: ) assert ( - url - == "https://generativelanguage.googleapis.com/v1beta/files/..%2F..%2Fmodels%2Fgemini-pro%3Fx%3D1%23frag" + url == "https://generativelanguage.googleapis.com/v1beta/files/..%2F..%2Fmodels%2Fgemini-pro%3Fx%3D1%23frag" ) assert params == {} @@ -358,7 +351,49 @@ class TestGoogleAIStudioFilesTransformation: ) assert ( - url - == "https://generativelanguage.googleapis.com/v1beta/files/..%2F..%2Fmodels%2Fgemini-pro%3Fx%3D1%23frag" + url == "https://generativelanguage.googleapis.com/v1beta/files/..%2F..%2Fmodels%2Fgemini-pro%3Fx%3D1%23frag" ) assert params == {} + + +@pytest.mark.parametrize( + "file_id", + ["batch-la6", "files/batch-la6", "https://generativelanguage.googleapis.com/v1beta/files/batch-la6"], +) +def test_file_content_request_downloads_the_file_as_media(file_id): + handler = GoogleAIStudioFilesHandler() + url, params = handler.transform_file_content_request( + file_content_request={"file_id": file_id}, + optional_params={}, + litellm_params={"api_base": "https://proxy.example"}, + ) + assert (url, params) == ("https://proxy.example/download/v1beta/files/batch-la6:download", {"alt": "media"}) + + +def test_file_content_response_passes_the_jsonl_bytes_through(): + handler = GoogleAIStudioFilesHandler() + body = b'{"custom_id":"r0","response":{"status_code":200,"body":{"object":"chat.completion"}}}\n' + raw = httpx.Response(200, content=body, request=httpx.Request("GET", "https://generativelanguage.googleapis.com/x")) + assert ( + handler.transform_file_content_response(raw_response=raw, logging_obj=Mock(), litellm_params={}).content == body + ) + + +def test_retrieve_file_response_accepts_nanosecond_create_time(): + handler = GoogleAIStudioFilesHandler() + raw = httpx.Response( + 200, + json={ + "name": "files/batch-la6", + "displayName": "Results of batch la6", + "mimeType": "application/jsonl", + "sizeBytes": "1024", + "createTime": "2026-11-10T21:35:00.902340270Z", + "updateTime": "2026-11-10T21:35:00.902340270Z", + "uri": "https://generativelanguage.googleapis.com/v1beta/files/batch-la6", + "state": "ACTIVE", + }, + request=httpx.Request("GET", "https://generativelanguage.googleapis.com/v1beta/files/batch-la6"), + ) + file_object = handler.transform_retrieve_file_response(raw_response=raw, logging_obj=Mock(), litellm_params={}) + assert (file_object.created_at, file_object.bytes, file_object.status) == (1794346500, 1024, "processed")