mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
fix(files): keep an explicit target_storage on its old path and refuse litellm_db as a caller choice
An explicit target_storage=litellm_db upload was accepted for any model, so an OpenAI model's litellm_db://<uuid> id was sent to OpenAI as input_file_id and a model-less upload left a content row nothing can read; it now answers 400 on target_storage. An explicit target_storage skips the files api probe and the purpose and single-target gates, which only decide whether LiteLLM keeps the file itself, so an azure_storage user_data upload for a vLLM model reaches the storage path again as it did before this branch. cancel_batch authorizes the model of every LiteLLM-managed batch id before it branches, the way retrieve_batch already does, so the LiteLLM-executed branch gets the check its provider sibling had. Restores the test_afile_delete_passes_trusted_model_credentials_to_router definition line an earlier commit dropped
This commit is contained in:
parent
020cbba4dd
commit
549548de62
5 changed files with 102 additions and 8 deletions
|
|
@ -1093,6 +1093,17 @@ async def cancel_batch(
|
|||
proxy_config=proxy_config,
|
||||
)
|
||||
|
||||
unified_model_id: Final = get_model_id_from_unified_batch_id(unified_batch_id) if unified_batch_id else None
|
||||
if unified_model_id is not None:
|
||||
resolved_unified_model: Final = (
|
||||
llm_router.resolve_model_name_from_model_id(unified_model_id) if llm_router is not None else None
|
||||
)
|
||||
await authorize_model_for_key(
|
||||
model_id=resolved_unified_model or unified_model_id,
|
||||
llm_router=llm_router,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
# SCENARIO 1: Batch ID is encoded with model info
|
||||
if model_from_id is not None:
|
||||
credentials: Final = await get_authorized_credentials_for_model(
|
||||
|
|
@ -1143,11 +1154,6 @@ async def cancel_batch(
|
|||
status_code=400,
|
||||
detail={"error": "Invalid LiteLLM managed batch ID. Missing model_id."},
|
||||
)
|
||||
await authorize_model_for_key(
|
||||
model_id=llm_router.resolve_model_name_from_model_id(model_id_from_batch) or model_id_from_batch,
|
||||
llm_router=llm_router,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
data["model"] = model_id_from_batch
|
||||
data["batch_id"] = get_batch_id_from_unified_batch_id(unified_batch_id)
|
||||
response = await llm_router.acancel_batch(**data)
|
||||
|
|
|
|||
|
|
@ -117,6 +117,7 @@ async def _litellm_executed_batch_input_model(
|
|||
model: str | None,
|
||||
target_model_names_list: Sequence[str],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
explicit_storage: str | None,
|
||||
) -> str | None:
|
||||
if llm_router is None:
|
||||
return None
|
||||
|
|
@ -129,6 +130,8 @@ async def _litellm_executed_batch_input_model(
|
|||
if _names_a_litellm_executed_provider(llm_router, candidate, team_id)
|
||||
)
|
||||
)
|
||||
if explicit_storage is not None:
|
||||
return None
|
||||
providers: Final = await asyncio.gather(
|
||||
*(resolve_litellm_executed_provider(llm_router, candidate, team_id) for candidate in candidates)
|
||||
)
|
||||
|
|
@ -305,10 +308,21 @@ async def route_create_file(
|
|||
5. Else -> use custom_llm_provider with files_settings
|
||||
"""
|
||||
|
||||
executed_model: Final = await _litellm_executed_batch_input_model(
|
||||
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
|
||||
if explicit_storage == LITELLM_DB_STORAGE_BACKEND_NAME:
|
||||
raise ProxyException(
|
||||
message=(
|
||||
f"target_storage={LITELLM_DB_STORAGE_BACKEND_NAME} is not a storage a caller can pick: LiteLLM "
|
||||
"chooses it on its own for the batch input files of a model whose batches it runs itself, so "
|
||||
"upload with purpose=batch and name that model instead of target_storage"
|
||||
),
|
||||
type="invalid_request_error",
|
||||
param="target_storage",
|
||||
code=400,
|
||||
)
|
||||
executed_model: Final = await _litellm_executed_batch_input_model(
|
||||
llm_router, purpose, model, target_model_names_list, user_api_key_dict, explicit_storage
|
||||
)
|
||||
storage: Final = explicit_storage or (LITELLM_DB_STORAGE_BACKEND_NAME if executed_model is not None else None)
|
||||
if storage is not None:
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
|
|
|
|||
|
|
@ -1764,6 +1764,10 @@ async def test_post_call_hook_leaves_litellm_executed_batches_untouched(llm_batc
|
|||
assert managed_files.store_unified_object_id.await_count == (1 if stores else 0)
|
||||
if not stores:
|
||||
assert response.id == original_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_delete_passes_trusted_model_credentials_to_router():
|
||||
"""
|
||||
afile_delete must hand the deployment's credential snapshot to the router
|
||||
call, since Bedrock validates the s3:// file id against the bucket in it.
|
||||
|
|
|
|||
|
|
@ -3220,3 +3220,16 @@ async def test_cancel__unified_batch_id_rejects_key_without_model_grant(cancel_h
|
|||
|
||||
assert exc_info.value.code == "403"
|
||||
cancel_harness.router_acancel.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel__executed_batch_rejects_key_without_model_grant(cancel_harness, executed_runner):
|
||||
runner, factory = executed_runner
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await call_cancel(cancel_harness, EXECUTED_BATCH_B64, user=_key_restricted_to("vertex-model"))
|
||||
|
||||
assert exc_info.value.code == "403"
|
||||
factory.assert_not_called()
|
||||
runner.cancel.assert_not_called()
|
||||
cancel_harness.router_acancel.assert_not_called()
|
||||
|
|
|
|||
|
|
@ -781,6 +781,63 @@ def test_upload_for_a_litellm_executed_model_goes_to_the_provider_unless_the_ser
|
|||
assert provider_upload.call_args.kwargs["api_base"] == "http://vllm.test/v1"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"form",
|
||||
[{}, {"target_model_names": "my-vllm"}, {"target_model_names": "gemini-2.0-flash"}],
|
||||
ids=["no model", "litellm-executed model", "provider model"],
|
||||
)
|
||||
def test_upload_naming_litellm_db_as_target_storage_is_rejected(batch_upload_seams, form: dict[str, str]):
|
||||
stored, provider_upload, upstream_files_route = batch_upload_seams
|
||||
|
||||
response = _upload_batch_file({}, {**form, "target_storage": "litellm_db"})
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
error = response.json()["error"]
|
||||
assert error["type"] == "invalid_request_error"
|
||||
assert error["param"] == "target_storage"
|
||||
assert "litellm_db" in error["message"]
|
||||
assert upstream_files_route.call_count == 0
|
||||
stored.assert_not_awaited()
|
||||
provider_upload.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("purpose", ["user_data", "batch"])
|
||||
def test_upload_with_an_explicit_target_storage_goes_where_the_caller_said_without_probing_the_server(
|
||||
batch_upload_seams, purpose: str
|
||||
):
|
||||
stored, provider_upload, upstream_files_route = batch_upload_seams
|
||||
|
||||
response = _upload_batch_file(
|
||||
{}, {"purpose": purpose, "target_model_names": "my-vllm", "target_storage": "azure_storage"}
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert upstream_files_route.call_count == 0
|
||||
provider_upload.assert_not_awaited()
|
||||
stored.assert_awaited_once()
|
||||
kwargs = stored.call_args.kwargs
|
||||
assert kwargs["target_storage"] == "azure_storage"
|
||||
assert tuple(kwargs["target_model_names"]) == ("my-vllm",)
|
||||
assert kwargs["purpose"] == purpose
|
||||
|
||||
|
||||
def test_upload_with_an_explicit_target_storage_still_refuses_a_key_without_the_executed_model(batch_upload_seams):
|
||||
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({}, {"target_model_names": "my-vllm", "target_storage": "azure_storage"})
|
||||
|
||||
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_for_a_provider_model_still_goes_to_the_provider(batch_upload_seams):
|
||||
stored, provider_upload, _ = batch_upload_seams
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue