mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
feat(gemini): file-backed batch jobs on /v1/batches with output download and spend tracking (#51032)
* feat(gemini): file-backed batch jobs on /v1/batches with output download
Adds GeminiBatchesConfig (create + retrieve) behind custom_llm_provider=gemini,
registers it in ProviderConfigManager, and implements file content retrieval on
GoogleAIStudioFilesHandler so GET /v1/files/{output_file_id}/content returns the
batch output. Gemini runs the OpenAI-shaped JSONL the proxy uploads today unchanged.
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(gemini): type the file content request params so no pyright rule grows
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* feat(gemini): price completed batch output rows so spend lands on the key
Gemini writes the usage of an OpenAI-shaped output row as camelCase
(promptTokens, completionTokens, totalTokens) with thinking tokens only in
totalTokens. Map it onto Usage with the thinking tokens folded into
completion_tokens, so the shared batch cost path bills them at the gemini
_batches rates.
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(gemini): parse nanosecond file timestamps
Gemini stamps files with nine fractional digits, which strptime's %f
rejects, so every retrieve of a batch output file logged a ValueError.
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(gemini): parse file timestamps on Python 3.10 too
fromisoformat on 3.10 takes neither a Z suffix nor more than six fractional digits.
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(gemini): price native batch output rows, stream output downloads, reject unsupported batch metadata
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(gemini): detect native batch rows by key, anchor the streaming file download handler
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* refactor(httpx): make anchor_handler_to public so the file content stream can use it
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---------
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu
GitHub
parent
c9d3c8b52f
commit
a0f2562088
14 changed files with 878 additions and 88 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
0
litellm/llms/gemini/batches/__init__.py
Normal file
0
litellm/llms/gemini/batches/__init__.py
Normal file
301
litellm/llms/gemini/batches/transformation.py
Normal file
301
litellm/llms/gemini/batches/transformation.py
Normal file
|
|
@ -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/<id>`` is the batch id ``<id>`` and its
|
||||
``files/<id>`` output file is the file id ``<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))
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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})
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
0
tests/unit/llms/gemini/batches/__init__.py
Normal file
0
tests/unit/llms/gemini/batches/__init__.py
Normal file
|
|
@ -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/<id>`` and ``files/<id>`` 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,
|
||||
)
|
||||
|
|
@ -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")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue