mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
fix(xai): name a completed xAI batch's results file and stop looking it up on xAI's files API (#45559)
* fix(xai): report a completed xAI batch's results as its output file details * fix(xai): keep a short results download and per-call timeouts out of the reported file details * fix(xai): describe a batch's results file from the batch alone and let the batch job fill its size Reading an xAI batch id as a file now costs one batch read and never downloads the results: the file reports a 0 byte placeholder, the results filename and the registration time. The managed proxy fills in the size the batch cost job measures when it downloads the results, keeps a known size when a provider reports none, and drops the per-process results download cache. * fix(managed_files): keep a measured batch output size when a refresh races the batch job A refresh that read the row before the batch cost job stored the measured size could save the provider's unsized details over it. The refresh now reads the row back from the writer and saves only while the row still holds what it read, retrying when another write landed in between. The new tests drive the real xAI and OpenAI retrieve paths through a fake HTTP upstream. * fix(managed_files): guard the refresh save with the row the request already read The conditional save first compares against the row the request read, so a refresh adds no database read unless another write landed in between; only then does it read the row back from the writer and try again. * fix(managed_files): save a refreshed file object with one guarded write and no extra database reads A refresh that loses its guarded write to a newer one now leaves that row alone and drops this worker's cached copy, so the next read loads the newer row instead of retrying against the database on the request path * fix(managed_files): retry the batch cost job's measured size after a read refreshes the row first The batch cost job re-reads the row from the writer and saves again when its guarded write loses to a concurrent refresh, so its measured size is never dropped. Reads without a measured size still make one write and no re-read * fix(managed_files): size a concurrently refreshed row instead of restoring the batch job's placeholder copy When the batch cost job loses its guarded write and the row it re-reads holds refreshed details, it applies its measured size to those details, so the results filename and any size the provider reported are kept * fix(managed_files): keep a batch output's earlier created_at when a refresh replaces its placeholder A refresh saved over a stored row now keeps the earlier of the two creation times, so an xAI placeholder keeps the time the proxy registered it instead of the read time, while a provider's own earlier creation time still wins * test(integration): an xAI placeholder row keeps its registration time after the proxy refreshes it * test(integration): drop the cached output row so the placeholder cell reads the seeded database row * test(integration): drop the cached output row without requiring it to still exist * fix(managed_files): limit the batch output size fill, guarded save and kept created_at to xAI rows * fix(xai): keep the batch id as the batch output filename Managed xAI batch output files keep the filename main already returns (the batch id) instead of renaming them to <batch id>_results.jsonl. --------- Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com>
This commit is contained in:
yucheng-berri
GitHub
parent
c0663d5202
commit
d0fa05076d
8 changed files with 1242 additions and 9 deletions
|
|
@ -24,7 +24,7 @@ from uuid import NAMESPACE_URL, uuid5
|
|||
|
||||
import httpx
|
||||
from fastapi import HTTPException
|
||||
from pydantic import ValidationError
|
||||
from pydantic import BaseModel, ValidationError
|
||||
from typing_extensions import ReadOnly, Unpack
|
||||
|
||||
import litellm
|
||||
|
|
@ -54,6 +54,7 @@ from litellm.llms.base_llm.managed_resources.isolation import (
|
|||
can_access_resource,
|
||||
resolve_resource_owner_id,
|
||||
)
|
||||
from litellm.llms.xai.batches.transformation import is_xai_batch_results_id
|
||||
from litellm.proxy._types import (
|
||||
CallTypes,
|
||||
LiteLLM_ManagedFileTable,
|
||||
|
|
@ -102,6 +103,7 @@ from litellm.types.utils import (
|
|||
CallTypesLiteral,
|
||||
LiteLLMBatch,
|
||||
LiteLLMFineTuningJob,
|
||||
LlmProviders,
|
||||
LLMResponseTypes,
|
||||
SpecialEnums,
|
||||
)
|
||||
|
|
@ -194,6 +196,32 @@ def _is_transient_file_retrieve_error(error: Exception) -> bool:
|
|||
return isinstance(error, (httpx.TransportError, APIConnectionError, asyncio.TimeoutError))
|
||||
|
||||
|
||||
_MEASURED_SIZE_SAVE_ATTEMPTS: Final = 3
|
||||
|
||||
|
||||
def _stored_json_text(raw_file_object: object) -> str:
|
||||
"""JSON text of a stored file object exactly as read, for a write that must not race a newer one."""
|
||||
if isinstance(raw_file_object, BaseModel):
|
||||
return raw_file_object.model_dump_json()
|
||||
return json.dumps(raw_file_object)
|
||||
|
||||
|
||||
def _with_measured_size(file_object: OpenAIFileObject, size_bytes: int | None) -> OpenAIFileObject:
|
||||
"""Fill a file object whose provider reported no size (0 bytes) with a size measured elsewhere."""
|
||||
if not size_bytes or file_object.bytes:
|
||||
return file_object
|
||||
return file_object.model_copy(update={"bytes": size_bytes})
|
||||
|
||||
|
||||
def _over_stored_object(file_object: OpenAIFileObject, stored_object: OpenAIFileObject | None) -> OpenAIFileObject:
|
||||
"""Refreshed details for a stored row, keeping its size when none is reported and its stored created_at."""
|
||||
if stored_object is None:
|
||||
return file_object
|
||||
return _with_measured_size(file_object, stored_object.bytes).model_copy(
|
||||
update={"created_at": stored_object.created_at}
|
||||
)
|
||||
|
||||
|
||||
def _proxy_llm_router() -> Router | None:
|
||||
import litellm.proxy.proxy_server as proxy_server_module
|
||||
|
||||
|
|
@ -252,6 +280,23 @@ def _has_provider_file_retrieve_route(
|
|||
return bool(router_credentials) or _model_name_file_retrieve_provider(model_name) is not None
|
||||
|
||||
|
||||
def _is_xai_batch_output(
|
||||
*,
|
||||
llm_router: Router | None,
|
||||
model_id: str | None,
|
||||
model_name: str | None,
|
||||
provider_file_id: str,
|
||||
) -> bool:
|
||||
"""Whether a batch output row holds an xAI batch's results, from its batch id and its deployment's provider."""
|
||||
if not is_xai_batch_results_id(provider_file_id):
|
||||
return False
|
||||
router_credentials: Final = _provider_file_retrieve_credentials(llm_router=llm_router, model_id=model_id)
|
||||
provider: Final = (router_credentials or {}).get("custom_llm_provider") or _model_name_file_retrieve_provider(
|
||||
model_name
|
||||
)
|
||||
return provider == LlmProviders.XAI.value
|
||||
|
||||
|
||||
class _ManagedFileRow(Protocol):
|
||||
unified_file_id: str
|
||||
file_object: OpenAIFileObject
|
||||
|
|
@ -845,6 +890,57 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
litellm_parent_otel_span=None,
|
||||
)
|
||||
|
||||
async def _save_xai_refreshed_file_object(
|
||||
self,
|
||||
stored: LiteLLM_ManagedFileTable,
|
||||
file_object: OpenAIFileObject,
|
||||
) -> OpenAIFileObject | None:
|
||||
"""Save an xAI batch output's refreshed details, keeping a known size, only over the row as read; drop the cached copy if a newer write won."""
|
||||
stored_object: Final = _parse_managed_file_object(stored.file_object, stored.unified_file_id)
|
||||
sized_file_object: Final = _over_stored_object(file_object, stored_object)
|
||||
if not await ManagedFileRepository(self.prisma_client).update_file_object(
|
||||
stored.unified_file_id,
|
||||
sized_file_object,
|
||||
if_file_object=_stored_json_text(stored.file_object) if stored.file_object is not None else None,
|
||||
):
|
||||
await self.internal_usage_cache.async_set_cache(
|
||||
key=stored.unified_file_id, value=None, litellm_parent_otel_span=None
|
||||
)
|
||||
return None
|
||||
await self.internal_usage_cache.async_set_cache(
|
||||
key=stored.unified_file_id,
|
||||
value=stored.model_copy(update={"file_object": sized_file_object}).model_dump(),
|
||||
litellm_parent_otel_span=None,
|
||||
)
|
||||
return sized_file_object
|
||||
|
||||
async def _save_batch_output_file_object(
|
||||
self,
|
||||
stored: LiteLLM_ManagedFileTable,
|
||||
file_object: OpenAIFileObject,
|
||||
size_bytes: int | None,
|
||||
) -> None:
|
||||
"""Save an xAI batch output's details; after losing to a refresh, size the refreshed details instead of replacing them."""
|
||||
current_row: LiteLLM_ManagedFileTable | None = stored # rebind-ok: the row as last read
|
||||
details: OpenAIFileObject = file_object # rebind-ok: a concurrent refresh's details replace the job's copy
|
||||
for _ in range(_MEASURED_SIZE_SAVE_ATTEMPTS):
|
||||
if current_row is None:
|
||||
return
|
||||
saved = await self._save_xai_refreshed_file_object( # rebind-ok: per attempt
|
||||
current_row, _with_measured_size(details, size_bytes)
|
||||
)
|
||||
if saved is not None or not size_bytes:
|
||||
return
|
||||
latest_row = await ManagedFileRepository( # rebind-ok: per attempt
|
||||
self.prisma_client, use_writer=True
|
||||
).table.find_first(where={"unified_file_id": stored.unified_file_id})
|
||||
current_row = (
|
||||
None if latest_row is None else LiteLLM_ManagedFileTable.model_validate(latest_row.model_dump())
|
||||
)
|
||||
latest_object = None if current_row is None else current_row.file_object # rebind-ok: per attempt
|
||||
if latest_object is not None and not latest_object.litellm_details_fallback:
|
||||
details = latest_object
|
||||
|
||||
async def store_batch_output_file(
|
||||
self,
|
||||
*,
|
||||
|
|
@ -885,6 +981,18 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
)
|
||||
return
|
||||
if stored_object is not None and not stored_object.litellm_details_fallback:
|
||||
if (
|
||||
stored_file is not None
|
||||
and size_bytes
|
||||
and not stored_object.bytes
|
||||
and _is_xai_batch_output(
|
||||
llm_router=_proxy_llm_router(),
|
||||
model_id=model_id,
|
||||
model_name=model_name,
|
||||
provider_file_id=provider_file_id,
|
||||
)
|
||||
):
|
||||
await self._save_batch_output_file_object(stored_file, stored_object, size_bytes)
|
||||
return
|
||||
|
||||
fallback_written_recently: Final = (
|
||||
|
|
@ -924,13 +1032,22 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
)
|
||||
)
|
||||
|
||||
xai_batch_output: Final = _is_xai_batch_output(
|
||||
llm_router=_proxy_llm_router(),
|
||||
model_id=model_id,
|
||||
model_name=model_name,
|
||||
provider_file_id=provider_file_id,
|
||||
)
|
||||
if stored_file is not None and xai_batch_output:
|
||||
await self._save_batch_output_file_object(stored_file, file_object, size_bytes)
|
||||
return
|
||||
if stored_file is not None:
|
||||
await self._save_refreshed_file_object(stored_file, file_object)
|
||||
return
|
||||
|
||||
await self.store_unified_file_id(
|
||||
file_id=unified_file_id,
|
||||
file_object=file_object,
|
||||
file_object=_with_measured_size(file_object, size_bytes) if xai_batch_output else file_object,
|
||||
litellm_parent_otel_span=litellm_parent_otel_span,
|
||||
model_mappings={model_id: provider_file_id} if model_id else {},
|
||||
user_api_key_dict=owner,
|
||||
|
|
@ -1943,8 +2060,20 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
)
|
||||
if refreshed_file_object is None:
|
||||
return _public_file_object(file_object, file_id)
|
||||
await self._save_refreshed_file_object(stored_file_object, refreshed_file_object)
|
||||
return _public_file_object(refreshed_file_object, file_id)
|
||||
if not _is_xai_batch_output(
|
||||
llm_router=llm_router if llm_router is not None else _proxy_llm_router(),
|
||||
model_id=model_id,
|
||||
model_name=model_id,
|
||||
provider_file_id=provider_file_id,
|
||||
):
|
||||
await self._save_refreshed_file_object(stored_file_object, refreshed_file_object)
|
||||
return _public_file_object(refreshed_file_object, file_id)
|
||||
saved_file_object: Final = await self._save_xai_refreshed_file_object(
|
||||
stored_file_object, refreshed_file_object
|
||||
)
|
||||
return _public_file_object(
|
||||
saved_file_object or _over_stored_object(refreshed_file_object, file_object), file_id
|
||||
)
|
||||
except Exception as error:
|
||||
verbose_logger.warning(
|
||||
"Failed to refresh batch file object for "
|
||||
|
|
|
|||
|
|
@ -366,7 +366,15 @@ def file_retrieve(
|
|||
|
||||
_is_async: Final = kwargs.pop("is_async", False) is True
|
||||
|
||||
if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
|
||||
if custom_llm_provider == LlmProviders.XAI.value and is_xai_batch_results_id(file_id):
|
||||
response = xai_batch_results_instance.batch_results_file( # rebind-ok: shares the result slot with the provider branches below
|
||||
_is_async=_is_async,
|
||||
batch_id=file_id,
|
||||
api_base=optional_params.api_base,
|
||||
api_key=optional_params.api_key,
|
||||
timeout=timeout,
|
||||
)
|
||||
elif custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
|
||||
openai_creds: Final = get_openai_credentials(
|
||||
api_base=optional_params.api_base,
|
||||
api_key=optional_params.api_key,
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
HTTPHandler,
|
||||
get_async_httpx_client,
|
||||
)
|
||||
from litellm.types.llms.openai import CreateBatchRequest, HttpxBinaryResponseContent
|
||||
from litellm.types.llms.openai import CreateBatchRequest, HttpxBinaryResponseContent, OpenAIFileObject
|
||||
from litellm.types.utils import LiteLLMBatch, LlmProviders
|
||||
|
||||
from .transformation import (
|
||||
|
|
@ -26,6 +26,7 @@ from .transformation import (
|
|||
to_create_batch_body,
|
||||
to_litellm_batch,
|
||||
to_openai_batch_list,
|
||||
xai_batch_results_file_object,
|
||||
xai_batches_url,
|
||||
)
|
||||
|
||||
|
|
@ -114,6 +115,24 @@ class XAIBatchesHandler:
|
|||
response: Final = self._sync(timeout).get(url, headers=headers, timeout=timeout)
|
||||
return to_litellm_batch(XAIBatch.model_validate(raise_for_xai_status(response).json()))
|
||||
|
||||
def batch_results_file(
|
||||
self,
|
||||
_is_async: bool,
|
||||
batch_id: str,
|
||||
api_base: str | None,
|
||||
api_key: str | None,
|
||||
timeout: float | httpx.Timeout,
|
||||
) -> OpenAIFileObject | Coroutine[None, None, OpenAIFileObject]:
|
||||
"""Describe a completed batch's results as a file from one batch read, never downloading the results."""
|
||||
retrieved: Final = self.retrieve_batch(_is_async, batch_id, api_base, api_key, timeout)
|
||||
if isinstance(retrieved, LiteLLMBatch):
|
||||
return xai_batch_results_file_object(retrieved)
|
||||
|
||||
async def _afile() -> OpenAIFileObject:
|
||||
return xai_batch_results_file_object(await retrieved)
|
||||
|
||||
return _afile()
|
||||
|
||||
def cancel_batch(
|
||||
self,
|
||||
_is_async: bool,
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ xAI batches carry request counters, not a status, and no output file: results ar
|
|||
"""
|
||||
|
||||
import json
|
||||
import time
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime, timezone
|
||||
from types import MappingProxyType
|
||||
|
|
@ -24,7 +25,7 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
|||
from litellm.llms.xai.common_utils import XAIModelInfo
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
from litellm.types.llms.openai import CreateBatchRequest
|
||||
from litellm.types.llms.openai import CreateBatchRequest, OpenAIFileObject
|
||||
from litellm.types.utils import LiteLLMBatch
|
||||
|
||||
OpenAIBatchStatus: TypeAlias = Literal[
|
||||
|
|
@ -85,6 +86,23 @@ def is_xai_batch_results_id(file_id: str) -> bool:
|
|||
return file_id.startswith(XAI_BATCH_ID_PREFIX)
|
||||
|
||||
|
||||
def xai_batch_results_file_object(batch: LiteLLMBatch) -> OpenAIFileObject:
|
||||
"""Describe a completed batch's results as a file without downloading them, so the size is a 0 placeholder."""
|
||||
if batch.output_file_id is None:
|
||||
raise xai_batches_error(
|
||||
f"xAI batch {batch.id} has no results file until it completes (status={batch.status})", 404, _EMPTY_HEADERS
|
||||
)
|
||||
return OpenAIFileObject(
|
||||
id=batch.output_file_id,
|
||||
object="file",
|
||||
purpose="batch_output",
|
||||
filename=batch.output_file_id,
|
||||
bytes=0,
|
||||
created_at=int(time.time()),
|
||||
status="processed",
|
||||
)
|
||||
|
||||
|
||||
class XAICreateBatchRequest(TypedDict):
|
||||
name: ReadOnly[str]
|
||||
input_file_id: NotRequired[ReadOnly[str]]
|
||||
|
|
|
|||
|
|
@ -10,9 +10,14 @@ if TYPE_CHECKING:
|
|||
class ManagedFileRepository(PrismaTableRepository["prisma_models.LiteLLM_ManagedFileTable"]):
|
||||
table_name = "litellm_managedfiletable"
|
||||
|
||||
async def update_file_object(self, unified_file_id: str, file_object: OpenAIFileObject) -> bool:
|
||||
async def update_file_object(
|
||||
self, unified_file_id: str, file_object: OpenAIFileObject, *, if_file_object: str | None = None
|
||||
) -> bool:
|
||||
"""Replace a row's file object, only while it still holds ``if_file_object`` (JSON text) when given."""
|
||||
updated_rows: Final = await self.table.update_many(
|
||||
where={"unified_file_id": unified_file_id},
|
||||
where={"unified_file_id": unified_file_id}
|
||||
if if_file_object is None
|
||||
else {"unified_file_id": unified_file_id, "file_object": {"equals": if_file_object}},
|
||||
data={"file_object": file_object.model_dump_json()},
|
||||
)
|
||||
return updated_rows > 0
|
||||
|
|
|
|||
334
tests/integration/providers/test_xai_batch_output_file_wire.py
Normal file
334
tests/integration/providers/test_xai_batch_output_file_wire.py
Normal file
|
|
@ -0,0 +1,334 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value
|
||||
from integration._support.database import read_rows, write_rows
|
||||
from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario
|
||||
from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse
|
||||
from pydantic import JsonValue
|
||||
from redis import Redis
|
||||
|
||||
XAI_MODEL: Final = "xai/grok-4.20-0309-non-reasoning"
|
||||
XAI_BATCH_ID: Final = "batch_$REQUEST_ID"
|
||||
REGISTERED_AT: Final = 1791400000
|
||||
_INPUT: Final = "".join(
|
||||
json.dumps(
|
||||
{
|
||||
"custom_id": custom_id,
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {"model": "grok-4.20-0309-non-reasoning", "messages": [{"role": "user", "content": "ping"}]},
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
for custom_id in ("req-1", "req-2")
|
||||
)
|
||||
|
||||
|
||||
def _completion(custom_id: str, content: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"batch_request_id": custom_id,
|
||||
"batch_result": {
|
||||
"response": {
|
||||
"chat_get_completion": {
|
||||
"id": f"chatcmpl-{custom_id}",
|
||||
"object": "chat.completion",
|
||||
"model": "grok-4.20-0309-non-reasoning",
|
||||
"choices": [
|
||||
{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}
|
||||
],
|
||||
"usage": {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9},
|
||||
}
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _xai_batch(*, completed: bool) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"batch_id": XAI_BATCH_ID,
|
||||
"name": "litellm-batch",
|
||||
"create_time": "2026-10-09T00:00:00Z",
|
||||
"expire_time": "2026-10-10T00:00:00Z",
|
||||
"state": {
|
||||
"num_requests": 2 if completed else 0,
|
||||
"num_pending": 0,
|
||||
"num_success": 2 if completed else 0,
|
||||
"num_error": 0,
|
||||
"num_cancelled": 0,
|
||||
},
|
||||
"input_file_id": "file_in_$REQUEST_ID",
|
||||
}
|
||||
|
||||
|
||||
def _xai_routes(*, results_status: int = 200) -> RoutedResponse:
|
||||
results: Final = (
|
||||
JsonResponse(
|
||||
content_type="application/json",
|
||||
body={
|
||||
"results": [_completion("req-1", "first"), _completion("req-2", "second")],
|
||||
"pagination_token": None,
|
||||
},
|
||||
)
|
||||
if results_status == 200
|
||||
else JsonResponse(
|
||||
content_type="application/json",
|
||||
body={"code": "internal", "error": "results unavailable"},
|
||||
status=results_status,
|
||||
)
|
||||
)
|
||||
return RoutedResponse(
|
||||
content_type="application/x-routed",
|
||||
routes={
|
||||
"POST /v1/files": JsonResponse(
|
||||
content_type="application/json",
|
||||
body={
|
||||
"id": "file_in_$REQUEST_ID",
|
||||
"bytes": 240,
|
||||
"created_at": 1791500000,
|
||||
"filename": "input.jsonl",
|
||||
"purpose": "",
|
||||
},
|
||||
),
|
||||
"GET /v1/files/file_in_$REQUEST_ID/content": TextResponse(content_type="application/jsonl", body=_INPUT),
|
||||
"POST /v1/batches": JsonResponse(content_type="application/json", body=_xai_batch(completed=False)),
|
||||
f"GET /v1/batches/{XAI_BATCH_ID}": JsonResponse(
|
||||
content_type="application/json", body=_xai_batch(completed=True)
|
||||
),
|
||||
f"GET /v1/batches/{XAI_BATCH_ID}/results": results,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _XAIBatch:
|
||||
owner_key: str
|
||||
batch_id: str
|
||||
upstream: ScenarioHandle
|
||||
|
||||
|
||||
def _create_xai_batch(scenario: Scenario, *, results_status: int = 200) -> _XAIBatch:
|
||||
handle: Final = register_scenario(
|
||||
f"xai-batch-output-{uuid.uuid4().hex}", _xai_routes(results_status=results_status)
|
||||
)
|
||||
scenario.cleanups.callback(delete_scenario, handle)
|
||||
model: Final = scenario.model(model=XAI_MODEL, api_key="integration-xai-key", api_base=handle.api_base())
|
||||
owner_key: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), models=[model])
|
||||
uploaded: Final = scenario.gateway.request_multipart(
|
||||
"/v1/files",
|
||||
{"purpose": "batch", "target_model_names": model},
|
||||
{"file": ("input.jsonl", _INPUT.encode(), "application/jsonl")},
|
||||
key=owner_key,
|
||||
)
|
||||
assert uploaded.status_code == 200, uploaded.text
|
||||
created: Final = scenario.gateway.request(
|
||||
"POST",
|
||||
"/v1/batches",
|
||||
{
|
||||
"input_file_id": string_value(JSON_OBJECT.validate_json(uploaded.content)["id"]),
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"completion_window": "24h",
|
||||
"model": model,
|
||||
},
|
||||
key=owner_key,
|
||||
)
|
||||
assert created.status_code == 200, created.text
|
||||
return _XAIBatch(owner_key, string_value(JSON_OBJECT.validate_json(created.content)["id"]), handle)
|
||||
|
||||
|
||||
def _upstream_paths(gateway: Gateway, scenario_id: str) -> tuple[str, ...]:
|
||||
response: Final = httpx.get(f"{gateway.upstream_url}/__observations", timeout=15, trust_env=False)
|
||||
assert response.status_code == 200, response.text
|
||||
requests: Final = JSON_OBJECT.validate_json(response.content)["requests"]
|
||||
assert isinstance(requests, list), response.text
|
||||
return tuple(
|
||||
f"{request.get('method')} {request.get('path')}"
|
||||
for request in map(object_value, requests)
|
||||
if scenario_id in str(request.get("path"))
|
||||
)
|
||||
|
||||
|
||||
def _results_reads(gateway: Gateway, batch: _XAIBatch, raw_batch_id: str) -> int:
|
||||
return sum(path.endswith(f"/{raw_batch_id}/results") for path in _upstream_paths(gateway, batch.upstream.scenario_id))
|
||||
|
||||
|
||||
def _stable_results_reads(gateway: Gateway, batch: _XAIBatch, raw_batch_id: str) -> int:
|
||||
def _read_twice() -> tuple[int, int]:
|
||||
first: Final = _results_reads(gateway, batch, raw_batch_id)
|
||||
time.sleep(3)
|
||||
return first, _results_reads(gateway, batch, raw_batch_id)
|
||||
|
||||
return eventually(_read_twice, lambda reads: reads[0] == reads[1], seconds=30)[1]
|
||||
|
||||
|
||||
def test_xai_batch_output_file_is_sized_by_the_batch_job_and_never_downloaded_or_looked_up_on_read(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
batch: Final = _create_xai_batch(scenario)
|
||||
raw_batch_id: Final = XAI_BATCH_ID.replace("$REQUEST_ID", batch.upstream.scenario_id)
|
||||
retrieved_response: Final = gateway.request("GET", f"/v1/batches/{batch.batch_id}", key=batch.owner_key)
|
||||
assert retrieved_response.status_code == 200, retrieved_response.text
|
||||
retrieved: Final = JSON_OBJECT.validate_json(retrieved_response.content)
|
||||
assert retrieved["status"] == "completed", retrieved
|
||||
output_id: Final = string_value(retrieved["output_file_id"])
|
||||
|
||||
first_read: Final = gateway.request("GET", f"/v1/files/{output_id}", key=batch.owner_key)
|
||||
assert first_read.status_code == 200, first_read.text
|
||||
first_detail: Final = JSON_OBJECT.validate_json(first_read.content)
|
||||
assert (first_detail["id"], first_detail["purpose"], first_detail["filename"]) == (
|
||||
output_id,
|
||||
"batch_output",
|
||||
raw_batch_id,
|
||||
), first_detail
|
||||
|
||||
content: Final = gateway.request("GET", f"/v1/files/{output_id}/content", key=batch.owner_key)
|
||||
assert content.status_code == 200, content.text
|
||||
lines: Final = tuple(json.loads(line) for line in content.content.decode().splitlines())
|
||||
assert [
|
||||
(line["custom_id"], line["response"]["body"]["choices"][0]["message"]["content"]) for line in lines
|
||||
] == [
|
||||
("req-1", "first"),
|
||||
("req-2", "second"),
|
||||
], lines
|
||||
|
||||
sized: Final = eventually(
|
||||
lambda: JSON_OBJECT.validate_json(
|
||||
gateway.request("GET", f"/v1/files/{output_id}", key=batch.owner_key).content
|
||||
),
|
||||
lambda detail: detail.get("bytes") == len(content.content),
|
||||
seconds=60,
|
||||
)
|
||||
assert (sized["filename"], "litellm_details_fallback" in sized) == (raw_batch_id, False)
|
||||
|
||||
settled_reads: Final = _stable_results_reads(gateway, batch, raw_batch_id)
|
||||
reads: Final = (
|
||||
*(gateway.request("GET", f"/v1/files/{output_id}", key=batch.owner_key) for _ in range(3)),
|
||||
gateway.request("GET", f"/v1/batches/{batch.batch_id}", key=batch.owner_key),
|
||||
gateway.request("GET", "/v1/files", key=batch.owner_key, params={"purpose": "batch_output"}),
|
||||
)
|
||||
assert all(response.status_code == 200 for response in reads), [response.text for response in reads]
|
||||
listed_output: Final = next(value for value in reads[-1].json()["data"] if value.get("id") == output_id)
|
||||
assert (listed_output["filename"], listed_output["bytes"]) == (
|
||||
raw_batch_id,
|
||||
len(content.content),
|
||||
), listed_output
|
||||
assert _results_reads(gateway, batch, raw_batch_id) == settled_reads
|
||||
|
||||
file_lookups: Final = tuple(
|
||||
path for path in _upstream_paths(gateway, batch.upstream.scenario_id) if "/v1/files/batch_" in path
|
||||
)
|
||||
assert file_lookups == (), f"xAI batch id was looked up as a file: {file_lookups}"
|
||||
|
||||
|
||||
def test_xai_batch_output_file_answers_at_once_while_results_are_unavailable_and_is_sized_once_they_return(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
batch: Final = _create_xai_batch(scenario, results_status=503)
|
||||
raw_batch_id: Final = XAI_BATCH_ID.replace("$REQUEST_ID", batch.upstream.scenario_id)
|
||||
retrieved_response: Final = gateway.request("GET", f"/v1/batches/{batch.batch_id}", key=batch.owner_key)
|
||||
assert retrieved_response.status_code == 200, retrieved_response.text
|
||||
output_id: Final = string_value(JSON_OBJECT.validate_json(retrieved_response.content)["output_file_id"])
|
||||
|
||||
started: Final = time.monotonic()
|
||||
unavailable: Final = gateway.request("GET", f"/v1/files/{output_id}", key=batch.owner_key)
|
||||
elapsed: Final = time.monotonic() - started
|
||||
assert unavailable.status_code == 200, unavailable.text
|
||||
unavailable_detail: Final = JSON_OBJECT.validate_json(unavailable.content)
|
||||
assert (unavailable_detail["filename"], unavailable_detail["bytes"]) == (
|
||||
raw_batch_id,
|
||||
0,
|
||||
), unavailable_detail
|
||||
assert elapsed < 5, elapsed
|
||||
|
||||
register_scenario(batch.upstream.scenario_id, _xai_routes())
|
||||
content: Final = gateway.request("GET", f"/v1/files/{output_id}/content", key=batch.owner_key)
|
||||
assert content.status_code == 200, content.text
|
||||
recovered: Final = eventually(
|
||||
lambda: JSON_OBJECT.validate_json(
|
||||
gateway.request("GET", f"/v1/files/{output_id}", key=batch.owner_key).content
|
||||
),
|
||||
lambda detail: detail.get("bytes") == len(content.content),
|
||||
seconds=60,
|
||||
)
|
||||
assert (recovered["purpose"], recovered["filename"]) == ("batch_output", raw_batch_id)
|
||||
|
||||
file_lookups: Final = tuple(
|
||||
path for path in _upstream_paths(gateway, batch.upstream.scenario_id) if "/v1/files/batch_" in path
|
||||
)
|
||||
assert file_lookups == (), f"xAI batch id was looked up as a file: {file_lookups}"
|
||||
|
||||
|
||||
def test_xai_batch_output_file_registered_by_the_batch_job_first_carries_the_downloaded_size(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
batch: Final = _create_xai_batch(scenario)
|
||||
raw_batch_id: Final = XAI_BATCH_ID.replace("$REQUEST_ID", batch.upstream.scenario_id)
|
||||
|
||||
listed: Final = eventually(
|
||||
lambda: gateway.request("GET", "/v1/files", key=batch.owner_key, params={"purpose": "batch_output"}).json()[
|
||||
"data"
|
||||
],
|
||||
lambda values: any(value.get("filename") == raw_batch_id for value in values),
|
||||
seconds=60,
|
||||
)
|
||||
listed_output: Final = next(
|
||||
value for value in listed if value.get("filename") == raw_batch_id
|
||||
)
|
||||
content: Final = gateway.request("GET", f"/v1/files/{listed_output['id']}/content", key=batch.owner_key)
|
||||
assert content.status_code == 200, content.text
|
||||
assert (listed_output["purpose"], listed_output["bytes"]) == ("batch_output", len(content.content)), (
|
||||
listed_output
|
||||
)
|
||||
|
||||
upstream_paths: Final = _upstream_paths(gateway, batch.upstream.scenario_id)
|
||||
file_lookups: Final = tuple(path for path in upstream_paths if "/v1/files/batch_" in path)
|
||||
assert file_lookups == (), f"xAI batch id was looked up as a file: {file_lookups}"
|
||||
results_reads: Final = sum(path.endswith(f"/{raw_batch_id}/results") for path in upstream_paths)
|
||||
state_reads: Final = sum(path.endswith(f"/v1/batches/{raw_batch_id}") for path in upstream_paths)
|
||||
assert (results_reads, state_reads) == (2, 2), upstream_paths
|
||||
|
||||
|
||||
def test_xai_batch_output_file_left_as_a_placeholder_keeps_its_registration_time_when_the_proxy_refreshes_it(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
batch: Final = _create_xai_batch(scenario)
|
||||
raw_batch_id: Final = XAI_BATCH_ID.replace("$REQUEST_ID", batch.upstream.scenario_id)
|
||||
retrieved_response: Final = gateway.request("GET", f"/v1/batches/{batch.batch_id}", key=batch.owner_key)
|
||||
assert retrieved_response.status_code == 200, retrieved_response.text
|
||||
output_id: Final = string_value(JSON_OBJECT.validate_json(retrieved_response.content)["output_file_id"])
|
||||
write_rows(
|
||||
'UPDATE "LiteLLM_ManagedFileTable" SET file_object = file_object || jsonb_build_object('
|
||||
"'bytes', 0, 'filename', %s::text, 'created_at', %s::bigint, 'litellm_details_fallback', true) "
|
||||
"WHERE unified_file_id = %s",
|
||||
(raw_batch_id, str(REGISTERED_AT), output_id),
|
||||
)
|
||||
with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache:
|
||||
cache.delete(output_id)
|
||||
|
||||
refreshed: Final = eventually(
|
||||
lambda: JSON_OBJECT.validate_json(
|
||||
gateway.request("GET", f"/v1/files/{output_id}", key=batch.owner_key).content
|
||||
),
|
||||
lambda detail: detail.get("created_at") == REGISTERED_AT,
|
||||
seconds=30,
|
||||
)
|
||||
assert refreshed["filename"] == raw_batch_id, refreshed
|
||||
stored: Final = object_value(
|
||||
read_rows('SELECT file_object FROM "LiteLLM_ManagedFileTable" WHERE unified_file_id=%s', (output_id,))[0][
|
||||
"file_object"
|
||||
]
|
||||
)
|
||||
assert (stored["created_at"], stored["filename"], stored.get("litellm_details_fallback")) == (
|
||||
REGISTERED_AT,
|
||||
raw_batch_id,
|
||||
None,
|
||||
), stored
|
||||
|
|
@ -10,6 +10,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
from fastapi import HTTPException
|
||||
from openai import APIConnectionError
|
||||
|
||||
|
|
@ -109,6 +110,11 @@ class _InMemoryManagedFileTable:
|
|||
existing_row: Final = self.rows.get(unified_file_id)
|
||||
if existing_row is None:
|
||||
return 0
|
||||
expected_file_object: Final = where.get("file_object")
|
||||
if isinstance(expected_file_object, Mapping) and json.loads(cast(str, expected_file_object["equals"])) != (
|
||||
None if existing_row.file_object is None else existing_row.file_object.model_dump(mode="json")
|
||||
):
|
||||
return 0
|
||||
file_object: Final = OpenAIFileObject.model_validate(
|
||||
json.loads(cast(str, data["file_object"]))
|
||||
)
|
||||
|
|
@ -156,6 +162,7 @@ class _InMemoryManagedFilesDatabase:
|
|||
class _InMemoryManagedFilesPrismaClient:
|
||||
def __init__(self, managed_file_table: _InMemoryManagedFileTable) -> None:
|
||||
self.db: Final = _InMemoryManagedFilesDatabase(managed_file_table)
|
||||
self.writer_db: Final = self.db
|
||||
|
||||
|
||||
def _managed_files_with_fake_prisma(
|
||||
|
|
@ -2389,6 +2396,673 @@ async def test_store_batch_output_file_ignores_unmarked_row_with_provider_route(
|
|||
assert managed_file_table.upsert_calls == []
|
||||
|
||||
|
||||
_XAI_API_BASE: Final = "https://api.x.ai"
|
||||
_XAI_BATCH_ID: Final = "batch_lit9385"
|
||||
_OPENAI_API_BASE: Final = "https://api.openai.com/v1"
|
||||
_OPENAI_FILE_ID: Final = "file-lit9385"
|
||||
|
||||
|
||||
def _provider_router(provider: str) -> MagicMock:
|
||||
router: Final = MagicMock()
|
||||
router.get_deployment_credentials_with_provider.return_value = (
|
||||
{"custom_llm_provider": "xai", "api_key": "xai-test-key", "api_base": _XAI_API_BASE}
|
||||
if provider == "xai"
|
||||
else {"custom_llm_provider": "openai", "api_key": "sk-test-key", "api_base": _OPENAI_API_BASE}
|
||||
)
|
||||
return router
|
||||
|
||||
|
||||
def _provider_file_id(provider: str) -> str:
|
||||
return _XAI_BATCH_ID if provider == "xai" else _OPENAI_FILE_ID
|
||||
|
||||
|
||||
def _xai_completed_batch() -> httpx.Response:
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"batch_id": _XAI_BATCH_ID,
|
||||
"name": "litellm-batch",
|
||||
"create_time": "2026-10-09",
|
||||
"expire_time": "2026-10-10",
|
||||
"state": {"num_requests": 3, "num_pending": 0, "num_success": 3, "num_error": 0, "num_cancelled": 0},
|
||||
"input_file_id": "file_in",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _provider_file_route(provider: str) -> respx.Route:
|
||||
if provider == "xai":
|
||||
return respx.get(f"{_XAI_API_BASE}/v1/batches/{_XAI_BATCH_ID}").mock(return_value=_xai_completed_batch())
|
||||
return respx.get(f"{_OPENAI_API_BASE}/files/{_OPENAI_FILE_ID}").respond(
|
||||
200,
|
||||
json={
|
||||
"id": _OPENAI_FILE_ID,
|
||||
"object": "file",
|
||||
"bytes": 836,
|
||||
"created_at": 456,
|
||||
"filename": "output.jsonl",
|
||||
"purpose": "batch_output",
|
||||
"status": "processed",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _batch_output_row(provider: str, *, size_bytes: int, fallback: bool) -> LiteLLM_ManagedFileTable:
|
||||
provider_file_id: Final = _provider_file_id(provider)
|
||||
return LiteLLM_ManagedFileTable(
|
||||
unified_file_id="unified-output",
|
||||
file_object=OpenAIFileObject(
|
||||
id="unified-output",
|
||||
object="file",
|
||||
bytes=size_bytes,
|
||||
created_at=123,
|
||||
filename=provider_file_id,
|
||||
purpose="batch_output",
|
||||
status="processed",
|
||||
litellm_details_fallback=True if fallback else None,
|
||||
),
|
||||
model_mappings={"model-123": provider_file_id},
|
||||
flat_model_file_ids=[provider_file_id],
|
||||
created_by="user-123",
|
||||
)
|
||||
|
||||
|
||||
def _size_written_by_batch_job(table: _InMemoryManagedFileTable) -> None:
|
||||
stored: Final = table.rows["unified-output"]
|
||||
assert stored.file_object is not None
|
||||
table.rows["unified-output"] = stored.model_copy(
|
||||
update={"file_object": stored.file_object.model_copy(update={"bytes": 2565})}
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _respx_sees_provider_calls(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
import litellm
|
||||
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("_respx_sees_provider_calls")
|
||||
@pytest.mark.parametrize(("stored_bytes", "expected_bytes"), [(0, 2565), (1, 1)])
|
||||
@respx.mock
|
||||
async def test_store_batch_output_file_fills_only_an_unknown_size_on_an_unmarked_row(
|
||||
stored_bytes: int, expected_bytes: int
|
||||
):
|
||||
import litellm.proxy.proxy_server as proxy_server_module
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
row: Final = _batch_output_row("xai", size_bytes=stored_bytes, fallback=False)
|
||||
proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(row)
|
||||
provider_route: Final = _provider_file_route("xai")
|
||||
with patch.object(proxy_server_module, "llm_router", _provider_router("xai")):
|
||||
await proxy_managed_files.store_batch_output_file(
|
||||
unified_file_id="unified-output",
|
||||
provider_file_id=_XAI_BATCH_ID,
|
||||
model_id="model-123",
|
||||
owner=UserAPIKeyAuth(user_id="user-123"),
|
||||
litellm_parent_otel_span=None,
|
||||
size_bytes=2565,
|
||||
)
|
||||
|
||||
assert row.file_object is not None
|
||||
assert provider_route.call_count == 0
|
||||
assert managed_file_table.rows["unified-output"].file_object == row.file_object.model_copy(
|
||||
update={"bytes": expected_bytes}
|
||||
)
|
||||
assert managed_file_table.upsert_calls == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("_respx_sees_provider_calls")
|
||||
@pytest.mark.parametrize(
|
||||
("provider", "expected_filename", "expected_bytes"),
|
||||
[("xai", _XAI_BATCH_ID, 2565), ("openai", "output.jsonl", 836)],
|
||||
)
|
||||
@pytest.mark.parametrize("has_marked_row", [False, True])
|
||||
@respx.mock
|
||||
async def test_store_batch_output_file_sizes_only_an_unsized_provider_file_from_the_downloaded_results(
|
||||
provider: str, expected_filename: str, expected_bytes: int, has_marked_row: bool
|
||||
):
|
||||
import litellm.proxy.proxy_server as proxy_server_module
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(
|
||||
*([_batch_output_row(provider, size_bytes=0, fallback=True)] if has_marked_row else [])
|
||||
)
|
||||
provider_route: Final = _provider_file_route(provider)
|
||||
with patch.object(proxy_server_module, "llm_router", _provider_router(provider)):
|
||||
await proxy_managed_files.store_batch_output_file(
|
||||
unified_file_id="unified-output",
|
||||
provider_file_id=_provider_file_id(provider),
|
||||
model_id="model-123",
|
||||
owner=UserAPIKeyAuth(user_id="user-123"),
|
||||
litellm_parent_otel_span=None,
|
||||
size_bytes=2565,
|
||||
)
|
||||
|
||||
stored: Final = managed_file_table.rows["unified-output"].file_object
|
||||
assert stored is not None
|
||||
assert provider_route.call_count == 1
|
||||
assert (stored.filename, stored.bytes, stored.litellm_details_fallback) == (
|
||||
expected_filename,
|
||||
expected_bytes,
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("_respx_sees_provider_calls")
|
||||
@respx.mock
|
||||
async def test_store_batch_output_file_keeps_its_measured_size_when_a_read_refreshes_the_row_first():
|
||||
import litellm.proxy.proxy_server as proxy_server_module
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(
|
||||
_batch_output_row("xai", size_bytes=0, fallback=True)
|
||||
)
|
||||
|
||||
def a_read_refreshes_the_row_then_xai_answers(_request: httpx.Request) -> httpx.Response:
|
||||
placeholder: Final = managed_file_table.rows["unified-output"]
|
||||
assert placeholder.file_object is not None
|
||||
managed_file_table.rows["unified-output"] = placeholder.model_copy(
|
||||
update={
|
||||
"file_object": placeholder.file_object.model_copy(
|
||||
update={"filename": _XAI_BATCH_ID, "litellm_details_fallback": None}
|
||||
)
|
||||
}
|
||||
)
|
||||
return _xai_completed_batch()
|
||||
|
||||
respx.get(f"{_XAI_API_BASE}/v1/batches/{_XAI_BATCH_ID}").mock(side_effect=a_read_refreshes_the_row_then_xai_answers)
|
||||
with patch.object(proxy_server_module, "llm_router", _provider_router("xai")):
|
||||
await proxy_managed_files.store_batch_output_file(
|
||||
unified_file_id="unified-output",
|
||||
provider_file_id=_provider_file_id("xai"),
|
||||
model_id="model-123",
|
||||
owner=UserAPIKeyAuth(user_id="user-123"),
|
||||
litellm_parent_otel_span=None,
|
||||
size_bytes=2565,
|
||||
)
|
||||
|
||||
stored: Final = managed_file_table.rows["unified-output"].file_object
|
||||
assert stored is not None
|
||||
assert (stored.filename, stored.bytes, stored.litellm_details_fallback) == (
|
||||
_XAI_BATCH_ID,
|
||||
2565,
|
||||
None,
|
||||
)
|
||||
assert len(managed_file_table.update_many_calls) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("_respx_sees_provider_calls")
|
||||
@respx.mock
|
||||
async def test_store_batch_output_file_from_a_batch_read_leaves_a_newer_row_alone_without_rereading_it():
|
||||
import litellm.proxy.proxy_server as proxy_server_module
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(
|
||||
_batch_output_row("xai", size_bytes=0, fallback=True)
|
||||
)
|
||||
|
||||
def the_batch_job_writes_then_xai_answers(_request: httpx.Request) -> httpx.Response:
|
||||
placeholder: Final = managed_file_table.rows["unified-output"]
|
||||
assert placeholder.file_object is not None
|
||||
managed_file_table.rows["unified-output"] = placeholder.model_copy(
|
||||
update={
|
||||
"file_object": placeholder.file_object.model_copy(
|
||||
update={
|
||||
"filename": _XAI_BATCH_ID,
|
||||
"bytes": 2565,
|
||||
"litellm_details_fallback": None,
|
||||
}
|
||||
)
|
||||
}
|
||||
)
|
||||
return _xai_completed_batch()
|
||||
|
||||
respx.get(f"{_XAI_API_BASE}/v1/batches/{_XAI_BATCH_ID}").mock(side_effect=the_batch_job_writes_then_xai_answers)
|
||||
reads_before_store: Final = len(managed_file_table.find_first_calls)
|
||||
with patch.object(proxy_server_module, "llm_router", _provider_router("xai")):
|
||||
await proxy_managed_files.store_batch_output_file(
|
||||
unified_file_id="unified-output",
|
||||
provider_file_id=_provider_file_id("xai"),
|
||||
model_id="model-123",
|
||||
owner=UserAPIKeyAuth(user_id="user-123"),
|
||||
litellm_parent_otel_span=None,
|
||||
)
|
||||
|
||||
stored: Final = managed_file_table.rows["unified-output"].file_object
|
||||
assert stored is not None
|
||||
assert (stored.filename, stored.bytes) == (_XAI_BATCH_ID, 2565)
|
||||
assert len(managed_file_table.update_many_calls) == 1
|
||||
assert len(managed_file_table.find_first_calls) - reads_before_store == 1
|
||||
|
||||
|
||||
class _RefreshLandsBeforeTheFirstWriteTable(_InMemoryManagedFileTable):
|
||||
def __init__(self, row: LiteLLM_ManagedFileTable, refreshed_file_object: OpenAIFileObject) -> None:
|
||||
super().__init__(row)
|
||||
self.pending_refreshes: Final = [refreshed_file_object]
|
||||
|
||||
async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int:
|
||||
if self.pending_refreshes:
|
||||
refreshed_row: Final = self.rows["unified-output"]
|
||||
self.rows["unified-output"] = refreshed_row.model_copy(
|
||||
update={"file_object": self.pending_refreshes.pop()}
|
||||
)
|
||||
return await super().update_many(where, data)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("refreshed_bytes", "expected_bytes"), [(0, 2565), (836, 836)])
|
||||
async def test_store_batch_output_file_sizes_the_refreshed_details_instead_of_restoring_its_placeholder(
|
||||
refreshed_bytes: int, expected_bytes: int
|
||||
):
|
||||
import litellm.proxy.proxy_server as proxy_server_module
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
placeholder_row: Final = _batch_output_row("xai", size_bytes=0, fallback=True)
|
||||
assert placeholder_row.file_object is not None
|
||||
recent_placeholder: Final = placeholder_row.file_object.model_copy(update={"created_at": int(time.time())})
|
||||
table: Final = _RefreshLandsBeforeTheFirstWriteTable(
|
||||
placeholder_row.model_copy(update={"file_object": recent_placeholder}),
|
||||
recent_placeholder.model_copy(
|
||||
update={
|
||||
"filename": _XAI_BATCH_ID,
|
||||
"bytes": refreshed_bytes,
|
||||
"litellm_details_fallback": None,
|
||||
}
|
||||
),
|
||||
)
|
||||
proxy_managed_files: Final = PROXY_LiteLLMManagedFiles(
|
||||
DualCache(), prisma_client=_InMemoryManagedFilesPrismaClient(table)
|
||||
)
|
||||
with patch.object(proxy_server_module, "llm_router", _provider_router("xai")):
|
||||
await proxy_managed_files.store_batch_output_file(
|
||||
unified_file_id="unified-output",
|
||||
provider_file_id=_provider_file_id("xai"),
|
||||
model_id="model-123",
|
||||
owner=UserAPIKeyAuth(user_id="user-123"),
|
||||
litellm_parent_otel_span=None,
|
||||
size_bytes=2565,
|
||||
)
|
||||
|
||||
stored: Final = table.rows["unified-output"].file_object
|
||||
assert stored is not None
|
||||
assert (stored.filename, stored.bytes, stored.litellm_details_fallback) == (
|
||||
_XAI_BATCH_ID,
|
||||
expected_bytes,
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
class _LaggingReplicaManagedFileTable:
|
||||
def __init__(self, replica: _InMemoryManagedFileTable, writer: _InMemoryManagedFileTable) -> None:
|
||||
self.replica: Final = replica
|
||||
self.writer: Final = writer
|
||||
|
||||
async def find_first(self, where: Mapping[str, object]) -> LiteLLM_ManagedFileTable | None:
|
||||
return await self.replica.find_first(where)
|
||||
|
||||
async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int:
|
||||
return await self.writer.update_many(where, data)
|
||||
|
||||
|
||||
class _LaggingReplicaPrismaClient:
|
||||
def __init__(self, replica: _InMemoryManagedFileTable, writer: _InMemoryManagedFileTable) -> None:
|
||||
self.db: Final = SimpleNamespace(
|
||||
litellm_managedfiletable=_LaggingReplicaManagedFileTable(replica, writer),
|
||||
litellm_managedobjecttable=_InMemoryManagedObjectTable(),
|
||||
)
|
||||
self.writer_db: Final = _InMemoryManagedFilesDatabase(writer)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("_respx_sees_provider_calls")
|
||||
@respx.mock
|
||||
async def test_store_batch_output_file_rereads_a_lost_row_from_the_writer_not_a_lagging_replica():
|
||||
import litellm.proxy.proxy_server as proxy_server_module
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
placeholder_row: Final = _batch_output_row("xai", size_bytes=0, fallback=True)
|
||||
assert placeholder_row.file_object is not None
|
||||
writer: Final = _InMemoryManagedFileTable(
|
||||
placeholder_row.model_copy(
|
||||
update={
|
||||
"file_object": placeholder_row.file_object.model_copy(
|
||||
update={"filename": _XAI_BATCH_ID, "litellm_details_fallback": None}
|
||||
)
|
||||
}
|
||||
)
|
||||
)
|
||||
proxy_managed_files: Final = PROXY_LiteLLMManagedFiles(
|
||||
DualCache(),
|
||||
prisma_client=_LaggingReplicaPrismaClient(_InMemoryManagedFileTable(placeholder_row), writer),
|
||||
)
|
||||
_provider_file_route("xai")
|
||||
with patch.object(proxy_server_module, "llm_router", _provider_router("xai")):
|
||||
await proxy_managed_files.store_batch_output_file(
|
||||
unified_file_id="unified-output",
|
||||
provider_file_id=_provider_file_id("xai"),
|
||||
model_id="model-123",
|
||||
owner=UserAPIKeyAuth(user_id="user-123"),
|
||||
litellm_parent_otel_span=None,
|
||||
size_bytes=2565,
|
||||
)
|
||||
|
||||
stored: Final = writer.rows["unified-output"].file_object
|
||||
assert stored is not None
|
||||
assert (stored.filename, stored.bytes) == (_XAI_BATCH_ID, 2565)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("_respx_sees_provider_calls")
|
||||
@pytest.mark.parametrize("stored_bytes", [0, 2565])
|
||||
@respx.mock
|
||||
async def test_afile_retrieve_refresh_keeps_the_registration_time_of_an_xai_placeholder(stored_bytes: int):
|
||||
proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(
|
||||
_batch_output_row("xai", size_bytes=stored_bytes, fallback=True)
|
||||
)
|
||||
_provider_file_route("xai")
|
||||
|
||||
response: Final = await proxy_managed_files.afile_retrieve(
|
||||
file_id="unified-output",
|
||||
litellm_parent_otel_span=None,
|
||||
llm_router=_provider_router("xai"),
|
||||
)
|
||||
|
||||
stored: Final = managed_file_table.rows["unified-output"].file_object
|
||||
assert stored is not None
|
||||
assert (response.created_at, stored.created_at, stored.filename) == (
|
||||
123,
|
||||
123,
|
||||
_XAI_BATCH_ID,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("_respx_sees_provider_calls")
|
||||
@respx.mock
|
||||
async def test_store_batch_output_file_keeps_the_registration_time_of_an_xai_placeholder():
|
||||
import litellm.proxy.proxy_server as proxy_server_module
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(
|
||||
_batch_output_row("xai", size_bytes=0, fallback=True)
|
||||
)
|
||||
_provider_file_route("xai")
|
||||
with patch.object(proxy_server_module, "llm_router", _provider_router("xai")):
|
||||
await proxy_managed_files.store_batch_output_file(
|
||||
unified_file_id="unified-output",
|
||||
provider_file_id=_provider_file_id("xai"),
|
||||
model_id="model-123",
|
||||
owner=UserAPIKeyAuth(user_id="user-123"),
|
||||
litellm_parent_otel_span=None,
|
||||
size_bytes=2565,
|
||||
)
|
||||
|
||||
stored: Final = managed_file_table.rows["unified-output"].file_object
|
||||
assert stored is not None
|
||||
assert (stored.created_at, stored.bytes, stored.filename) == (123, 2565, _XAI_BATCH_ID)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("_respx_sees_provider_calls")
|
||||
@respx.mock
|
||||
async def test_afile_retrieve_refresh_keeps_the_known_size_when_the_provider_reports_none():
|
||||
proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(
|
||||
_batch_output_row("xai", size_bytes=2565, fallback=True)
|
||||
)
|
||||
_provider_file_route("xai")
|
||||
|
||||
response: Final = await proxy_managed_files.afile_retrieve(
|
||||
file_id="unified-output",
|
||||
litellm_parent_otel_span=None,
|
||||
llm_router=_provider_router("xai"),
|
||||
)
|
||||
|
||||
stored: Final = managed_file_table.rows["unified-output"].file_object
|
||||
assert stored is not None
|
||||
assert (response.id, response.filename, response.bytes) == (
|
||||
"unified-output",
|
||||
_XAI_BATCH_ID,
|
||||
2565,
|
||||
)
|
||||
assert (stored.filename, stored.bytes, stored.litellm_details_fallback) == (
|
||||
_XAI_BATCH_ID,
|
||||
2565,
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("_respx_sees_provider_calls")
|
||||
@respx.mock
|
||||
async def test_afile_retrieve_refresh_keeps_a_size_the_batch_job_writes_while_it_runs():
|
||||
placeholder_row: Final = _batch_output_row("xai", size_bytes=0, fallback=True)
|
||||
proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(placeholder_row)
|
||||
await proxy_managed_files.internal_usage_cache.async_set_cache(
|
||||
key="unified-output", value=placeholder_row.model_dump(), litellm_parent_otel_span=None
|
||||
)
|
||||
|
||||
def batch_job_writes_then_xai_answers(_request: httpx.Request) -> httpx.Response:
|
||||
_size_written_by_batch_job(managed_file_table)
|
||||
return _xai_completed_batch()
|
||||
|
||||
respx.get(f"{_XAI_API_BASE}/v1/batches/{_XAI_BATCH_ID}").mock(side_effect=batch_job_writes_then_xai_answers)
|
||||
|
||||
raced_response: Final = await proxy_managed_files.afile_retrieve(
|
||||
file_id="unified-output",
|
||||
litellm_parent_otel_span=None,
|
||||
llm_router=_provider_router("xai"),
|
||||
)
|
||||
raced_row: Final = managed_file_table.rows["unified-output"].file_object
|
||||
assert raced_row is not None
|
||||
assert (raced_response.filename, raced_row.bytes, raced_row.filename) == (
|
||||
_XAI_BATCH_ID,
|
||||
2565,
|
||||
_XAI_BATCH_ID,
|
||||
)
|
||||
|
||||
next_response: Final = await proxy_managed_files.afile_retrieve(
|
||||
file_id="unified-output",
|
||||
litellm_parent_otel_span=None,
|
||||
llm_router=_provider_router("xai"),
|
||||
)
|
||||
stored: Final = managed_file_table.rows["unified-output"].file_object
|
||||
assert stored is not None
|
||||
assert (next_response.bytes, stored.bytes, stored.filename, stored.litellm_details_fallback) == (
|
||||
2565,
|
||||
2565,
|
||||
_XAI_BATCH_ID,
|
||||
None,
|
||||
)
|
||||
assert len(managed_file_table.update_many_calls) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("_respx_sees_provider_calls")
|
||||
@respx.mock
|
||||
async def test_afile_retrieve_refresh_that_loses_its_write_still_answers_with_the_known_size():
|
||||
proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(
|
||||
_batch_output_row("xai", size_bytes=2565, fallback=True)
|
||||
)
|
||||
read_row: Final = managed_file_table.rows["unified-output"]
|
||||
assert read_row.file_object is not None
|
||||
newer_file_object: Final = read_row.file_object.model_copy(update={"created_at": 456})
|
||||
|
||||
def another_write_lands_then_xai_answers(_request: httpx.Request) -> httpx.Response:
|
||||
managed_file_table.rows["unified-output"] = read_row.model_copy(update={"file_object": newer_file_object})
|
||||
return _xai_completed_batch()
|
||||
|
||||
respx.get(f"{_XAI_API_BASE}/v1/batches/{_XAI_BATCH_ID}").mock(side_effect=another_write_lands_then_xai_answers)
|
||||
|
||||
response: Final = await proxy_managed_files.afile_retrieve(
|
||||
file_id="unified-output",
|
||||
litellm_parent_otel_span=None,
|
||||
llm_router=_provider_router("xai"),
|
||||
)
|
||||
|
||||
assert (response.bytes, response.created_at, response.filename) == (
|
||||
2565,
|
||||
123,
|
||||
_XAI_BATCH_ID,
|
||||
)
|
||||
assert managed_file_table.rows["unified-output"].file_object == newer_file_object
|
||||
|
||||
|
||||
_NON_XAI_BATCH_OUTPUTS: Final = pytest.mark.parametrize(
|
||||
("provider", "provider_file_id"),
|
||||
[("openai", _OPENAI_FILE_ID), ("openai", _XAI_BATCH_ID), ("xai", "file_lit9385")],
|
||||
ids=["openai-file", "openai-batch-shaped-id", "xai-input-file"],
|
||||
)
|
||||
|
||||
|
||||
def _non_xai_batch_output_row(
|
||||
provider_file_id: str, *, size_bytes: int, fallback: bool
|
||||
) -> LiteLLM_ManagedFileTable:
|
||||
return _batch_output_row("xai", size_bytes=size_bytes, fallback=fallback).model_copy(
|
||||
update={"model_mappings": {"model-123": provider_file_id}, "flat_model_file_ids": [provider_file_id]}
|
||||
)
|
||||
|
||||
|
||||
def _provider_file_details_route(provider_file_id: str, *, size_bytes: int) -> respx.Route:
|
||||
return respx.get(url__regex=rf".*/files/{provider_file_id}$").respond(
|
||||
200,
|
||||
json={
|
||||
"id": provider_file_id,
|
||||
"object": "file",
|
||||
"bytes": size_bytes,
|
||||
"created_at": 456,
|
||||
"filename": "output.jsonl",
|
||||
"purpose": "batch_output",
|
||||
"status": "processed",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("_respx_sees_provider_calls")
|
||||
@_NON_XAI_BATCH_OUTPUTS
|
||||
@respx.mock
|
||||
async def test_store_batch_output_file_leaves_an_unsized_non_xai_row_alone(provider: str, provider_file_id: str):
|
||||
import litellm.proxy.proxy_server as proxy_server_module
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
row: Final = _non_xai_batch_output_row(provider_file_id, size_bytes=0, fallback=False)
|
||||
proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(row)
|
||||
provider_route: Final = _provider_file_details_route(provider_file_id, size_bytes=836)
|
||||
with patch.object(proxy_server_module, "llm_router", _provider_router(provider)):
|
||||
await proxy_managed_files.store_batch_output_file(
|
||||
unified_file_id="unified-output",
|
||||
provider_file_id=provider_file_id,
|
||||
model_id="model-123",
|
||||
owner=UserAPIKeyAuth(user_id="user-123"),
|
||||
litellm_parent_otel_span=None,
|
||||
size_bytes=2565,
|
||||
)
|
||||
|
||||
assert (provider_route.call_count, managed_file_table.update_many_calls, managed_file_table.upsert_calls) == (
|
||||
0,
|
||||
[],
|
||||
[],
|
||||
)
|
||||
assert managed_file_table.rows["unified-output"].file_object == row.file_object
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("_respx_sees_provider_calls")
|
||||
@_NON_XAI_BATCH_OUTPUTS
|
||||
@respx.mock
|
||||
async def test_store_batch_output_file_saves_non_xai_provider_details_unguarded_and_as_reported(
|
||||
provider: str, provider_file_id: str
|
||||
):
|
||||
import litellm.proxy.proxy_server as proxy_server_module
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(
|
||||
_non_xai_batch_output_row(provider_file_id, size_bytes=0, fallback=True)
|
||||
)
|
||||
_provider_file_details_route(provider_file_id, size_bytes=0)
|
||||
with patch.object(proxy_server_module, "llm_router", _provider_router(provider)):
|
||||
await proxy_managed_files.store_batch_output_file(
|
||||
unified_file_id="unified-output",
|
||||
provider_file_id=provider_file_id,
|
||||
model_id="model-123",
|
||||
owner=UserAPIKeyAuth(user_id="user-123"),
|
||||
litellm_parent_otel_span=None,
|
||||
size_bytes=2565,
|
||||
)
|
||||
|
||||
stored: Final = managed_file_table.rows["unified-output"].file_object
|
||||
assert stored is not None
|
||||
assert (stored.bytes, stored.created_at, stored.filename) == (0, 456, "output.jsonl")
|
||||
assert managed_file_table.update_many_calls == [
|
||||
({"unified_file_id": "unified-output"}, {"file_object": stored.model_dump_json()})
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("_respx_sees_provider_calls")
|
||||
@_NON_XAI_BATCH_OUTPUTS
|
||||
@respx.mock
|
||||
async def test_store_batch_output_file_registers_non_xai_provider_details_without_the_measured_size(
|
||||
provider: str, provider_file_id: str
|
||||
):
|
||||
import litellm.proxy.proxy_server as proxy_server_module
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma()
|
||||
_provider_file_details_route(provider_file_id, size_bytes=0)
|
||||
with patch.object(proxy_server_module, "llm_router", _provider_router(provider)):
|
||||
await proxy_managed_files.store_batch_output_file(
|
||||
unified_file_id="unified-output",
|
||||
provider_file_id=provider_file_id,
|
||||
model_id="model-123",
|
||||
owner=UserAPIKeyAuth(user_id="user-123"),
|
||||
litellm_parent_otel_span=None,
|
||||
size_bytes=2565,
|
||||
)
|
||||
|
||||
stored: Final = managed_file_table.rows["unified-output"].file_object
|
||||
assert stored is not None
|
||||
assert (stored.bytes, stored.created_at, stored.filename) == (0, 456, "output.jsonl")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("_respx_sees_provider_calls")
|
||||
@_NON_XAI_BATCH_OUTPUTS
|
||||
@respx.mock
|
||||
async def test_afile_retrieve_refresh_of_a_non_xai_row_answers_and_saves_the_provider_details_as_reported(
|
||||
provider: str, provider_file_id: str
|
||||
):
|
||||
proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(
|
||||
_non_xai_batch_output_row(provider_file_id, size_bytes=2565, fallback=True)
|
||||
)
|
||||
_provider_file_details_route(provider_file_id, size_bytes=0)
|
||||
|
||||
response: Final = await proxy_managed_files.afile_retrieve(
|
||||
file_id="unified-output",
|
||||
litellm_parent_otel_span=None,
|
||||
llm_router=_provider_router(provider),
|
||||
)
|
||||
next_response: Final = await proxy_managed_files.afile_retrieve(
|
||||
file_id="unified-output",
|
||||
litellm_parent_otel_span=None,
|
||||
llm_router=_provider_router(provider),
|
||||
)
|
||||
|
||||
stored: Final = managed_file_table.rows["unified-output"].file_object
|
||||
assert stored is not None
|
||||
assert (response.bytes, response.created_at, response.filename) == (0, 456, "output.jsonl")
|
||||
assert next_response.model_dump() == response.model_dump()
|
||||
assert (stored.bytes, stored.created_at) == (0, 456)
|
||||
assert managed_file_table.update_many_calls == [
|
||||
({"unified_file_id": "unified-output"}, {"file_object": stored.model_dump_json()})
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_batch_output_file_rejects_non_provider_model_prefix():
|
||||
import litellm.proxy.proxy_server as proxy_server_module
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import json
|
||||
import time
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
|
@ -278,6 +279,51 @@ async def test_file_content_of_a_batch_id_walks_every_results_page(sync_mode: bo
|
|||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@respx.mock
|
||||
async def test_file_retrieve_of_a_batch_id_describes_its_results_from_one_batch_read(sync_mode: bool) -> None:
|
||||
batch_id: Final = f"batch_retrieve_{sync_mode}"
|
||||
batch: Final = respx.get(f"{API_BASE}/v1/batches/{batch_id}").respond(200, json={**_XAI_BATCH, "batch_id": batch_id})
|
||||
results: Final = respx.get(f"{API_BASE}/v1/batches/{batch_id}/results").respond(
|
||||
200, json={"results": [], "pagination_token": None}
|
||||
)
|
||||
files_lookup: Final = respx.get(f"{API_BASE}/v1/files/{batch_id}").respond(404, json={"code": "not-found"})
|
||||
|
||||
kwargs: Final = {"file_id": batch_id, "custom_llm_provider": "xai", "api_key": KEY, "api_base": API_BASE}
|
||||
before: Final = int(time.time())
|
||||
file_object: Final = litellm.file_retrieve(**kwargs) if sync_mode else await litellm.afile_retrieve(**kwargs)
|
||||
after: Final = int(time.time())
|
||||
|
||||
assert (batch.call_count, results.call_count, files_lookup.call_count) == (1, 0, 0)
|
||||
assert (file_object.id, file_object.purpose, file_object.filename, file_object.bytes, file_object.status) == (
|
||||
batch_id,
|
||||
"batch_output",
|
||||
batch_id,
|
||||
0,
|
||||
"processed",
|
||||
)
|
||||
assert before <= file_object.created_at <= after
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@respx.mock
|
||||
async def test_file_retrieve_of_an_unfinished_batch_is_a_404_without_reading_results(sync_mode: bool) -> None:
|
||||
batch_id: Final = f"batch_unfinished_{sync_mode}"
|
||||
respx.get(f"{API_BASE}/v1/batches/{batch_id}").respond(
|
||||
200,
|
||||
json={**_XAI_BATCH, "batch_id": batch_id, "state": {**_XAI_BATCH["state"], "num_pending": 1}},
|
||||
)
|
||||
results: Final = respx.get(f"{API_BASE}/v1/batches/{batch_id}/results").respond(
|
||||
200, json={"results": [], "pagination_token": None}
|
||||
)
|
||||
|
||||
kwargs: Final = {"file_id": batch_id, "custom_llm_provider": "xai", "api_key": KEY, "api_base": API_BASE}
|
||||
with pytest.raises(XAIBatchesError, match="has no results file until it completes") as raised:
|
||||
litellm.file_retrieve(**kwargs) if sync_mode else await litellm.afile_retrieve(**kwargs)
|
||||
|
||||
assert (raised.value.status_code, results.call_count) == (404, 0)
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_file_content_unwraps_image_and_video_result_bodies() -> None:
|
||||
respx.get(f"{API_BASE}/v1/batches/batch_1/results").respond(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue