mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-26 06:03:19 +00:00
fbc2e358bc
* fix(responses): bill a background response once, from the poll job
Reading a completed background response with GET /v1/responses/{id} wrote a
full-cost spend row on every read, on top of the row the background cost poll
writes. Reads of a Responses object (GET, DELETE, cancel) now never bill, and
the poll job is the single biller.
The poll job now bills with the creating request's attribution (key hash, team,
org, user, aliases, tags) read off the managed object row, through a helper it
shares with the batch cost poll. The background create persists the creating
key and request tags so that attribution survives until the poll runs.
* test(proxy): a background response read reports zero cost
* fix(responses): bill a background response once, on the first read or poll that sees it finished
A conditional update moves the stored row out of queued or in_progress so only one worker or pod bills it, through one re-read stamped with the creator's attribution. A failed billing read puts the row back for the next poll
* fix(responses): keep the caller's response id out of the billing failure warning
* fix(responses): keep the creating client's User-Agent tags on the one background charge
* fix(responses): keep managed background reads free and bill each response once
The managed files hook marks reads and cancels of proxy-issued response ids, so only those stay free. SDK and database-less reads of a background response still bill, extra_body background creates are tracked, and the poll's status read no longer cools down a deployment on a 404.
* fix(responses): bill finished background responses from the fetched body and store rows Postgres accepts
* fix(responses): give the background billing event its own spend log row and bill responses the create left untracked
* fix(responses): narrow the managed files hook with a protocol instead of an enterprise import
* fix(responses): leave a read the caller retargets with client-side credentials to the cost poll
* fix(responses): bill background responses on the creator's end user and report the charge on the settling read
The read that settles the bill returns the charge in x-litellm-response-cost, the billed event
carries the creator's end user and the key, team, and user budget fields, and its spend row is
named after the id the client got. Override reads, cancels, and deletes settle from the
configured deployment, a streamed create marks its row billed, stored rows keep no model
output, and model insights counts the billed row's spend without a request.
* test(responses): cover background response billing across proxies, reads and chaos
* fix(responses): settle background billing only for callers allowed to address the response
A delete now settles the bill only after the response id access check passes, a streamed
create is marked billed when it sets background inside extra_body or its completed event
arrives after the last chunk, the billing event keeps the deployment's litellm_params
base_model, a key with no user and no team bills no user instead of its owner scope, and the
enterprise cost check import is optional
* fix(responses): record the bound on the pending status IN filters
* fix(responses): send the billed cost split on the settling read and keep the batch cost event as it was
The read that settles a background response now sends the whole cost header
family from the billed event, so the breakdown headers add up to the charge.
Batch cost events go back to main's attribution (creator id as the user, no
budget fields); only background responses get the end user and budgets.
Integration tests expect the charge on the settling read (or none when the
poll billed first), one charge for a streamed create, and SDK flows that poll
with the id each read returns.
* fix(responses): bill each background response once across upgrades, streams, and failed settles
Rows stored before this billing model are found by the issued id and rekeyed by the
provider id, and the cost poll retires one a read already billed. A streamed background
create is free on its own and bills the finished response through the same managed row
before the client sees response.completed, so a racing read cannot bill it again. A
settle that fails retries in the background, a retargeted read with no row tracks and
bills from the configured deployment, and a delete keeps its billed cost header.
Reverts the unrelated json.dumps narrowing in the completion transformation.
* fix(responses): bill cache-polled background streams and match poll rows by the issued id
* test(responses): cover deletes, failed settles, budgets, base models, and cache-polled background billing
Adds integration cells for the background response billing paths the audit
inventory lists: a DELETE that settles the bill and refuses a stranger, a
settle the database refuses that retries once it recovers, the Prometheus
remaining budget, base_model pricing from model_info and litellm_params, the
daily model usage rollup, the end user, the settling read's cost split, stored
rows without model output, a forged managed marker on the create, streamed
creates with background in extra_body, and the cache-polled create
---------
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
2146 lines
77 KiB
Python
2146 lines
77 KiB
Python
"""
|
|
Tests for enterprise/litellm_enterprise/proxy/hooks/managed_files.py
|
|
|
|
Regression test for afile_retrieve called without credentials in
|
|
async_post_call_success_hook when processing completed batch responses.
|
|
"""
|
|
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
import logging
|
|
from types import MappingProxyType
|
|
|
|
import pytest
|
|
from typing import Final, Optional
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth
|
|
from litellm.proxy.openai_files_endpoints.common_utils import BATCH_CREATE_HIDDEN_PARAM
|
|
from litellm.types.llms.openai import FileListPage, OpenAIFileObject
|
|
from litellm.types.utils import LiteLLMBatch
|
|
|
|
|
|
def _make_file_object(file_id: str = "file-output-abc") -> OpenAIFileObject:
|
|
return OpenAIFileObject(
|
|
id=file_id,
|
|
bytes=100,
|
|
created_at=1700000000,
|
|
filename="output.jsonl",
|
|
object="file",
|
|
purpose="batch_output",
|
|
status="processed",
|
|
)
|
|
|
|
|
|
def _make_batch_response(
|
|
batch_id: str = "batch-123",
|
|
output_file_id: Optional[str] = "file-output-abc",
|
|
status: str = "completed",
|
|
model_id: str = "model-deploy-xyz",
|
|
model_name: str = "azure/gpt-4",
|
|
) -> LiteLLMBatch:
|
|
"""Create a LiteLLMBatch response with hidden params set as the router would."""
|
|
batch = LiteLLMBatch(
|
|
id=batch_id,
|
|
completion_window="24h",
|
|
created_at=1700000000,
|
|
endpoint="/v1/chat/completions",
|
|
input_file_id="file-input-abc",
|
|
object="batch",
|
|
status=status,
|
|
output_file_id=output_file_id,
|
|
)
|
|
batch._hidden_params = {
|
|
"unified_file_id": "some-unified-id",
|
|
"unified_batch_id": "some-unified-batch-id",
|
|
"model_id": model_id,
|
|
"model_name": model_name,
|
|
}
|
|
return batch
|
|
|
|
|
|
def _make_user_api_key_dict() -> UserAPIKeyAuth:
|
|
return UserAPIKeyAuth(
|
|
api_key="sk-test",
|
|
user_id="test-user",
|
|
parent_otel_span=None,
|
|
)
|
|
|
|
|
|
def _make_team_member_api_key_dict() -> UserAPIKeyAuth:
|
|
"""The shape most real virtual keys carry: a user_id and a team_id."""
|
|
return UserAPIKeyAuth(
|
|
api_key="sk-test",
|
|
user_id="test-user",
|
|
team_id="test-team",
|
|
parent_otel_span=None,
|
|
)
|
|
|
|
|
|
def _make_service_account_api_key_dict() -> UserAPIKeyAuth:
|
|
return UserAPIKeyAuth(
|
|
api_key="sk-service",
|
|
team_id="test-team",
|
|
parent_otel_span=None,
|
|
)
|
|
|
|
|
|
def _make_admin_api_key_dict() -> UserAPIKeyAuth:
|
|
return UserAPIKeyAuth(
|
|
api_key="sk-admin",
|
|
user_id="admin-user",
|
|
user_role=LitellmUserRoles.PROXY_ADMIN,
|
|
parent_otel_span=None,
|
|
)
|
|
|
|
|
|
def _make_managed_file_row(
|
|
unified_file_id: str,
|
|
purpose: str = "batch_output",
|
|
created_by: str = "test-user",
|
|
team_id: Optional[str] = None,
|
|
) -> MagicMock:
|
|
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(),
|
|
created_by=created_by,
|
|
team_id=team_id,
|
|
)
|
|
|
|
|
|
def _make_unparseable_managed_file_row(
|
|
unified_file_id: str,
|
|
created_by: str = "test-user",
|
|
team_id: Optional[str] = None,
|
|
) -> MagicMock:
|
|
"""A row whose stored blob cannot be parsed back into a file object."""
|
|
return MagicMock(
|
|
unified_file_id=unified_file_id,
|
|
file_object=None,
|
|
created_by=created_by,
|
|
team_id=team_id,
|
|
)
|
|
|
|
|
|
def _row_matches_where(row, where) -> bool:
|
|
"""Apply the Prisma ``where`` shapes build_owner_filter actually emits:
|
|
``{}``, a single equality, and the ``OR`` of equalities a key carrying
|
|
both a user_id and a team_id produces."""
|
|
for field, expected in where.items():
|
|
if field == "OR":
|
|
if not any(_row_matches_where(row, clause) for clause in expected):
|
|
return False
|
|
elif getattr(row, field) != expected:
|
|
return False
|
|
return True
|
|
|
|
|
|
class _FakeManagedFileTable:
|
|
"""In-memory stand-in for the managed file table, newest row first."""
|
|
|
|
def __init__(self, rows):
|
|
self.rows = list(rows)
|
|
self.find_many_calls = []
|
|
self.find_first_calls = []
|
|
|
|
def _owned_rows(self, where):
|
|
return [row for row in self.rows if _row_matches_where(row, where)]
|
|
|
|
async def find_first(self, where):
|
|
self.find_first_calls.append(where)
|
|
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}
|
|
)
|
|
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"]
|
|
)
|
|
rows = rows[start + skip :]
|
|
return rows if take is None else rows[:take]
|
|
|
|
|
|
def _make_managed_files_over_rows(rows):
|
|
managed_files = _make_managed_files_instance()
|
|
table = _FakeManagedFileTable(rows)
|
|
managed_files.prisma_client.db.litellm_managedfiletable = table
|
|
return managed_files, table
|
|
|
|
|
|
def _make_managed_files_instance():
|
|
"""Create a PROXY_LiteLLMManagedFiles with storage methods mocked out."""
|
|
from litellm_enterprise.proxy.hooks.managed_files import (
|
|
PROXY_LiteLLMManagedFiles,
|
|
)
|
|
|
|
mock_cache = MagicMock()
|
|
mock_cache.async_get_cache = AsyncMock(return_value=None)
|
|
mock_prisma = MagicMock()
|
|
mock_prisma.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None)
|
|
|
|
instance = PROXY_LiteLLMManagedFiles(
|
|
internal_usage_cache=mock_cache,
|
|
prisma_client=mock_prisma,
|
|
)
|
|
instance.store_unified_file_id = AsyncMock()
|
|
instance.store_unified_object_id = AsyncMock()
|
|
return instance
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_pass_credentials_to_afile_retrieve():
|
|
"""
|
|
When async_post_call_success_hook processes a completed batch with an output_file_id,
|
|
it calls afile_retrieve to fetch file metadata. It must pass credentials from the
|
|
router deployment, not just custom_llm_provider and file_id.
|
|
|
|
Regression test for: managed_files.py:919 calling afile_retrieve without api_key/api_base.
|
|
"""
|
|
managed_files = _make_managed_files_instance()
|
|
batch_response = _make_batch_response(
|
|
model_id="model-deploy-xyz",
|
|
model_name="azure/gpt-4",
|
|
output_file_id="file-output-abc",
|
|
)
|
|
user_api_key_dict = _make_user_api_key_dict()
|
|
|
|
mock_credentials = {
|
|
"api_key": "test-azure-key",
|
|
"api_base": "https://my-azure.openai.azure.com/",
|
|
"api_version": "2025-03-01-preview",
|
|
"custom_llm_provider": "azure",
|
|
}
|
|
|
|
mock_router = MagicMock()
|
|
mock_router.get_deployment_credentials_with_provider = MagicMock(
|
|
return_value=mock_credentials
|
|
)
|
|
|
|
mock_afile_retrieve = AsyncMock(return_value=_make_file_object("file-output-abc"))
|
|
|
|
with (
|
|
patch("litellm.afile_retrieve", mock_afile_retrieve),
|
|
patch("litellm.proxy.proxy_server.llm_router", mock_router),
|
|
):
|
|
await managed_files.async_post_call_success_hook(
|
|
data={},
|
|
user_api_key_dict=user_api_key_dict,
|
|
response=batch_response,
|
|
)
|
|
|
|
mock_afile_retrieve.assert_called()
|
|
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}"
|
|
)
|
|
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}"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_user_created_file_ids_skips_rows_without_file_object():
|
|
managed_files = _make_managed_files_instance()
|
|
managed_files.prisma_client.db.litellm_managedfiletable.find_many = AsyncMock(
|
|
return_value=[
|
|
MagicMock(
|
|
file_object=_make_file_object().model_dump(),
|
|
unified_file_id="unified-id-1",
|
|
),
|
|
MagicMock(file_object=None, unified_file_id="unified-id-2"),
|
|
]
|
|
)
|
|
|
|
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"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_user_created_file_ids_remaps_stored_raw_provider_id_to_unified_id():
|
|
"""
|
|
Rows registered from batch outputs store the provider's file object, whose
|
|
id is the raw provider id (e.g. file-abc). Listing must return the row's
|
|
unified_file_id so callers get ids that work on the managed routes.
|
|
|
|
Regression test for https://github.com/BerriAI/litellm/issues/35362.
|
|
"""
|
|
unified_id = "bGl0ZWxsbV9wcm94eTt1bmlmaWVkX2lkLGRlYWRiZWVm"
|
|
raw_provider_object = _make_file_object("file-raw-provider-123")
|
|
managed_files = _make_managed_files_instance()
|
|
managed_files.prisma_client.db.litellm_managedfiletable.find_many = AsyncMock(
|
|
return_value=[
|
|
MagicMock(
|
|
file_object=raw_provider_object.model_dump(),
|
|
unified_file_id=unified_id,
|
|
),
|
|
]
|
|
)
|
|
|
|
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
|
|
assert files[0].purpose == raw_provider_object.purpose
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
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",
|
|
flat_model_file_ids=["file-provider-1", "file-provider-2"],
|
|
)
|
|
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(),
|
|
)
|
|
)
|
|
|
|
assert unified_file_ids == {
|
|
"file-provider-1": "unified-file-id",
|
|
"file-provider-2": "unified-file-id",
|
|
}
|
|
assert isinstance(unified_file_ids, MappingProxyType)
|
|
find_many.assert_awaited_once_with(
|
|
where={
|
|
"OR": [{"created_by": "test-user"}, {"team_id": "test-team"}],
|
|
"flat_model_file_ids": {
|
|
"hasSome": ["file-provider-1", "file-provider-2", "file-unmanaged-2"],
|
|
},
|
|
}
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
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
|
|
no_owner: Final = UserAPIKeyAuth(
|
|
api_key=None,
|
|
token=None,
|
|
user_id=None,
|
|
team_id=None,
|
|
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,
|
|
)
|
|
)
|
|
|
|
assert unified_file_ids == {}
|
|
assert isinstance(unified_file_ids, MappingProxyType)
|
|
find_many.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_provider_file_id_resolver_skips_database_query_for_empty_input() -> None:
|
|
managed_files: Final = _make_managed_files_instance()
|
|
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(),
|
|
)
|
|
)
|
|
|
|
assert unified_file_ids == {}
|
|
assert isinstance(unified_file_ids, MappingProxyType)
|
|
find_many.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_list_returns_owner_scoped_managed_files():
|
|
managed_files = _make_managed_files_instance()
|
|
managed_files.prisma_client.db.litellm_managedfiletable.find_many = AsyncMock(
|
|
return_value=[
|
|
MagicMock(
|
|
file_object=_make_file_object("file-provider-id").model_dump(),
|
|
unified_file_id="unified-file-id",
|
|
),
|
|
MagicMock(
|
|
file_object=_make_file_object("file-other-purpose").model_copy(
|
|
update={"purpose": "batch"}
|
|
).model_dump(),
|
|
unified_file_id="unified-other-purpose",
|
|
),
|
|
]
|
|
)
|
|
|
|
response = await managed_files.afile_list(
|
|
purpose="batch_output",
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=_make_user_api_key_dict(),
|
|
)
|
|
|
|
managed_files.prisma_client.db.litellm_managedfiletable.find_many.assert_awaited_once_with(
|
|
where={"created_by": "test-user"},
|
|
take=10001,
|
|
order=[{"created_at": "desc"}, {"unified_file_id": "desc"}],
|
|
)
|
|
assert [file.id for file in response.data] == ["unified-file-id"]
|
|
assert response.first_id == "unified-file-id"
|
|
assert response.last_id == "unified-file-id"
|
|
assert response.has_more is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_list_returns_a_page_object_callbacks_can_read():
|
|
"""Post-call hooks receive the listing and read ``.data`` off it, the way the
|
|
provider SDK's page lets them. The body on the wire stays a plain list page."""
|
|
from fastapi.encoders import jsonable_encoder
|
|
|
|
managed_files, _ = _make_managed_files_over_rows([_make_managed_file_row("unified-file-id")])
|
|
|
|
page = await managed_files.afile_list(
|
|
purpose=None,
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=_make_user_api_key_dict(),
|
|
)
|
|
|
|
assert isinstance(page, FileListPage)
|
|
assert [file.id for file in page.data] == ["unified-file-id"]
|
|
|
|
body = jsonable_encoder(page)
|
|
assert list(body) == ["object", "data", "first_id", "last_id", "has_more"]
|
|
assert body["object"] == "list"
|
|
assert [file["id"] for file in body["data"]] == ["unified-file-id"]
|
|
assert body["first_id"] == "unified-file-id"
|
|
assert body["last_id"] == "unified-file-id"
|
|
assert body["has_more"] is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("purpose", ["nonexistent_purpose", "EVALS", "batch "])
|
|
async def test_afile_list_rejects_a_purpose_the_files_api_never_accepts(purpose):
|
|
"""No stored file can carry an undocumented purpose, so filtering on one is a
|
|
bad request rather than a legitimately empty page."""
|
|
managed_files, table = _make_managed_files_over_rows([_make_managed_file_row("unified-file-id")])
|
|
|
|
with pytest.raises(ProxyException) as exc_info:
|
|
await managed_files.afile_list(
|
|
purpose=purpose,
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=_make_user_api_key_dict(),
|
|
)
|
|
|
|
assert exc_info.value.code == "400"
|
|
assert exc_info.value.type == "invalid_request_error"
|
|
assert exc_info.value.param == "purpose"
|
|
assert table.find_many_calls == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("purpose", ["batch", "assistants", "fine-tune", "evals", None])
|
|
async def test_afile_list_accepts_every_documented_purpose(purpose):
|
|
managed_files, _ = _make_managed_files_over_rows([_make_managed_file_row("unified-file-id")])
|
|
|
|
page = await managed_files.afile_list(
|
|
purpose=purpose,
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=_make_user_api_key_dict(),
|
|
)
|
|
|
|
assert isinstance(page, FileListPage)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_list_does_not_leak_another_callers_files():
|
|
managed_files, table = _make_managed_files_over_rows(
|
|
[
|
|
_make_managed_file_row("unified-mine-2"),
|
|
_make_managed_file_row("unified-theirs", created_by="other-user"),
|
|
_make_managed_file_row("unified-mine-1"),
|
|
]
|
|
)
|
|
|
|
response = await managed_files.afile_list(
|
|
purpose=None,
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=_make_user_api_key_dict(),
|
|
)
|
|
|
|
assert [file.id for file in response.data] == ["unified-mine-2", "unified-mine-1"]
|
|
assert table.find_many_calls[0]["where"] == {"created_by": "test-user"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_list_returns_own_and_team_files_for_a_key_carrying_both_ids():
|
|
managed_files, table = _make_managed_files_over_rows(
|
|
[
|
|
_make_managed_file_row("unified-mine"),
|
|
_make_managed_file_row("unified-teammates", created_by="other-user", team_id="test-team"),
|
|
_make_managed_file_row("unified-outsiders", created_by="outsider", team_id="other-team"),
|
|
]
|
|
)
|
|
|
|
response = await managed_files.afile_list(
|
|
purpose=None,
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=_make_team_member_api_key_dict(),
|
|
)
|
|
|
|
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"}]
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_list_scopes_a_service_account_key_to_its_team():
|
|
managed_files, table = _make_managed_files_over_rows(
|
|
[
|
|
_make_managed_file_row("unified-teams", created_by="other-user", team_id="test-team"),
|
|
_make_managed_file_row("unified-outsiders", created_by="outsider", team_id="other-team"),
|
|
]
|
|
)
|
|
|
|
response = await managed_files.afile_list(
|
|
purpose=None,
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=_make_service_account_api_key_dict(),
|
|
)
|
|
|
|
assert [file.id for file in response.data] == ["unified-teams"]
|
|
assert table.find_many_calls[0]["where"] == {"team_id": "test-team"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_list_returns_every_callers_files_for_a_proxy_admin():
|
|
managed_files, table = _make_managed_files_over_rows(
|
|
[
|
|
_make_managed_file_row("unified-mine"),
|
|
_make_managed_file_row("unified-theirs", created_by="other-user", team_id="other-team"),
|
|
]
|
|
)
|
|
|
|
response = await managed_files.afile_list(
|
|
purpose=None,
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=_make_admin_api_key_dict(),
|
|
)
|
|
|
|
assert [file.id for file in response.data] == ["unified-mine", "unified-theirs"]
|
|
assert table.find_many_calls[0]["where"] == {}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_list_pages_a_team_key_across_both_halves_of_its_filter():
|
|
"""Keyset pagination has to walk an OR filter as one ordered set, without
|
|
repeating a row across pages or dropping one between them."""
|
|
managed_files, table = _make_managed_files_over_rows(
|
|
[
|
|
_make_managed_file_row("unified-0"),
|
|
_make_managed_file_row("unified-1", created_by="other-user", team_id="test-team"),
|
|
_make_managed_file_row("unified-2"),
|
|
_make_managed_file_row("unified-3", created_by="outsider", team_id="other-team"),
|
|
_make_managed_file_row("unified-4", created_by="other-user", team_id="test-team"),
|
|
]
|
|
)
|
|
user_api_key_dict = _make_team_member_api_key_dict()
|
|
|
|
seen = []
|
|
cursor = None
|
|
for _ in range(4):
|
|
response = await managed_files.afile_list(
|
|
purpose=None,
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=user_api_key_dict,
|
|
limit=2,
|
|
after=cursor,
|
|
)
|
|
seen.extend(file.id for file in response.data)
|
|
if not response.has_more:
|
|
break
|
|
cursor = response.last_id
|
|
|
|
assert seen == ["unified-0", "unified-1", "unified-2", "unified-4"]
|
|
assert all(
|
|
call["where"] == {"OR": [{"created_by": "test-user"}, {"team_id": "test-team"}]}
|
|
for call in table.find_many_calls
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_list_orders_newest_first_and_breaks_ties_on_the_cursor_column():
|
|
managed_files, table = _make_managed_files_over_rows([_make_managed_file_row("unified-mine")])
|
|
|
|
await managed_files.afile_list(
|
|
purpose=None,
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=_make_user_api_key_dict(),
|
|
)
|
|
|
|
assert table.find_many_calls[0]["order"] == [
|
|
{"created_at": "desc"},
|
|
{"unified_file_id": "desc"},
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_list_scopes_a_keyless_key_to_its_own_hashed_token():
|
|
caller = UserAPIKeyAuth(api_key="sk-test", parent_otel_span=None)
|
|
managed_files, table = _make_managed_files_over_rows(
|
|
[
|
|
_make_managed_file_row("unified-mine", created_by=f"key:{caller.token}"),
|
|
_make_managed_file_row("unified-theirs", created_by="other-user"),
|
|
]
|
|
)
|
|
|
|
response = await managed_files.afile_list(
|
|
purpose=None,
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=caller,
|
|
)
|
|
|
|
assert [file.id for file in response.data] == ["unified-mine"]
|
|
assert table.find_many_calls[0]["where"] == {"created_by": f"key:{caller.token}"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_list_denies_a_caller_with_no_identity_at_all():
|
|
managed_files, table = _make_managed_files_over_rows([_make_managed_file_row("unified-mine")])
|
|
|
|
response = await managed_files.afile_list(
|
|
purpose=None,
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=UserAPIKeyAuth(parent_otel_span=None),
|
|
)
|
|
|
|
assert response.data == []
|
|
assert response.has_more is False
|
|
assert table.find_many_calls == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_list_filters_by_purpose():
|
|
managed_files, _ = _make_managed_files_over_rows(
|
|
[
|
|
_make_managed_file_row("unified-batch-output"),
|
|
_make_managed_file_row("unified-batch", purpose="batch"),
|
|
]
|
|
)
|
|
|
|
response = await managed_files.afile_list(
|
|
purpose="batch",
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=_make_user_api_key_dict(),
|
|
)
|
|
|
|
assert [file.id for file in response.data] == ["unified-batch"]
|
|
|
|
|
|
async def _walk_afile_list(managed_files, user_api_key_dict, purpose, limit):
|
|
"""Page through the listing the way the official SDK does, off ``data[-1].id``."""
|
|
seen = []
|
|
after = None
|
|
while True:
|
|
page = await managed_files.afile_list(
|
|
purpose=purpose,
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=user_api_key_dict,
|
|
limit=limit,
|
|
after=after,
|
|
)
|
|
page_ids = [file.id for file in page.data]
|
|
assert not set(page_ids) & set(seen)
|
|
seen.extend(page_ids)
|
|
if not page.has_more:
|
|
return seen
|
|
assert page_ids, "an SDK stops paging on an empty page, so has_more must never ride one"
|
|
after = page_ids[-1]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_list_fills_a_page_past_rows_the_purpose_filter_drops():
|
|
"""The newest rows do not match, so the page must reach past them rather than come back empty."""
|
|
managed_files, _ = _make_managed_files_over_rows(
|
|
[
|
|
_make_managed_file_row("unified-0"),
|
|
_make_managed_file_row("unified-1"),
|
|
_make_managed_file_row("unified-2", purpose="batch"),
|
|
_make_managed_file_row("unified-3"),
|
|
_make_managed_file_row("unified-4", purpose="batch"),
|
|
]
|
|
)
|
|
user_api_key_dict = _make_user_api_key_dict()
|
|
|
|
first_page = await managed_files.afile_list(
|
|
purpose="batch",
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=user_api_key_dict,
|
|
limit=1,
|
|
)
|
|
|
|
assert [file.id for file in first_page.data] == ["unified-2"]
|
|
assert first_page.has_more is True
|
|
assert first_page.last_id == "unified-2"
|
|
|
|
second_page = await managed_files.afile_list(
|
|
purpose="batch",
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=user_api_key_dict,
|
|
limit=1,
|
|
after=first_page.last_id,
|
|
)
|
|
|
|
assert [file.id for file in second_page.data] == ["unified-4"]
|
|
assert second_page.has_more is False
|
|
|
|
|
|
@pytest.mark.parametrize("limit", [1, 2, 3])
|
|
@pytest.mark.asyncio
|
|
async def test_afile_list_walks_every_purpose_match_at_any_limit(limit):
|
|
managed_files, _ = _make_managed_files_over_rows(
|
|
[
|
|
_make_managed_file_row("unified-0"),
|
|
_make_managed_file_row("unified-1"),
|
|
_make_managed_file_row("unified-2", purpose="batch"),
|
|
_make_managed_file_row("unified-3"),
|
|
_make_managed_file_row("unified-4", purpose="batch"),
|
|
_make_managed_file_row("unified-5", purpose="batch"),
|
|
_make_managed_file_row("unified-6"),
|
|
]
|
|
)
|
|
|
|
seen = await _walk_afile_list(managed_files, _make_user_api_key_dict(), "batch", limit)
|
|
|
|
assert seen == ["unified-2", "unified-4", "unified-5"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_list_fills_a_page_past_rows_that_do_not_parse():
|
|
managed_files, _ = _make_managed_files_over_rows(
|
|
[
|
|
_make_unparseable_managed_file_row("unified-0"),
|
|
_make_unparseable_managed_file_row("unified-1"),
|
|
_make_managed_file_row("unified-2"),
|
|
]
|
|
)
|
|
|
|
page = await managed_files.afile_list(
|
|
purpose=None,
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=_make_user_api_key_dict(),
|
|
limit=1,
|
|
)
|
|
|
|
assert [file.id for file in page.data] == ["unified-2"]
|
|
assert page.has_more is False
|
|
|
|
|
|
_DEEP_SCAN_ROW_COUNT = 2000
|
|
_DEEP_SCAN_QUERY_BUDGET = 10
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_list_bounds_the_queries_a_deep_purpose_match_costs():
|
|
"""A tiny limit over rows the filter drops must not turn one request into thousands of queries."""
|
|
managed_files, table = _make_managed_files_over_rows(
|
|
[_make_managed_file_row(f"unified-{index:05d}") for index in range(_DEEP_SCAN_ROW_COUNT)]
|
|
+ [_make_managed_file_row("unified-match", purpose="batch")]
|
|
)
|
|
|
|
page = await managed_files.afile_list(
|
|
purpose="batch",
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=_make_user_api_key_dict(),
|
|
limit=1,
|
|
)
|
|
|
|
assert [file.id for file in page.data] == ["unified-match"]
|
|
assert page.has_more is False
|
|
assert len(table.find_many_calls) <= _DEEP_SCAN_QUERY_BUDGET
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_list_bounds_the_queries_a_deep_unparseable_run_costs():
|
|
"""Rows that will not parse drop out like a filter does, so they get the same bound."""
|
|
managed_files, table = _make_managed_files_over_rows(
|
|
[_make_unparseable_managed_file_row(f"unified-{index:05d}") for index in range(_DEEP_SCAN_ROW_COUNT)]
|
|
+ [_make_managed_file_row("unified-parses")]
|
|
)
|
|
|
|
page = await managed_files.afile_list(
|
|
purpose=None,
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=_make_user_api_key_dict(),
|
|
limit=1,
|
|
)
|
|
|
|
assert [file.id for file in page.data] == ["unified-parses"]
|
|
assert page.has_more is False
|
|
assert len(table.find_many_calls) <= _DEEP_SCAN_QUERY_BUDGET
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_list_reads_one_chunk_when_the_first_one_fills_the_page():
|
|
"""The widened chunk must stay off the common path, where the newest rows already fill the page."""
|
|
managed_files, table = _make_managed_files_over_rows(
|
|
[_make_managed_file_row(f"unified-{index:05d}") for index in range(_DEEP_SCAN_ROW_COUNT)]
|
|
)
|
|
|
|
page = await managed_files.afile_list(
|
|
purpose=None,
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=_make_user_api_key_dict(),
|
|
limit=2,
|
|
)
|
|
|
|
assert [file.id for file in page.data] == ["unified-00000", "unified-00001"]
|
|
assert page.has_more is True
|
|
assert [call["take"] for call in table.find_many_calls] == [3]
|
|
|
|
|
|
@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)]
|
|
)
|
|
|
|
page = await managed_files.afile_list(
|
|
purpose="batch",
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=_make_user_api_key_dict(),
|
|
limit=2,
|
|
)
|
|
|
|
assert page.data == []
|
|
assert page.has_more is False
|
|
assert page.first_id is None
|
|
assert page.last_id is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_list_honors_limit_and_reports_more_pages():
|
|
managed_files, table = _make_managed_files_over_rows(
|
|
[_make_managed_file_row(f"unified-{index}") for index in range(5)]
|
|
)
|
|
|
|
response = await managed_files.afile_list(
|
|
purpose=None,
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=_make_user_api_key_dict(),
|
|
limit=2,
|
|
)
|
|
|
|
assert [file.id for file in response.data] == ["unified-0", "unified-1"]
|
|
assert response.has_more is True
|
|
assert table.find_many_calls[0]["take"] == 3
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_list_pages_through_every_file_without_overlap():
|
|
managed_files, table = _make_managed_files_over_rows(
|
|
[_make_managed_file_row(f"unified-{index}") for index in range(5)]
|
|
)
|
|
user_api_key_dict = _make_user_api_key_dict()
|
|
|
|
seen = []
|
|
after = None
|
|
while True:
|
|
page = await managed_files.afile_list(
|
|
purpose=None,
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=user_api_key_dict,
|
|
limit=2,
|
|
after=after,
|
|
)
|
|
page_ids = [file.id for file in page.data]
|
|
assert not set(page_ids) & set(seen)
|
|
seen.extend(page_ids)
|
|
if not page.has_more:
|
|
break
|
|
after = page.last_id
|
|
|
|
assert seen == [f"unified-{index}" for index in range(5)]
|
|
assert table.find_many_calls[1]["cursor"] == {"unified_file_id": "unified-1"}
|
|
assert table.find_many_calls[1]["skip"] == 1
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"unknown_cursor",
|
|
["unified-theirs", "unified-nowhere"],
|
|
ids=["another-users-file", "no-such-file"],
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_afile_list_rejects_an_after_cursor_outside_the_callers_files(unknown_cursor):
|
|
from litellm.proxy._types import ProxyException
|
|
|
|
managed_files, table = _make_managed_files_over_rows(
|
|
[
|
|
_make_managed_file_row("unified-mine"),
|
|
_make_managed_file_row("unified-theirs", created_by="other-user"),
|
|
]
|
|
)
|
|
|
|
with pytest.raises(ProxyException) as exc_info:
|
|
await managed_files.afile_list(
|
|
purpose=None,
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=_make_user_api_key_dict(),
|
|
after=unknown_cursor,
|
|
)
|
|
|
|
assert exc_info.value.code == "400"
|
|
assert exc_info.value.type == "invalid_request_error"
|
|
assert exc_info.value.param == "after"
|
|
assert exc_info.value.message == f"Invalid 'after' cursor: no file found with id '{unknown_cursor}'."
|
|
assert table.find_first_calls[0] == {
|
|
"created_by": "test-user",
|
|
"unified_file_id": unknown_cursor,
|
|
}
|
|
assert table.find_many_calls == []
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"limit, bound, expected_range",
|
|
[
|
|
(0, "below minimum", ">= 1"),
|
|
(-1, "below minimum", ">= 1"),
|
|
(10001, "above maximum", "<= 10000"),
|
|
],
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_afile_list_rejects_a_limit_outside_the_openai_range(limit, bound, expected_range):
|
|
from litellm.proxy._types import ProxyException
|
|
|
|
managed_files, table = _make_managed_files_over_rows([_make_managed_file_row("unified-mine")])
|
|
|
|
with pytest.raises(ProxyException) as exc_info:
|
|
await managed_files.afile_list(
|
|
purpose=None,
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=_make_user_api_key_dict(),
|
|
limit=limit,
|
|
)
|
|
|
|
assert exc_info.value.code == "400"
|
|
assert exc_info.value.type == "invalid_request_error"
|
|
assert exc_info.value.param == "limit"
|
|
assert exc_info.value.message == (
|
|
f"Invalid 'limit': integer {bound} value. Expected a value {expected_range}, but got {limit} instead."
|
|
)
|
|
assert table.find_many_calls == []
|
|
|
|
|
|
@pytest.mark.parametrize("limit", [1, 10000])
|
|
@pytest.mark.asyncio
|
|
async def test_afile_list_accepts_the_ends_of_the_openai_limit_range(limit):
|
|
managed_files, table = _make_managed_files_over_rows([_make_managed_file_row("unified-mine")])
|
|
|
|
response = await managed_files.afile_list(
|
|
purpose=None,
|
|
litellm_parent_otel_span=None,
|
|
user_api_key_dict=_make_user_api_key_dict(),
|
|
limit=limit,
|
|
)
|
|
|
|
assert [file.id for file in response.data] == ["unified-mine"]
|
|
assert response.has_more is False
|
|
assert table.find_many_calls[0]["take"] == limit + 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parse_managed_file_object_warning_omits_rejected_values(caplog):
|
|
from litellm_enterprise.proxy.hooks.managed_files import (
|
|
_parse_managed_file_object,
|
|
)
|
|
|
|
with caplog.at_level(logging.WARNING):
|
|
parsed = _parse_managed_file_object(
|
|
{"id": "file-corrupt", "object": "file", "filename": "confidential.jsonl"},
|
|
"unified-corrupt",
|
|
)
|
|
|
|
assert parsed is None
|
|
assert "unified-corrupt" in caplog.text
|
|
assert "bytes" in caplog.text
|
|
assert "confidential.jsonl" not in caplog.text
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_user_created_file_ids_skips_unparseable_rows():
|
|
managed_files = _make_managed_files_instance()
|
|
managed_files.prisma_client.db.litellm_managedfiletable.find_many = AsyncMock(
|
|
return_value=[
|
|
MagicMock(
|
|
file_object={"id": "file-corrupt", "object": "file"},
|
|
unified_file_id="unified-corrupt",
|
|
),
|
|
MagicMock(
|
|
file_object=_make_file_object().model_dump(),
|
|
unified_file_id="unified-valid",
|
|
),
|
|
]
|
|
)
|
|
|
|
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"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_fallback_when_no_router():
|
|
"""
|
|
When llm_router is not available, afile_retrieve should still be called
|
|
with the fallback behavior (custom_llm_provider extracted from model_name).
|
|
"""
|
|
managed_files = _make_managed_files_instance()
|
|
batch_response = _make_batch_response(
|
|
model_id="model-deploy-xyz",
|
|
model_name="azure/gpt-4",
|
|
output_file_id="file-output-abc",
|
|
)
|
|
user_api_key_dict = _make_user_api_key_dict()
|
|
|
|
mock_afile_retrieve = AsyncMock(return_value=_make_file_object("file-output-abc"))
|
|
|
|
with (
|
|
patch("litellm.afile_retrieve", mock_afile_retrieve),
|
|
patch("litellm.proxy.proxy_server.llm_router", None),
|
|
):
|
|
await managed_files.async_post_call_success_hook(
|
|
data={},
|
|
user_api_key_dict=user_api_key_dict,
|
|
response=batch_response,
|
|
)
|
|
|
|
mock_afile_retrieve.assert_called()
|
|
call_kwargs = mock_afile_retrieve.call_args
|
|
assert call_kwargs.kwargs.get("custom_llm_provider") == "azure"
|
|
assert call_kwargs.kwargs.get("file_id") == "file-output-abc"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_not_double_wrap_already_unified_output_file_id():
|
|
"""After ensure_batch_response_managed_file_ids, retrieve must not re-wrap
|
|
output_file_id or store a nested unified id as the provider mapping."""
|
|
import base64
|
|
|
|
managed_files = _make_managed_files_instance()
|
|
provider_file_id = "file-WXWt9R4LzmU5WpeKzjCfLR"
|
|
model_id = "openai/openai/gpt-5.5-batch"
|
|
already_unified = managed_files.get_unified_output_file_id(
|
|
output_file_id=provider_file_id,
|
|
model_id=model_id,
|
|
model_name="openai/openai/gpt-5.5-batch",
|
|
)
|
|
|
|
batch_response = _make_batch_response(
|
|
model_id=model_id,
|
|
model_name="openai/openai/gpt-5.5-batch",
|
|
output_file_id=already_unified,
|
|
)
|
|
user_api_key_dict = _make_user_api_key_dict()
|
|
|
|
mock_credentials = {
|
|
"api_key": "test-key",
|
|
"api_base": "https://api.openai.com/v1",
|
|
"custom_llm_provider": "openai",
|
|
}
|
|
mock_router = MagicMock()
|
|
mock_router.get_deployment_credentials_with_provider = MagicMock(
|
|
return_value=mock_credentials
|
|
)
|
|
mock_afile_retrieve = AsyncMock(return_value=_make_file_object(provider_file_id))
|
|
|
|
with (
|
|
patch("litellm.afile_retrieve", mock_afile_retrieve),
|
|
patch("litellm.proxy.proxy_server.llm_router", mock_router),
|
|
):
|
|
await managed_files.async_post_call_success_hook(
|
|
data={},
|
|
user_api_key_dict=user_api_key_dict,
|
|
response=batch_response,
|
|
)
|
|
|
|
assert batch_response.output_file_id == already_unified
|
|
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
|
|
}
|
|
|
|
decoded = base64.urlsafe_b64decode(
|
|
already_unified + "=" * (-len(already_unified) % 4)
|
|
).decode()
|
|
assert decoded.count(f"llm_output_file_id,{provider_file_id}") == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_skip_non_file_unified_id_on_output_file_id():
|
|
"""Batch-style unified ids lack llm_output_file_id; must not IndexError or re-wrap."""
|
|
import base64
|
|
|
|
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"
|
|
)
|
|
.decode()
|
|
.rstrip("=")
|
|
)
|
|
|
|
batch_response = _make_batch_response(
|
|
model_id="openai/openai/gpt-5.5-batch",
|
|
model_name="openai/openai/gpt-5.5-batch",
|
|
output_file_id=batch_unified,
|
|
)
|
|
user_api_key_dict = _make_user_api_key_dict()
|
|
|
|
with patch("litellm.afile_retrieve", AsyncMock()) as mock_afile_retrieve:
|
|
await managed_files.async_post_call_success_hook(
|
|
data={},
|
|
user_api_key_dict=user_api_key_dict,
|
|
response=batch_response,
|
|
)
|
|
|
|
assert batch_response.output_file_id == batch_unified
|
|
mock_afile_retrieve.assert_not_called()
|
|
managed_files.store_unified_file_id.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_content_passes_trusted_model_credentials_to_router():
|
|
"""
|
|
afile_content must hand the deployment's credential snapshot to the router
|
|
call as an immutable server-side mapping. Cloud-storage providers (Bedrock
|
|
S3) validate file ids against the bucket in that snapshot, so without it
|
|
unified-id content retrieval only works when AWS_S3_BUCKET_NAME is set.
|
|
"""
|
|
from types import MappingProxyType
|
|
|
|
managed_files = _make_managed_files_instance()
|
|
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}}
|
|
)
|
|
|
|
mock_router = MagicMock()
|
|
mock_router.get_deployment_credentials_with_provider = MagicMock(
|
|
return_value={
|
|
"custom_llm_provider": "bedrock",
|
|
"s3_bucket_name": "my-bucket",
|
|
"aws_region_name": "us-west-2",
|
|
}
|
|
)
|
|
mock_router.afile_content = AsyncMock(return_value=MagicMock())
|
|
|
|
await managed_files.afile_content(
|
|
file_id=unified_file_id,
|
|
litellm_parent_otel_span=None,
|
|
llm_router=mock_router,
|
|
)
|
|
|
|
call_kwargs = mock_router.afile_content.call_args.kwargs
|
|
assert call_kwargs["model"] == "model-123"
|
|
assert call_kwargs["file_id"] == s3_uri
|
|
trusted_credentials = call_kwargs["_litellm_internal_model_credentials"]
|
|
assert isinstance(trusted_credentials, MappingProxyType)
|
|
assert trusted_credentials["s3_bucket_name"] == "my-bucket"
|
|
|
|
|
|
def _managed_deletion_file_id(provider_file_id):
|
|
from litellm.types.utils import SpecialEnums
|
|
|
|
value = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format(
|
|
"application/json", "test-file", "batch-model", provider_file_id, "model-123"
|
|
)
|
|
return base64.urlsafe_b64encode(value.encode()).decode().rstrip("=")
|
|
|
|
|
|
def _managed_files_with_deletion_row(unified_file_id, provider_file_id, file_object):
|
|
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,
|
|
model_mappings={"model-123": provider_file_id},
|
|
flat_model_file_ids=[provider_file_id],
|
|
file_object=file_object,
|
|
)
|
|
table = MagicMock(
|
|
find_first=AsyncMock(return_value=row),
|
|
delete=AsyncMock(),
|
|
)
|
|
return PROXY_LiteLLMManagedFiles(
|
|
internal_usage_cache=DualCache(),
|
|
prisma_client=MagicMock(db=MagicMock(litellm_managedfiletable=table)),
|
|
), table
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_delete_bedrock_uses_deployment_bucket_and_signed_s3_delete(monkeypatch):
|
|
import httpx
|
|
import respx
|
|
|
|
from litellm import Router
|
|
|
|
monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False)
|
|
monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False)
|
|
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "bedrock-batch",
|
|
"litellm_params": {
|
|
"model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
"aws_access_key_id": "AKIAEXAMPLE",
|
|
"aws_secret_access_key": "secret",
|
|
"aws_region_name": "us-west-2",
|
|
"s3_bucket_name": "my-bucket",
|
|
},
|
|
"model_info": {"id": "model-123"},
|
|
}
|
|
],
|
|
num_retries=0,
|
|
)
|
|
s3_uri = "s3://my-bucket/litellm-bedrock-files/input.jsonl"
|
|
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))
|
|
response = await managed_files.afile_delete(
|
|
file_id=unified_file_id,
|
|
litellm_parent_otel_span=None,
|
|
llm_router=router,
|
|
_litellm_internal_model_credentials={"s3_bucket_name": "request-bucket"},
|
|
)
|
|
|
|
assert len(route.calls) == 1
|
|
assert route.calls[0].request.headers["Authorization"].startswith("AWS4-HMAC-SHA256")
|
|
assert response.id == unified_file_id
|
|
assert response.deleted is True
|
|
table.delete.assert_awaited_once_with(where={"unified_file_id": unified_file_id})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_delete_returns_managed_id_for_stored_provider_output():
|
|
from openai.types import FileDeleted
|
|
|
|
provider_file_id = "file-error-output"
|
|
unified_file_id = _managed_deletion_file_id(provider_file_id)
|
|
stored_file = _make_file_object(provider_file_id)
|
|
managed_files, table = _managed_files_with_deletion_row(unified_file_id, provider_file_id, stored_file)
|
|
router = MagicMock(
|
|
get_deployment_credentials_with_provider=MagicMock(return_value=None),
|
|
afile_delete=AsyncMock(return_value=FileDeleted(id=provider_file_id, object="file", deleted=True)),
|
|
)
|
|
response = await managed_files.afile_delete(
|
|
file_id=unified_file_id,
|
|
litellm_parent_otel_span=None,
|
|
llm_router=router,
|
|
_litellm_internal_model_credentials={"s3_bucket_name": "request-bucket"},
|
|
)
|
|
|
|
assert response.id == unified_file_id
|
|
assert response.object == "file"
|
|
assert response.deleted is True
|
|
assert stored_file.id == provider_file_id
|
|
router.afile_delete.assert_awaited_once_with(model="model-123", file_id=provider_file_id)
|
|
table.delete.assert_awaited_once_with(where={"unified_file_id": unified_file_id})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_content_bedrock_unified_id_end_to_end(monkeypatch):
|
|
"""
|
|
Proxy repro for Bedrock batch output retrieval: a unified file id that
|
|
resolves to an s3:// output object must be fetched via a SigV4-signed S3
|
|
GET using the deployment's s3_bucket_name (no AWS_S3_BUCKET_NAME env).
|
|
|
|
Regression test for "BedrockFilesConfig does not support file content
|
|
retrieval" raised on this path.
|
|
"""
|
|
import httpx
|
|
import respx
|
|
|
|
import litellm
|
|
from litellm import Router
|
|
|
|
monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False)
|
|
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
|
litellm.in_memory_llm_clients_cache.flush_cache()
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "bedrock-claude",
|
|
"litellm_params": {
|
|
"model": "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
|
|
"aws_access_key_id": "AKIAEXAMPLE",
|
|
"aws_secret_access_key": "secret",
|
|
"aws_region_name": "us-west-2",
|
|
"s3_bucket_name": "my-bucket",
|
|
},
|
|
"model_info": {"id": "model-123"},
|
|
}
|
|
]
|
|
)
|
|
|
|
managed_files = _make_managed_files_instance()
|
|
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}}
|
|
)
|
|
|
|
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"}')
|
|
)
|
|
|
|
response = await managed_files.afile_content(
|
|
file_id=unified_file_id,
|
|
litellm_parent_otel_span=None,
|
|
llm_router=router,
|
|
)
|
|
|
|
assert route.called
|
|
assert (
|
|
route.calls[0].request.headers["Authorization"].startswith("AWS4-HMAC-SHA256")
|
|
)
|
|
assert response.content == b'{"recordId": "x"}'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_content_error_reports_unified_id_not_provider_uri():
|
|
"""When every model attempt fails, the error must name the caller's unified
|
|
file id, never the resolved internal s3:// URI (no internal-path leak)."""
|
|
managed_files = _make_managed_files_instance()
|
|
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}}
|
|
)
|
|
|
|
mock_router = MagicMock()
|
|
mock_router.get_deployment_credentials_with_provider = 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:
|
|
await managed_files.afile_content(
|
|
file_id=unified_file_id,
|
|
litellm_parent_otel_span=None,
|
|
llm_router=mock_router,
|
|
)
|
|
|
|
message = str(exc_info.value)
|
|
assert unified_file_id in message
|
|
assert s3_uri not in message
|
|
|
|
|
|
def _make_real_managed_files_instance():
|
|
"""Create a PROXY_LiteLLMManagedFiles with a real store_unified_file_id but
|
|
an AsyncMock prisma client, so the DB write path itself can be asserted."""
|
|
from litellm_enterprise.proxy.hooks.managed_files import (
|
|
PROXY_LiteLLMManagedFiles,
|
|
)
|
|
|
|
mock_cache = MagicMock()
|
|
mock_cache.async_get_cache = AsyncMock(return_value=None)
|
|
mock_cache.async_set_cache = AsyncMock()
|
|
|
|
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"
|
|
)
|
|
)
|
|
|
|
return (
|
|
PROXY_LiteLLMManagedFiles(
|
|
internal_usage_cache=mock_cache,
|
|
prisma_client=mock_prisma,
|
|
),
|
|
mock_prisma,
|
|
)
|
|
|
|
|
|
def _make_object_store_instance():
|
|
"""A real store_unified_object_id over an AsyncMock prisma client, so both the
|
|
upsert and the update-only write path can be asserted."""
|
|
from litellm_enterprise.proxy.hooks.managed_files import (
|
|
PROXY_LiteLLMManagedFiles,
|
|
)
|
|
|
|
mock_cache = MagicMock()
|
|
mock_cache.async_set_cache = AsyncMock()
|
|
|
|
mock_prisma = MagicMock()
|
|
mock_prisma.db.litellm_managedobjecttable.upsert = AsyncMock()
|
|
mock_prisma.db.litellm_managedobjecttable.update_many = AsyncMock()
|
|
|
|
return (
|
|
PROXY_LiteLLMManagedFiles(
|
|
internal_usage_cache=mock_cache,
|
|
prisma_client=mock_prisma,
|
|
),
|
|
mock_prisma,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poll_refreshes_batch_state_without_claiming_the_row():
|
|
"""Regression (stale batch state): a poll observes a batch it did not create, so it
|
|
must still refresh status and file_object -- otherwise GET /v1/batches serves the
|
|
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
|
|
)
|
|
|
|
await managed_files.store_unified_object_id(
|
|
unified_object_id="uoi-1",
|
|
file_object=_make_batch_response(status="completed"),
|
|
litellm_parent_otel_span=None,
|
|
model_object_id="batch-123",
|
|
file_purpose="batch",
|
|
user_api_key_dict=poller,
|
|
request_tags=("poller:tag",),
|
|
persist_attribution=False,
|
|
create_if_missing=False,
|
|
)
|
|
|
|
# the row is refreshed in place, and cannot be conjured by a poll
|
|
mock_prisma.db.litellm_managedobjecttable.upsert.assert_not_awaited()
|
|
update_many = mock_prisma.db.litellm_managedobjecttable.update_many
|
|
update_many.assert_awaited_once()
|
|
call = update_many.await_args
|
|
assert call.kwargs["where"] == {"unified_object_id": "uoi-1"}
|
|
|
|
written = call.kwargs["data"]
|
|
assert written["status"] == "completed"
|
|
assert json.loads(written["file_object"])["output_file_id"] == "file-output-abc"
|
|
# nothing the poller could be billed for
|
|
for owned in ("api_key", "request_tags", "created_by", "team_id"):
|
|
assert owned not in written
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
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
|
|
)
|
|
|
|
await managed_files.store_unified_object_id(
|
|
unified_object_id="uoi-2",
|
|
file_object=_make_batch_response(status="validating"),
|
|
litellm_parent_otel_span=None,
|
|
model_object_id="batch-456",
|
|
file_purpose="batch",
|
|
user_api_key_dict=creator,
|
|
request_tags=("env:prod",),
|
|
persist_attribution=True,
|
|
)
|
|
|
|
mock_prisma.db.litellm_managedobjecttable.update_many.assert_not_awaited()
|
|
upsert = mock_prisma.db.litellm_managedobjecttable.upsert
|
|
upsert.assert_awaited_once()
|
|
created = upsert.await_args.kwargs["data"]["create"]
|
|
# UserAPIKeyAuth hashes an sk- token on construction; the hash is what is billed
|
|
assert created["api_key"] == creator.api_key
|
|
assert created["api_key"] != "sk-the-creator"
|
|
assert created["created_by"] == "alice"
|
|
assert created["team_id"] == "team-alpha"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_stored_background_response_names_its_creators_end_user_and_holds_none_of_its_output():
|
|
"""Billing reads only a stored response's status and model, so neither the row nor its cache
|
|
entry keeps what the model generated, and the row names the end user the create billed to."""
|
|
from litellm.types.llms.openai import ResponsesAPIResponse
|
|
|
|
managed_files, mock_prisma = _make_object_store_instance()
|
|
generated_text = "the generated answer"
|
|
response = ResponsesAPIResponse(
|
|
id="resp_unified",
|
|
object="response",
|
|
status="completed",
|
|
created_at=1,
|
|
model="gpt-5.6",
|
|
output=[
|
|
{
|
|
"type": "message",
|
|
"id": "msg_1",
|
|
"role": "assistant",
|
|
"status": "completed",
|
|
"content": [{"type": "output_text", "text": generated_text, "annotations": []}],
|
|
}
|
|
],
|
|
background=True,
|
|
)
|
|
|
|
await managed_files.store_unified_object_id(
|
|
unified_object_id="resp_unified",
|
|
file_object=response,
|
|
litellm_parent_otel_span=None,
|
|
model_object_id="resp_provider",
|
|
file_purpose="response",
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-the-creator", end_user_id="customer-42", parent_otel_span=None),
|
|
persist_attribution=True,
|
|
)
|
|
|
|
written = mock_prisma.db.litellm_managedobjecttable.upsert.await_args.kwargs["data"]
|
|
cached = managed_files.internal_usage_cache.async_set_cache.await_args.kwargs["value"]
|
|
assert generated_text not in json.dumps(written, default=str)
|
|
assert generated_text not in json.dumps(cached, default=str)
|
|
assert json.loads(written["create"]["file_object"])["status"] == "completed"
|
|
assert written["create"]["end_user_id"] == "customer-42"
|
|
assert generated_text in response.model_dump_json()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_default_callers_still_create_their_rows():
|
|
"""create_if_missing defaults to True, so the fine-tune, Responses and managed
|
|
/v1/batches callers, none of which passes it, keep upserting exactly as before."""
|
|
managed_files, mock_prisma = _make_object_store_instance()
|
|
|
|
await managed_files.store_unified_object_id(
|
|
unified_object_id="uoi-3",
|
|
file_object=_make_batch_response(),
|
|
litellm_parent_otel_span=None,
|
|
model_object_id="ft-789",
|
|
file_purpose="fine-tune",
|
|
user_api_key_dict=_make_user_api_key_dict(),
|
|
)
|
|
|
|
mock_prisma.db.litellm_managedobjecttable.upsert.assert_awaited_once()
|
|
mock_prisma.db.litellm_managedobjecttable.update_many.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_store_unified_file_id_is_idempotent_via_upsert():
|
|
"""Regression test for the managed-batch retrieve 500 (UniqueViolationError on
|
|
unified_file_id): re-registering an already-stored output file id must upsert on
|
|
unified_file_id, never do an unconditional create that raises on conflict."""
|
|
managed_files, mock_prisma = _make_real_managed_files_instance()
|
|
file_id = "litellm_proxy_unified_output_id_abc"
|
|
model_mappings = {"model-deploy-xyz": "file-output-abc"}
|
|
|
|
for _ in range(2):
|
|
await managed_files.store_unified_file_id(
|
|
file_id=file_id,
|
|
file_object=_make_file_object(),
|
|
litellm_parent_otel_span=None,
|
|
model_mappings=model_mappings,
|
|
user_api_key_dict=_make_user_api_key_dict(),
|
|
)
|
|
|
|
mock_prisma.db.litellm_managedfiletable.create.assert_not_awaited()
|
|
upsert_mock = mock_prisma.db.litellm_managedfiletable.upsert
|
|
assert upsert_mock.await_count == 2
|
|
for upsert_call in upsert_mock.await_args_list:
|
|
assert upsert_call.kwargs["where"] == {"unified_file_id": file_id}
|
|
upsert_data = upsert_call.kwargs["data"]
|
|
assert upsert_data["create"]["unified_file_id"] == file_id
|
|
assert json.loads(upsert_data["create"]["model_mappings"]) == model_mappings
|
|
assert json.loads(upsert_data["update"]["model_mappings"]) == model_mappings
|
|
|
|
|
|
def test_get_unified_output_file_id_is_deterministic_per_output_file():
|
|
managed_files, _ = _make_real_managed_files_instance()
|
|
|
|
first = managed_files.get_unified_output_file_id(
|
|
output_file_id="file-output-abc",
|
|
model_id="model-deploy-xyz",
|
|
model_name="azure/gpt-4",
|
|
)
|
|
repeat = managed_files.get_unified_output_file_id(
|
|
output_file_id="file-output-abc",
|
|
model_id="model-deploy-xyz",
|
|
model_name="azure/gpt-4",
|
|
)
|
|
other_file = managed_files.get_unified_output_file_id(
|
|
output_file_id="file-output-def",
|
|
model_id="model-deploy-xyz",
|
|
model_name="azure/gpt-4",
|
|
)
|
|
other_model = managed_files.get_unified_output_file_id(
|
|
output_file_id="file-output-abc",
|
|
model_id="model-deploy-other",
|
|
model_name="azure/gpt-4",
|
|
)
|
|
|
|
assert first == repeat
|
|
assert len({first, other_file, other_model}) == 3
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_first_registrations_converge_on_one_row():
|
|
managed_files, mock_prisma = _make_real_managed_files_instance()
|
|
|
|
minted_ids = tuple(
|
|
managed_files.get_unified_output_file_id(
|
|
output_file_id="file-output-abc",
|
|
model_id="model-deploy-xyz",
|
|
model_name=None,
|
|
)
|
|
for _ in range(2)
|
|
)
|
|
await asyncio.gather(
|
|
*(
|
|
managed_files.store_unified_file_id(
|
|
file_id=unified_id,
|
|
file_object=None,
|
|
litellm_parent_otel_span=None,
|
|
model_mappings={"model-deploy-xyz": "file-output-abc"},
|
|
user_api_key_dict=_make_user_api_key_dict(),
|
|
)
|
|
for unified_id in minted_ids
|
|
)
|
|
)
|
|
|
|
upserted_row_keys = {
|
|
upsert_call.kwargs["where"]["unified_file_id"]
|
|
for upsert_call in mock_prisma.db.litellm_managedfiletable.upsert.await_args_list
|
|
}
|
|
assert minted_ids[0] == minted_ids[1]
|
|
assert upserted_row_keys == {minted_ids[0]}
|
|
|
|
|
|
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}"
|
|
)
|
|
return base64.urlsafe_b64encode(unified_input_file_id.encode()).decode().rstrip("=")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
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"
|
|
)
|
|
|
|
mock_router = MagicMock()
|
|
mock_router.get_deployment_credentials_with_provider = MagicMock(return_value={})
|
|
|
|
with (
|
|
patch("litellm.afile_retrieve", AsyncMock(return_value=_make_file_object())),
|
|
patch("litellm.proxy.proxy_server.llm_router", mock_router),
|
|
):
|
|
await managed_files.async_post_call_success_hook(
|
|
data={},
|
|
user_api_key_dict=_make_user_api_key_dict(),
|
|
response=batch_response,
|
|
)
|
|
|
|
assert batch_response.output_file_id == managed_files.get_unified_output_file_id(
|
|
output_file_id="file-output-abc",
|
|
model_id="model-deploy-xyz",
|
|
model_name="model-a,model-b",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_hook_mint_falls_back_to_response_input_file_id_target_models():
|
|
managed_files = _make_managed_files_instance()
|
|
batch_response = _make_batch_response()
|
|
batch_response.input_file_id = _b64_unified_input_file_id("model-a,model-b")
|
|
batch_response._hidden_params = {
|
|
"unified_batch_id": "some-unified-batch-id",
|
|
"model_id": "model-deploy-xyz",
|
|
}
|
|
|
|
mock_router = MagicMock()
|
|
mock_router.get_deployment_credentials_with_provider = MagicMock(return_value={})
|
|
|
|
with (
|
|
patch("litellm.afile_retrieve", AsyncMock(return_value=_make_file_object())),
|
|
patch("litellm.proxy.proxy_server.llm_router", mock_router),
|
|
):
|
|
await managed_files.async_post_call_success_hook(
|
|
data={},
|
|
user_api_key_dict=_make_user_api_key_dict(),
|
|
response=batch_response,
|
|
)
|
|
|
|
assert batch_response.output_file_id == managed_files.get_unified_output_file_id(
|
|
output_file_id="file-output-abc",
|
|
model_id="model-deploy-xyz",
|
|
model_name="model-a,model-b",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cost_job_and_retrieve_paths_mint_identical_unified_output_file_ids():
|
|
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)
|
|
unified_input_file_id = _b64_unified_input_file_id("model-a")
|
|
|
|
retrieve_response = LiteLLMBatch(
|
|
id="batch-123",
|
|
completion_window="24h",
|
|
created_at=1700000000,
|
|
endpoint="/v1/chat/completions",
|
|
input_file_id=unified_input_file_id,
|
|
object="batch",
|
|
status="completed",
|
|
output_file_id="file-output-abc",
|
|
)
|
|
retrieve_response._hidden_params = {"model_id": "model-deploy-xyz"}
|
|
|
|
await ensure_batch_response_managed_file_ids(
|
|
response=retrieve_response,
|
|
managed_files_obj=managed_files,
|
|
prisma_client=mock_prisma,
|
|
verbose_proxy_logger=MagicMock(),
|
|
user_api_key_dict=_make_user_api_key_dict(),
|
|
)
|
|
|
|
job = MagicMock()
|
|
job.file_object = {
|
|
"id": "batch-123",
|
|
"completion_window": "24h",
|
|
"created_at": 1700000000,
|
|
"endpoint": "/v1/chat/completions",
|
|
"input_file_id": unified_input_file_id,
|
|
"object": "batch",
|
|
"status": "completed",
|
|
}
|
|
cost_job_model_name = CheckBatchCost._get_managed_file_model_name(
|
|
job=job, deployment_info=MagicMock(model_name="vertex_ai/gemini-3-pro")
|
|
)
|
|
|
|
assert cost_job_model_name == "model-a"
|
|
assert retrieve_response.output_file_id == managed_files.get_unified_output_file_id(
|
|
output_file_id="file-output-abc",
|
|
model_id="model-deploy-xyz",
|
|
model_name=cost_job_model_name,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_batch_create_hook_persists_creating_key_and_tags():
|
|
"""Regression: the /v1/batches create hook must persist the creating key and the
|
|
request's tags on the managed object row. CheckBatchCost, which owns the batch's
|
|
accounting once the retrieve path defers to it, bills whatever the row carries, and
|
|
without these columns the cost lands on the user alone and the key's spend and
|
|
budget never see it."""
|
|
managed_files = _make_managed_files_instance()
|
|
creator = UserAPIKeyAuth(api_key="sk-the-creator", user_id="alice", parent_otel_span=None)
|
|
create_response = _make_batch_response(status="validating", output_file_id=None)
|
|
create_response._hidden_params = {
|
|
BATCH_CREATE_HIDDEN_PARAM: True,
|
|
"model_id": "model-deploy-xyz",
|
|
"model_name": "azure/gpt-4",
|
|
}
|
|
|
|
await managed_files.async_post_call_success_hook(
|
|
data={"litellm_metadata": {"tags": ["env:prod", "team:ml"], "user_api_key": creator.api_key}},
|
|
user_api_key_dict=creator,
|
|
response=create_response,
|
|
)
|
|
|
|
managed_files.store_unified_object_id.assert_awaited_once()
|
|
stored = managed_files.store_unified_object_id.await_args.kwargs
|
|
assert stored["persist_attribution"] is True
|
|
assert stored["request_tags"] == ("env:prod", "team:ml")
|
|
assert stored["user_api_key_dict"] is creator
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_batch_create_hook_records_created_metric_once():
|
|
managed_files = _make_managed_files_instance()
|
|
prometheus_logger = MagicMock()
|
|
managed_files._get_prometheus_logger = MagicMock(return_value=prometheus_logger)
|
|
create_response = _make_batch_response(status="validating", output_file_id=None)
|
|
create_response._hidden_params = {
|
|
BATCH_CREATE_HIDDEN_PARAM: True,
|
|
"model_id": "model-deploy-xyz",
|
|
"model_name": "azure/gpt-4",
|
|
}
|
|
|
|
await managed_files.async_post_call_success_hook(
|
|
data={},
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-the-creator", user_id="alice", parent_otel_span=None),
|
|
response=create_response,
|
|
)
|
|
|
|
prometheus_logger.record_managed_batch_created.assert_called_once()
|
|
recorded = prometheus_logger.record_managed_batch_created.call_args.kwargs
|
|
assert recorded["model"] == "azure/gpt-4"
|
|
assert recorded["api_provider"] == "azure"
|
|
assert recorded["user"] == "alice"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_batch_retrieve_hook_does_not_record_created_metric():
|
|
managed_files = _make_managed_files_instance()
|
|
prometheus_logger = MagicMock()
|
|
managed_files._get_prometheus_logger = MagicMock(return_value=prometheus_logger)
|
|
retrieve_response = _make_batch_response(status="in_progress", output_file_id=None)
|
|
retrieve_response._hidden_params = {
|
|
"unified_batch_id": "some-unified-batch-id",
|
|
"model_id": "model-deploy-xyz",
|
|
"model_name": "azure/gpt-4",
|
|
}
|
|
|
|
await managed_files.async_post_call_success_hook(
|
|
data={},
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-the-poller", user_id="bob", parent_otel_span=None),
|
|
response=retrieve_response,
|
|
)
|
|
|
|
prometheus_logger.record_managed_batch_created.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_batch_retrieve_hook_does_not_claim_attribution():
|
|
"""A retrieve carries unified_batch_id but no unified_file_id, so it must not rewrite
|
|
the row's paying key to whoever happens to poll the batch."""
|
|
managed_files = _make_managed_files_instance()
|
|
retrieve_response = _make_batch_response(status="in_progress", output_file_id=None)
|
|
retrieve_response._hidden_params = {
|
|
"unified_batch_id": "some-unified-batch-id",
|
|
"model_id": "model-deploy-xyz",
|
|
"model_name": "azure/gpt-4",
|
|
}
|
|
|
|
await managed_files.async_post_call_success_hook(
|
|
data={"litellm_metadata": {"tags": ["poller:tag"]}},
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-the-poller", user_id="bob", parent_otel_span=None),
|
|
response=retrieve_response,
|
|
)
|
|
|
|
managed_files.store_unified_object_id.assert_awaited_once()
|
|
assert managed_files.store_unified_object_id.await_args.kwargs["persist_attribution"] is False
|
|
|
|
|
|
def _unified_batch_id(llm_batch_id: str) -> str:
|
|
decoded = f"litellm_proxy;model_id:my-vllm;llm_batch_id:{llm_batch_id}"
|
|
return base64.urlsafe_b64encode(decoded.encode()).decode().rstrip("=")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"llm_batch_id, stores",
|
|
[("litellm_batch_abc", False), ("batch_abc", True)],
|
|
ids=["litellm-executed batch is left alone", "provider batch is still stored"],
|
|
)
|
|
async def test_post_call_hook_leaves_litellm_executed_batches_untouched(llm_batch_id: str, stores: bool):
|
|
managed_files = _make_managed_files_instance()
|
|
response = _make_batch_response(status="in_progress", output_file_id=None)
|
|
response.id = _unified_batch_id(llm_batch_id)
|
|
response._hidden_params = {
|
|
"unified_batch_id": response.id,
|
|
"model_id": "my-vllm",
|
|
"model_name": "hosted_vllm/qwen",
|
|
}
|
|
original_id = response.id
|
|
|
|
returned = await managed_files.async_post_call_success_hook(
|
|
data={},
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-the-poller", user_id="bob", parent_otel_span=None),
|
|
response=response,
|
|
)
|
|
|
|
assert returned is response
|
|
assert managed_files.store_unified_object_id.await_count == (1 if stores else 0)
|
|
if not stores:
|
|
assert response.id == original_id
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_delete_passes_trusted_model_credentials_to_router():
|
|
"""
|
|
afile_delete must hand the deployment's credential snapshot to the router
|
|
call, since Bedrock validates the s3:// file id against the bucket in it.
|
|
"""
|
|
from types import MappingProxyType
|
|
|
|
managed_files = _make_managed_files_instance()
|
|
unified_file_id = "unified-file-id"
|
|
s3_uri = "s3://my-bucket/litellm-bedrock-files/job-123/input.jsonl"
|
|
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.delete_unified_file_id = AsyncMock(return_value=_make_file_object(unified_file_id))
|
|
|
|
mock_router = MagicMock()
|
|
mock_router.get_deployment_credentials_with_provider = MagicMock(
|
|
return_value={
|
|
"custom_llm_provider": "bedrock",
|
|
"s3_bucket_name": "my-bucket",
|
|
"aws_region_name": "us-west-2",
|
|
}
|
|
)
|
|
mock_router.afile_delete = AsyncMock(return_value=MagicMock())
|
|
|
|
await managed_files.afile_delete(
|
|
file_id=unified_file_id,
|
|
litellm_parent_otel_span=None,
|
|
llm_router=mock_router,
|
|
)
|
|
|
|
call_kwargs = mock_router.afile_delete.call_args.kwargs
|
|
assert call_kwargs["model"] == "model-123"
|
|
assert call_kwargs["file_id"] == s3_uri
|
|
trusted_credentials = call_kwargs["_litellm_internal_model_credentials"]
|
|
assert isinstance(trusted_credentials, MappingProxyType)
|
|
assert trusted_credentials["s3_bucket_name"] == "my-bucket"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_delete_bedrock_unified_id_end_to_end(monkeypatch):
|
|
"""
|
|
Proxy repro for deleting a Bedrock batch input file by unified id: the
|
|
s3:// object must be removed via a SigV4-signed S3 DELETE using the
|
|
deployment's s3_bucket_name (no AWS_S3_BUCKET_NAME env).
|
|
|
|
Regression test for "BedrockFilesConfig does not support file deletion"
|
|
raised on this path.
|
|
"""
|
|
import httpx
|
|
import respx
|
|
|
|
import litellm
|
|
from litellm import Router
|
|
|
|
monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False)
|
|
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
|
litellm.in_memory_llm_clients_cache.flush_cache()
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "bedrock-claude",
|
|
"litellm_params": {
|
|
"model": "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
|
|
"aws_access_key_id": "AKIAEXAMPLE",
|
|
"aws_secret_access_key": "secret",
|
|
"aws_region_name": "us-west-2",
|
|
"s3_bucket_name": "my-bucket",
|
|
},
|
|
"model_info": {"id": "model-123"},
|
|
}
|
|
]
|
|
)
|
|
|
|
managed_files = _make_managed_files_instance()
|
|
unified_file_id = "unified-file-id"
|
|
s3_uri = "s3://my-bucket/litellm-bedrock-files/job-123/input.jsonl"
|
|
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.delete_unified_file_id = AsyncMock(return_value=_make_file_object(unified_file_id))
|
|
|
|
expected_url = "https://s3.us-west-2.amazonaws.com/my-bucket/litellm-bedrock-files/job-123/input.jsonl"
|
|
with respx.mock:
|
|
route = respx.delete(expected_url).mock(return_value=httpx.Response(204))
|
|
|
|
response = await managed_files.afile_delete(
|
|
file_id=unified_file_id,
|
|
litellm_parent_otel_span=None,
|
|
llm_router=router,
|
|
)
|
|
|
|
assert route.called
|
|
assert route.calls[0].request.headers["Authorization"].startswith("AWS4-HMAC-SHA256")
|
|
assert response.id == unified_file_id
|
|
assert response.model_dump() == {"id": unified_file_id, "object": "file", "deleted": True}
|
|
managed_files.delete_unified_file_id.assert_awaited_once_with(unified_file_id, None)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_delete_storage_backed_row_deletes_stored_content_not_provider_files():
|
|
from litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles
|
|
from openai.types import FileDeleted
|
|
|
|
from litellm.caching import DualCache
|
|
from litellm.models.managed_files import LiteLLM_ManagedFileTable
|
|
|
|
storage_url = "litellm_db://content-row-1"
|
|
unified_file_id = _managed_deletion_file_id(storage_url)
|
|
row = LiteLLM_ManagedFileTable(
|
|
unified_file_id=unified_file_id,
|
|
model_mappings={"vllm-batch": storage_url},
|
|
flat_model_file_ids=[storage_url],
|
|
file_object=_make_file_object(unified_file_id),
|
|
storage_backend="litellm_db",
|
|
storage_url=storage_url,
|
|
)
|
|
file_table = MagicMock(find_first=AsyncMock(return_value=row), delete=AsyncMock())
|
|
content_table = MagicMock(delete=AsyncMock())
|
|
managed_files = PROXY_LiteLLMManagedFiles(
|
|
internal_usage_cache=DualCache(),
|
|
prisma_client=MagicMock(
|
|
db=MagicMock(litellm_managedfiletable=file_table, litellm_managedfilecontenttable=content_table)
|
|
),
|
|
)
|
|
router = MagicMock(
|
|
get_deployment_credentials_with_provider=MagicMock(return_value=None),
|
|
afile_delete=AsyncMock(),
|
|
)
|
|
|
|
response = await managed_files.afile_delete(
|
|
file_id=unified_file_id,
|
|
litellm_parent_otel_span=None,
|
|
llm_router=router,
|
|
)
|
|
|
|
content_table.delete.assert_awaited_once_with(where={"id": "content-row-1"})
|
|
router.afile_delete.assert_not_awaited()
|
|
file_table.delete.assert_awaited_once_with(where={"unified_file_id": unified_file_id})
|
|
assert response == FileDeleted(id=unified_file_id, object="file", deleted=True)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_afile_content_storage_backed_row_returns_stored_bytes_not_provider_content():
|
|
from litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles
|
|
from prisma import Base64
|
|
|
|
from litellm.caching import DualCache
|
|
from litellm.models.managed_files import LiteLLM_ManagedFileTable
|
|
|
|
storage_url = "litellm_db://content-row-1"
|
|
unified_file_id = _managed_deletion_file_id(storage_url)
|
|
stored_bytes = b'{"custom_id": "line-1", "method": "POST", "url": "/v1/chat/completions", "body": {}}\n'
|
|
row = LiteLLM_ManagedFileTable(
|
|
unified_file_id=unified_file_id,
|
|
model_mappings={"vllm-batch": storage_url},
|
|
flat_model_file_ids=[storage_url],
|
|
file_object=_make_file_object(unified_file_id),
|
|
storage_backend="litellm_db",
|
|
storage_url=storage_url,
|
|
)
|
|
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(),
|
|
prisma_client=MagicMock(
|
|
db=MagicMock(litellm_managedfiletable=file_table, litellm_managedfilecontenttable=content_table)
|
|
),
|
|
)
|
|
router = MagicMock(
|
|
get_deployment_credentials_with_provider=MagicMock(return_value=None),
|
|
afile_content=AsyncMock(),
|
|
)
|
|
|
|
response = await managed_files.afile_content(
|
|
file_id=unified_file_id,
|
|
litellm_parent_otel_span=None,
|
|
llm_router=router,
|
|
)
|
|
|
|
assert response.content == stored_bytes
|
|
content_table.find_unique.assert_awaited_once_with(where={"id": "content-row-1"})
|
|
router.afile_content.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_store_unified_object_id_batch_processed_is_written_only_when_asked():
|
|
managed_files, mock_prisma = _make_object_store_instance()
|
|
upsert = mock_prisma.db.litellm_managedobjecttable.upsert
|
|
creator = UserAPIKeyAuth(api_key="sk-creator", user_id="alice", team_id="team-alpha", parent_otel_span=None)
|
|
|
|
await managed_files.store_unified_object_id(
|
|
unified_object_id="uoi-processed",
|
|
file_object=_make_batch_response(status="completed"),
|
|
litellm_parent_otel_span=None,
|
|
model_object_id="batch-processed",
|
|
file_purpose="batch",
|
|
user_api_key_dict=creator,
|
|
batch_processed=True,
|
|
)
|
|
await managed_files.store_unified_object_id(
|
|
unified_object_id="uoi-default",
|
|
file_object=_make_batch_response(status="completed"),
|
|
litellm_parent_otel_span=None,
|
|
model_object_id="batch-default",
|
|
file_purpose="batch",
|
|
user_api_key_dict=creator,
|
|
)
|
|
|
|
processed_create, default_create = (call.kwargs["data"]["create"] for call in upsert.await_args_list)
|
|
assert processed_create["batch_processed"] is True
|
|
assert default_create["batch_processed"] is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_store_unified_file_id_caches_the_storage_location_the_db_row_gets():
|
|
from litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles
|
|
|
|
from litellm.caching import DualCache
|
|
|
|
file_table = MagicMock(upsert=AsyncMock(), find_first=AsyncMock(side_effect=AssertionError("cache miss")))
|
|
managed_files = PROXY_LiteLLMManagedFiles(
|
|
internal_usage_cache=DualCache(),
|
|
prisma_client=MagicMock(db=MagicMock(litellm_managedfiletable=file_table)),
|
|
)
|
|
stored = _make_file_object("file-kept").model_copy(update={"purpose": "batch"})
|
|
stored._hidden_params = {"storage_backend": "litellm_db", "storage_url": "litellm_db://content-row-1"}
|
|
|
|
await managed_files.store_unified_file_id(
|
|
file_id="unified-kept",
|
|
file_object=stored,
|
|
litellm_parent_otel_span=None,
|
|
model_mappings={"vllm-batch": "litellm_db://content-row-1"},
|
|
user_api_key_dict=_make_user_api_key_dict(),
|
|
)
|
|
cached = await managed_files.get_unified_file_id("unified-kept")
|
|
|
|
assert cached is not None
|
|
assert (cached.storage_backend, cached.storage_url) == ("litellm_db", "litellm_db://content-row-1")
|
|
create_data = file_table.upsert.await_args.kwargs["data"]["create"]
|
|
assert (create_data["storage_backend"], create_data["storage_url"]) == ("litellm_db", "litellm_db://content-row-1")
|