scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 100
litellm/tests/unit/enterprise/proxy/test_managed_files_hook.py
devin-ai-integration[bot] bf0187072b
ci: move caching, proxy-extras, gateway and enterprise tests into tests/unit and run them from litellm-tests (#42902) 
* ci: fix the litellm-tests unit job with sysmon coverage, an env allowlist and coverage upload on failure

* test: replace key-dependent proxy, enterprise and mcp unit tests with synthetic values and integration and e2e coverage

* test: drop key reads at the legacy proxy, enterprise and mcp paths and wire the gemini pass-through split

* ci: move caching, proxy-extras, gateway and enterprise tests into tests/unit and run them from litellm-tests under their legacy flags

* ci: move caching, proxy-extras, gateway and enterprise tests into tests/unit and run them from litellm-tests under their legacy flags

* ci: fail the unit shard when circleci tests split errors

* test: drop restating comments from the gemini pass-through split

* ci: exit the unit shard cleanly when circleci tests split assigns it no files

---------

Co-authored-by: yuneng <yuneng@berri.ai>
2026-09-24 15:49:59 -07:00

2012 lines
72 KiB
Python
Raw Blame History

"""
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
import pytest
from typing import 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_prisma = MagicMock()
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_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_set_cache = AsyncMock()
mock_prisma = MagicMock()
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_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")
Reference in a new issue View git blame Copy permalink