scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-05 07:37:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 233

refactor(redis): isolate managed file metadata store

Browse source
This commit is contained in:
Devin AI 2026-10-30 19:48:23 +00:00
parent f624f3402f
commit 069ded7b67
5 changed files with 845 additions and 938 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

354
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -4,19 +4,18 @@
import asyncio
import base64
import json
import math
import time
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timedelta
from itertools import chain
from types import MappingProxyType
from typing import (
TYPE_CHECKING,
Any,
Dict,
Final,
List,
Literal,
Optional,
Protocol,
TypedDict,
Union,
@ -36,6 +35,8 @@ from litellm import Router, verbose_logger
from litellm._internal_context import with_service_target
from litellm._uuid import uuid
from litellm.caching.caching import DualCache
from litellm.caching.redis_batch import active_request_redis_batch
from litellm.caching.redis_cache import RedisCache
from litellm.constants import (
BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS,
BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS,
@ -93,6 +94,8 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import (
request_tags_from_metadata,
)
from litellm.redis.stores.managed_files import ManagedFileStore, TransportManagedFileStore, file_key
from litellm.redis.transport.contract import RedisFailure
from litellm.repositories.managed_file_repository import ManagedFileRepository
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import ( # pyright: ignore[reportAttributeAccessIssue]
@ -158,7 +161,7 @@ def _decode_json_blob(blob: object) -> object:
return json.loads(blob) if isinstance(blob, str) else blob
def _parse_managed_batch_row(row: "PrismaManagedObjectRow") -> Optional[LiteLLMBatch]:
def _parse_managed_batch_row(row: "PrismaManagedObjectRow") -> LiteLLMBatch | None:
try:
batch_obj: Final = LiteLLMBatch.model_validate(_decode_json_blob(row.file_object))
except Exception as e:
@ -168,7 +171,7 @@ def _parse_managed_batch_row(row: "PrismaManagedObjectRow") -> Optional[LiteLLMB
return batch_obj
def _parse_managed_file_object(raw_file_object: object, unified_file_id: str) -> Optional[OpenAIFileObject]:
def _parse_managed_file_object(raw_file_object: object, unified_file_id: str) -> OpenAIFileObject | None:
if raw_file_object is None:
return None
try:
@ -311,16 +314,16 @@ class _ManagedFileRow(Protocol):
file_object: OpenAIFileObject
model_mappings: Mapping[str, str]
flat_model_file_ids: Sequence[str]
storage_backend: Optional[str]
storage_url: Optional[str]
created_by: Optional[str]
team_id: Optional[str]
storage_backend: str | None
storage_url: str | None
created_by: str | None
team_id: str | None
def model_dump(self) -> Mapping[str, object]: ...
class _ManagedFileTableActions(Protocol):
async def find_first(self, where: Mapping[str, object]) -> Optional[_ManagedFileRow]: ...
async def find_first(self, where: Mapping[str, object]) -> _ManagedFileRow | None: ...
async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ...
@ -328,24 +331,24 @@ class _ManagedFileTableActions(Protocol):
self,
where: Mapping[str, object],
take: int = ...,
order: Union[Mapping[str, str], Sequence[Mapping[str, str]]] = ...,
order: Mapping[str, str] | Sequence[Mapping[str, str]] = ...,
cursor: Mapping[str, str] = ...,
skip: int = ...,
) -> Sequence[_ManagedFileRow]: ...
async def upsert(self, where: Mapping[str, str], data: Mapping[str, Mapping[str, object]]) -> _ManagedFileRow: ...
async def delete(self, where: Mapping[str, str]) -> Optional[_ManagedFileRow]: ...
async def delete(self, where: Mapping[str, str]) -> _ManagedFileRow | None: ...
class _ManagedObjectTableActions(Protocol):
async def find_first(self, where: Mapping[str, object]) -> "Optional[PrismaManagedObjectRow]": ...
async def find_first(self, where: Mapping[str, object]) -> "PrismaManagedObjectRow | None": ...
async def find_many(
self,
where: Mapping[str, object],
take: int,
order: Union[Mapping[str, str], Sequence[Mapping[str, str]]],
order: Mapping[str, str] | Sequence[Mapping[str, str]],
cursor: Mapping[str, str] = ...,
skip: int = ...,
) -> "Sequence[PrismaManagedObjectRow]": ...
@ -517,11 +520,7 @@ def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, s
"Mapping[str, object]", getattr(file_object, "_hidden_params", None) or {}
)
return MappingProxyType(
{
key: value
for key in ("storage_backend", "storage_url")
if isinstance(value := hidden_params.get(key), str)
}
{key: value for key in ("storage_backend", "storage_url") if isinstance(value := hidden_params.get(key), str)}
)
@ -578,6 +577,19 @@ def _keyed_get(source: object, key: str, default: object = None) -> object:
return require_attributes(source, _SupportsKeyedGet, "get").get(key, default)
class _LocalMetadataCache(Protocol):
async def async_get_cache(self, key: str, *, local_only: bool) -> object: ...
async def async_set_cache(self, key: str, value: object, *, local_only: bool) -> None: ...
@dataclass(frozen=True, slots=True)
class _MetadataCaches:
local: _LocalMetadataCache
redis: RedisCache | None
in_memory_ttl: float | None
class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Class variables or attributes
def __init__(
@ -593,6 +605,56 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
self._sleep = sleep
self._clock = clock
def _metadata_cache(self) -> _MetadataCaches:
cache: Final = self.internal_usage_cache
dual: Final = cache if isinstance(cache, DualCache) else getattr(cache, "dual_cache", None)
if isinstance(dual, DualCache):
return _MetadataCaches(local=dual, redis=dual.redis_cache, in_memory_ttl=dual.default_in_memory_ttl)
return _MetadataCaches(
local=require_attributes(cache, _LocalMetadataCache, "async_get_cache", "async_set_cache"),
redis=None,
in_memory_ttl=None,
)
async def _cache_file(self, file_id: str, record: Mapping[str, object] | None) -> None:
cache: Final = self._metadata_cache()
local_cache: Final = cache.local
await local_cache.async_set_cache(file_key(file_id), record, local_only=True)
if cache.redis is None:
return
store: Final[ManagedFileStore] = TransportManagedFileStore(cache.redis.transport())
try:
seconds: Final = (
int(cache.in_memory_ttl)
if cache.in_memory_ttl is not None and not math.isnan(cache.in_memory_ttl)
else cache.redis.default_ttl
)
ttl: Final = timedelta(seconds=seconds)
except (OverflowError, ValueError):
return
if record is None:
await store.invalidate_file(file_id, ttl=ttl)
else:
await store.write_file(file_id, record, ttl=ttl)
async def _cached_file(self, file_id: str) -> object:
cache: Final = self._metadata_cache()
local_cache: Final = cache.local
local: Final = await local_cache.async_get_cache(file_key(file_id), local_only=True)
if local is not None or cache.redis is None:
return local
request_batch: Final = active_request_redis_batch(cache.redis)
if request_batch is not None and request_batch.read_as_missing(file_key(file_id)):
return None
store: Final[ManagedFileStore] = TransportManagedFileStore(cache.redis.transport())
result: Final = await store.read_file(file_id)
if result is None or isinstance(result, RedisFailure):
return None
value: Final = result.value
if value is not None:
await local_cache.async_set_cache(file_key(file_id), value, local_only=True)
return value
@staticmethod
def _get_prometheus_logger():
"""Find PrometheusLogger from litellm.callbacks, if registered."""
@ -604,9 +666,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
async def store_unified_file_id(
self,
file_id: str,
file_object: Optional[OpenAIFileObject],
litellm_parent_otel_span: Optional[Span],
model_mappings: Dict[str, str],
file_object: OpenAIFileObject | None,
litellm_parent_otel_span: Span | None,
model_mappings: dict[str, str],
user_api_key_dict: UserAPIKeyAuth,
) -> None:
verbose_logger.info(f"Storing LiteLLM Managed File object with id={file_id} in cache")
@ -623,11 +685,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
storage_backend=storage_metadata.get("storage_backend"),
storage_url=storage_metadata.get("storage_url"),
)
await self.internal_usage_cache.async_set_cache(
key=file_id,
value=litellm_managed_file_object.model_dump(),
litellm_parent_otel_span=litellm_parent_otel_span,
)
await self._cache_file(file_id, litellm_managed_file_object.model_dump())
## STORE MODEL MAPPINGS IN DB
@ -663,7 +721,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
verbose_logger.debug(f"LiteLLM Managed File object with id={file_id} stored in db: {result}")
async def _resolve_creator_org_id(self, user_api_key_dict: UserAPIKeyAuth) -> Optional[str]:
async def _resolve_creator_org_id(self, user_api_key_dict: UserAPIKeyAuth) -> str | None:
if user_api_key_dict.org_id:
return user_api_key_dict.org_id
if not user_api_key_dict.team_id:
@ -689,7 +747,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
self,
unified_object_id: str,
file_object: Union[LiteLLMBatch, LiteLLMFineTuningJob, "ResponsesAPIResponse"],
litellm_parent_otel_span: Optional[Span],
litellm_parent_otel_span: Span | None,
model_object_id: str,
file_purpose: Literal["batch", "fine-tune", "response"],
user_api_key_dict: UserAPIKeyAuth,
@ -817,16 +875,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
@with_service_target(_MANAGED_FILES_TARGET)
async def get_unified_file_id(
self, file_id: str, litellm_parent_otel_span: Optional[Span] = None
) -> Optional[LiteLLM_ManagedFileTable]:
self, file_id: str, litellm_parent_otel_span: Span | None = None
) -> LiteLLM_ManagedFileTable | None:
## CHECK CACHE
result = cast(
Optional[dict],
await self.internal_usage_cache.async_get_cache(
key=file_id,
litellm_parent_otel_span=litellm_parent_otel_span,
),
)
result: Final = await self._cached_file(file_id)
if result:
return LiteLLM_ManagedFileTable.model_validate(result)
@ -840,18 +892,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
@with_service_target(_MANAGED_FILES_TARGET)
async def delete_unified_file_id(
self, file_id: str, litellm_parent_otel_span: Optional[Span] = None
self, file_id: str, litellm_parent_otel_span: Span | None = None
) -> OpenAIFileObject:
## get old value
initial_value = await _managed_file_table(self.prisma_client).find_first(where={"unified_file_id": file_id})
if initial_value is None:
raise Exception(f"LiteLLM Managed File object with id={file_id} not found")
## delete old value
await self.internal_usage_cache.async_set_cache(
key=file_id,
value=None,
litellm_parent_otel_span=litellm_parent_otel_span,
)
await self._cache_file(file_id, None)
await _managed_file_table(self.prisma_client).delete(where={"unified_file_id": file_id})
return initial_value.file_object
@ -887,9 +935,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
detail=f"Object not found: {unified_object_id}",
)
async def enforce_batch_object_access(
self, object_id: str, user_api_key_dict: UserAPIKeyAuth
) -> None:
async def enforce_batch_object_access(self, object_id: str, user_api_key_dict: UserAPIKeyAuth) -> None:
"""Deny access to a provider-format batch id owned by another caller.
Ids with no ownership row (batches created before ownership tracking,
@ -913,9 +959,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
detail=f"User {user_api_key_dict.user_id} does not have access to the object {object_id}",
)
async def enforce_provider_file_access(
self, file_id: str, user_api_key_dict: UserAPIKeyAuth
) -> None:
async def enforce_provider_file_access(self, file_id: str, user_api_key_dict: UserAPIKeyAuth) -> None:
"""Deny access to a provider-format file id owned by another caller.
Ownership rows for provider-format ids are written when a managed
@ -940,7 +984,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
async def store_batch_output_file_ownership(
self, response: LiteLLMBatch, litellm_parent_otel_span: Optional[Span]
self, response: LiteLLMBatch, litellm_parent_otel_span: Span | None
) -> None:
"""Record ownership rows for a batch's provider-format output/error
file ids, inherited from the owning batch row (never the caller), so
@ -957,16 +1001,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
return
if self.prisma_client is None:
return
batch_row = await _managed_object_table(self.prisma_client).find_first(
where={"unified_object_id": response.id}
)
if batch_row is None or (
batch_row.created_by is None and batch_row.team_id is None
):
batch_row = await _managed_object_table(self.prisma_client).find_first(where={"unified_object_id": response.id})
if batch_row is None or (batch_row.created_by is None and batch_row.team_id is None):
return
owner_identity = UserAPIKeyAuth(
user_id=batch_row.created_by, team_id=batch_row.team_id
)
owner_identity = UserAPIKeyAuth(user_id=batch_row.created_by, team_id=batch_row.team_id)
for file_id in provider_file_ids:
model_name = decode_model_from_file_id(file_id)
raw_file_id = get_original_file_id(file_id)
@ -1084,8 +1122,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if raise_on_failure:
if isinstance(error, TimeoutError) and not str(error):
raise TimeoutError(
"Provider file retrieve timed out "
f"after {BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS} seconds"
f"Provider file retrieve timed out after {BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS} seconds"
) from error
raise
return None, True
@ -1098,11 +1135,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if not await ManagedFileRepository(self.prisma_client).update_file_object(stored.unified_file_id, file_object):
return
refreshed_row: Final = stored.model_copy(update={"file_object": file_object})
await self.internal_usage_cache.async_set_cache(
key=stored.unified_file_id,
value=refreshed_row.model_dump(),
litellm_parent_otel_span=None,
)
await self._cache_file(stored.unified_file_id, refreshed_row.model_dump())
async def _save_xai_refreshed_file_object(
self,
@ -1117,14 +1150,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
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
)
await self._cache_file(stored.unified_file_id, 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,
await self._cache_file(
stored.unified_file_id, stored.model_copy(update={"file_object": sized_file_object}).model_dump()
)
return sized_file_object
@ -1270,12 +1299,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
async def list_user_batches(
self,
user_api_key_dict: UserAPIKeyAuth,
limit: Optional[int] = None,
after: Optional[str] = None,
provider: Optional[str] = None,
target_model_names: Optional[str] = None,
llm_router: Optional[Router] = None,
) -> Dict[str, object]:
limit: int | None = None,
after: str | None = None,
provider: str | None = None,
target_model_names: str | None = None,
llm_router: Router | None = None,
) -> dict[str, object]:
# Provider filtering is not supported for managed batches
# This is because the encoded object ids stored in the managed objects table do not contain the provider information
# To support provider filtering, we would need to store the provider information in the encoded object ids
@ -1305,7 +1334,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if owner_filter is None:
return build_list_page([])
where_clause: Dict[str, object] = {"file_purpose": "batch", **owner_filter}
where_clause: dict[str, object] = {"file_purpose": "batch", **owner_filter}
if after:
cursor_row = await _managed_object_table(self.prisma_client).find_first(
@ -1329,7 +1358,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
async def _collect_listed_batches(
self,
where_clause: Mapping[str, object],
after: Optional[str],
after: str | None,
wanted: int,
user_api_key_dict: UserAPIKeyAuth,
) -> tuple[LiteLLMBatch, ...]:
@ -1341,7 +1370,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
and every chunk advances the keyset cursor, so the walk ends once the
caller's rows are exhausted."""
matches: tuple[LiteLLMBatch, ...] = () # rebind-ok: accumulates survivors across chunks
cursor_id: Optional[str] = after # rebind-ok: keyset cursor advances to each chunk's last row
cursor_id: str | None = after # rebind-ok: keyset cursor advances to each chunk's last row
chunk_size: int = wanted # rebind-ok: widens once a scan has to continue past the first chunk
while len(matches) < wanted:
cursor_args: _CursorPageArgs = {"cursor": {"unified_object_id": cursor_id}, "skip": 1} if cursor_id else {}
@ -1398,7 +1427,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
batch_obj: LiteLLMBatch,
unified_id_by_raw_id: Mapping[str, str],
user_api_key_dict: UserAPIKeyAuth,
) -> Optional[LiteLLMBatch]:
) -> LiteLLMBatch | None:
apply_unified_file_ids(batch_obj, unified_id_by_raw_id)
try:
await ensure_batch_response_managed_file_ids(
@ -1450,8 +1479,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
async def get_user_created_file_ids(
self, user_api_key_dict: UserAPIKeyAuth, model_object_ids: List[str]
) -> List[OpenAIFileObject]:
self, user_api_key_dict: UserAPIKeyAuth, model_object_ids: list[str]
) -> list[OpenAIFileObject]:
"""
Get all file ids the caller is allowed to see for a list of model
object ids. Service-account keys (no user_id) are scoped to their
@ -1476,8 +1505,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None
]
async def check_managed_file_id_access(self, data: Dict, user_api_key_dict: UserAPIKeyAuth) -> bool:
retrieve_file_id = cast(Optional[str], data.get("file_id"))
async def check_managed_file_id_access(self, data: dict, user_api_key_dict: UserAPIKeyAuth) -> bool:
retrieve_file_id = cast(str | None, data.get("file_id"))
potential_file_id = _is_base64_encoded_unified_file_id(retrieve_file_id) if retrieve_file_id else False
if potential_file_id and retrieve_file_id:
if await self.can_user_call_unified_file_id(retrieve_file_id, user_api_key_dict):
@ -1488,9 +1517,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
detail=f"User {user_api_key_dict.user_id} does not have access to the file {retrieve_file_id}",
)
if retrieve_file_id:
await self.enforce_provider_file_access(
retrieve_file_id, user_api_key_dict
)
await self.enforce_provider_file_access(retrieve_file_id, user_api_key_dict)
return False
async def check_file_ids_access(self, file_ids: Sequence[str], user_api_key_dict: UserAPIKeyAuth) -> None:
@ -1518,9 +1545,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: Dict,
data: dict,
call_type: CallTypesLiteral,
) -> Union[Exception, str, Dict, None]:
) -> Exception | str | dict | None:
"""
- Detect litellm_proxy/ file_id
- add dictionary of mappings of litellm_proxy/ file_id -> provider_file_id => {litellm_proxy/file_id: {"model_id": id, "file_id": provider_file_id}}
@ -1593,7 +1620,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if _is_streamed_background_create(create_request):
return _with_managed_background_marker(create_request)
elif call_type == CallTypes.afile_content.value:
retrieve_file_id = cast(Optional[str], data.get("file_id"))
retrieve_file_id = cast(str | None, data.get("file_id"))
potential_file_id = _is_base64_encoded_unified_file_id(retrieve_file_id) if retrieve_file_id else False
if potential_file_id and "llm_output_file_id," in potential_file_id:
model_id = self.get_model_id_from_unified_file_id(potential_file_id)
@ -1601,7 +1628,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
data["model"] = model_id
data["file_id"] = self.get_output_file_id_from_unified_file_id(potential_file_id)
elif call_type == CallTypes.acreate_batch.value:
input_file_id = cast(Optional[str], data.get("input_file_id"))
input_file_id = cast(str | None, data.get("input_file_id"))
if input_file_id:
await self.check_file_ids_access((input_file_id,), user_api_key_dict)
model_file_id_mapping = await self.get_model_file_id_mapping(
@ -1615,8 +1642,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
or call_type == CallTypes.acancel_fine_tuning_job.value
or call_type == CallTypes.aretrieve_fine_tuning_job.value
):
accessor_key: Optional[str] = None
retrieve_object_id: Optional[str] = None
accessor_key: str | None = None
retrieve_object_id: str | None = None
if call_type == CallTypes.aretrieve_batch.value or call_type == CallTypes.acancel_batch.value:
accessor_key = "batch_id"
elif (
@ -1626,7 +1653,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
accessor_key = "fine_tuning_job_id"
if accessor_key:
retrieve_object_id = cast(Optional[str], data.get(accessor_key))
retrieve_object_id = cast(str | None, data.get(accessor_key))
potential_llm_object_id = (
_is_base64_encoded_unified_file_id(retrieve_object_id) if retrieve_object_id else False
@ -1670,11 +1697,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
async def async_filter_deployments(
self,
model: str,
healthy_deployments: List,
messages: Optional[List[AllMessageValues]],
request_kwargs: Optional[Dict] = None,
parent_otel_span: Optional[Span] = None,
) -> List[Dict]:
healthy_deployments: list,
messages: list[AllMessageValues] | None,
request_kwargs: dict | None = None,
parent_otel_span: Span | None = None,
) -> list[dict]:
if request_kwargs is None:
return healthy_deployments
@ -1701,8 +1728,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
return holders if holders or "input_file_id" in request else healthy_deployments
async def async_pre_call_deployment_hook(
self, kwargs: dict[str, object], call_type: Optional[CallTypes]
) -> Optional[dict]:
self, kwargs: dict[str, object], call_type: CallTypes | None
) -> dict | None:
"""
Allow modifying the request just before it's sent to the deployment.
"""
@ -1710,18 +1737,18 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if accessor_keys is None:
return kwargs
model_file_id_mapping = cast(Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping"))
model_file_id_mapping = cast(dict[str, dict[str, str]] | None, kwargs.get("model_file_id_mapping"))
# model_info may be at top-level or nested under litellm_metadata
# (batch/file operations use litellm_metadata)
model_id = cast(Optional[str], _keyed_get(kwargs.get("model_info", {}), "id"))
model_id = cast(str | None, _keyed_get(kwargs.get("model_info", {}), "id"))
if model_id is None:
model_id = cast(
Optional[str],
str | None,
_keyed_get(_keyed_get(kwargs.get("litellm_metadata", {}), "model_info", {}), "id"),
)
for accessor_key in accessor_keys:
input_file_id = cast(Optional[str], kwargs.get(accessor_key))
mapped_file_id: Optional[str] = None
input_file_id = cast(str | None, kwargs.get(accessor_key))
mapped_file_id: str | None = None
if input_file_id and model_file_id_mapping and model_id:
mapped_file_id = model_file_id_mapping.get(input_file_id, {}).get(model_id, None)
if mapped_file_id:
@ -1729,7 +1756,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
return kwargs
def get_file_ids_from_messages(self, messages: List[AllMessageValues]) -> List[str]:
def get_file_ids_from_messages(self, messages: list[AllMessageValues]) -> list[str]:
"""
Gets file ids from messages
"""
@ -1749,7 +1776,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
file_ids.append(file_id)
return file_ids
def get_file_ids_from_responses_input(self, input: Union[str, List[Dict[str, object]]]) -> List[str]:
def get_file_ids_from_responses_input(self, input: str | list[dict[str, object]]) -> list[str]:
"""
Gets file ids from responses API input.
@ -1759,7 +1786,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
- type: "input_file" with file_id
- content: a list that can contain items with type: "input_file" and file_id
"""
file_ids: List[str] = []
file_ids: list[str] = []
if isinstance(input, str):
return file_ids
@ -1788,7 +1815,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
return file_ids
def get_file_ids_from_responses_tools(self, tools: List[Dict[str, object]]) -> List[str]:
def get_file_ids_from_responses_tools(self, tools: list[dict[str, object]]) -> list[str]:
"""
Gets file ids from responses API tools parameter.
@ -1800,7 +1827,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
}
]
"""
file_ids: List[str] = []
file_ids: list[str] = []
if not isinstance(tools, list):
return file_ids
@ -1821,7 +1848,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
return file_ids
def get_vector_store_ids_from_file_search_tools(self, tools: List[Dict[str, object]]) -> List[str]:
def get_vector_store_ids_from_file_search_tools(self, tools: list[dict[str, object]]) -> list[str]:
"""
Extract unified vector_store_ids from file_search tools.
@ -1832,7 +1859,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
is_base64_encoded_unified_id,
)
vs_ids: List[str] = []
vs_ids: list[str] = []
if not isinstance(tools, list):
return vs_ids
@ -1850,7 +1877,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
async def check_vector_store_ids_access(
self,
vector_store_ids: List[str],
vector_store_ids: list[str],
user_api_key_dict: UserAPIKeyAuth,
) -> None:
"""
@ -1876,7 +1903,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
return
# Map each unified ID to its internal UUID for a single batch DB fetch
uuid_to_unified: Dict[str, str] = {}
uuid_to_unified: dict[str, str] = {}
for vs_id in vector_store_ids:
uuid = extract_unified_uuid_from_unified_id(vs_id)
if uuid:
@ -1931,7 +1958,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
}
"""
file_id_mapping: Dict[str, Dict[str, str]] = {}
file_id_mapping: dict[str, dict[str, str]] = {}
litellm_managed_file_ids = []
for file_id in file_ids:
@ -1953,11 +1980,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
async def create_file_for_each_model(
self,
llm_router: Optional[Router],
llm_router: Router | None,
_create_file_request: CreateFileRequest,
target_model_names_list: List[str],
target_model_names_list: list[str],
litellm_parent_otel_span: Span | None,
) -> List[OpenAIFileObject]:
) -> list[OpenAIFileObject]:
if llm_router is None:
raise Exception("LLM Router not initialized. Ensure models added to proxy.")
responses = []
@ -1971,7 +1998,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
self,
create_file_request: CreateFileRequest,
llm_router: Router,
target_model_names_list: List[str],
target_model_names_list: list[str],
litellm_parent_otel_span: Span | None,
user_api_key_dict: UserAPIKeyAuth,
) -> OpenAIFileObject:
@ -1990,7 +2017,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
## STORE MODEL MAPPINGS IN DB
model_mappings: Dict[str, str] = {}
model_mappings: dict[str, str] = {}
for file_object in responses:
file_hidden_params = cast( # cast-ok: preserve mapping operations on dynamic file metadata
@ -2036,11 +2063,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
@staticmethod
async def return_unified_file_id(
file_objects: List[OpenAIFileObject],
file_objects: list[OpenAIFileObject],
create_file_request: CreateFileRequest,
internal_usage_cache: InternalUsageCache,
litellm_parent_otel_span: Span | None,
target_model_names_list: List[str],
target_model_names_list: list[str],
) -> OpenAIFileObject:
## GET THE FILE TYPE FROM THE CREATE FILE REQUEST
_, file_type = extract_file_metadata(create_file_request["file"])
@ -2085,7 +2112,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
unified_batch_id = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format(model_id, batch_id)
return base64.urlsafe_b64encode(unified_batch_id.encode()).decode().rstrip("=")
def get_unified_output_file_id(self, output_file_id: str, model_id: str, model_name: Optional[str]) -> str:
def get_unified_output_file_id(self, output_file_id: str, model_id: str, model_name: str | None) -> str:
deterministic_uuid: Final = uuid5(uuid5(NAMESPACE_URL, model_id), output_file_id)
unified_output_file_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format(
"application/json",
@ -2106,7 +2133,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
return file_id.split(marker, 1)[1].split(";")[0]
async def async_post_call_success_hook(
self, data: Dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes
self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes
) -> LLMResponseTypes:
if isinstance(response, LiteLLMBatch):
decoded_batch_id: Final = _is_base64_encoded_unified_file_id(response.id)
@ -2117,8 +2144,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
unified_file_id = response_hidden_params.get("unified_file_id")
unified_batch_id = response_hidden_params.get("unified_batch_id")
is_batch_create: Final = response_hidden_params.get(BATCH_CREATE_HIDDEN_PARAM) is True
model_id = cast(Optional[str], response_hidden_params.get("model_id"))
model_name = cast(Optional[str], response_hidden_params.get("model_name"))
model_id = cast(str | None, response_hidden_params.get("model_id"))
model_name = cast(str | None, response_hidden_params.get("model_name"))
resolved_model_name = resolve_managed_output_file_model_name(
unified_input_file_id=unified_file_id if isinstance(unified_file_id, str) else response.input_file_id,
@ -2206,7 +2233,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
unified_file_id = finetuning_response_hidden_params.get("unified_file_id")
unified_finetuning_job_id = finetuning_response_hidden_params.get("unified_finetuning_job_id")
model_id = cast(Optional[str], finetuning_response_hidden_params.get("model_id"))
model_id = cast(str | None, finetuning_response_hidden_params.get("model_id"))
original_response_id = response.id
if (unified_file_id or unified_finetuning_job_id) and model_id:
response.id = self.get_unified_generic_response_id(model_id=model_id, generic_response_id=response.id)
@ -2224,15 +2251,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"""
## check if file object
if hasattr(response, "data") and isinstance(response.data, list):
if all(isinstance(file_object, FileObject) for file_object in response.data):
## Get all file id's
## Check which file id's were created by the user
## Filter the response to only include the files created by the user
## Return the filtered response
file_ids = [
file_object.id
for file_object in cast(List[FileObject], response.data) # type: ignore
]
page_items: Final = TypeAdapter(tuple[object, ...]).validate_python(response.data)
file_objects: Final = tuple(item for item in page_items if isinstance(item, FileObject))
if len(file_objects) == len(page_items):
file_ids = [file_object.id for file_object in file_objects]
user_created_file_ids = await self.get_user_created_file_ids(user_api_key_dict, file_ids)
## Filter the response to only include the files created by the user
response.data = user_created_file_ids # type: ignore
@ -2242,7 +2264,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
return response
@staticmethod
def _scope_list_page_cursors(response: AsyncCursorPage, data: List[OpenAIFileObject]) -> None:
def _scope_list_page_cursors(response: AsyncCursorPage, data: list[OpenAIFileObject]) -> None:
"""Rebuild ``first_id`` / ``last_id`` from the caller-scoped page.
The upstream cursors point at rows that were just filtered out, so
@ -2258,7 +2280,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
response.has_more = False
async def afile_retrieve(
self, file_id: str, litellm_parent_otel_span: Optional[Span], llm_router: Optional[Router] = None
self, file_id: str, litellm_parent_otel_span: Span | None, llm_router: Router | None = None
) -> OpenAIFileObject:
"""Return public details for a managed file ID, refreshing a basic entry when possible."""
stored_file_object = await self.get_unified_file_id(file_id, litellm_parent_otel_span)
@ -2325,7 +2347,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
allow_default_provider=True,
)
except Exception as e:
raise Exception(f"Failed to retrieve file {file_id} from provider: {str(e)}") from e
raise Exception(f"Failed to retrieve file {file_id} from provider: {e!s}") from e
if not fetch_route_available:
raise Exception(
f"LiteLLM Managed File object with id={file_id} has no file_object and no provider route to fetch it"
@ -2336,14 +2358,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
async def afile_list(
self,
purpose: Optional[str],
litellm_parent_otel_span: Optional[Span],
purpose: str | None,
litellm_parent_otel_span: Span | None,
user_api_key_dict: UserAPIKeyAuth,
limit: Optional[int] = None,
after: Optional[str] = None,
model_scope: Optional[FileListModelScope] = None,
before: Optional[str] = None,
**data: Dict,
limit: int | None = None,
after: str | None = None,
model_scope: FileListModelScope | None = None,
before: str | None = None,
**data: dict,
) -> FileListPage:
"""List the managed files the caller owns, newest first.
@ -2404,7 +2426,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
return FileListPage.model_validate(build_list_page([]))
page_size: Final = min(limit or MAX_FILE_LIST_LIMIT, MAX_FILE_LIST_LIMIT)
matches: Final[List[OpenAIFileObject]] = []
matches: Final[list[OpenAIFileObject]] = []
cursor_id = cursor
chunk_size = page_size + 1
direction: Final[Literal["asc", "desc"]] = "asc" if before else "desc"
@ -2471,7 +2493,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
verbose_logger.warning(f"Error checking batch polling configuration: {e}. Assuming disabled.")
return False
async def _get_batches_referencing_file(self, file_id: str) -> List[Dict[str, object]]:
async def _get_batches_referencing_file(self, file_id: str) -> list[dict[str, object]]:
"""
Find batches that reference this file and still need cost tracking.
Find batches that are in non-terminal state and have not yet been processed by CheckBatchCost.
@ -2607,9 +2629,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
async def afile_delete(
self,
file_id: str,
litellm_parent_otel_span: Optional[Span],
litellm_parent_otel_span: Span | None,
llm_router: Router,
**data: Dict,
**data: dict,
) -> FileDeleted:
# Check if file deletion should be blocked due to batch references
@ -2668,9 +2690,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
async def afile_content(
self,
file_id: str,
litellm_parent_otel_span: Optional[Span],
litellm_parent_otel_span: Span | None,
llm_router: Router,
**data: Dict,
**data: dict,
) -> HttpxBinaryResponseContent:
"""
Get the content of a file from first model that has it
@ -2696,7 +2718,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# request params.
credentials = llm_router.registry().credentials(model_id, form="with_provider")
if credentials is not None:
data["_litellm_internal_model_credentials"] = cast(Dict, MappingProxyType(dict(credentials)))
data["_litellm_internal_model_credentials"] = cast(dict, MappingProxyType(dict(credentials)))
else:
data.pop("_litellm_internal_model_credentials", None)
return cast(
@ -2722,9 +2744,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
async def _convert_storage_files_to_base64(
self,
messages: List[AllMessageValues],
file_ids: List[str],
litellm_parent_otel_span: Optional[Span],
messages: list[AllMessageValues],
file_ids: list[str],
litellm_parent_otel_span: Span | None,
) -> None:
"""
Convert files stored in storage backends to base64 format for Vertex AI/Gemini.
@ -2757,9 +2779,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
try:
storage_backend = get_storage_backend(storage_backend_name, prisma_client=self.prisma_client)
except ValueError as e:
verbose_logger.warning(
f"Storage backend '{storage_backend_name}' error for file {file_id}: {str(e)}"
)
verbose_logger.warning(f"Storage backend '{storage_backend_name}' error for file {file_id}: {e!s}")
continue
file_content = await storage_backend.download_file(storage_url)
@ -2774,7 +2794,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Update messages to use base64 instead of file_id
self._update_messages_with_base64_data(messages, file_id, base64_data_uri, content_type)
except Exception as e:
verbose_logger.exception(f"Error converting file {file_id} from storage backend to base64: {str(e)}")
verbose_logger.exception(f"Error converting file {file_id} from storage backend to base64: {e!s}")
# Continue with other files even if one fails
continue
@ -2800,7 +2820,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
def _update_messages_with_base64_data(
self,
messages: List[AllMessageValues],
messages: list[AllMessageValues],
file_id: str,
base64_data_uri: str,
content_type: str,
@ -2834,4 +2854,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
verbose_logger.debug(
f"Converted file {file_id} from storage backend to base64 with format {content_type}"
)
PROXY_LiteLLMManagedFiles = _PROXY_LiteLLMManagedFiles

65
litellm/redis/stores/managed_files.py Normal file
View file

@ -0,0 +1,65 @@
import ast
import json
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import timedelta
from typing import Final, Protocol
from pydantic import TypeAdapter
from litellm.redis.transport.contract import RedisFailure, RedisTransport
@dataclass(frozen=True, slots=True)
class ManagedFileValue:
value: object
_VALUE: Final[TypeAdapter[ManagedFileValue]] = TypeAdapter(ManagedFileValue)
def file_key(unified_file_id: str) -> str:
return unified_file_id
class ManagedFileStore(Protocol):
async def read_file(self, unified_file_id: str) -> ManagedFileValue | None | RedisFailure: ...
async def write_file(
self, unified_file_id: str, record: Mapping[str, object], *, ttl: timedelta | None
) -> None | RedisFailure: ...
async def invalidate_file(self, unified_file_id: str, *, ttl: timedelta | None) -> None | RedisFailure: ...
@dataclass(frozen=True, slots=True)
class TransportManagedFileStore:
transport: RedisTransport
async def read_file(self, unified_file_id: str) -> ManagedFileValue | None | RedisFailure:
raw: Final = await self.transport.get(file_key(unified_file_id))
if raw is None or isinstance(raw, RedisFailure):
return raw
try:
return _VALUE.validate_python({"value": json.loads(raw.decode("utf-8"))})
except (ValueError, UnicodeError, RecursionError):
try:
return _VALUE.validate_python({"value": ast.literal_eval(raw.decode("utf-8"))})
except (ValueError, SyntaxError, UnicodeError, RecursionError) as exc:
return RedisFailure(kind="decode", message=str(exc))
async def write_file(
self, unified_file_id: str, record: Mapping[str, object], *, ttl: timedelta | None
) -> None | RedisFailure:
try:
encoded: Final = json.dumps(record)
except (TypeError, ValueError) as exc:
return RedisFailure(kind="decode", message=str(exc))
return await self._write(unified_file_id, encoded, ttl)
async def invalidate_file(self, unified_file_id: str, *, ttl: timedelta | None) -> None | RedisFailure:
return await self._write(unified_file_id, "null", ttl)
async def _write(self, unified_file_id: str, encoded: str, ttl: timedelta | None) -> None | RedisFailure:
result: Final = await self.transport.set_seconds(file_key(unified_file_id), encoded, ttl=ttl)
return result if isinstance(result, RedisFailure) else None

903
tests/unit/enterprise/proxy/hooks/test_managed_files.py
View file

File diff suppressed because it is too large Load diff

197
tests/unit/enterprise/proxy/test_managed_files_hook.py
View file

@ -10,14 +10,15 @@ import base64
import json
import logging
from types import MappingProxyType
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from typing import Final, Optional
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.llms.base_llm.files.transformation import FileListModelScope
from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth
from litellm.proxy.openai_files_endpoints.common_utils import BATCH_CREATE_HIDDEN_PARAM
from litellm.proxy.utils import InternalUsageCache
from litellm.types.llms.openai import FileListPage, OpenAIFileObject
from litellm.types.utils import LiteLLMBatch
@ -36,7 +37,7 @@ def _make_file_object(file_id: str = "file-output-abc") -> OpenAIFileObject:
def _make_batch_response(
batch_id: str = "batch-123",
output_file_id: Optional[str] = "file-output-abc",
output_file_id: str | None = "file-output-abc",
status: str = "completed",
model_id: str = "model-deploy-xyz",
model_name: str = "azure/gpt-4",
@ -100,11 +101,9 @@ def _make_managed_file_row(
unified_file_id: str,
purpose: str = "batch_output",
created_by: str = "test-user",
team_id: Optional[str] = None,
team_id: str | None = None,
) -> MagicMock:
file_object = _make_file_object(f"file-provider-{unified_file_id}").model_copy(
update={"purpose": purpose}
)
file_object = _make_file_object(f"file-provider-{unified_file_id}").model_copy(update={"purpose": purpose})
return MagicMock(
unified_file_id=unified_file_id,
file_object=file_object.model_dump(),
@ -116,7 +115,7 @@ def _make_managed_file_row(
def _make_unparseable_managed_file_row(
unified_file_id: str,
created_by: str = "test-user",
team_id: Optional[str] = None,
team_id: str | None = None,
) -> MagicMock:
"""A row whose stored blob cannot be parsed back into a file object."""
return MagicMock(
@ -156,16 +155,10 @@ class _FakeManagedFileTable:
return next(iter(self._owned_rows(where)), None)
async def find_many(self, where, take=None, order=None, cursor=None, skip=0):
self.find_many_calls.append(
{"where": where, "take": take, "order": order, "cursor": cursor, "skip": skip}
)
self.find_many_calls.append({"where": where, "take": take, "order": order, "cursor": cursor, "skip": skip})
rows = self._owned_rows(where)
if cursor is not None:
start = next(
index
for index, row in enumerate(rows)
if row.unified_file_id == cursor["unified_file_id"]
)
start = next(index for index, row in enumerate(rows) if row.unified_file_id == cursor["unified_file_id"])
rows = rows[start + skip :]
return rows if take is None else rows[:take]
@ -183,13 +176,14 @@ def _make_managed_files_instance():
PROXY_LiteLLMManagedFiles,
)
mock_cache = MagicMock()
mock_cache.async_get_cache = AsyncMock(return_value=None)
from litellm.caching.dual_cache import DualCache
cache: Final = DualCache()
mock_prisma = MagicMock()
mock_prisma.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None)
instance = PROXY_LiteLLMManagedFiles(
internal_usage_cache=mock_cache,
internal_usage_cache=cache,
prisma_client=mock_prisma,
)
instance.store_unified_file_id = AsyncMock()
@ -222,9 +216,7 @@ async def test_should_pass_credentials_to_afile_retrieve():
}
mock_router = MagicMock()
mock_router.registry.return_value.credentials = MagicMock(
return_value=mock_credentials
)
mock_router.registry.return_value.credentials = MagicMock(return_value=mock_credentials)
mock_afile_retrieve = AsyncMock(return_value=_make_file_object("file-output-abc"))
@ -242,14 +234,10 @@ async def test_should_pass_credentials_to_afile_retrieve():
call_kwargs = mock_afile_retrieve.call_args
assert call_kwargs.kwargs.get("api_key") == "test-azure-key", (
f"afile_retrieve must receive api_key from router credentials. "
f"Got kwargs: {call_kwargs.kwargs}"
f"afile_retrieve must receive api_key from router credentials. Got kwargs: {call_kwargs.kwargs}"
)
assert (
call_kwargs.kwargs.get("api_base") == "https://my-azure.openai.azure.com/"
), (
f"afile_retrieve must receive api_base from router credentials. "
f"Got kwargs: {call_kwargs.kwargs}"
assert call_kwargs.kwargs.get("api_base") == "https://my-azure.openai.azure.com/", (
f"afile_retrieve must receive api_base from router credentials. Got kwargs: {call_kwargs.kwargs}"
)
@ -266,9 +254,7 @@ async def test_get_user_created_file_ids_skips_rows_without_file_object():
]
)
files = await managed_files.get_user_created_file_ids(
_make_user_api_key_dict(), ["file-output-abc"]
)
files = await managed_files.get_user_created_file_ids(_make_user_api_key_dict(), ["file-output-abc"])
assert [file.id for file in files] == ["unified-id-1"]
@ -294,9 +280,7 @@ async def test_get_user_created_file_ids_remaps_stored_raw_provider_id_to_unifie
]
)
files = await managed_files.get_user_created_file_ids(
_make_user_api_key_dict(), ["file-raw-provider-123"]
)
files = await managed_files.get_user_created_file_ids(_make_user_api_key_dict(), ["file-raw-provider-123"])
assert [file.id for file in files] == [unified_id]
assert files[0].filename == raw_provider_object.filename
@ -304,9 +288,7 @@ async def test_get_user_created_file_ids_remaps_stored_raw_provider_id_to_unifie
@pytest.mark.asyncio
async def test_provider_file_id_resolver_returns_owned_mappings_with_owner_scoped_filter() -> (
None
):
async def test_provider_file_id_resolver_returns_owned_mappings_with_owner_scoped_filter() -> None:
managed_files: Final = _make_managed_files_instance()
managed_row: Final = MagicMock(
unified_file_id="unified-file-id",
@ -315,16 +297,14 @@ async def test_provider_file_id_resolver_returns_owned_mappings_with_owner_scope
find_many: Final = AsyncMock(return_value=[managed_row])
managed_files.prisma_client.db.litellm_managedfiletable.find_many = find_many
unified_file_ids: Final = (
await managed_files.get_unified_file_ids_for_provider_file_ids(
provider_file_ids=(
"file-provider-1",
"file-provider-2",
"file-unmanaged-2",
"file-provider-1",
),
user_api_key_dict=_make_team_member_api_key_dict(),
)
unified_file_ids: Final = await managed_files.get_unified_file_ids_for_provider_file_ids(
provider_file_ids=(
"file-provider-1",
"file-provider-2",
"file-unmanaged-2",
"file-provider-1",
),
user_api_key_dict=_make_team_member_api_key_dict(),
)
assert unified_file_ids == {
@ -343,9 +323,7 @@ async def test_provider_file_id_resolver_returns_owned_mappings_with_owner_scope
@pytest.mark.asyncio
async def test_provider_file_id_resolver_denies_unowned_callers_without_database_query() -> (
None
):
async def test_provider_file_id_resolver_denies_unowned_callers_without_database_query() -> None:
managed_files: Final = _make_managed_files_instance()
find_many: Final = AsyncMock()
managed_files.prisma_client.db.litellm_managedfiletable.find_many = find_many
@ -357,11 +335,9 @@ async def test_provider_file_id_resolver_denies_unowned_callers_without_database
parent_otel_span=None,
)
unified_file_ids: Final = (
await managed_files.get_unified_file_ids_for_provider_file_ids(
provider_file_ids=("file-provider-1",),
user_api_key_dict=no_owner,
)
unified_file_ids: Final = await managed_files.get_unified_file_ids_for_provider_file_ids(
provider_file_ids=("file-provider-1",),
user_api_key_dict=no_owner,
)
assert unified_file_ids == {}
@ -375,11 +351,9 @@ async def test_provider_file_id_resolver_skips_database_query_for_empty_input()
find_many: Final = AsyncMock()
managed_files.prisma_client.db.litellm_managedfiletable.find_many = find_many
unified_file_ids: Final = (
await managed_files.get_unified_file_ids_for_provider_file_ids(
provider_file_ids=(),
user_api_key_dict=_make_user_api_key_dict(),
)
unified_file_ids: Final = await managed_files.get_unified_file_ids_for_provider_file_ids(
provider_file_ids=(),
user_api_key_dict=_make_user_api_key_dict(),
)
assert unified_file_ids == {}
@ -397,9 +371,9 @@ async def test_afile_list_returns_owner_scoped_managed_files():
unified_file_id="unified-file-id",
),
MagicMock(
file_object=_make_file_object("file-other-purpose").model_copy(
update={"purpose": "batch"}
).model_dump(),
file_object=_make_file_object("file-other-purpose")
.model_copy(update={"purpose": "batch"})
.model_dump(),
unified_file_id="unified-other-purpose",
),
]
@ -519,9 +493,7 @@ async def test_afile_list_returns_own_and_team_files_for_a_key_carrying_both_ids
)
assert [file.id for file in response.data] == ["unified-mine", "unified-teammates"]
assert table.find_many_calls[0]["where"] == {
"OR": [{"created_by": "test-user"}, {"team_id": "test-team"}]
}
assert table.find_many_calls[0]["where"] == {"OR": [{"created_by": "test-user"}, {"team_id": "test-team"}]}
@pytest.mark.asyncio
@ -848,9 +820,7 @@ async def test_afile_list_reads_one_chunk_when_the_first_one_fills_the_page():
@pytest.mark.asyncio
async def test_afile_list_reports_no_more_pages_when_nothing_matches():
managed_files, _ = _make_managed_files_over_rows(
[_make_managed_file_row(f"unified-{index}") for index in range(5)]
)
managed_files, _ = _make_managed_files_over_rows([_make_managed_file_row(f"unified-{index}") for index in range(5)])
page = await managed_files.afile_list(
purpose="batch",
@ -1029,9 +999,7 @@ async def test_get_user_created_file_ids_skips_unparseable_rows():
]
)
files = await managed_files.get_user_created_file_ids(
_make_user_api_key_dict(), ["file-output-abc"]
)
files = await managed_files.get_user_created_file_ids(_make_user_api_key_dict(), ["file-output-abc"])
assert [file.id for file in files] == ["unified-valid"]
@ -1096,9 +1064,7 @@ async def test_should_not_double_wrap_already_unified_output_file_id():
"custom_llm_provider": "openai",
}
mock_router = MagicMock()
mock_router.registry.return_value.credentials = MagicMock(
return_value=mock_credentials
)
mock_router.registry.return_value.credentials = MagicMock(return_value=mock_credentials)
mock_afile_retrieve = AsyncMock(return_value=_make_file_object(provider_file_id))
with (
@ -1115,13 +1081,9 @@ async def test_should_not_double_wrap_already_unified_output_file_id():
mock_afile_retrieve.assert_called_once()
assert mock_afile_retrieve.call_args.kwargs["file_id"] == provider_file_id
managed_files.store_unified_file_id.assert_awaited_once()
assert managed_files.store_unified_file_id.await_args.kwargs["model_mappings"] == {
model_id: provider_file_id
}
assert managed_files.store_unified_file_id.await_args.kwargs["model_mappings"] == {model_id: provider_file_id}
decoded = base64.urlsafe_b64decode(
already_unified + "=" * (-len(already_unified) % 4)
).decode()
decoded = base64.urlsafe_b64decode(already_unified + "=" * (-len(already_unified) % 4)).decode()
assert decoded.count(f"llm_output_file_id,{provider_file_id}") == 1
@ -1132,9 +1094,7 @@ async def test_should_skip_non_file_unified_id_on_output_file_id():
managed_files = _make_managed_files_instance()
batch_unified = (
base64.urlsafe_b64encode(
b"litellm_proxy;model_id:openai/openai/gpt-5.5-batch;llm_batch_id:batch_abc"
)
base64.urlsafe_b64encode(b"litellm_proxy;model_id:openai/openai/gpt-5.5-batch;llm_batch_id:batch_abc")
.decode()
.rstrip("=")
)
@ -1172,9 +1132,7 @@ async def test_afile_content_passes_trusted_model_credentials_to_router():
unified_file_id = "unified-file-id"
s3_uri = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out"
managed_files.get_unified_file_id = AsyncMock(return_value=None)
managed_files.get_model_file_id_mapping = AsyncMock(
return_value={unified_file_id: {"model-123": s3_uri}}
)
managed_files.get_model_file_id_mapping = AsyncMock(return_value={unified_file_id: {"model-123": s3_uri}})
mock_router = MagicMock()
mock_router.registry.return_value.credentials = MagicMock(
@ -1210,9 +1168,10 @@ def _managed_deletion_file_id(provider_file_id):
def _managed_files_with_deletion_row(unified_file_id, provider_file_id, file_object):
from litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles
from litellm.caching import DualCache
from litellm.models.managed_files import LiteLLM_ManagedFileTable
from litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles
row = LiteLLM_ManagedFileTable(
unified_file_id=unified_file_id,
@ -1225,7 +1184,7 @@ def _managed_files_with_deletion_row(unified_file_id, provider_file_id, file_obj
delete=AsyncMock(),
)
return PROXY_LiteLLMManagedFiles(
internal_usage_cache=DualCache(),
internal_usage_cache=InternalUsageCache(DualCache()),
prisma_client=MagicMock(db=MagicMock(litellm_managedfiletable=table)),
), table
@ -1260,9 +1219,9 @@ async def test_afile_delete_bedrock_uses_deployment_bucket_and_signed_s3_delete(
unified_file_id = _managed_deletion_file_id(s3_uri)
managed_files, table = _managed_files_with_deletion_row(unified_file_id, s3_uri, None)
with respx.mock:
route = respx.delete(
"https://s3.us-west-2.amazonaws.com/my-bucket/litellm-bedrock-files/input.jsonl"
).mock(return_value=httpx.Response(204))
route = respx.delete("https://s3.us-west-2.amazonaws.com/my-bucket/litellm-bedrock-files/input.jsonl").mock(
return_value=httpx.Response(204)
)
response = await managed_files.afile_delete(
file_id=unified_file_id,
litellm_parent_otel_span=None,
@ -1344,15 +1303,11 @@ async def test_afile_content_bedrock_unified_id_end_to_end(monkeypatch):
unified_file_id = "unified-file-id"
s3_uri = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out"
managed_files.get_unified_file_id = AsyncMock(return_value=None)
managed_files.get_model_file_id_mapping = AsyncMock(
return_value={unified_file_id: {"model-123": s3_uri}}
)
managed_files.get_model_file_id_mapping = AsyncMock(return_value={unified_file_id: {"model-123": s3_uri}})
expected_url = "https://s3.us-west-2.amazonaws.com/my-bucket/litellm-batch-outputs/job-123/input.jsonl.out"
with respx.mock:
route = respx.get(expected_url).mock(
return_value=httpx.Response(200, content=b'{"recordId": "x"}')
)
route = respx.get(expected_url).mock(return_value=httpx.Response(200, content=b'{"recordId": "x"}'))
response = await managed_files.afile_content(
file_id=unified_file_id,
@ -1361,9 +1316,7 @@ async def test_afile_content_bedrock_unified_id_end_to_end(monkeypatch):
)
assert route.called
assert (
route.calls[0].request.headers["Authorization"].startswith("AWS4-HMAC-SHA256")
)
assert route.calls[0].request.headers["Authorization"].startswith("AWS4-HMAC-SHA256")
assert response.content == b'{"recordId": "x"}'
@ -1375,15 +1328,13 @@ async def test_afile_content_error_reports_unified_id_not_provider_uri():
unified_file_id = "litellm_proxy_unified_id_abc"
s3_uri = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out"
managed_files.get_unified_file_id = AsyncMock(return_value=None)
managed_files.get_model_file_id_mapping = AsyncMock(
return_value={unified_file_id: {"model-123": s3_uri}}
)
managed_files.get_model_file_id_mapping = AsyncMock(return_value={unified_file_id: {"model-123": s3_uri}})
mock_router = MagicMock()
mock_router.registry.return_value.credentials = MagicMock(return_value=None)
mock_router.afile_content = AsyncMock(side_effect=Exception("deployment failed"))
with pytest.raises(Exception, match='LiteLLM Managed File object with') as exc_info:
with pytest.raises(Exception, match="LiteLLM Managed File object with") as exc_info:
await managed_files.afile_content(
file_id=unified_file_id,
litellm_parent_otel_span=None,
@ -1402,17 +1353,15 @@ def _make_real_managed_files_instance():
PROXY_LiteLLMManagedFiles,
)
mock_cache = MagicMock()
mock_cache.async_get_cache = AsyncMock(return_value=None)
mock_cache.async_set_cache = AsyncMock()
from litellm.caching.dual_cache import DualCache
mock_cache = DualCache()
mock_prisma = MagicMock()
mock_prisma.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None)
mock_prisma.db.litellm_managedfiletable.upsert = AsyncMock()
mock_prisma.db.litellm_managedfiletable.create = AsyncMock(
side_effect=AssertionError(
"store_unified_file_id must upsert, not create, on the retrieve path"
)
side_effect=AssertionError("store_unified_file_id must upsert, not create, on the retrieve path")
)
return (
@ -1454,9 +1403,7 @@ async def test_poll_refreshes_batch_state_without_claiming_the_row():
create-time snapshot forever -- while writing none of the attribution columns and
never creating a row it would then own."""
managed_files, mock_prisma = _make_object_store_instance()
poller = UserAPIKeyAuth(
api_key="sk-the-poller", user_id="bob", team_id="team-bravo", parent_otel_span=None
)
poller = UserAPIKeyAuth(api_key="sk-the-poller", user_id="bob", team_id="team-bravo", parent_otel_span=None)
await managed_files.store_unified_object_id(
unified_object_id="uoi-1",
@ -1490,9 +1437,7 @@ async def test_create_still_upserts_and_claims_attribution():
"""The create is the one caller that can speak for the batch, so it keeps the upsert
(creating the row when absent) and writes the attribution columns."""
managed_files, mock_prisma = _make_object_store_instance()
creator = UserAPIKeyAuth(
api_key="sk-the-creator", user_id="alice", team_id="team-alpha", parent_otel_span=None
)
creator = UserAPIKeyAuth(api_key="sk-the-creator", user_id="alice", team_id="team-alpha", parent_otel_span=None)
await managed_files.store_unified_object_id(
unified_object_id="uoi-2",
@ -1672,8 +1617,7 @@ async def test_concurrent_first_registrations_converge_on_one_row():
def _b64_unified_input_file_id(target_model_names: str) -> str:
unified_input_file_id = (
"litellm_proxy:application/octet-stream;unified_id,input-uuid;"
f"target_model_names,{target_model_names}"
f"litellm_proxy:application/octet-stream;unified_id,input-uuid;target_model_names,{target_model_names}"
)
return base64.urlsafe_b64encode(unified_input_file_id.encode()).decode().rstrip("=")
@ -1682,9 +1626,7 @@ def _b64_unified_input_file_id(target_model_names: str) -> str:
async def test_hook_mint_prefers_input_file_target_model_names():
managed_files = _make_managed_files_instance()
batch_response = _make_batch_response(model_name="model-a")
batch_response._hidden_params["unified_file_id"] = _b64_unified_input_file_id(
"model-a,model-b"
)
batch_response._hidden_params["unified_file_id"] = _b64_unified_input_file_id("model-a,model-b")
mock_router = MagicMock()
mock_router.registry.return_value.credentials = MagicMock(return_value={})
@ -1738,10 +1680,11 @@ async def test_hook_mint_falls_back_to_response_input_file_id_target_models():
@pytest.mark.asyncio
async def test_cost_job_and_retrieve_paths_mint_identical_unified_output_file_ids():
from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
from litellm.proxy.openai_files_endpoints.common_utils import (
ensure_batch_response_managed_file_ids,
)
from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
managed_files, mock_prisma = _make_real_managed_files_instance()
mock_prisma.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None)
@ -2040,7 +1983,7 @@ async def test_afile_delete_storage_backed_row_deletes_stored_content_not_provid
file_table = MagicMock(find_first=AsyncMock(return_value=row), delete=AsyncMock())
content_table = MagicMock(delete=AsyncMock())
managed_files = PROXY_LiteLLMManagedFiles(
internal_usage_cache=DualCache(),
internal_usage_cache=InternalUsageCache(DualCache()),
prisma_client=MagicMock(
db=MagicMock(litellm_managedfiletable=file_table, litellm_managedfilecontenttable=content_table)
),
@ -2084,7 +2027,7 @@ async def test_afile_content_storage_backed_row_returns_stored_bytes_not_provide
file_table = MagicMock(find_first=AsyncMock(return_value=row))
content_table = MagicMock(find_unique=AsyncMock(return_value=MagicMock(content=Base64.encode(stored_bytes))))
managed_files = PROXY_LiteLLMManagedFiles(
internal_usage_cache=DualCache(),
internal_usage_cache=InternalUsageCache(DualCache()),
prisma_client=MagicMock(
db=MagicMock(litellm_managedfiletable=file_table, litellm_managedfilecontenttable=content_table)
),
@ -2142,7 +2085,7 @@ async def test_store_unified_file_id_caches_the_storage_location_the_db_row_gets
file_table = MagicMock(upsert=AsyncMock(), find_first=AsyncMock(side_effect=AssertionError("cache miss")))
managed_files = PROXY_LiteLLMManagedFiles(
internal_usage_cache=DualCache(),
internal_usage_cache=InternalUsageCache(DualCache()),
prisma_client=MagicMock(db=MagicMock(litellm_managedfiletable=file_table)),
)
stored = _make_file_object("file-kept").model_copy(update={"purpose": "batch"})

264
tests/unit/redis/stores/test_managed_files.py Normal file
View file

@ -0,0 +1,264 @@
import json
import sys
import time
from collections.abc import Mapping
from datetime import timedelta
from typing import Final, cast
from unittest.mock import patch
import fakeredis
import pytest
from litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles
from pydantic import ValidationError
from redis.asyncio import Redis
from redis.exceptions import TimeoutError as RedisTimeoutError
from litellm.caching.dual_cache import DualCache
from litellm.models.managed_files import LiteLLM_ManagedFileTable
from litellm.proxy.utils import InternalUsageCache, PrismaClient
from litellm.types.llms.openai import OpenAIFileObject
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from tests._support.fake_redis import fake_redis_cache
from tests.unit.enterprise.proxy.hooks.test_managed_files import (
_InMemoryManagedFilesPrismaClient, # pyright: ignore[reportPrivateUsage] # in-memory doubles shared with the hook suite
_InMemoryManagedFileTable, # pyright: ignore[reportPrivateUsage] # in-memory doubles shared with the hook suite
)
async def _stored(raw: "Redis[bytes]", key: str) -> bytes | None:
return await raw.get(key)
def _row(file_id: str = "unified-é") -> LiteLLM_ManagedFileTable:
return LiteLLM_ManagedFileTable(
unified_file_id=file_id,
file_object=OpenAIFileObject(
id=file_id, object="file", bytes=123, created_at=456, filename="文.jsonl", purpose="batch"
),
model_mappings={"model": "provider-file"},
flat_model_file_ids=["provider-file"],
created_by="owner",
updated_by="owner",
)
def _hook(cache: DualCache, table: _InMemoryManagedFileTable) -> PROXY_LiteLLMManagedFiles:
return PROXY_LiteLLMManagedFiles(
InternalUsageCache(cache), cast(PrismaClient, _InMemoryManagedFilesPrismaClient(table))
)
@pytest.mark.asyncio
@pytest.mark.parametrize("file_id", ["unified-é", "tenant:unified-é"])
@pytest.mark.parametrize("ttl", [1, 17.9, None, 0.5, 0, -1])
async def test_file_metadata_matches_legacy_wire_and_ttl(
server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, file_id: str, ttl: float | None
) -> None:
cache: Final = fake_redis_cache(server, namespace="tenant")
dual: Final = DualCache(redis_cache=cache, default_in_memory_ttl=ttl, default_redis_ttl=999)
dual.update_cache_ttl(ttl, None)
hook: Final = _hook(dual, _InMemoryManagedFileTable())
row: Final = _row(file_id)
key: Final = file_id if file_id.startswith("tenant:") else f"tenant:{file_id}"
with patch("time.time", return_value=time.time()):
await hook.store_unified_file_id(
file_id, row.file_object, None, row.model_mappings, UserAPIKeyAuth(user_id="owner")
)
encoded: Final = await _stored(raw, key)
expiry: Final = await raw.pttl(key)
await cache.async_set_cache("legacy", row.model_dump(), **({"ttl": ttl} if ttl is not None else {}))
assert encoded == await raw.get("tenant:legacy")
assert expiry == await raw.pttl("tenant:legacy")
assert encoded == (None if ttl is not None and ttl < 1 else json.dumps(row.model_dump()).encode())
assert await cache.async_get_cache(file_id) == (None if ttl is not None and ttl < 1 else row.model_dump())
assert await raw.get(f"tenant:{key}") is None
@pytest.mark.asyncio
async def test_old_writer_new_reader_backfills_only_redis_hits(
server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis
) -> None:
cache: Final = fake_redis_cache(server, namespace="tenant")
dual: Final = DualCache(redis_cache=cache, default_in_memory_ttl=1)
row: Final = _row()
table: Final = _InMemoryManagedFileTable(row)
hook: Final = _hook(dual, table)
await cache.async_set_cache(row.unified_file_id, row.model_dump(), ttl=60)
assert await hook.get_unified_file_id(row.unified_file_id) == row
assert table.find_first_calls == []
server.connected = False
assert await hook.get_unified_file_id(row.unified_file_id) == row
server.connected = True
dual.in_memory_cache.flush_cache()
await raw.delete(f"tenant:{row.unified_file_id}")
assert await hook.get_unified_file_id(row.unified_file_id) == row
assert await raw.get(f"tenant:{row.unified_file_id}") is None
assert len(table.find_first_calls) == 1
@pytest.mark.asyncio
async def test_file_expiry_and_legacy_literal_record(
server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis
) -> None:
row: Final = _row()
with patch("time.time", return_value=1_800_000_000):
cache: Final = fake_redis_cache(server)
table: Final = _InMemoryManagedFileTable(row)
hook: Final = _hook(DualCache(redis_cache=cache, default_in_memory_ttl=1), table)
await raw.set(row.unified_file_id, repr(row.model_dump()), ex=1)
assert await hook.get_unified_file_id(row.unified_file_id) == row
assert table.find_first_calls == []
with patch("time.time", return_value=1_800_000_002):
assert await hook.get_unified_file_id(row.unified_file_id) == row
assert len(table.find_first_calls) == 1
assert await raw.get(row.unified_file_id) is None
@pytest.mark.asyncio
@pytest.mark.parametrize("encoded", [b"null", b"{}", b"broken", b"\xff"])
async def test_file_misses_and_malformed_bytes_fall_back_to_database(
server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, encoded: bytes
) -> None:
row: Final = _row()
table: Final = _InMemoryManagedFileTable(row)
hook: Final = _hook(DualCache(redis_cache=fake_redis_cache(server)), table)
await raw.set(row.unified_file_id, encoded)
assert await hook.get_unified_file_id(row.unified_file_id) == row
assert len(table.find_first_calls) == 1
assert await raw.get(row.unified_file_id) == encoded
@pytest.mark.asyncio
async def test_recursive_file_metadata_retains_legacy_database_fallback(
server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis
) -> None:
row: Final = _row()
table: Final = _InMemoryManagedFileTable(row)
cache: Final = fake_redis_cache(server)
hook: Final = _hook(DualCache(redis_cache=cache), table)
await raw.set(row.unified_file_id, b"[" * 10000 + b"0" + b"]" * 10000)
limit: Final = sys.getrecursionlimit()
try:
sys.setrecursionlimit(1000)
assert await cache.async_get_cache(row.unified_file_id) is None
assert await hook.get_unified_file_id(row.unified_file_id) == row
assert len(table.find_first_calls) == 1
finally:
sys.setrecursionlimit(limit)
@pytest.mark.asyncio
@pytest.mark.parametrize("encoded", [b'{"invalid": true}', b"(1, 2)", b"{'item'}", b"{1: 'value'}"])
async def test_truthy_invalid_file_record_still_raises_validation_error(
server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, encoded: bytes
) -> None:
table: Final = _InMemoryManagedFileTable(_row())
hook: Final = _hook(DualCache(redis_cache=fake_redis_cache(server)), table)
await raw.set("unified-é", encoded)
with pytest.raises(ValidationError):
await hook.get_unified_file_id("unified-é")
assert table.find_first_calls == []
class _DeleteObservingTable(_InMemoryManagedFileTable):
def __init__(self, raw: fakeredis.FakeAsyncRedis, row: LiteLLM_ManagedFileTable) -> None:
super().__init__(row)
self.raw = raw
async def delete(self, where: Mapping[str, object]) -> LiteLLM_ManagedFileTable | None:
assert await self.raw.get(f"tenant:{where['unified_file_id']}") == b"null"
assert await self.raw.pttl(f"tenant:{where['unified_file_id']}") == 1000
return await super().delete(where)
@pytest.mark.asyncio
async def test_file_delete_writes_null_before_database_delete(
server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis
) -> None:
row: Final = _row()
table: Final = _DeleteObservingTable(raw, row)
hook: Final = _hook(
DualCache(redis_cache=fake_redis_cache(server, namespace="tenant"), default_in_memory_ttl=1), table
)
with patch("time.time", return_value=time.time()):
assert await hook.delete_unified_file_id(row.unified_file_id) == row.file_object
assert table.rows == {}
assert await hook.get_unified_file_id(row.unified_file_id) is None
assert len(table.find_first_calls) == 2
@pytest.mark.asyncio
async def test_file_outage_keeps_local_write_and_database_policy(server: fakeredis.FakeServer) -> None:
row: Final = _row()
cache: Final = fake_redis_cache(server)
table: Final = _InMemoryManagedFileTable()
dual: Final = DualCache(redis_cache=cache, default_in_memory_ttl=1)
hook: Final = _hook(dual, table)
server.connected = False
await hook.store_unified_file_id(
row.unified_file_id, row.file_object, None, row.model_mappings, UserAPIKeyAuth(user_id="owner")
)
assert await hook.get_unified_file_id(row.unified_file_id) == row
assert table.rows[row.unified_file_id] == row
dual.in_memory_cache.flush_cache()
for _ in range(8):
assert await hook.get_unified_file_id(row.unified_file_id) == row
assert cache._circuit_breaker.state == "open" # pyright: ignore[reportPrivateUsage] # asserts the breaker opened after the outage
@pytest.mark.asyncio
@pytest.mark.parametrize("ttl", [float("nan"), float("inf"), float("-inf")])
async def test_invalid_file_cache_ttl_does_not_block_database_write(
server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, ttl: float
) -> None:
row: Final = _row()
table: Final = _InMemoryManagedFileTable()
dual: Final = DualCache(redis_cache=fake_redis_cache(server), default_in_memory_ttl=ttl)
hook: Final = _hook(dual, table)
with patch("time.time", return_value=1_800_000_000):
await dual.async_set_cache("legacy", row.model_dump())
await hook.store_unified_file_id(
row.unified_file_id, row.file_object, None, row.model_mappings, UserAPIKeyAuth(user_id="owner")
)
assert table.rows[row.unified_file_id] == row
assert await raw.get(row.unified_file_id) == await raw.get("legacy")
assert await raw.pttl(row.unified_file_id) == await raw.pttl("legacy")
@pytest.mark.asyncio
async def test_unknown_applied_file_write_is_not_rolled_back(
server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis
) -> None:
cache: Final = fake_redis_cache(server, namespace="tenant")
row: Final = _row()
table: Final = _InMemoryManagedFileTable()
hook: Final = _hook(DualCache(redis_cache=cache, default_in_memory_ttl=1), table)
async def applied_then_timeout(name: str, value: str, *, ex: timedelta | None, nx: bool) -> bool:
await raw.set(name, value, ex=ex, nx=nx)
raise RedisTimeoutError("reply lost after apply")
with patch.object(cache.init_async_client(), "set", side_effect=applied_then_timeout):
await hook.store_unified_file_id(
row.unified_file_id, row.file_object, None, row.model_mappings, UserAPIKeyAuth(user_id="owner")
)
assert await raw.get(f"tenant:{row.unified_file_id}") == json.dumps(row.model_dump()).encode()
assert table.rows[row.unified_file_id] == row
@pytest.mark.asyncio
async def test_file_refresh_and_lost_compare_and_swap_preserve_cache_order(
server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis
) -> None:
row: Final = _row()
table: Final = _InMemoryManagedFileTable(row)
hook: Final = _hook(DualCache(redis_cache=fake_redis_cache(server), default_in_memory_ttl=1), table)
assert row.file_object is not None
updated: Final = row.file_object.model_copy(update={"bytes": 987})
await hook._save_refreshed_file_object(row, updated) # pyright: ignore[reportPrivateUsage] # exercises the refresh write path
assert await raw.get(row.unified_file_id) == json.dumps(table.rows[row.unified_file_id].model_dump()).encode()
assert await hook._save_xai_refreshed_file_object(row, updated) is None # pyright: ignore[reportPrivateUsage] # exercises the refresh write path
assert await raw.get(row.unified_file_id) == b"null"
assert table.rows[row.unified_file_id].file_object == updated
assert await hook._save_xai_refreshed_file_object(table.rows[row.unified_file_id], updated) == updated # pyright: ignore[reportPrivateUsage] # exercises the refresh write path
assert await raw.get(row.unified_file_id) == json.dumps(table.rows[row.unified_file_id].model_dump()).encode()