mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-13 03:59:12 +00:00
Merge cf5832bf84 into 145c53bedf
This commit is contained in:
mrinal-berri
GitHub
commit
48bf44522b
8 changed files with 537 additions and 24 deletions
|
|
@ -1162,7 +1162,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
)
|
||||
return False
|
||||
|
||||
async def check_file_ids_access(self, file_ids: List[str], user_api_key_dict: UserAPIKeyAuth) -> None:
|
||||
async def check_file_ids_access(self, file_ids: Sequence[str], user_api_key_dict: UserAPIKeyAuth) -> None:
|
||||
"""
|
||||
Check if the user has access to a list of file IDs.
|
||||
Only checks managed (unified) file IDs.
|
||||
|
|
@ -1269,6 +1269,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
elif call_type == CallTypes.acreate_batch.value:
|
||||
input_file_id = cast(Optional[str], data.get("input_file_id"))
|
||||
if input_file_id:
|
||||
await self.check_file_ids_access((input_file_id,), user_api_key_dict)
|
||||
model_file_id_mapping = await self.get_model_file_id_mapping(
|
||||
[input_file_id], user_api_key_dict.parent_otel_span
|
||||
)
|
||||
|
|
@ -1315,11 +1316,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
elif retrieve_object_id and accessor_key == "batch_id":
|
||||
await self.enforce_batch_object_access(retrieve_object_id, user_api_key_dict)
|
||||
elif call_type == CallTypes.acreate_fine_tuning_job.value:
|
||||
input_file_id = cast(Optional[str], data.get("training_file"))
|
||||
if input_file_id:
|
||||
model_file_id_mapping = await self.get_model_file_id_mapping(
|
||||
[input_file_id], user_api_key_dict.parent_otel_span
|
||||
)
|
||||
fine_tuning_file_ids: Final = tuple(
|
||||
file_id
|
||||
for file_id in (data.get("training_file"), data.get("validation_file"))
|
||||
if isinstance(file_id, str) and file_id
|
||||
)
|
||||
await self.check_file_ids_access(fine_tuning_file_ids, user_api_key_dict)
|
||||
|
||||
return data
|
||||
|
||||
|
|
|
|||
|
|
@ -1128,6 +1128,24 @@ class _PROXY_BatchRateLimiter(CustomLogger):
|
|||
},
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _check_managed_file_ownership(
|
||||
managed_files_obj: object,
|
||||
file_id: str,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> None:
|
||||
"""Deny before the input file is read, so a caller cannot make the proxy
|
||||
download, count, or charge its own limits for another tenant's file."""
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import ManagedResourceAccessChecker
|
||||
|
||||
if not isinstance(managed_files_obj, ManagedResourceAccessChecker):
|
||||
raise HTTPException(status_code=500, detail="Managed resource ownership validation is unavailable.")
|
||||
if not await managed_files_obj.can_user_call_unified_file_id(file_id, user_api_key_dict):
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=f"User {user_api_key_dict.user_id} does not have access to the file {file_id}",
|
||||
)
|
||||
|
||||
async def _fetch_managed_file_content(
|
||||
self,
|
||||
file_id: str,
|
||||
|
|
@ -1170,6 +1188,8 @@ class _PROXY_BatchRateLimiter(CustomLogger):
|
|||
if llm_router is None:
|
||||
raise ValueError("llm_router not available. Cannot access managed files.")
|
||||
|
||||
await self._check_managed_file_ownership(managed_files_obj, file_id, user_api_key_dict)
|
||||
|
||||
# Use the managed files hook to get file content
|
||||
# This properly handles user permissions and file ownership
|
||||
file_content: Final = await managed_files_obj.afile_content(
|
||||
|
|
|
|||
|
|
@ -122,6 +122,7 @@ from litellm.proxy.management_helpers.access_group_key_sync import (
|
|||
sync_key_update_access_group_membership,
|
||||
)
|
||||
from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at
|
||||
from litellm.proxy.management_helpers.managed_resource_owner_sync import repoint_key_owned_managed_resources
|
||||
from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: F401 # legacy module exports
|
||||
ObjectPermissionUpsert,
|
||||
_set_object_permission, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
|
|
@ -5957,6 +5958,13 @@ async def _execute_virtual_key_regeneration(
|
|||
|
||||
await evict_and_broadcast(cache_keys=jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache)
|
||||
|
||||
await repoint_key_owned_managed_resources(
|
||||
prisma_client=prisma_client,
|
||||
existing_key_row=key_in_db,
|
||||
previous_key_token=hashed_api_key,
|
||||
new_key_token=new_token_hash,
|
||||
)
|
||||
|
||||
# After credential invalidation, so a failure here can never keep the old key alive.
|
||||
await sync_key_regeneration_access_group_membership(
|
||||
prisma_client=prisma_client,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,44 @@
|
|||
"""Carry key-owned managed resources over to a regenerated key's new hash."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Final, Protocol
|
||||
|
||||
from litellm.proxy._types import LiteLLM_VerificationToken
|
||||
from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper
|
||||
from litellm.repositories.table_repositories import ManagedObjectRepository
|
||||
|
||||
_KEY_OWNED_TABLES: Final = (
|
||||
"LiteLLM_ManagedFileTable",
|
||||
"LiteLLM_ManagedObjectTable",
|
||||
"LiteLLM_ManagedVectorStoreTable",
|
||||
)
|
||||
|
||||
|
||||
class _RawWriter(Protocol):
|
||||
async def execute_raw(self, query: str, *args: str) -> int: ...
|
||||
|
||||
|
||||
def _raw_writer(prisma_client: object) -> _RawWriter:
|
||||
db: Final = ManagedObjectRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client
|
||||
return writer_wrapper(db) # pyright: ignore[reportAny, reportReturnType] # untyped Prisma client behind the pin
|
||||
|
||||
|
||||
async def repoint_key_owned_managed_resources(
|
||||
prisma_client: object,
|
||||
existing_key_row: LiteLLM_VerificationToken,
|
||||
previous_key_token: str,
|
||||
new_key_token: str,
|
||||
) -> None:
|
||||
"""A key with neither a user_id nor a team_id owns its managed files, batches and
|
||||
vector stores as ``key:<hashed token>`` (``resolve_resource_owner_id``), so a
|
||||
regenerated key would otherwise be denied its own resources."""
|
||||
if existing_key_row.user_id is not None or existing_key_row.team_id is not None:
|
||||
return
|
||||
writer: Final = _raw_writer(prisma_client)
|
||||
for table in _KEY_OWNED_TABLES:
|
||||
await writer.execute_raw(
|
||||
f'UPDATE "{table}" SET "created_by" = $2 WHERE "created_by" = $1',
|
||||
f"key:{previous_key_token}",
|
||||
f"key:{new_key_token}",
|
||||
)
|
||||
|
|
@ -0,0 +1,271 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import JSON_OBJECT, Gateway, Scenario, string_value
|
||||
from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario
|
||||
from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse
|
||||
from pydantic import JsonValue
|
||||
|
||||
_FILE: Final = {
|
||||
"id": "file-in-$REQUEST_ID",
|
||||
"object": "file",
|
||||
"purpose": "batch",
|
||||
"bytes": 100,
|
||||
"created_at": 1,
|
||||
"filename": "input.jsonl",
|
||||
"status": "processed",
|
||||
}
|
||||
_BATCH: Final = {
|
||||
"id": "batch-$REQUEST_ID",
|
||||
"object": "batch",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"errors": None,
|
||||
"input_file_id": "file-in-$REQUEST_ID",
|
||||
"completion_window": "24h",
|
||||
"status": "validating",
|
||||
"output_file_id": None,
|
||||
"error_file_id": None,
|
||||
"created_at": 1,
|
||||
"expires_at": 1,
|
||||
"request_counts": {"total": 1, "completed": 0, "failed": 0},
|
||||
"metadata": None,
|
||||
}
|
||||
_JOB: Final = {
|
||||
"id": "ftjob-$REQUEST_ID",
|
||||
"object": "fine_tuning.job",
|
||||
"model": "gpt-4o-mini-2024-07-18",
|
||||
"created_at": 1,
|
||||
"fine_tuned_model": None,
|
||||
"organization_id": "org-integration",
|
||||
"result_files": [],
|
||||
"status": "validating_files",
|
||||
"validation_file": None,
|
||||
"training_file": "file-in-$REQUEST_ID",
|
||||
"hyperparameters": {"n_epochs": "auto"},
|
||||
"seed": 1,
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Tenants:
|
||||
model: str
|
||||
handle: ScenarioHandle
|
||||
owner_key: str
|
||||
teammate_key: str
|
||||
outsider_key: str
|
||||
owner_file_id: str
|
||||
|
||||
|
||||
def _routes() -> RoutedResponse:
|
||||
return RoutedResponse(
|
||||
content_type="application/x-routed",
|
||||
routes={
|
||||
"POST /files": JsonResponse(content_type="application/json", body=_FILE),
|
||||
"POST /batches": JsonResponse(content_type="application/json", body=_BATCH),
|
||||
"POST /fine_tuning/jobs": JsonResponse(content_type="application/json", body=_JOB),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _upload(scenario: Scenario, key: str, model: str, purpose: str) -> str:
|
||||
line: Final = json.dumps(
|
||||
{
|
||||
"custom_id": "req-1",
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {"model": model, "messages": [{"role": "user", "content": "ping"}]},
|
||||
}
|
||||
).encode()
|
||||
uploaded: Final = scenario.gateway.request_multipart(
|
||||
"/v1/files",
|
||||
{"purpose": purpose, "target_model_names": model},
|
||||
{"file": ("input.jsonl", line + b"\n", "application/jsonl")},
|
||||
key=key,
|
||||
)
|
||||
assert uploaded.status_code == 200, uploaded.text
|
||||
return string_value(JSON_OBJECT.validate_json(uploaded.content)["id"])
|
||||
|
||||
|
||||
def _tenants(scenario: Scenario, purpose: str) -> _Tenants:
|
||||
handle: Final = register_scenario(f"managed-file-create-ownership-{uuid.uuid4().hex}", _routes())
|
||||
scenario.cleanups.callback(delete_scenario, handle)
|
||||
model: Final = scenario.model(api_base=handle.api_base())
|
||||
team_a: Final = scenario.team(models=[model])
|
||||
team_b: Final = scenario.team(models=[model])
|
||||
owner_key: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), team_id=team_a, models=[model])
|
||||
teammate_key: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), team_id=team_a, models=[model])
|
||||
outsider_key: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), team_id=team_b, models=[model])
|
||||
return _Tenants(model, handle, owner_key, teammate_key, outsider_key, _upload(scenario, owner_key, model, purpose))
|
||||
|
||||
|
||||
def _upstream_creates(gateway: Gateway, handle: ScenarioHandle, path: str) -> int:
|
||||
response: Final = httpx.get(f"{gateway.upstream_url}/__observations", timeout=15, trust_env=False)
|
||||
assert response.status_code == 200, response.text
|
||||
requests: Final = JSON_OBJECT.validate_json(response.content)["requests"]
|
||||
assert isinstance(requests, list), response.text
|
||||
return sum(
|
||||
1
|
||||
for request in requests
|
||||
if isinstance(request, dict)
|
||||
and request.get("method") == "POST"
|
||||
and handle.scenario_id in str(request.get("path"))
|
||||
and str(request.get("path")).endswith(path)
|
||||
)
|
||||
|
||||
|
||||
def _create_batch(gateway: Gateway, key: str, model: str, input_file_id: JsonValue) -> httpx.Response:
|
||||
return gateway.request(
|
||||
"POST",
|
||||
"/v1/batches",
|
||||
{
|
||||
"input_file_id": input_file_id,
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"completion_window": "24h",
|
||||
"model": model,
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
|
||||
|
||||
def _create_job(gateway: Gateway, key: str, model: str, files: dict[str, JsonValue]) -> httpx.Response:
|
||||
return gateway.request("POST", "/v1/fine_tuning/jobs", {"model": model, **files}, key=key)
|
||||
|
||||
|
||||
def test_batch_create_from_another_teams_managed_file_is_403_and_never_reaches_the_provider(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
tenants: Final = _tenants(scenario, "batch")
|
||||
denied: Final = _create_batch(gateway, tenants.outsider_key, tenants.model, tenants.owner_file_id)
|
||||
assert denied.status_code == 403, denied.text
|
||||
assert tenants.owner_file_id in denied.text, denied.text
|
||||
assert _upstream_creates(gateway, tenants.handle, "/batches") == 0
|
||||
|
||||
|
||||
def test_batch_create_from_own_or_teammates_managed_file_still_dispatches(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
tenants: Final = _tenants(scenario, "batch")
|
||||
owner: Final = _create_batch(gateway, tenants.owner_key, tenants.model, tenants.owner_file_id)
|
||||
teammate: Final = _create_batch(gateway, tenants.teammate_key, tenants.model, tenants.owner_file_id)
|
||||
admin: Final = _create_batch(gateway, gateway.key, tenants.model, tenants.owner_file_id)
|
||||
assert (owner.status_code, teammate.status_code, admin.status_code) == (200, 200, 200), (
|
||||
owner.text,
|
||||
teammate.text,
|
||||
admin.text,
|
||||
)
|
||||
assert JSON_OBJECT.validate_json(owner.content)["input_file_id"] == tenants.owner_file_id, owner.text
|
||||
assert _upstream_creates(gateway, tenants.handle, "/batches") == 3
|
||||
|
||||
|
||||
def test_batch_create_from_a_managed_file_id_with_no_owner_row_is_404(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
tenants: Final = _tenants(scenario, "batch")
|
||||
forged_raw: Final = (
|
||||
f"litellm_proxy:application/octet-stream;unified_id,{uuid.uuid4()};target_model_names,{tenants.model};"
|
||||
"llm_output_file_id,file-forged;llm_output_file_model_id,model-forged"
|
||||
)
|
||||
forged: Final = base64.urlsafe_b64encode(forged_raw.encode()).decode().rstrip("=")
|
||||
missing: Final = _create_batch(gateway, tenants.outsider_key, tenants.model, forged)
|
||||
assert missing.status_code == 404, missing.text
|
||||
assert _upstream_creates(gateway, tenants.handle, "/batches") == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["training_file", "validation_file"])
|
||||
def test_fine_tuning_create_from_another_teams_managed_file_is_403(gateway: Gateway, field: str) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
tenants: Final = _tenants(scenario, "fine-tune")
|
||||
own_training: Final = _upload(scenario, tenants.outsider_key, tenants.model, "fine-tune")
|
||||
files: Final[dict[str, JsonValue]] = {"training_file": own_training, field: tenants.owner_file_id}
|
||||
denied: Final = _create_job(gateway, tenants.outsider_key, tenants.model, files)
|
||||
assert denied.status_code == 403, denied.text
|
||||
assert _upstream_creates(gateway, tenants.handle, "/fine_tuning/jobs") == 0
|
||||
|
||||
|
||||
def test_fine_tuning_create_from_own_managed_files_still_dispatches(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
tenants: Final = _tenants(scenario, "fine-tune")
|
||||
validation: Final = _upload(scenario, tenants.teammate_key, tenants.model, "fine-tune")
|
||||
created: Final = _create_job(
|
||||
gateway,
|
||||
tenants.owner_key,
|
||||
tenants.model,
|
||||
{"training_file": tenants.owner_file_id, "validation_file": validation},
|
||||
)
|
||||
assert created.status_code == 200, created.text
|
||||
assert _upstream_creates(gateway, tenants.handle, "/fine_tuning/jobs") == 1
|
||||
|
||||
|
||||
def _upstream_calls(gateway: Gateway, handle: ScenarioHandle, method: str, suffix: str) -> int:
|
||||
response: Final = httpx.get(f"{gateway.upstream_url}/__observations", timeout=15, trust_env=False)
|
||||
assert response.status_code == 200, response.text
|
||||
requests: Final = JSON_OBJECT.validate_json(response.content)["requests"]
|
||||
assert isinstance(requests, list), response.text
|
||||
return sum(
|
||||
1
|
||||
for request in requests
|
||||
if isinstance(request, dict)
|
||||
and request.get("method") == method
|
||||
and handle.scenario_id in str(request.get("path"))
|
||||
and str(request.get("path")).endswith(suffix)
|
||||
)
|
||||
|
||||
|
||||
def test_batch_create_from_another_teams_file_is_403_before_the_rate_limiter_reads_it(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
scenario_id: Final = f"managed-file-create-ownership-{uuid.uuid4().hex}"
|
||||
handle: Final = register_scenario(scenario_id, _routes())
|
||||
scenario.cleanups.callback(delete_scenario, handle)
|
||||
model: Final = scenario.model(api_base=handle.api_base())
|
||||
line: Final = json.dumps(
|
||||
{
|
||||
"custom_id": "req-1",
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {"model": model, "messages": [{"role": "user", "content": "ping"}], "max_tokens": 5},
|
||||
}
|
||||
)
|
||||
register_scenario(
|
||||
scenario_id,
|
||||
RoutedResponse(
|
||||
content_type="application/x-routed",
|
||||
routes={
|
||||
**_routes().routes,
|
||||
"GET /files/file-in-$REQUEST_ID/content": TextResponse(
|
||||
content_type="application/jsonl", body=line + "\n"
|
||||
),
|
||||
},
|
||||
),
|
||||
)
|
||||
team_a: Final = scenario.team(models=[model])
|
||||
team_b: Final = scenario.team(models=[model])
|
||||
owner: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), team_id=team_a)
|
||||
outsider: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), team_id=team_b, tpm_limit=1)
|
||||
owner_file_id: Final = _upload(scenario, owner, model, "batch")
|
||||
denied: Final = _create_batch(gateway, outsider, model, owner_file_id)
|
||||
assert denied.status_code == 403, denied.text
|
||||
assert owner_file_id in denied.text, denied.text
|
||||
assert _upstream_calls(gateway, handle, "GET", "/content") == 0
|
||||
|
||||
|
||||
def test_regenerated_key_without_user_or_team_keeps_its_managed_files(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
tenants: Final = _tenants(scenario, "batch")
|
||||
original: Final = string_value(gateway.post("/key/generate", {})["key"])
|
||||
file_id: Final = _upload(scenario, original, tenants.model, "batch")
|
||||
regenerated: Final = string_value(gateway.post("/key/regenerate", {"key": original})["key"])
|
||||
scenario.cleanups.callback(scenario.delete_key, regenerated)
|
||||
batch: Final = _create_batch(gateway, regenerated, tenants.model, file_id)
|
||||
retrieved: Final = gateway.request("GET", f"/v1/files/{file_id}", key=regenerated)
|
||||
stranger: Final = _create_batch(
|
||||
gateway, string_value(gateway.post("/key/generate", {})["key"]), tenants.model, file_id
|
||||
)
|
||||
assert (batch.status_code, retrieved.status_code, stranger.status_code) == (200, 200, 403), (
|
||||
batch.text,
|
||||
retrieved.text,
|
||||
stranger.text,
|
||||
)
|
||||
|
|
@ -5264,3 +5264,113 @@ async def test_list_user_batches_target_model_names_filter_rejected_with_400():
|
|||
assert exc.value.type == "invalid_request_error"
|
||||
assert exc.value.param == "target_model_names"
|
||||
assert exc.value.message == "Filtering by 'target_model_names' is not supported when using managed batches."
|
||||
|
||||
|
||||
def _tenant_a_input_file_id() -> str:
|
||||
raw: Final = (
|
||||
"litellm_proxy:application/octet-stream;unified_id,7d1c4b0e-5a7f-4a43-9a0e-3f8e2a1b6c11;"
|
||||
"target_model_names,gpt-4o-mini;llm_output_file_id,file-tenant-a;llm_output_file_model_id,model-a"
|
||||
)
|
||||
return base64.urlsafe_b64encode(raw.encode()).decode().rstrip("=")
|
||||
|
||||
|
||||
def _tenant_a_input_file_row(unified_file_id: str) -> LiteLLM_ManagedFileTable:
|
||||
return LiteLLM_ManagedFileTable(
|
||||
unified_file_id=unified_file_id,
|
||||
file_object=None,
|
||||
model_mappings={"model-a": "file-tenant-a"},
|
||||
flat_model_file_ids=["file-tenant-a"],
|
||||
created_by="user_a",
|
||||
team_id="team_a",
|
||||
)
|
||||
|
||||
|
||||
_INPUT_FILE_CREATE_CALLS: Final = (
|
||||
("acreate_batch", "input_file_id"),
|
||||
("acreate_fine_tuning_job", "training_file"),
|
||||
("acreate_fine_tuning_job", "validation_file"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("call_type", "file_field"), _INPUT_FILE_CREATE_CALLS)
|
||||
async def test_create_from_another_tenants_managed_file_is_denied(call_type: str, file_field: str) -> None:
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
unified_file_id: Final = _tenant_a_input_file_id()
|
||||
managed_files, _ = _managed_files_with_fake_prisma(_tenant_a_input_file_row(unified_file_id))
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await managed_files.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="user_b", team_id="team_b", parent_otel_span=MagicMock()),
|
||||
cache=MagicMock(),
|
||||
data={"model": "gpt-4o-mini", "training_file": "file-own-training", file_field: unified_file_id},
|
||||
call_type=call_type,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 403, f"{call_type}.{file_field} let tenant B use tenant A's file"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("call_type", "file_field"), _INPUT_FILE_CREATE_CALLS)
|
||||
async def test_create_from_managed_file_id_with_no_db_row_is_404(call_type: str, file_field: str) -> None:
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
managed_files, _ = _managed_files_with_fake_prisma()
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await managed_files.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="user_a", team_id="team_a", parent_otel_span=MagicMock()),
|
||||
cache=MagicMock(),
|
||||
data={"model": "gpt-4o-mini", file_field: _tenant_a_input_file_id()},
|
||||
call_type=call_type,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 404, f"{call_type}.{file_field} dispatched a managed id with no owner row"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("call_type", "file_field"), _INPUT_FILE_CREATE_CALLS)
|
||||
@pytest.mark.parametrize(
|
||||
"caller_kwargs",
|
||||
[
|
||||
{"user_id": "user_a", "team_id": "team_a"},
|
||||
{"user_id": "teammate_of_a", "team_id": "team_a"},
|
||||
{"user_id": "admin", "team_id": "team_b", "user_role": "proxy_admin"},
|
||||
],
|
||||
ids=["owner", "same_team", "proxy_admin"],
|
||||
)
|
||||
async def test_create_from_accessible_managed_file_is_allowed(
|
||||
call_type: str, file_field: str, caller_kwargs: dict[str, str]
|
||||
) -> None:
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
unified_file_id: Final = _tenant_a_input_file_id()
|
||||
managed_files, _ = _managed_files_with_fake_prisma(_tenant_a_input_file_row(unified_file_id))
|
||||
|
||||
result: Final = await managed_files.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(**caller_kwargs, parent_otel_span=MagicMock()),
|
||||
cache=MagicMock(),
|
||||
data={"model": "gpt-4o-mini", file_field: unified_file_id},
|
||||
call_type=call_type,
|
||||
)
|
||||
|
||||
assert isinstance(result, dict) and result[file_field] == unified_file_id, result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("call_type", "file_field"), _INPUT_FILE_CREATE_CALLS)
|
||||
async def test_create_from_raw_provider_file_id_skips_managed_ownership_lookup(call_type: str, file_field: str) -> None:
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
managed_files, table = _managed_files_with_fake_prisma()
|
||||
|
||||
result: Final = await managed_files.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="user_b", team_id="team_b", parent_otel_span=MagicMock()),
|
||||
cache=MagicMock(),
|
||||
data={"model": "gpt-4o-mini", file_field: "file-raw-provider-id"},
|
||||
call_type=call_type,
|
||||
)
|
||||
|
||||
assert isinstance(result, dict) and result[file_field] == "file-raw-provider-id", result
|
||||
assert not any("unified_file_id" in where for where in table.find_first_calls), table.find_first_calls
|
||||
|
|
|
|||
|
|
@ -57,7 +57,7 @@ from litellm.types.llms.openai import BatchJobStatus
|
|||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
|
||||
from litellm.types.utils import CredentialItem, LiteLLMBatch, SpecialEnums
|
||||
|
||||
from fastapi import Request, Response
|
||||
from fastapi import HTTPException, Request, Response
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Fixtures: distinguishable credentials per model so a wrong/hardcoded model_id
|
||||
|
|
@ -1164,6 +1164,34 @@ async def test_create__uses_acreate_batch_route_type(harness, openai_env_creds):
|
|||
assert harness.pre_call.call_args.kwargs["route_type"] == "acreate_batch"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("loadbalancing", [True, False], ids=["load_balanced", "unified_branch"])
|
||||
async def test_create__pre_call_ownership_denial_stops_every_dispatch_branch(harness, loadbalancing: bool):
|
||||
set_body(
|
||||
harness,
|
||||
{
|
||||
"input_file_id": "litellm_proxy_unified_id",
|
||||
"model": "vertex-model",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"completion_window": "24h",
|
||||
},
|
||||
)
|
||||
harness.is_known_model.return_value = True
|
||||
harness.pre_call.side_effect = HTTPException(status_code=403, detail="does not have access to the file")
|
||||
|
||||
with (
|
||||
patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", loadbalancing),
|
||||
patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value="unified-xyz"),
|
||||
patch.object(endpoints, "get_models_from_unified_file_id", return_value=["vertex-model"]),
|
||||
pytest.raises(ProxyException) as exc,
|
||||
):
|
||||
await call_create(harness)
|
||||
|
||||
assert exc.value.code == "403", exc.value
|
||||
harness.router_acreate.assert_not_called()
|
||||
harness.litellm_acreate.assert_not_called()
|
||||
|
||||
|
||||
def install_managed_files_hook(harness: Harness) -> AsyncMock:
|
||||
prisma_client = AsyncMock()
|
||||
managed_files = PROXY_LiteLLMManagedFiles(MagicMock(async_set_cache=AsyncMock()), prisma_client=prisma_client)
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import time
|
|||
from collections.abc import Iterator, Sequence
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
|
@ -481,10 +482,9 @@ async def test_batch_rate_limiter_without_user_context(openai_files: respx.MockR
|
|||
assert usage_without_context.total_tokens == usage_with_context.total_tokens == _token_counter_total(rows)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_managed_file_is_read_through_the_managed_files_hook_with_user_context(
|
||||
openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
async def _install_managed_file_hook(
|
||||
openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rows: str
|
||||
) -> tuple[respx.Route, respx.Route]:
|
||||
from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles
|
||||
|
||||
from litellm import Router
|
||||
|
|
@ -494,24 +494,22 @@ async def test_managed_file_is_read_through_the_managed_files_hook_with_user_con
|
|||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
assert is_base64_encoded_unified_file_id(_MANAGED_FILE_ID)
|
||||
rows: Final = _batch_rows(("Test message for regression",))
|
||||
provider_route: Final = _serve_file(openai_files, "file-provider-1", rows)
|
||||
standard_route: Final = _serve_file(openai_files, "file-abc123", rows)
|
||||
unrouted_managed_read: Final = _serve_file(openai_files, _MANAGED_FILE_ID, rows)
|
||||
file_cache: Final = InternalUsageCache(dual_cache=DualCache())
|
||||
await file_cache.async_set_cache(
|
||||
key=_MANAGED_FILE_ID,
|
||||
value=LiteLLM_ManagedFileTable(
|
||||
unified_file_id=_MANAGED_FILE_ID,
|
||||
model_mappings={"deployment-1": "file-provider-1"},
|
||||
flat_model_file_ids=["file-provider-1"],
|
||||
created_by="test-user-regression",
|
||||
).model_dump(),
|
||||
litellm_parent_otel_span=None,
|
||||
row: Final = LiteLLM_ManagedFileTable(
|
||||
unified_file_id=_MANAGED_FILE_ID,
|
||||
model_mappings={"deployment-1": "file-provider-1"},
|
||||
flat_model_file_ids=["file-provider-1"],
|
||||
created_by="test-user-regression",
|
||||
team_id="team-owner",
|
||||
)
|
||||
file_cache: Final = InternalUsageCache(dual_cache=DualCache())
|
||||
await file_cache.async_set_cache(key=_MANAGED_FILE_ID, value=row.model_dump(), litellm_parent_otel_span=None)
|
||||
prisma_client: Final = MagicMock()
|
||||
prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(return_value=row)
|
||||
proxy_logging: Final = ProxyLogging(user_api_key_cache=DualCache())
|
||||
proxy_logging.proxy_hook_mapping["managed_files"] = _PROXY_LiteLLMManagedFiles(
|
||||
internal_usage_cache=file_cache, prisma_client=None
|
||||
internal_usage_cache=file_cache, prisma_client=prisma_client
|
||||
)
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
|
|
@ -524,6 +522,16 @@ async def test_managed_file_is_read_through_the_managed_files_hook_with_user_con
|
|||
)
|
||||
monkeypatch.setattr(proxy_server, "proxy_logging_obj", proxy_logging)
|
||||
monkeypatch.setattr(proxy_server, "llm_router", router)
|
||||
return provider_route, unrouted_managed_read
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_managed_file_is_read_through_the_managed_files_hook_with_user_context(
|
||||
openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
rows: Final = _batch_rows(("Test message for regression",))
|
||||
provider_route, unrouted_managed_read = await _install_managed_file_hook(openai_files, monkeypatch, rows)
|
||||
standard_route: Final = _serve_file(openai_files, "file-abc123", rows)
|
||||
_, _, batch_limiter = _make_limiters()
|
||||
user_api_key_dict: Final = UserAPIKeyAuth(
|
||||
api_key="test-key-regression", user_id="test-user-regression", tpm_limit=1000, rpm_limit=10
|
||||
|
|
@ -541,3 +549,25 @@ async def test_managed_file_is_read_through_the_managed_files_hook_with_user_con
|
|||
assert not unrouted_managed_read.called
|
||||
assert managed_usage.request_count == standard_usage.request_count == 1
|
||||
assert managed_usage.total_tokens == standard_usage.total_tokens == _token_counter_total(rows)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_managed_file_of_another_tenant_is_denied_before_it_is_read(
|
||||
openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
rows: Final = _batch_rows(("Test message for regression",))
|
||||
provider_route, unrouted_managed_read = await _install_managed_file_hook(openai_files, monkeypatch, rows)
|
||||
_, _, batch_limiter = _make_limiters()
|
||||
outsider: Final = UserAPIKeyAuth(api_key="test-key-outsider", user_id="outsider", team_id="team-other", tpm_limit=1)
|
||||
|
||||
with pytest.raises(HTTPException) as denied:
|
||||
await batch_limiter.async_pre_call_hook(
|
||||
user_api_key_dict=outsider,
|
||||
cache=DualCache(),
|
||||
data=dict(_create_batch_data(_MANAGED_FILE_ID)),
|
||||
call_type="acreate_batch",
|
||||
)
|
||||
|
||||
assert denied.value.status_code == 403
|
||||
assert not provider_route.called
|
||||
assert not unrouted_managed_read.called
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue