mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-03 07:23:33 +00:00
fix(managed_files): keep the upload the managed id encodes when two target models share a deployment (#49420)
* fix(managed_files): keep the upload the managed id encodes when two target models share a deployment * fix(proxy): skip cost tracking for a file delete routed to a deployment * fix(managed_files): delete repeat uploads after the upload answers Run the repeat-copy deletes together in a background task so a slow provider delete no longer holds the upload's response, log the unbilled file-delete cost skip at debug, and type the new test helpers * fix(proxy): run managed file repeat-upload cleanup as a drained internal call Delete repeat uploads after the record is saved, through the shared background task list the shutdown drain waits on, and as an internal call so logging, spend tracking and alerts skip it. Leave litellm_proxy deployments on the old keep-last behavior, since their per-target copies are distinct files. Revert the cost tracker skip change, which the internal call makes unnecessary. * fix(managed_files): delete the copy a managed id names on DELETE and isolate cleanup metadata * fix(managed_files): read the copy DELETE removes from the id's gateway-written fields * fix(managed_files): delete the copy a managed id carries as an internal call * fix(managed_files): pass the parent call id to the cleanup deletes * fix(managed_files): file each repeat-upload delete under the upload that made its copy POST /v1/files sets no request call id, so the cleanup delete passed none and its usage ledger row had no parent. Each copy now keeps the call id of the router upload that made it, and its delete runs under that id. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
quietjules
GitHub
parent
acd084ab38
commit
4187e1ae89
3 changed files with 557 additions and 28 deletions
|
|
@ -3,9 +3,11 @@
|
|||
|
||||
import asyncio
|
||||
import base64
|
||||
import copy
|
||||
import json
|
||||
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
|
||||
|
|
@ -34,7 +36,7 @@ from typing_extensions import ReadOnly, Unpack
|
|||
|
||||
import litellm
|
||||
from litellm import Router, verbose_logger
|
||||
from litellm._internal_context import with_service_target
|
||||
from litellm._internal_context import internal_sub_call, post_response_phase, with_service_target
|
||||
from litellm._uuid import uuid
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.constants import (
|
||||
|
|
@ -95,6 +97,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio
|
|||
request_tags_from_metadata,
|
||||
)
|
||||
from litellm.repositories.managed_file_repository import ManagedFileRepository
|
||||
from litellm.responses.streaming_iterator import spawn_logging_task
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.types.llms.openai import ( # pyright: ignore[reportAttributeAccessIssue]
|
||||
AllMessageValues,
|
||||
|
|
@ -395,6 +398,134 @@ class _RouterFileCallKwargs(TypedDict, total=False):
|
|||
custom_llm_provider: ReadOnly[str | None]
|
||||
|
||||
|
||||
_MODEL_FILE_ID_MAPPING: Final = TypeAdapter(dict[str, str])
|
||||
_FILE_UPLOAD_ONLY_FIELDS: Final = frozenset({"file", "purpose", "expires_after", "extra_body"})
|
||||
_FILE_DELETE_RESERVED_FIELDS: Final = frozenset({"model", "file_id", "_litellm_internal_model_credentials"})
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ProviderUpload:
|
||||
model_id: str
|
||||
provider_file_id: str
|
||||
upload_call_id: str | None = None
|
||||
|
||||
|
||||
def _provider_uploads(file_object: OpenAIFileObject) -> tuple[_ProviderUpload, ...]:
|
||||
file_hidden_params: Final = cast( # cast-ok: preserve mapping operations on dynamic file metadata
|
||||
dict[str, object], getattr(file_object, HIDDEN_PARAMS_ATTR)
|
||||
)
|
||||
model_file_id_mapping: Final = _MODEL_FILE_ID_MAPPING.validate_python(
|
||||
file_hidden_params.get("model_file_id_mapping") or {}
|
||||
)
|
||||
upload_call_id: Final = file_hidden_params.get("litellm_call_id")
|
||||
return tuple(
|
||||
_ProviderUpload(
|
||||
model_id=model_id,
|
||||
provider_file_id=provider_file_id,
|
||||
upload_call_id=upload_call_id if isinstance(upload_call_id, str) else None,
|
||||
)
|
||||
for model_id, provider_file_id in model_file_id_mapping.items()
|
||||
)
|
||||
|
||||
|
||||
def _keeps_one_upload_per_target(llm_router: Router, model_id: str) -> bool:
|
||||
credentials: Final = _provider_file_retrieve_credentials(llm_router=llm_router, model_id=model_id)
|
||||
return (credentials or {}).get("custom_llm_provider") == LlmProviders.LITELLM_PROXY.value
|
||||
|
||||
|
||||
def _is_kept_upload(uploads: Sequence[_ProviderUpload], index: int, per_target_model_ids: frozenset[str]) -> bool:
|
||||
upload: Final = uploads[index]
|
||||
other_uploads: Final = uploads[index + 1 :] if upload.model_id in per_target_model_ids else uploads[:index]
|
||||
return all(other.model_id != upload.model_id for other in other_uploads)
|
||||
|
||||
|
||||
def _upload_from_trailing_id_fields(decoded_unified_file_id: str) -> _ProviderUpload:
|
||||
trailing_fields: Final = decoded_unified_file_id.rpartition(";llm_output_file_id,")[2]
|
||||
provider_file_id, _, model_id = trailing_fields.partition(";llm_output_file_model_id,")
|
||||
return _ProviderUpload(model_id=model_id, provider_file_id=provider_file_id)
|
||||
|
||||
|
||||
def _carried_upload_missing_from(unified_file_id: str, model_file_ids: Mapping[str, str]) -> _ProviderUpload | None:
|
||||
decoded: Final = _is_base64_encoded_unified_file_id(unified_file_id)
|
||||
carried: Final = _upload_from_trailing_id_fields(decoded) if decoded else None
|
||||
if carried is None or model_file_ids.get(carried.model_id) in (None, carried.provider_file_id):
|
||||
return None
|
||||
return carried
|
||||
|
||||
|
||||
async def _delete_provider_file(
|
||||
llm_router: Router,
|
||||
model_id: str,
|
||||
provider_file_id: str,
|
||||
request_data: Mapping[str, object],
|
||||
) -> None:
|
||||
credentials: Final = llm_router.registry().credentials(model_id, form="with_provider")
|
||||
delete_data: Final = {
|
||||
**{k: v for k, v in request_data.items() if k not in _FILE_DELETE_RESERVED_FIELDS},
|
||||
**(
|
||||
{"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
|
||||
if credentials is not None
|
||||
else {}
|
||||
),
|
||||
}
|
||||
await llm_router.afile_delete(
|
||||
model=model_id,
|
||||
file_id=provider_file_id,
|
||||
**cast(_RouterFileCallKwargs, delete_data),
|
||||
)
|
||||
|
||||
|
||||
def _with_own_nested_dicts(request_data: Mapping[str, object]) -> MappingProxyType[str, object]:
|
||||
return MappingProxyType({k: copy.copy(v) if isinstance(v, dict) else v for k, v in request_data.items()})
|
||||
|
||||
|
||||
async def _delete_redundant_upload(
|
||||
llm_router: Router,
|
||||
upload: _ProviderUpload,
|
||||
delete_request: Mapping[str, object],
|
||||
parent_call_id: str | None,
|
||||
) -> None:
|
||||
try:
|
||||
with internal_sub_call(parent_call_id):
|
||||
await _delete_provider_file(
|
||||
llm_router, upload.model_id, upload.provider_file_id, _with_own_nested_dicts(delete_request)
|
||||
)
|
||||
except Exception as error:
|
||||
verbose_logger.warning(
|
||||
f"Could not delete the redundant upload provider_file_id={sanitize_for_log(upload.provider_file_id)} "
|
||||
f"at deployment model_id={sanitize_for_log(upload.model_id)}: "
|
||||
f"{type(error).__name__} {sanitize_for_log(error)}"
|
||||
)
|
||||
|
||||
|
||||
async def _delete_redundant_uploads(
|
||||
llm_router: Router,
|
||||
uploads: Sequence[_ProviderUpload],
|
||||
delete_request: Mapping[str, object],
|
||||
) -> None:
|
||||
await asyncio.gather(
|
||||
*(_delete_redundant_upload(llm_router, upload, delete_request, upload.upload_call_id) for upload in uploads)
|
||||
)
|
||||
|
||||
|
||||
def _start_deleting_redundant_uploads(
|
||||
llm_router: Router,
|
||||
uploads: Sequence[_ProviderUpload],
|
||||
create_file_request: CreateFileRequest,
|
||||
) -> None:
|
||||
if not uploads:
|
||||
return
|
||||
delete_request: Final = MappingProxyType(
|
||||
{k: v for k, v in create_file_request.items() if k not in _FILE_UPLOAD_ONLY_FIELDS}
|
||||
)
|
||||
with post_response_phase():
|
||||
spawn_logging_task(
|
||||
asyncio.get_running_loop(),
|
||||
_delete_redundant_uploads(llm_router, uploads, delete_request),
|
||||
task_name="managed file repeat upload cleanup",
|
||||
)
|
||||
|
||||
|
||||
def _managed_database(prisma_client: _ManagedResourcePrismaClient) -> _ManagedResourceDatabase:
|
||||
return prisma_client.db
|
||||
|
||||
|
|
@ -2034,16 +2165,16 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
target_model_names_list=target_model_names_list,
|
||||
)
|
||||
|
||||
## STORE MODEL MAPPINGS IN DB
|
||||
model_mappings: Dict[str, str] = {}
|
||||
|
||||
for file_object in responses:
|
||||
file_hidden_params = cast( # cast-ok: preserve mapping operations on dynamic file metadata
|
||||
dict[str, object], getattr(file_object, HIDDEN_PARAMS_ATTR)
|
||||
)
|
||||
model_file_id_mapping = file_hidden_params.get("model_file_id_mapping")
|
||||
if model_file_id_mapping and isinstance(model_file_id_mapping, dict):
|
||||
model_mappings.update(model_file_id_mapping)
|
||||
uploads: Final = tuple(chain.from_iterable(map(_provider_uploads, responses)))
|
||||
per_target_model_ids: Final = frozenset(
|
||||
model_id
|
||||
for model_id in {upload.model_id for upload in uploads}
|
||||
if _keeps_one_upload_per_target(llm_router, model_id)
|
||||
)
|
||||
kept_uploads: Final = tuple(
|
||||
upload for index, upload in enumerate(uploads) if _is_kept_upload(uploads, index, per_target_model_ids)
|
||||
)
|
||||
model_mappings: Final = {upload.model_id: upload.provider_file_id for upload in kept_uploads}
|
||||
|
||||
await self.store_unified_file_id(
|
||||
file_id=response.id,
|
||||
|
|
@ -2052,6 +2183,15 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
model_mappings=model_mappings,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
_start_deleting_redundant_uploads(
|
||||
llm_router,
|
||||
tuple(
|
||||
upload
|
||||
for upload in uploads
|
||||
if upload not in kept_uploads and upload.model_id not in per_target_model_ids
|
||||
),
|
||||
create_file_request,
|
||||
)
|
||||
|
||||
# Emit Prometheus metrics for managed file creation
|
||||
prom_logger = self._get_prometheus_logger()
|
||||
|
|
@ -2692,23 +2832,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
specific_model_file_id_mapping: Final = model_file_id_mapping.get(file_id)
|
||||
if not specific_model_file_id_mapping:
|
||||
return
|
||||
filtered_data: Final = {
|
||||
k: v for k, v in data.items() if k not in ("model", "file_id", "_litellm_internal_model_credentials")
|
||||
}
|
||||
for model_id, model_file_id in specific_model_file_id_mapping.items():
|
||||
credentials = llm_router.registry().credentials(model_id, form="with_provider")
|
||||
delete_data = {
|
||||
**filtered_data,
|
||||
**(
|
||||
{"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
|
||||
if credentials is not None
|
||||
else {}
|
||||
),
|
||||
}
|
||||
await llm_router.afile_delete(
|
||||
model=model_id,
|
||||
file_id=model_file_id,
|
||||
**cast(_RouterFileCallKwargs, delete_data),
|
||||
await _delete_provider_file(llm_router, model_id, model_file_id, data)
|
||||
unrecorded_upload: Final = _carried_upload_missing_from(file_id, specific_model_file_id_mapping)
|
||||
delete_call_id: Final = data.get("litellm_call_id")
|
||||
if unrecorded_upload is not None:
|
||||
await _delete_redundant_upload(
|
||||
llm_router, unrecorded_upload, data, delete_call_id if isinstance(delete_call_id, str) else None
|
||||
)
|
||||
|
||||
async def afile_content(
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ from concurrent.futures import ThreadPoolExecutor
|
|||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from itertools import count
|
||||
from pathlib import Path
|
||||
from queue import SimpleQueue
|
||||
from types import MappingProxyType
|
||||
|
|
@ -1042,6 +1043,77 @@ def test_multi_model_upload_maps_only_the_provider_id_the_managed_id_carries(gat
|
|||
assert second_ids == ((managed_a,) if carried == second.file_id("a.txt") else (second.file_id("a.txt"),))
|
||||
|
||||
|
||||
def _team_wildcard(scenario: Scenario, pattern: str, api_base: str, bearer: str, team: str) -> str:
|
||||
identity: Final = f"wildcard-{uuid.uuid4().hex}"
|
||||
scenario.gateway.post(
|
||||
"/model/new",
|
||||
{
|
||||
"model_name": pattern,
|
||||
"litellm_params": {"model": "openai/*", "api_key": bearer, "api_base": api_base},
|
||||
"model_info": {"id": identity, "team_id": team},
|
||||
},
|
||||
)
|
||||
scenario.cleanups.callback(scenario.delete_model, identity)
|
||||
eventually(lambda: _every_worker_has_deployment(scenario.gateway, identity), lambda served: served, seconds=90)
|
||||
return identity
|
||||
|
||||
|
||||
def test_two_targets_served_by_one_team_wildcard_keep_one_upload_that_delete_removes(gateway: Gateway) -> None:
|
||||
store: Final = "vs_" + uuid.uuid4().hex
|
||||
bearer: Final = "provider-key-" + uuid.uuid4().hex[:8]
|
||||
prefix: Final = f"wc-{uuid.uuid4().hex}"
|
||||
upload_numbers: Final = count(1)
|
||||
first_upload: Final = f"file-{prefix}-1"
|
||||
second_upload: Final = f"file-{prefix}-2"
|
||||
served: Final = _provider(store, _constant_listing(store, first_upload))
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
if (request.method, urlsplit(request.target).path) != ("POST", "/v1/files"):
|
||||
return served(request)
|
||||
_assert_provider_upload(request, bearer, "a.txt")
|
||||
return _json_reply(_file_object(f"file-{prefix}-{next(upload_numbers)}"))
|
||||
|
||||
with gateway.scenario() as scenario, wire_server(respond) as wire:
|
||||
team: Final = scenario.team(models=[])
|
||||
_team_wildcard(scenario, f"{prefix}/*", f"{wire.url}/v1", bearer, team)
|
||||
key: Final = scenario.key(team_id=team, user_id=scenario.member(team))
|
||||
wire.drain()
|
||||
|
||||
managed: Final = _upload(gateway, key, f"{prefix}/gpt-one,{prefix}/gpt-two", "a.txt")
|
||||
assert _carried_provider_file_id(managed) == first_upload, _decoded(managed)
|
||||
assert [row["flat_model_file_ids"] for row in read_rows(MANAGED_FILE_ROW, (managed,))] == [[first_upload]]
|
||||
eventually(lambda: wire.received.qsize(), lambda received: received >= 3)
|
||||
create_requests: Final = wire.drain()
|
||||
assert [(request.method, request.target) for request in create_requests] == [
|
||||
("POST", "/v1/files"),
|
||||
("POST", "/v1/files"),
|
||||
("DELETE", f"/v1/files/{second_upload}"),
|
||||
], create_requests
|
||||
assert create_requests[2].headers["authorization"] == f"Bearer {bearer}", create_requests[2].headers
|
||||
|
||||
attached: Final = gateway.request(
|
||||
"POST", f"/v1/vector_stores/{store}/files", {"file_id": managed}, key=key, params={"model": f"{prefix}/*"}
|
||||
)
|
||||
assert attached.status_code == 200, attached.text
|
||||
listed: Final = gateway.request(
|
||||
"GET", f"/v1/vector_stores/{store}/files", key=key, params={"model": f"{prefix}/*"}
|
||||
)
|
||||
assert _ids(_listed(listed)) == (managed,), listed.text
|
||||
attach_bodies: Final = [
|
||||
JSON_OBJECT.validate_json(request.body) for request in wire.drain() if request.method == "POST"
|
||||
]
|
||||
assert attach_bodies == [{"file_id": first_upload}], attach_bodies
|
||||
|
||||
deleted: Final = gateway.request("DELETE", f"/v1/files/{managed}", key=key)
|
||||
assert deleted.status_code == 200, deleted.text
|
||||
assert read_rows(MANAGED_FILE_ROW, (managed,)) == []
|
||||
delete_requests: Final = wire.drain()
|
||||
assert [(request.method, request.target) for request in delete_requests] == [
|
||||
("DELETE", f"/v1/files/{first_upload}")
|
||||
], delete_requests
|
||||
assert delete_requests[0].headers["authorization"] == f"Bearer {bearer}", delete_requests[0].headers
|
||||
|
||||
|
||||
def test_deleting_the_managed_file_makes_its_listing_raw_again(gateway: Gateway) -> None:
|
||||
with _rig(gateway, "a.txt") as rig:
|
||||
member: Final = _member(rig.scenario, rig.model)
|
||||
|
|
|
|||
|
|
@ -9,18 +9,24 @@ import asyncio
|
|||
import base64
|
||||
import json
|
||||
import logging
|
||||
from collections.abc import Callable, Mapping
|
||||
from types import MappingProxyType
|
||||
|
||||
import pytest
|
||||
from typing import Final, Optional
|
||||
from typing import TYPE_CHECKING, Final, Optional
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from litellm._internal_context import internal_parent_request_id, is_internal_call
|
||||
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.responses.streaming_iterator import drain_logging_tasks
|
||||
from litellm.types.llms.openai import FileListPage, OpenAIFileObject
|
||||
from litellm.types.utils import LiteLLMBatch
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles
|
||||
|
||||
|
||||
def _make_file_object(file_id: str = "file-output-abc") -> OpenAIFileObject:
|
||||
return OpenAIFileObject(
|
||||
|
|
@ -2268,3 +2274,324 @@ async def test_batch_create_hook_passes_the_router_deployment_as_creating_model_
|
|||
first, second = managed_files.store_unified_object_id.await_args_list
|
||||
assert first.kwargs["creating_model_id"] == "deploy-team-b"
|
||||
assert second.kwargs["creating_model_id"] is None
|
||||
|
||||
|
||||
class _FakeFilesRouter:
|
||||
def __init__(
|
||||
self,
|
||||
deployments_by_model: Mapping[str, tuple[str, ...]],
|
||||
delete_error: Exception | None = None,
|
||||
delete_gate: asyncio.Event | None = None,
|
||||
provider: str = "openai",
|
||||
gone_file_ids: frozenset[str] = frozenset(),
|
||||
) -> None:
|
||||
self.deployments_by_model = deployments_by_model
|
||||
self.delete_error = delete_error
|
||||
self.gone_file_ids = gone_file_ids
|
||||
self.delete_gate = delete_gate
|
||||
self.provider = provider
|
||||
self.upload_count = 0
|
||||
self.deletes: tuple[tuple[str, str, Mapping[str, object], Mapping[str, str]], ...] = ()
|
||||
self.deletes_marked_internal: tuple[bool, ...] = ()
|
||||
self.delete_parent_request_ids: tuple[str | None, ...] = ()
|
||||
|
||||
async def acreate_file(
|
||||
self,
|
||||
model: str,
|
||||
file: tuple[str, bytes, str],
|
||||
purpose: str,
|
||||
litellm_metadata: Mapping[str, object],
|
||||
) -> OpenAIFileObject:
|
||||
deployment_ids: Final = self.deployments_by_model[model]
|
||||
provider_file_ids: Final = tuple(
|
||||
f"file-upload-{self.upload_count + offset + 1}" for offset in range(len(deployment_ids))
|
||||
)
|
||||
self.upload_count += len(deployment_ids)
|
||||
response: Final = _make_file_object(provider_file_ids[0]).model_copy(update={"purpose": purpose})
|
||||
response._hidden_params = {
|
||||
"model_id": deployment_ids[0],
|
||||
"custom_llm_provider": self.provider,
|
||||
"model_file_id_mapping": dict(zip(deployment_ids, provider_file_ids)),
|
||||
"litellm_call_id": f"call-of-{provider_file_ids[0]}",
|
||||
}
|
||||
return response
|
||||
|
||||
def registry(self) -> "_FakeFilesRouter":
|
||||
return self
|
||||
|
||||
def credentials(self, model_id: str, form: str) -> dict[str, str]:
|
||||
assert form == "with_provider"
|
||||
return {"custom_llm_provider": self.provider, "api_key": f"key-of-{model_id}"}
|
||||
|
||||
async def afile_delete(
|
||||
self,
|
||||
model: str,
|
||||
file_id: str,
|
||||
litellm_metadata: Mapping[str, object],
|
||||
_litellm_internal_model_credentials: Mapping[str, str],
|
||||
litellm_call_id: str | None = None,
|
||||
) -> None:
|
||||
if self.delete_gate is not None:
|
||||
await self.delete_gate.wait()
|
||||
if isinstance(litellm_metadata, dict):
|
||||
litellm_metadata.update({"model_info": {"id": model}})
|
||||
self.deletes = (*self.deletes, (model, file_id, litellm_metadata, _litellm_internal_model_credentials))
|
||||
self.deletes_marked_internal = (*self.deletes_marked_internal, is_internal_call.get())
|
||||
self.delete_parent_request_ids = (*self.delete_parent_request_ids, internal_parent_request_id.get())
|
||||
if file_id in self.gone_file_ids:
|
||||
raise RuntimeError(f"No such File object: {file_id}")
|
||||
if self.delete_error is not None:
|
||||
raise self.delete_error
|
||||
|
||||
|
||||
_REQUEST_METADATA: Final = MappingProxyType({"user_api_key_team_id": "test-team"})
|
||||
|
||||
|
||||
async def _eventually(condition: Callable[[], bool]) -> None:
|
||||
async def _poll() -> None:
|
||||
while not condition():
|
||||
await asyncio.sleep(0)
|
||||
|
||||
await asyncio.wait_for(_poll(), timeout=1)
|
||||
|
||||
|
||||
async def _upload(
|
||||
managed_files: "PROXY_LiteLLMManagedFiles",
|
||||
router: _FakeFilesRouter,
|
||||
target_model_names: list[str],
|
||||
litellm_metadata: Mapping[str, object] = _REQUEST_METADATA,
|
||||
content_type: str = "text/plain",
|
||||
) -> OpenAIFileObject:
|
||||
return await managed_files.acreate_file(
|
||||
create_file_request={
|
||||
"file": ("notes.txt", b"the launch code is ORCHID-42\n", content_type),
|
||||
"purpose": "user_data",
|
||||
"litellm_metadata": litellm_metadata,
|
||||
},
|
||||
llm_router=router,
|
||||
target_model_names_list=target_model_names,
|
||||
litellm_parent_otel_span=None,
|
||||
user_api_key_dict=_make_team_member_api_key_dict(),
|
||||
)
|
||||
|
||||
|
||||
async def _create_managed_file(
|
||||
router: _FakeFilesRouter, target_model_names: list[str]
|
||||
) -> tuple[dict[str, str], list[str], tuple[str, str]]:
|
||||
managed_files, mock_prisma = _make_real_managed_files_instance()
|
||||
response: Final = await _upload(managed_files, router, target_model_names)
|
||||
stored: Final = mock_prisma.db.litellm_managedfiletable.upsert.await_args.kwargs["data"]["create"]
|
||||
decoded: Final = base64.urlsafe_b64decode(response.id + "=" * (-len(response.id) % 4)).decode()
|
||||
carried: Final = (
|
||||
managed_files.get_model_id_from_unified_file_id(decoded),
|
||||
managed_files.get_output_file_id_from_unified_file_id(decoded),
|
||||
)
|
||||
return json.loads(stored["model_mappings"]), stored["flat_model_file_ids"], carried
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acreate_file_keeps_the_upload_the_managed_id_carries_when_two_targets_share_a_deployment():
|
||||
router: Final = _FakeFilesRouter({"gpt-a": ("wildcard-dep",), "gpt-b": ("wildcard-dep",)})
|
||||
|
||||
model_mappings, flat_ids, carried = await _create_managed_file(router, ["gpt-a", "gpt-b"])
|
||||
await _eventually(lambda: len(router.deletes) == 1)
|
||||
|
||||
assert carried == ("wildcard-dep", "file-upload-1")
|
||||
assert model_mappings == {"wildcard-dep": "file-upload-1"}
|
||||
assert flat_ids == ["file-upload-1"]
|
||||
assert router.deletes == (
|
||||
(
|
||||
"wildcard-dep",
|
||||
"file-upload-2",
|
||||
_REQUEST_METADATA,
|
||||
{"custom_llm_provider": "openai", "api_key": "key-of-wildcard-dep"},
|
||||
),
|
||||
)
|
||||
assert isinstance(router.deletes[0][3], MappingProxyType)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acreate_file_deletes_only_the_repeat_upload_and_keeps_every_deployment_reached_once():
|
||||
router: Final = _FakeFilesRouter({"gpt-a": ("dep-1", "dep-2"), "gpt-b": ("dep-2",), "gpt-c": ("dep-3",)})
|
||||
|
||||
model_mappings, flat_ids, carried = await _create_managed_file(router, ["gpt-a", "gpt-b", "gpt-c"])
|
||||
await _eventually(lambda: len(router.deletes) == 1)
|
||||
|
||||
assert carried == ("dep-1", "file-upload-1")
|
||||
assert list(model_mappings.items()) == [
|
||||
("dep-1", "file-upload-1"),
|
||||
("dep-2", "file-upload-2"),
|
||||
("dep-3", "file-upload-4"),
|
||||
]
|
||||
assert flat_ids == ["file-upload-1", "file-upload-2", "file-upload-4"]
|
||||
assert [delete[:2] for delete in router.deletes] == [("dep-2", "file-upload-3")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acreate_file_still_returns_the_managed_file_when_deleting_the_repeat_upload_fails(caplog: pytest.LogCaptureFixture):
|
||||
router: Final = _FakeFilesRouter(
|
||||
{"gpt-a": ("wildcard-dep",), "gpt-b": ("wildcard-dep",)},
|
||||
delete_error=RuntimeError("provider unavailable"),
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
model_mappings, flat_ids, carried = await _create_managed_file(router, ["gpt-a", "gpt-b"])
|
||||
await _eventually(lambda: "provider unavailable" in caplog.text)
|
||||
|
||||
assert carried == ("wildcard-dep", "file-upload-1")
|
||||
assert model_mappings == {"wildcard-dep": "file-upload-1"}
|
||||
assert flat_ids == ["file-upload-1"]
|
||||
assert [delete[:2] for delete in router.deletes] == [("wildcard-dep", "file-upload-2")]
|
||||
assert "file-upload-2" in caplog.text and "wildcard-dep" in caplog.text
|
||||
assert "provider unavailable" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acreate_file_answers_before_deleting_the_repeat_upload_finishes():
|
||||
delete_gate: Final = asyncio.Event()
|
||||
router: Final = _FakeFilesRouter({"gpt-a": ("wildcard-dep",), "gpt-b": ("wildcard-dep",)}, delete_gate=delete_gate)
|
||||
|
||||
model_mappings, flat_ids, carried = await asyncio.wait_for(
|
||||
_create_managed_file(router, ["gpt-a", "gpt-b"]), timeout=1
|
||||
)
|
||||
|
||||
assert carried == ("wildcard-dep", "file-upload-1")
|
||||
assert model_mappings == {"wildcard-dep": "file-upload-1"}
|
||||
assert flat_ids == ["file-upload-1"]
|
||||
assert router.deletes == ()
|
||||
delete_gate.set()
|
||||
await _eventually(lambda: len(router.deletes) == 1)
|
||||
assert router.deletes[0][:2] == ("wildcard-dep", "file-upload-2")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shutdown_drain_waits_for_the_repeat_upload_delete():
|
||||
delete_gate: Final = asyncio.Event()
|
||||
router: Final = _FakeFilesRouter({"gpt-a": ("wildcard-dep",), "gpt-b": ("wildcard-dep",)}, delete_gate=delete_gate)
|
||||
|
||||
await _create_managed_file(router, ["gpt-a", "gpt-b"])
|
||||
asyncio.get_running_loop().call_later(0.05, delete_gate.set)
|
||||
await drain_logging_tasks(timeout=1)
|
||||
|
||||
assert [delete[:2] for delete in router.deletes] == [("wildcard-dep", "file-upload-2")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_each_repeat_upload_delete_runs_as_an_internal_call_of_the_upload_that_made_the_copy():
|
||||
router: Final = _FakeFilesRouter(
|
||||
{"gpt-a": ("wildcard-dep",), "gpt-b": ("wildcard-dep",), "gpt-c": ("wildcard-dep",)}
|
||||
)
|
||||
|
||||
await _create_managed_file(router, ["gpt-a", "gpt-b", "gpt-c"])
|
||||
await drain_logging_tasks(timeout=1)
|
||||
|
||||
assert router.deletes_marked_internal == (True, True)
|
||||
assert sorted(zip((delete[1] for delete in router.deletes), router.delete_parent_request_ids)) == [
|
||||
("file-upload-2", "call-of-file-upload-2"),
|
||||
("file-upload-3", "call-of-file-upload-3"),
|
||||
]
|
||||
assert is_internal_call.get() is False
|
||||
assert internal_parent_request_id.get() is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_litellm_proxy_wildcard_keeps_its_last_per_target_upload_and_deletes_nothing():
|
||||
router: Final = _FakeFilesRouter(
|
||||
{"gpt-a": ("proxy-wildcard-dep",), "gpt-b": ("proxy-wildcard-dep",)}, provider="litellm_proxy"
|
||||
)
|
||||
|
||||
model_mappings, flat_ids, _ = await _create_managed_file(router, ["gpt-a", "gpt-b"])
|
||||
await drain_logging_tasks(timeout=1)
|
||||
|
||||
assert model_mappings == {"proxy-wildcard-dep": "file-upload-2"}
|
||||
assert flat_ids == ["file-upload-2"]
|
||||
assert router.deletes == ()
|
||||
|
||||
|
||||
async def _upload_then_delete(
|
||||
router: _FakeFilesRouter, target_model_names: list[str], content_type: str = "text/plain"
|
||||
) -> tuple[str, ...]:
|
||||
managed_files, mock_prisma = _make_real_managed_files_instance()
|
||||
response: Final = await _upload(managed_files, router, target_model_names, content_type=content_type)
|
||||
await drain_logging_tasks(timeout=1)
|
||||
deletes_during_upload: Final = len(router.deletes)
|
||||
managed_files.internal_usage_cache.async_get_cache = AsyncMock(
|
||||
return_value=managed_files.internal_usage_cache.async_set_cache.await_args.kwargs["value"]
|
||||
)
|
||||
mock_prisma.db.litellm_managedfiletable.find_first = AsyncMock(return_value=MagicMock())
|
||||
mock_prisma.db.litellm_managedfiletable.delete = AsyncMock()
|
||||
|
||||
deleted: Final = await managed_files.afile_delete(
|
||||
file_id=response.id,
|
||||
litellm_parent_otel_span=None,
|
||||
llm_router=router,
|
||||
litellm_metadata=_REQUEST_METADATA,
|
||||
litellm_call_id="delete-call",
|
||||
)
|
||||
|
||||
assert deleted.deleted is True
|
||||
mock_prisma.db.litellm_managedfiletable.delete.assert_awaited_once_with(where={"unified_file_id": response.id})
|
||||
return tuple(delete[1] for delete in router.deletes[deletes_during_upload:])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deleting_a_litellm_proxy_wildcard_file_removes_the_upload_its_id_carries_too():
|
||||
router: Final = _FakeFilesRouter(
|
||||
{"gpt-a": ("proxy-wildcard-dep",), "gpt-b": ("proxy-wildcard-dep",)}, provider="litellm_proxy"
|
||||
)
|
||||
|
||||
assert await _upload_then_delete(router, ["gpt-a", "gpt-b"]) == ("file-upload-2", "file-upload-1")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deleting_a_wildcard_file_removes_its_kept_upload_once():
|
||||
router: Final = _FakeFilesRouter({"gpt-a": ("wildcard-dep",), "gpt-b": ("wildcard-dep",)})
|
||||
|
||||
assert await _upload_then_delete(router, ["gpt-a", "gpt-b"]) == ("file-upload-1",)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"content_type",
|
||||
[
|
||||
"text/plain;llm_output_file_id,file-of-another-team;llm_output_file_model_id,wildcard-dep",
|
||||
"text/plain;llm_output_file_id,file-of-another-team",
|
||||
],
|
||||
)
|
||||
async def test_deleting_a_file_removes_only_the_uploads_the_gateway_made_for_it(content_type: str):
|
||||
router: Final = _FakeFilesRouter({"gpt-a": ("wildcard-dep",), "gpt-b": ("wildcard-dep",)})
|
||||
|
||||
assert await _upload_then_delete(router, ["gpt-a", "gpt-b"], content_type) == ("file-upload-1",)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deleting_a_file_still_succeeds_when_the_upload_its_id_carries_is_already_gone(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
):
|
||||
router: Final = _FakeFilesRouter(
|
||||
{"gpt-a": ("proxy-wildcard-dep",), "gpt-b": ("proxy-wildcard-dep",)},
|
||||
provider="litellm_proxy",
|
||||
gone_file_ids=frozenset({"file-upload-1"}),
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
assert await _upload_then_delete(router, ["gpt-a", "gpt-b"]) == ("file-upload-2", "file-upload-1")
|
||||
|
||||
assert "file-upload-1" in caplog.text and "No such File object" in caplog.text
|
||||
assert router.deletes_marked_internal == (False, True)
|
||||
assert router.delete_parent_request_ids == (None, "delete-call")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_repeat_upload_delete_leaves_the_upload_request_metadata_unchanged():
|
||||
router: Final = _FakeFilesRouter({"gpt-a": ("wildcard-dep",), "gpt-b": ("wildcard-dep",)})
|
||||
upload_metadata: Final = {"user_api_key_team_id": "test-team"}
|
||||
managed_files, _ = _make_real_managed_files_instance()
|
||||
|
||||
await _upload(managed_files, router, ["gpt-a", "gpt-b"], upload_metadata)
|
||||
await drain_logging_tasks(timeout=1)
|
||||
|
||||
assert upload_metadata == {"user_api_key_team_id": "test-team"}
|
||||
assert [delete[2] for delete in router.deletes] == [
|
||||
{"user_api_key_team_id": "test-team", "model_info": {"id": "wildcard-dep"}}
|
||||
]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue