diff --git a/litellm/llms/base_llm/files/litellm_db_storage_backend.py b/litellm/llms/base_llm/files/litellm_db_storage_backend.py index bca4b8f4c6f..a686062b2f7 100644 --- a/litellm/llms/base_llm/files/litellm_db_storage_backend.py +++ b/litellm/llms/base_llm/files/litellm_db_storage_backend.py @@ -1,13 +1,9 @@ -from collections.abc import Mapping from typing import TYPE_CHECKING, Final from litellm.llms.base_llm.files.storage_backend import BaseFileStorageBackend -from litellm.repositories.prisma_protocols import TableActions -from litellm.repositories.table_repositories import PrismaTableRepository +from litellm.repositories.managed_file_content_repository import ManagedFileContentRepository if TYPE_CHECKING: - from prisma import models as prisma_models - from litellm.proxy.utils import PrismaClient LITELLM_DB_STORAGE_BACKEND_NAME: Final = "litellm_db" @@ -20,21 +16,9 @@ def storage_url_to_row_id(storage_url: str) -> str: return storage_url.removeprefix(LITELLM_DB_STORAGE_URL_PREFIX) -def _where_id(storage_url: str) -> Mapping[str, str]: - return {"id": storage_url_to_row_id(storage_url)} # mutable-ok: Prisma filter - - -class ManagedFileContentRepository(PrismaTableRepository["prisma_models.LiteLLM_ManagedFileContentTable"]): - table_name = "litellm_managedfilecontenttable" - - class LiteLLMDbStorageBackend(BaseFileStorageBackend): def __init__(self, prisma_client: "PrismaClient") -> None: - self._prisma_client = prisma_client - - @property - def _table(self) -> "TableActions[prisma_models.LiteLLM_ManagedFileContentTable]": - return ManagedFileContentRepository(self._prisma_client).table + self._contents = ManagedFileContentRepository(prisma_client) async def upload_file( self, @@ -44,22 +28,13 @@ class LiteLLMDbStorageBackend(BaseFileStorageBackend): path_prefix: str | None = None, file_naming_strategy: str = "uuid", ) -> str: - from prisma import Base64 - - data: Final = {"content": Base64.encode(file_content)} # mutable-ok: Prisma payload - row: Final = await self._table.create(data=data) - return f"{LITELLM_DB_STORAGE_URL_PREFIX}{row.id}" + return f"{LITELLM_DB_STORAGE_URL_PREFIX}{await self._contents.store(file_content)}" async def download_file(self, storage_url: str) -> bytes: - row: Final = await self._table.find_unique(where=_where_id(storage_url)) - if row is None: + content: Final = await self._contents.load(storage_url_to_row_id(storage_url)) + if content is None: raise ValueError(f"No stored file content for {storage_url}") - return row.content.decode() + return content async def delete_file(self, storage_url: str) -> None: - from prisma.errors import RecordNotFoundError - - try: - await self._table.delete(where=_where_id(storage_url)) - except RecordNotFoundError: - return + await self._contents.delete(storage_url_to_row_id(storage_url)) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 6d60bc3fda5..a5921cc6380 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -37,7 +37,10 @@ from litellm.llms.base_llm.files.transformation import BaseFileEndpoints from litellm.llms.base_llm.managed_resources.isolation import build_list_page from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.batches_endpoints.litellm_executed_batches import resolve_litellm_executed_provider +from litellm.proxy.batches_endpoints.litellm_executed_batches import ( + litellm_executed_provider_of, + resolve_litellm_executed_provider, +) from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, @@ -69,6 +72,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, add_internal_model_credentials, apply_team_provider_credentials, + authorize_model_for_key, encode_file_id_with_model, extract_file_creation_params, get_authorized_credentials_for_model, @@ -102,16 +106,29 @@ from litellm.types.llms.openai import ( router: Final = APIRouter() +def _names_a_litellm_executed_provider(llm_router: Router, candidate: str, team_id: str | None) -> bool: + credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=candidate, team_id=team_id) + return credentials is not None and litellm_executed_provider_of(credentials) is not None + + async def _litellm_executed_batch_input_model( llm_router: Router | None, purpose: OpenAIFilesPurpose, model: str | None, target_model_names_list: Sequence[str], - team_id: str | None, + user_api_key_dict: UserAPIKeyAuth, ) -> str | None: if llm_router is None: return None candidates: Final = (model,) if model is not None else tuple(target_model_names_list) + team_id: Final = user_api_key_dict.team_id + await asyncio.gather( + *( + authorize_model_for_key(model_id=candidate, llm_router=llm_router, user_api_key_dict=user_api_key_dict) + for candidate in candidates + if _names_a_litellm_executed_provider(llm_router, candidate, team_id) + ) + ) providers: Final = await asyncio.gather( *(resolve_litellm_executed_provider(llm_router, candidate, team_id) for candidate in candidates) ) @@ -289,7 +306,7 @@ async def route_create_file( """ executed_model: Final = await _litellm_executed_batch_input_model( - llm_router, purpose, model, target_model_names_list, user_api_key_dict.team_id + llm_router, purpose, model, target_model_names_list, user_api_key_dict ) explicit_storage: Final = target_storage if target_storage and target_storage != "default" else None storage: Final = explicit_storage or (LITELLM_DB_STORAGE_BACKEND_NAME if executed_model is not None else None) diff --git a/litellm/repositories/managed_file_content_repository.py b/litellm/repositories/managed_file_content_repository.py new file mode 100644 index 00000000000..8810269279c --- /dev/null +++ b/litellm/repositories/managed_file_content_repository.py @@ -0,0 +1,30 @@ +from typing import TYPE_CHECKING, Final + +from litellm.repositories.table_repositories import PrismaTableRepository + +if TYPE_CHECKING: + from prisma import models as prisma_models # noqa: F401 # used by the quoted base-class subscript + + +class ManagedFileContentRepository(PrismaTableRepository["prisma_models.LiteLLM_ManagedFileContentTable"]): + table_name = "litellm_managedfilecontenttable" + + async def store(self, content: bytes) -> str: + from prisma import Base64 + + row: Final = await self.table.create( + data={"content": Base64.encode(content)} # mutable-ok: prisma payloads are plain dicts + ) + return row.id + + async def load(self, row_id: str) -> bytes | None: + row: Final = await self.table.find_unique(where={"id": row_id}) # mutable-ok: prisma filters are plain dicts + return None if row is None else row.content.decode() + + async def delete(self, row_id: str) -> None: + from prisma.errors import RecordNotFoundError + + try: + await self.table.delete(where={"id": row_id}) # mutable-ok: prisma filters are plain dicts + except RecordNotFoundError: + return diff --git a/tests/test_litellm/llms/base_llm/files/test_storage_backend_factory.py b/tests/test_litellm/llms/base_llm/files/test_storage_backend_factory.py index 39b0adb56fc..945691c5b98 100644 --- a/tests/test_litellm/llms/base_llm/files/test_storage_backend_factory.py +++ b/tests/test_litellm/llms/base_llm/files/test_storage_backend_factory.py @@ -1,21 +1,27 @@ -from unittest.mock import MagicMock +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock import pytest from litellm.llms.base_llm.files.litellm_db_storage_backend import ( LITELLM_DB_STORAGE_BACKEND_NAME, + LITELLM_DB_STORAGE_URL_PREFIX, LiteLLMDbStorageBackend, ) from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend -def test_litellm_db_backend_is_built_on_the_given_prisma_client(): - prisma_client = MagicMock() +@pytest.mark.asyncio +async def test_litellm_db_backend_stores_through_the_given_prisma_client(): + table = MagicMock(create=AsyncMock(return_value=SimpleNamespace(id="row-1"))) + prisma_client = MagicMock(db=MagicMock(litellm_managedfilecontenttable=table)) backend = get_storage_backend(LITELLM_DB_STORAGE_BACKEND_NAME, prisma_client=prisma_client) assert isinstance(backend, LiteLLMDbStorageBackend) - assert backend._table is prisma_client.db.litellm_managedfilecontenttable + stored_at = await backend.upload_file(file_content=b"line\n", filename="input.jsonl", content_type="text/plain") + assert stored_at == f"{LITELLM_DB_STORAGE_URL_PREFIX}row-1" + table.create.assert_awaited_once() def test_litellm_db_backend_without_a_database_is_rejected(): diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 6c0f4c012ba..6787aaa3525 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -706,6 +706,30 @@ def test_batch_upload_for_a_litellm_executed_model_is_kept_by_litellm( assert kwargs["purpose"] == "batch" +@pytest.mark.parametrize( + "headers, form", + [({"x-litellm-model": "my-vllm"}, {}), ({}, {"target_model_names": "my-vllm"})], + ids=["x-litellm-model header", "target_model_names form field"], +) +def test_batch_upload_for_a_litellm_executed_model_the_key_cannot_call_is_refused_before_the_server_is_probed( + batch_upload_seams, headers: dict[str, str], form: dict[str, str] +): + import litellm.proxy.proxy_server as ps + + stored, provider_upload, upstream_files_route = batch_upload_seams + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="restricted-user", models=["gemini-2.0-flash"] + ) + + response = _upload_batch_file(headers, form) + + assert response.status_code == 403, response.text + assert "my-vllm" in response.text + assert upstream_files_route.call_count == 0 + stored.assert_not_awaited() + provider_upload.assert_not_awaited() + + def test_batch_upload_naming_an_executed_and_a_provider_model_is_rejected(batch_upload_seams): stored, provider_upload, _ = batch_upload_seams