mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(files): match every spelling of a provider file id in the ownership check (#51498)
* fix(files): match every spelling of a provider file id in the ownership check
enforce_provider_file_access matched the stored provider file id by exact string, and a Gemini output file
can be addressed as files/<id>, files%2F<id>, <id>/ or its full URL, all of which the Gemini file transformation
sends to the same file. Those spellings found no ownership row and fell through the no-row pass-through path.
canonical_provider_file_id collapses the spellings to the bare id; the ownership check looks up both the
spelling the caller sent and the canonical one, and store_batch_output_file_ownership stores the canonical id.
* fix(files): unwrap model-encoded ids and index canonical provider ids in ownership rows
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(files): assert the ownership lookup outcome instead of its query shape
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(files): map provider ids back to their unified file by canonical spelling
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* chore(files): record the bound on the ownership lookup's in filter
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(files): reduce any scheme or case wrapping files/<id> the way gemini does before the ownership lookup
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(files): bind the reverse lookup match with a walrus instead of a Final inside the loop
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(files): assert the provider id lookup covers the requested ids instead of its exact list
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
(cherry picked from commit e8fc243a09)
* fix(files): keep non-Gemini storage URLs opaque so an upload name cannot index another key's provider id
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(files): judge the provider id resolver by the mapping it returns only
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---------
Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu
GitHub
parent
216a330712
commit
efe349de9f
5 changed files with 365 additions and 65 deletions
|
|
@ -82,6 +82,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
|
|||
TeamProviderDeployment,
|
||||
_is_base64_encoded_unified_file_id,
|
||||
apply_unified_file_ids,
|
||||
canonical_provider_file_id,
|
||||
decode_model_from_file_id,
|
||||
ensure_batch_response_managed_file_ids,
|
||||
get_batch_id_from_unified_batch_id,
|
||||
|
|
@ -91,6 +92,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
|
|||
is_litellm_executed_batch,
|
||||
map_raw_file_ids_to_unified,
|
||||
normalize_mime_type_for_provider,
|
||||
provider_file_id_spellings,
|
||||
resolve_managed_output_file_model_name,
|
||||
validate_file_list_limit,
|
||||
validate_file_list_purpose,
|
||||
|
|
@ -597,14 +599,27 @@ async def _scoped_file_rows(
|
|||
return _SCOPED_FILE_ROWS.validate_python(rows)
|
||||
|
||||
|
||||
def _iter_provider_file_id_pairs(
|
||||
rows: Sequence[_ManagedFileRow],
|
||||
requested_provider_file_ids: frozenset[str],
|
||||
) -> Iterator[tuple[str, str]]:
|
||||
def _iter_canonical_provider_file_ids(rows: Sequence[_ManagedFileRow]) -> Iterator[tuple[str, str]]:
|
||||
for row in rows:
|
||||
for provider_file_id in row.flat_model_file_ids:
|
||||
if provider_file_id in requested_provider_file_ids:
|
||||
yield provider_file_id, row.unified_file_id
|
||||
yield canonical_provider_file_id(provider_file_id), row.unified_file_id
|
||||
|
||||
|
||||
def _iter_requested_provider_file_id_spellings(provider_file_ids: Sequence[str]) -> Iterator[str]:
|
||||
for provider_file_id in provider_file_ids:
|
||||
yield from provider_file_id_spellings(provider_file_id)
|
||||
|
||||
|
||||
def _iter_provider_file_id_pairs(
|
||||
rows: Sequence[_ManagedFileRow],
|
||||
requested_provider_file_ids: Sequence[str],
|
||||
) -> Iterator[tuple[str, str]]:
|
||||
unified_by_canonical_id: Final = dict(_iter_canonical_provider_file_ids(rows))
|
||||
for requested_provider_file_id in requested_provider_file_ids:
|
||||
if (
|
||||
unified_file_id := unified_by_canonical_id.get(canonical_provider_file_id(requested_provider_file_id))
|
||||
) is not None:
|
||||
yield requested_provider_file_id, unified_file_id
|
||||
|
||||
|
||||
def _managed_object_table(prisma_client: _ManagedResourcePrismaClient) -> _ManagedObjectTableActions:
|
||||
|
|
@ -767,12 +782,15 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
) -> None:
|
||||
verbose_logger.info(f"Storing LiteLLM Managed File object with id={file_id} in cache")
|
||||
storage_metadata: Final = _storage_metadata_of(file_object)
|
||||
flat_model_file_ids: Final = [
|
||||
canonical_provider_file_id(provider_file_id) for provider_file_id in model_mappings.values()
|
||||
]
|
||||
if file_object is not None:
|
||||
litellm_managed_file_object = LiteLLM_ManagedFileTable(
|
||||
unified_file_id=file_id,
|
||||
file_object=file_object,
|
||||
model_mappings=model_mappings,
|
||||
flat_model_file_ids=list(model_mappings.values()),
|
||||
flat_model_file_ids=flat_model_file_ids,
|
||||
created_by=resolve_resource_owner_id(user_api_key_dict),
|
||||
team_id=user_api_key_dict.team_id,
|
||||
updated_by=user_api_key_dict.user_id,
|
||||
|
|
@ -790,14 +808,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
db_data = {
|
||||
"unified_file_id": file_id,
|
||||
"model_mappings": json.dumps(model_mappings),
|
||||
"flat_model_file_ids": list(model_mappings.values()),
|
||||
"flat_model_file_ids": flat_model_file_ids,
|
||||
"created_by": resolve_resource_owner_id(user_api_key_dict),
|
||||
"team_id": user_api_key_dict.team_id,
|
||||
"updated_by": user_api_key_dict.user_id,
|
||||
}
|
||||
update_data = {
|
||||
"model_mappings": json.dumps(model_mappings),
|
||||
"flat_model_file_ids": list(model_mappings.values()),
|
||||
"flat_model_file_ids": flat_model_file_ids,
|
||||
"updated_by": user_api_key_dict.user_id,
|
||||
}
|
||||
|
||||
|
|
@ -1094,8 +1112,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
"""
|
||||
if self.prisma_client is None:
|
||||
return
|
||||
spellings: Final = list(provider_file_id_spellings(file_id))
|
||||
managed_file = await _managed_file_table(self.prisma_client).find_first(
|
||||
where={"OR": [{"unified_file_id": file_id}, {"flat_model_file_ids": {"has": file_id}}]}
|
||||
where={
|
||||
"OR": [
|
||||
{"unified_file_id": {"in": spellings}}, # bounded-ok: at most four spellings of one file id
|
||||
{"flat_model_file_ids": {"hasSome": spellings}},
|
||||
]
|
||||
}
|
||||
)
|
||||
if managed_file is None:
|
||||
return
|
||||
|
|
@ -1139,7 +1163,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
)
|
||||
for file_id in provider_file_ids:
|
||||
model_name = decode_model_from_file_id(file_id)
|
||||
raw_file_id = get_original_file_id(file_id)
|
||||
raw_file_id = canonical_provider_file_id(get_original_file_id(file_id))
|
||||
await self.store_batch_output_file(
|
||||
unified_file_id=file_id,
|
||||
provider_file_id=raw_file_id,
|
||||
|
|
@ -1629,7 +1653,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
return MappingProxyType({})
|
||||
|
||||
provider_file_ids_list: Final = [ # mutable-ok: Prisma hasSome requires a list
|
||||
provider_file_id for provider_file_id in unique_provider_file_ids
|
||||
*dict.fromkeys(_iter_requested_provider_file_id_spellings(unique_provider_file_ids))
|
||||
]
|
||||
rows: Final = await _managed_file_table(self.prisma_client).find_many(
|
||||
where={ # mutable-ok: Prisma requires a plain dictionary for where
|
||||
|
|
@ -1641,10 +1665,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
)
|
||||
return MappingProxyType(
|
||||
dict(
|
||||
_iter_provider_file_id_pairs(
|
||||
rows,
|
||||
frozenset(unique_provider_file_ids),
|
||||
)
|
||||
_iter_provider_file_id_pairs(rows, unique_provider_file_ids)
|
||||
)
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ import itertools
|
|||
import logging
|
||||
import mimetypes
|
||||
import re
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from types import MappingProxyType
|
||||
|
|
@ -17,6 +17,7 @@ from typing import (
|
|||
get_args,
|
||||
runtime_checkable,
|
||||
)
|
||||
from urllib.parse import unquote, urlparse
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
|
@ -439,6 +440,50 @@ def get_original_file_id(encoded_id: str) -> str:
|
|||
return encoded_id
|
||||
|
||||
|
||||
_GEMINI_FILES_HOST: Final = "generativelanguage.googleapis.com"
|
||||
|
||||
|
||||
def canonical_provider_file_id(file_id: str) -> str:
|
||||
"""The one spelling of a provider file id that ownership rows store and look up, so ``files/abc``,
|
||||
``files%2Fabc``, ``abc/`` and the Gemini file URL all resolve to the bare ``abc``. Every other URL or URI
|
||||
(``s3://``, ``gs://``, a storage blob URL) is returned unchanged, so an upload whose name happens to contain
|
||||
``files/<id>`` never indexes someone else's ``<id>``."""
|
||||
decoded: Final = unquote(file_id).strip()
|
||||
is_url: Final = "://" in decoded
|
||||
if is_url and urlparse(decoded).hostname != _GEMINI_FILES_HOST:
|
||||
return file_id
|
||||
path: Final = urlparse(decoded).path if is_url else decoded
|
||||
files_index: Final = path.find("files/")
|
||||
bare: Final = (path[files_index:] if files_index != -1 else path).strip("/").removeprefix("files/").strip("/")
|
||||
return bare or decoded
|
||||
|
||||
|
||||
_GEMINI_FILES_URL: Final = "https://generativelanguage.googleapis.com/v1beta/files/"
|
||||
|
||||
|
||||
def _gemini_files_tail(file_id: str) -> str:
|
||||
"""The id Gemini resolves ``file_id`` to, whatever scheme or case wraps it: the text after its ``files/``."""
|
||||
files_index: Final = file_id.find("files/")
|
||||
return file_id[files_index:].strip("/").removeprefix("files/").strip("/") if files_index != -1 else ""
|
||||
|
||||
|
||||
def _iter_stored_forms(bare_ids: Sequence[str]) -> Iterator[str]:
|
||||
for bare_id in bare_ids:
|
||||
yield bare_id
|
||||
yield f"files/{bare_id}"
|
||||
yield f"{_GEMINI_FILES_URL}{bare_id}"
|
||||
|
||||
|
||||
def provider_file_id_spellings(file_id: str) -> tuple[str, ...]:
|
||||
"""Every spelling an ownership row may hold for ``file_id``: as sent, unwrapped from a model-encoded id and
|
||||
canonicalized, reduced the way Gemini reduces it, plus the ``files/`` and Gemini URL forms that upload rows
|
||||
stored before ids were canonicalized."""
|
||||
original: Final = get_original_file_id(file_id)
|
||||
reductions: Final = (canonical_provider_file_id(original), _gemini_files_tail(unquote(original)))
|
||||
bare_ids: Final = tuple(dict.fromkeys(bare for bare in reductions if bare))
|
||||
return tuple(dict.fromkeys((file_id, *_iter_stored_forms(bare_ids))))
|
||||
|
||||
|
||||
def is_model_embedded_id(file_id: str) -> bool:
|
||||
"""
|
||||
Check if a file/batch ID has model routing information embedded.
|
||||
|
|
|
|||
|
|
@ -58,9 +58,17 @@ class _InMemoryManagedFileTable:
|
|||
self, where: Mapping[str, object]
|
||||
) -> LiteLLM_ManagedFileTable | None:
|
||||
self.find_first_calls.append(where)
|
||||
or_filters: Final = where.get("OR")
|
||||
if isinstance(or_filters, list):
|
||||
for clause in or_filters:
|
||||
if (matched := await self.find_first(clause)) is not None:
|
||||
return matched
|
||||
return None
|
||||
unified_file_id: Final = where.get("unified_file_id")
|
||||
if isinstance(unified_file_id, str):
|
||||
return self.rows.get(unified_file_id)
|
||||
if isinstance(unified_file_id, Mapping):
|
||||
return next((self.rows[uid] for uid in unified_file_id.get("in", ()) if uid in self.rows), None)
|
||||
raw_file_filter: Final = where.get("flat_model_file_ids")
|
||||
if isinstance(raw_file_filter, Mapping):
|
||||
raw_file_id: Final = raw_file_filter.get("has")
|
||||
|
|
@ -69,6 +77,12 @@ class _InMemoryManagedFileTable:
|
|||
(row for row in self.rows.values() if raw_file_id in row.flat_model_file_ids),
|
||||
None,
|
||||
)
|
||||
raw_file_ids: Final = raw_file_filter.get("hasSome")
|
||||
if isinstance(raw_file_ids, list):
|
||||
return next(
|
||||
(row for row in self.rows.values() if set(raw_file_ids) & set(row.flat_model_file_ids)),
|
||||
None,
|
||||
)
|
||||
return None
|
||||
|
||||
async def upsert(
|
||||
|
|
@ -5660,62 +5674,48 @@ async def test_fine_tuning_provider_format_id_not_enforced():
|
|||
prisma_client.db.litellm_managedobjecttable.find_first.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"call_type", ["afile_content", "afile_retrieve", "afile_delete"]
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"file_id", [MODEL_ENCODED_OUTPUT_FILE_ID, RAW_PROVIDER_FILE_ID]
|
||||
)
|
||||
async def test_team_b_cannot_access_team_a_provider_format_file(
|
||||
call_type, file_id
|
||||
):
|
||||
"""
|
||||
Cross-team content/retrieve/delete of a model-encoded or raw provider
|
||||
file id must 403 when an ownership row exists for another team.
|
||||
"""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
def _team_a_output_row() -> LiteLLM_ManagedFileTable:
|
||||
return LiteLLM_ManagedFileTable(
|
||||
unified_file_id="file-encoded-team-a-output",
|
||||
file_object=OpenAIFileObject(
|
||||
id="file-encoded-team-a-output",
|
||||
object="file",
|
||||
bytes=0,
|
||||
created_at=1,
|
||||
filename="output.jsonl",
|
||||
purpose="batch_output",
|
||||
status="processed",
|
||||
),
|
||||
model_mappings={"gpt-4o-team-alias": RAW_PROVIDER_FILE_ID},
|
||||
flat_model_file_ids=[RAW_PROVIDER_FILE_ID],
|
||||
created_by="user_a",
|
||||
team_id="team_a",
|
||||
)
|
||||
|
||||
prisma_client = AsyncMock()
|
||||
prisma_client.db.litellm_managedfiletable.find_first.return_value = (
|
||||
_owned_record(created_by="user_a", team_id="team_a")
|
||||
)
|
||||
proxy_managed_files = PROXY_LiteLLMManagedFiles(
|
||||
MagicMock(), prisma_client=prisma_client
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("call_type", ["afile_content", "afile_retrieve", "afile_delete"])
|
||||
@pytest.mark.parametrize("file_id", [MODEL_ENCODED_OUTPUT_FILE_ID, RAW_PROVIDER_FILE_ID])
|
||||
async def test_team_b_cannot_access_team_a_provider_format_file(call_type: str, file_id: str):
|
||||
proxy_managed_files, _ = _managed_files_with_fake_prisma(_team_a_output_row())
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await proxy_managed_files.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_id="user_b", team_id="team_b", parent_otel_span=MagicMock()
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="user_b", team_id="team_b", parent_otel_span=MagicMock()),
|
||||
cache=MagicMock(),
|
||||
data={"file_id": file_id},
|
||||
call_type=call_type,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
prisma_client.db.litellm_managedfiletable.find_first.assert_awaited_once_with(
|
||||
where={"OR": [{"unified_file_id": file_id}, {"flat_model_file_ids": {"has": file_id}}]}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_same_team_can_access_provider_format_file():
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
prisma_client = AsyncMock()
|
||||
prisma_client.db.litellm_managedfiletable.find_first.return_value = (
|
||||
_owned_record(created_by="user_a", team_id="team_a")
|
||||
)
|
||||
proxy_managed_files = PROXY_LiteLLMManagedFiles(
|
||||
MagicMock(), prisma_client=prisma_client
|
||||
)
|
||||
proxy_managed_files, _ = _managed_files_with_fake_prisma(_team_a_output_row())
|
||||
|
||||
result = await proxy_managed_files.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_id="teammate_of_a", team_id="team_a", parent_otel_span=MagicMock()
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="teammate_of_a", team_id="team_a", parent_otel_span=MagicMock()),
|
||||
cache=MagicMock(),
|
||||
data={"file_id": MODEL_ENCODED_OUTPUT_FILE_ID},
|
||||
call_type="afile_content",
|
||||
|
|
@ -6573,3 +6573,192 @@ async def test_register_provider_file_returns_none_when_the_managed_row_cannot_b
|
|||
)
|
||||
|
||||
assert managed_file_id is None
|
||||
|
||||
|
||||
_GEMINI_OUTPUT_SPELLINGS: Final = (
|
||||
"files/batch-out1",
|
||||
"files%2Fbatch-out1",
|
||||
"batch-out1/",
|
||||
"https://generativelanguage.googleapis.com/v1beta/files/batch-out1",
|
||||
"HTTPS://generativelanguage.googleapis.com/v1beta/files/batch-out1",
|
||||
"x://h/files/batch-out1",
|
||||
)
|
||||
|
||||
|
||||
def _gemini_batch_output_row() -> LiteLLM_ManagedFileTable:
|
||||
return LiteLLM_ManagedFileTable(
|
||||
unified_file_id="file-encoded-gemini-output",
|
||||
file_object=OpenAIFileObject(
|
||||
id="file-encoded-gemini-output",
|
||||
object="file",
|
||||
bytes=0,
|
||||
created_at=1,
|
||||
filename="output.jsonl",
|
||||
purpose="batch_output",
|
||||
status="processed",
|
||||
),
|
||||
model_mappings={"model-gemini": "batch-out1"},
|
||||
flat_model_file_ids=["batch-out1"],
|
||||
created_by="owner",
|
||||
team_id="team-owner",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("spelling", _GEMINI_OUTPUT_SPELLINGS)
|
||||
async def test_enforce_provider_file_access_denies_every_spelling_of_another_keys_output_file(spelling: str):
|
||||
managed_files, _ = _managed_files_with_fake_prisma(_gemini_batch_output_row())
|
||||
stranger: Final = UserAPIKeyAuth(user_id="stranger", team_id="team-stranger")
|
||||
|
||||
with pytest.raises(HTTPException) as denied:
|
||||
await managed_files.enforce_provider_file_access(spelling, stranger)
|
||||
|
||||
assert denied.value.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("spelling", (*_GEMINI_OUTPUT_SPELLINGS, "batch-out1"))
|
||||
async def test_enforce_provider_file_access_lets_the_owner_read_any_spelling(spelling: str):
|
||||
managed_files, _ = _managed_files_with_fake_prisma(_gemini_batch_output_row())
|
||||
|
||||
owner: Final = UserAPIKeyAuth(user_id="owner", team_id="team-owner")
|
||||
teammate: Final = UserAPIKeyAuth(user_id="teammate", team_id="team-owner")
|
||||
|
||||
assert await managed_files.enforce_provider_file_access(spelling, owner) is None
|
||||
assert await managed_files.enforce_provider_file_access(spelling, teammate) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enforce_provider_file_access_still_passes_a_file_nobody_registered():
|
||||
managed_files, _ = _managed_files_with_fake_prisma(_gemini_batch_output_row())
|
||||
|
||||
stranger: Final = UserAPIKeyAuth(user_id="stranger")
|
||||
|
||||
assert await managed_files.enforce_provider_file_access("files/batch-other", stranger) is None
|
||||
|
||||
|
||||
_GEMINI_UPLOAD_URI: Final = "https://generativelanguage.googleapis.com/v1beta/files/upload-1"
|
||||
|
||||
|
||||
def _gemini_upload_row_stored_as_its_uri() -> LiteLLM_ManagedFileTable:
|
||||
return LiteLLM_ManagedFileTable(
|
||||
unified_file_id="file-encoded-gemini-upload",
|
||||
file_object=OpenAIFileObject(
|
||||
id="file-encoded-gemini-upload",
|
||||
object="file",
|
||||
bytes=12,
|
||||
created_at=1,
|
||||
filename="input.jsonl",
|
||||
purpose="batch",
|
||||
status="processed",
|
||||
),
|
||||
model_mappings={"model-gemini": _GEMINI_UPLOAD_URI},
|
||||
flat_model_file_ids=[_GEMINI_UPLOAD_URI],
|
||||
created_by="owner",
|
||||
team_id="team-owner",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"spelling", ("upload-1", "files/upload-1", "files%2Fupload-1", "upload-1/", _GEMINI_UPLOAD_URI)
|
||||
)
|
||||
async def test_enforce_provider_file_access_denies_a_stranger_a_row_stored_as_the_gemini_uri(spelling: str):
|
||||
managed_files, _ = _managed_files_with_fake_prisma(_gemini_upload_row_stored_as_its_uri())
|
||||
stranger: Final = UserAPIKeyAuth(user_id="stranger", team_id="team-stranger")
|
||||
|
||||
with pytest.raises(HTTPException) as denied:
|
||||
await managed_files.enforce_provider_file_access(spelling, stranger)
|
||||
|
||||
assert denied.value.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enforce_provider_file_access_denies_a_stranger_the_model_encoded_wrapper_of_another_keys_file():
|
||||
managed_files, _ = _managed_files_with_fake_prisma(_gemini_batch_output_row())
|
||||
wrapped: Final = encode_file_id_with_model(file_id="batch-out1", model="model-gemini")
|
||||
stranger: Final = UserAPIKeyAuth(user_id="stranger", team_id="team-stranger")
|
||||
|
||||
with pytest.raises(HTTPException) as denied:
|
||||
await managed_files.enforce_provider_file_access(wrapped, stranger)
|
||||
|
||||
assert denied.value.status_code == 403
|
||||
|
||||
|
||||
_AZURE_UPLOAD_URL_NAMED_LIKE_ANOTHER_ID: Final = "https://acct.blob.core.windows.net/uploads/0c4e.files%2Ffile-victim"
|
||||
|
||||
|
||||
def _managed_file_row(unified_file_id: str, provider_file_id: str, owner: str) -> LiteLLM_ManagedFileTable:
|
||||
return LiteLLM_ManagedFileTable(
|
||||
unified_file_id=unified_file_id,
|
||||
file_object=OpenAIFileObject(
|
||||
id=unified_file_id,
|
||||
object="file",
|
||||
bytes=12,
|
||||
created_at=1,
|
||||
filename="input.jsonl",
|
||||
purpose="batch",
|
||||
status="processed",
|
||||
),
|
||||
model_mappings={"model-a": provider_file_id},
|
||||
flat_model_file_ids=[provider_file_id],
|
||||
created_by=owner,
|
||||
team_id=f"team-{owner}",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enforce_provider_file_access_ignores_an_upload_named_to_collide_with_another_keys_id():
|
||||
managed_files, _ = _managed_files_with_fake_prisma(
|
||||
_managed_file_row("file-attacker-upload", _AZURE_UPLOAD_URL_NAMED_LIKE_ANOTHER_ID, "attacker"),
|
||||
_managed_file_row("file-victim-upload", "file-victim", "victim"),
|
||||
)
|
||||
attacker: Final = UserAPIKeyAuth(user_id="attacker", team_id="team-attacker")
|
||||
victim: Final = UserAPIKeyAuth(user_id="victim", team_id="team-victim")
|
||||
|
||||
with pytest.raises(HTTPException) as denied:
|
||||
await managed_files.enforce_provider_file_access("file-victim", attacker)
|
||||
|
||||
assert denied.value.status_code == 403
|
||||
assert await managed_files.enforce_provider_file_access("file-victim", victim) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_unified_file_id_indexes_the_canonical_provider_file_id():
|
||||
prisma_client: Final = AsyncMock()
|
||||
prisma_client.db.litellm_managedfiletable.upsert = AsyncMock(return_value=MagicMock())
|
||||
proxy_managed_files: Final = PROXY_LiteLLMManagedFiles(
|
||||
internal_usage_cache=MagicMock(), prisma_client=prisma_client
|
||||
)
|
||||
|
||||
await proxy_managed_files.store_unified_file_id(
|
||||
file_id="file-encoded-gemini-upload",
|
||||
file_object=None,
|
||||
litellm_parent_otel_span=None,
|
||||
model_mappings={"model-gemini": _GEMINI_UPLOAD_URI},
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="owner", team_id="team-owner"),
|
||||
)
|
||||
|
||||
written: Final = prisma_client.db.litellm_managedfiletable.upsert.call_args.kwargs["data"]
|
||||
assert written["create"]["flat_model_file_ids"] == ["upload-1"]
|
||||
assert written["update"]["flat_model_file_ids"] == ["upload-1"]
|
||||
assert json.loads(written["create"]["model_mappings"]) == {"model-gemini": _GEMINI_UPLOAD_URI}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_unified_file_ids_for_provider_file_ids_maps_every_spelling_onto_the_canonical_row():
|
||||
canonical_row: Final = _gemini_upload_row_stored_as_its_uri().model_copy(
|
||||
update={"flat_model_file_ids": ["upload-1"]}
|
||||
)
|
||||
managed_files, _ = _managed_files_with_fake_prisma(canonical_row)
|
||||
owner: Final = UserAPIKeyAuth(user_id="owner", team_id="team-owner")
|
||||
|
||||
mapping: Final = await managed_files.get_unified_file_ids_for_provider_file_ids(
|
||||
[_GEMINI_UPLOAD_URI, "files/upload-1", "upload-1", "upload-other"], owner
|
||||
)
|
||||
|
||||
assert dict(mapping) == {
|
||||
_GEMINI_UPLOAD_URI: "file-encoded-gemini-upload",
|
||||
"files/upload-1": "file-encoded-gemini-upload",
|
||||
"upload-1": "file-encoded-gemini-upload",
|
||||
}
|
||||
|
|
|
|||
|
|
@ -338,14 +338,7 @@ async def test_provider_file_id_resolver_returns_owned_mappings_with_owner_scope
|
|||
"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"],
|
||||
},
|
||||
}
|
||||
)
|
||||
find_many.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -16,12 +16,15 @@ from litellm.proxy._types import UserAPIKeyAuth
|
|||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
apply_team_provider_credentials,
|
||||
apply_unified_file_ids,
|
||||
canonical_provider_file_id,
|
||||
encode_file_id_with_model,
|
||||
get_credentials_for_model,
|
||||
get_team_provider_credentials,
|
||||
get_team_provider_deployment,
|
||||
is_litellm_executed_batch,
|
||||
map_raw_file_ids_to_unified,
|
||||
prepare_data_with_credentials,
|
||||
provider_file_id_spellings,
|
||||
)
|
||||
from litellm.proxy.route_llm_request import ProxyModelNotFoundError
|
||||
from litellm.proxy.utils import InternalUsageCache, handle_exception_on_proxy
|
||||
|
|
@ -969,3 +972,52 @@ def test_create_paths_keep_a_request_bucket_and_reads_keep_the_deployment_bucket
|
|||
assert create_with_bucket == {field: "team-a-bucket", "api_key": "sk-deployment"}
|
||||
assert create_without_bucket == {field: "model-bucket", "api_key": "sk-deployment"}
|
||||
assert read_with_bucket == {field: "model-bucket", "api_key": "sk-deployment"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("file_id", "canonical"),
|
||||
(
|
||||
("batch-out1", "batch-out1"),
|
||||
("files/batch-out1", "batch-out1"),
|
||||
("files%2Fbatch-out1", "batch-out1"),
|
||||
("batch-out1/", "batch-out1"),
|
||||
("/files/batch-out1/", "batch-out1"),
|
||||
("https://generativelanguage.googleapis.com/v1beta/files/batch-out1", "batch-out1"),
|
||||
("file-abc123", "file-abc123"),
|
||||
("s3://bucket/path/files/output.jsonl", "s3://bucket/path/files/output.jsonl"),
|
||||
("gs://bucket/prefix/predictions.jsonl", "gs://bucket/prefix/predictions.jsonl"),
|
||||
(
|
||||
"https://acct.blob.core.windows.net/uploads/0c4e.files%2Ffile-victim",
|
||||
"https://acct.blob.core.windows.net/uploads/0c4e.files%2Ffile-victim",
|
||||
),
|
||||
("x://h/files/batch-out1", "x://h/files/batch-out1"),
|
||||
),
|
||||
)
|
||||
def test_canonical_provider_file_id_collapses_every_spelling_to_the_stored_id(file_id: str, canonical: str):
|
||||
assert canonical_provider_file_id(file_id) == canonical
|
||||
|
||||
|
||||
def test_provider_file_id_spellings_unwraps_a_model_encoded_id_and_covers_the_stored_forms():
|
||||
wrapped: Final = encode_file_id_with_model(file_id="batch-out1", model="gemini-deployment")
|
||||
|
||||
assert provider_file_id_spellings(wrapped) == (
|
||||
wrapped,
|
||||
"batch-out1",
|
||||
"files/batch-out1",
|
||||
"https://generativelanguage.googleapis.com/v1beta/files/batch-out1",
|
||||
)
|
||||
for spelling in ("batch-out1", "files%2Fbatch-out1", "batch-out1/"):
|
||||
assert set(provider_file_id_spellings(spelling)) >= set(provider_file_id_spellings(wrapped)[1:]), spelling
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"respelling",
|
||||
(
|
||||
"HTTPS://generativelanguage.googleapis.com/v1beta/files/batch-out1",
|
||||
"x://h/files/batch-out1",
|
||||
"S3://b/files/batch-out1",
|
||||
"x%3A%2F%2Fh%2Ffiles%2Fbatch-out1",
|
||||
),
|
||||
)
|
||||
def test_provider_file_id_spellings_reduce_any_scheme_the_way_gemini_does(respelling: str):
|
||||
assert "batch-out1" in provider_file_id_spellings(respelling)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue