diff --git a/litellm/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/litellm/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index bbb09e75b47..9c594c70634 100644 --- a/litellm/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/litellm/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -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) ) ) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index a81fced9021..222edfb7612 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -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/`` never indexes someone else's ````.""" + 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. diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py index d8a647fda27..d165b5fb06e 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_files.py +++ b/tests/unit/enterprise/proxy/hooks/test_managed_files.py @@ -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", + } diff --git a/tests/unit/enterprise/proxy/test_managed_files_hook.py b/tests/unit/enterprise/proxy/test_managed_files_hook.py index 0704076ccca..0da8ec9ebb7 100644 --- a/tests/unit/enterprise/proxy/test_managed_files_hook.py +++ b/tests/unit/enterprise/proxy/test_managed_files_hook.py @@ -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 diff --git a/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py index 2f469bb8684..3e748ed8cb8 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py @@ -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)