scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-17 04:34:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 172

fix(batches): read a restricted key's batch input through its team deployment

Browse source 
The batch access check downloaded the input JSONL with the proxy's OPENAI_API_KEY while the upload and the create used the team deployment, so a key restricted under a team openai/* deployment could not create a batch when the two keys belong to different OpenAI projects. The read now resolves the team deployment the way the create does when no model pins it

Adds the audit cells for keys restricted under the team wildcard on /rag, /files and /batches, and for the rag-vs-chat split of a key restricted to a global exact-name deployment
This commit is contained in:
mateo-berri 2026-10-12 16:41:03 -07:00
parent 641cf49a3d
commit e70cc0da43
3 changed files with 218 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

42
litellm/proxy/hooks/batch_rate_limiter.py
View file

@ -509,6 +509,7 @@ class _PROXY_BatchRateLimiter(CustomLogger):
file_id: str,
custom_llm_provider: str,
data: dict,
user_api_key_dict: UserAPIKeyAuth | None = None,
) -> tuple[str, dict[str, Any]]:
"""
Map proxy-facing file IDs to provider file IDs and credentials.
@ -521,6 +522,7 @@ class _PROXY_BatchRateLimiter(CustomLogger):
decode_model_from_file_id,
get_credentials_for_model,
get_original_file_id,
get_team_provider_credentials,
)
from litellm.proxy.proxy_server import llm_router
@ -547,21 +549,32 @@ class _PROXY_BatchRateLimiter(CustomLogger):
return get_original_file_id(file_id), fetch_kwargs
request_model: Final = data.get("model")
if isinstance(request_model, str) and request_model and llm_router is not None:
try:
credentials = get_credentials_for_model(
llm_router=llm_router,
model_id=request_model,
operation_context="batch input file read (rate limiting)",
)
fetch_kwargs.update(extract_file_access_credentials(credentials))
fetch_kwargs["model"] = request_model
provider = credentials.get("custom_llm_provider")
if provider:
fetch_kwargs["custom_llm_provider"] = provider
except HTTPException:
pass
if isinstance(request_model, str) and request_model:
if llm_router is not None:
try:
credentials = get_credentials_for_model(
llm_router=llm_router,
model_id=request_model,
operation_context="batch input file read (rate limiting)",
)
fetch_kwargs.update(extract_file_access_credentials(credentials))
fetch_kwargs["model"] = request_model
provider = credentials.get("custom_llm_provider")
if provider:
fetch_kwargs["custom_llm_provider"] = provider
except HTTPException:
pass
return file_id, fetch_kwargs
if user_api_key_dict is None:
return file_id, fetch_kwargs
team_credentials: Final = get_team_provider_credentials(
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
if team_credentials is not None:
fetch_kwargs.update(extract_file_access_credentials(team_credentials))
return file_id, fetch_kwargs
async def _reserve_batch_enqueued_tokens(
@ -912,6 +925,7 @@ class _PROXY_BatchRateLimiter(CustomLogger):
file_id=file_id,
custom_llm_provider=custom_llm_provider,
data=data or {},
user_api_key_dict=user_api_key_dict,
)
# For non-managed files, use the standard litellm.afile_content
file_content = await litellm.afile_content(

143
tests/integration/authorization/test_rag_team_provider_credentials.py
View file

@ -46,13 +46,24 @@ GEMINI_API_KEY: Final = "integration-rag-gemini-key"
AZURE_API_KEY: Final = "integration-rag-azure-key"
JWT_KEY_ID: Final = "integration-rag-team-provider-jwt"
WILDCARD_MODEL: Final = "openai/gpt-4o-mini"
BARE_WILDCARD_MODEL: Final = "gpt-4o-mini"
OPENAI: Final[Mapping[str, JsonValue]] = MappingProxyType({"custom_llm_provider": "openai"})
DOCUMENT: Final = ("document.txt", b"team provider credentials", "text/plain")
BATCH_INPUT_ROW: Final[Mapping[str, JsonValue]] = MappingProxyType(
{
"custom_id": "integration-rag-batch",
"method": "POST",
"url": "/v1/chat/completions",
"body": {"model": BARE_WILDCARD_MODEL, "messages": [{"role": "user", "content": "batch input"}]},
}
)
BATCH_INPUT: Final = ("batch.jsonl", (json.dumps(dict(BATCH_INPUT_ROW)) + "\n").encode(), "application/jsonl")
FAILING_STORE_NAMES: Final = MappingProxyType({"rag-upstream-401": 401, "rag-upstream-429": 429})
OWNED_PROXY_CELL_SECONDS: Final = 2 * graceful_stop_seconds() + 120
STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
PREFIXED: Final = re.compile(r"^/(team|wildcard|shared|environment|credential|caller|gemini|azure)(/[^?]*)")
STORE_FILES: Final = re.compile(r"^/v1/vector_stores/([^/]+)/files$")
FILE_CONTENT: Final = re.compile(r"^/v1/files/([^/]+)/content$")
STORE_SEARCH: Final = re.compile(r"^(?:/v1|/openai)/vector_stores/([^/]+)/search$")
GEMINI_UPLOAD: Final = re.compile(r"^/upload/v1beta/fileSearchStores/([^:]+):uploadToFileSearchStore$")
@ -78,13 +89,18 @@ class Call:
@dataclass(frozen=True, slots=True)
class Tenants:
shared_model: str
global_model: str
team_a: str
team_a_model: str
team_a_key: str
restricted_key: str
viewer_key: str
jwt_member: str
team_w: str
team_w_key: str
wildcard_bare_key: str
wildcard_prefixed_key: str
wildcard_global_key: str
team_s_key: str
teamless_key: str
credential_name: str
@ -180,6 +196,30 @@ def _file_reply() -> Reply:
)
def _batch_input_reply() -> Reply:
return Reply(status=200, body=BATCH_INPUT[1])
def _batch_reply(request: Request) -> Reply:
body: Final = _json_body(request)
return _json_reply(
{
"id": f"batch_{uuid.uuid4().hex}",
"object": "batch",
"endpoint": str(body.get("endpoint")),
"input_file_id": str(body.get("input_file_id")),
"completion_window": str(body.get("completion_window")),
"status": "validating",
"created_at": int(time.time()),
"errors": None,
"output_file_id": None,
"error_file_id": None,
"request_counts": {"total": 0, "completed": 0, "failed": 0},
"metadata": None,
}
)
def _attach_reply(request: Request, store_id: str) -> Reply:
return _json_reply(
{
@ -292,12 +332,16 @@ def _gemini_reply(request: Request, route: str) -> Reply:
def _openai_reply(request: Request, route: str, outage: threading.Event) -> Reply:
if request.method == "GET":
return _batch_input_reply() if FILE_CONTENT.match(route) is not None else _error_reply(404)
if request.method != "POST":
return _error_reply(404)
if route == "/v1/vector_stores":
return _vector_store_reply(request, outage)
if route == "/v1/files":
return _file_reply()
if route == "/v1/batches":
return _batch_reply(request)
if route == "/v1/chat/completions":
return _chat_reply(request)
if route == "/v1/responses":
@ -398,6 +442,16 @@ def _tenants(gateway: Gateway, scenario: Scenario, upstream: str) -> Tenants:
api_base=f"{upstream}/shared/v1",
team_id=None,
)
global_model: Final = f"openai/integration-rag-global-{uuid.uuid4().hex}"
_deployment(
gateway,
scenario,
name=global_model,
model="openai/gpt-4o-mini",
api_key=SHARED_API_KEY,
api_base=f"{upstream}/shared/v1",
team_id=None,
)
team_a_model: Final = f"integration-rag-team-{uuid.uuid4().hex}"
team_a: Final = scenario.team(models=[team_a_model, shared_model])
_deployment(
@ -424,13 +478,18 @@ def _tenants(gateway: Gateway, scenario: Scenario, upstream: str) -> Tenants:
team_s: Final = scenario.team(models=[shared_model])
return Tenants(
shared_model=shared_model,
global_model=global_model,
team_a=team_a,
team_a_model=team_a_model,
team_a_key=scenario.key(team_id=team_a),
restricted_key=scenario.key(team_id=team_a, models=[shared_model]),
viewer_key=scenario.key(user_id=viewer, team_id=team_a),
jwt_member=scenario.member(team_a),
team_w=team_w,
team_w_key=scenario.key(team_id=team_w),
wildcard_bare_key=scenario.key(team_id=team_w, models=[BARE_WILDCARD_MODEL]),
wildcard_prefixed_key=scenario.key(team_id=team_w, models=[WILDCARD_MODEL]),
wildcard_global_key=scenario.key(team_id=team_w, models=[global_model]),
team_s_key=scenario.key(team_id=team_s),
teamless_key=scenario.key(),
credential_name=_credential(gateway, scenario, f"{upstream}/credential/v1"),
@ -726,6 +785,90 @@ def test_files_upload_keeps_the_accessible_deployment_fallback(
assert rag_rig.signatures() == ((prefix, "POST", "/v1/files", f"Bearer {api_key}"),)
@pytest.mark.parametrize(
("caller", "model"),
(("wildcard_bare_key", BARE_WILDCARD_MODEL), ("wildcard_prefixed_key", WILDCARD_MODEL)),
ids=("bare_model", "prefixed_model"),
)
def test_key_restricted_to_a_model_under_the_team_wildcard_uses_the_team_deployment(
rag_rig: RagRig, caller: Literal["wildcard_bare_key", "wildcard_prefixed_key"], model: str
) -> None:
key: Final = getattr(rag_rig.tenants, caller)
store_id: Final = _completed_store(rag_rig, _ingest(rag_rig.gateway, key))
assert rag_rig.signatures() == _ingest_signatures("wildcard", WILDCARD_API_KEY, store_id)
rows: Final = _store_rows(store_id)
assert [row["team_id"] for row in rows] == [rag_rig.tenants.team_w], rows
assert WILDCARD_API_KEY not in json.dumps(rows), rows
response: Final = _query(rag_rig.gateway, key, model, {**OPENAI, "vector_store_id": store_id})
assert response.status_code == 200, response.text
assert rag_rig.signatures() == (
_search_signature("wildcard", WILDCARD_API_KEY, store_id),
_chat_signature("wildcard", WILDCARD_API_KEY),
)
def test_key_restricted_to_a_global_model_under_the_team_wildcard_splits_rag_and_chat(rag_rig: RagRig) -> None:
tenants: Final = rag_rig.tenants
store_id: Final = _completed_store(rag_rig, _ingest(rag_rig.gateway, tenants.wildcard_global_key))
assert rag_rig.signatures() == _ingest_signatures("wildcard", WILDCARD_API_KEY, store_id)
response: Final = _query(
rag_rig.gateway, tenants.wildcard_global_key, tenants.global_model, {**OPENAI, "vector_store_id": store_id}
)
assert response.status_code == 200, response.text
assert rag_rig.signatures() == (
_search_signature("wildcard", WILDCARD_API_KEY, store_id),
_chat_signature("shared", SHARED_API_KEY),
)
def test_files_upload_by_a_key_restricted_under_the_team_wildcard_uses_the_team_deployment(rag_rig: RagRig) -> None:
response: Final = rag_rig.gateway.request_multipart(
"/v1/files", {"purpose": "assistants"}, {"file": DOCUMENT}, key=rag_rig.tenants.wildcard_bare_key
)
assert response.status_code == 200, response.text
assert rag_rig.signatures() == (("wildcard", "POST", "/v1/files", f"Bearer {WILDCARD_API_KEY}"),)
@pytest.mark.parametrize(
("caller", "prefix", "api_key", "reads_input"),
(("team_a_key", "team", TEAM_API_KEY, False), ("wildcard_bare_key", "wildcard", WILDCARD_API_KEY, True)),
ids=("team_key", "restricted_wildcard_key"),
)
def test_batch_create_reads_the_input_and_creates_the_batch_through_the_team_deployment(
rag_rig: RagRig, caller: Literal["team_a_key", "wildcard_bare_key"], prefix: str, api_key: str, reads_input: bool
) -> None:
key: Final = getattr(rag_rig.tenants, caller)
uploaded: Final = rag_rig.gateway.request_multipart(
"/v1/files", {"purpose": "batch"}, {"file": BATCH_INPUT}, key=key
)
assert uploaded.status_code == 200, uploaded.text
file_id: Final = string_value(object_value(uploaded.json())["id"])
created: Final = rag_rig.gateway.request(
"POST",
"/v1/batches",
{"input_file_id": file_id, "endpoint": "/v1/chat/completions", "completion_window": "24h"},
key=key,
)
assert created.status_code == 200, created.text
assert object_value(created.json())["input_file_id"] == file_id, created.text
bearer: Final = f"Bearer {api_key}"
input_read: Final = ((prefix, "GET", f"/v1/files/{file_id}/content", bearer),) if reads_input else ()
assert rag_rig.signatures() == (
(prefix, "POST", "/v1/files", bearer),
*input_read,
(prefix, "POST", "/v1/batches", bearer),
)
def test_gemini_ingest_never_receives_the_team_openai_key(rag_rig: RagRig) -> None:
store_id: Final = _completed_store(
rag_rig, _ingest(rag_rig.gateway, rag_rig.tenants.team_a_key, {"custom_llm_provider": "gemini"})

47
tests/unit/proxy/hooks/test_batch_rate_limiter.py
View file

@ -481,6 +481,53 @@ 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_batch_input_without_a_model_is_read_through_the_team_deployment(
openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
):
from litellm import Router
from litellm.proxy import proxy_server
monkeypatch.setattr(
proxy_server,
"llm_router",
Router(
model_list=[
{
"model_name": "openai/*",
"litellm_params": {
"model": "openai/*",
"api_key": "sk-team",
"api_base": "https://team.example/v1",
},
"model_info": {"id": "team-openai", "team_id": "team-1"},
}
]
),
)
_, _, batch_limiter = _make_limiters()
rows: Final = _batch_rows(("Hello",))
environment_route: Final = _serve_file(openai_files, "file-abc123", rows)
team_route: Final = openai_files.get("https://team.example/v1/files/file-abc123/content").mock(
return_value=httpx.Response(200, content="\n".join(json.dumps(row) for row in rows).encode())
)
team_key: Final = UserAPIKeyAuth(api_key="sk", team_id="team-1", team_models=["openai/*"], models=[_BATCH_MODEL])
teamless_key: Final = UserAPIKeyAuth(api_key="sk")
team_usage: Final = await batch_limiter.count_input_file_usage(
file_id="file-abc123", custom_llm_provider="openai", user_api_key_dict=team_key, data={}
)
teamless_usage: Final = await batch_limiter.count_input_file_usage(
file_id="file-abc123", custom_llm_provider="openai", user_api_key_dict=teamless_key, data={}
)
assert team_route.call_count == 1
assert team_route.calls.last.request.headers["authorization"] == "Bearer sk-team"
assert environment_route.call_count == 1
assert environment_route.calls.last.request.headers["authorization"] == "Bearer sk-test"
assert team_usage.total_tokens == teamless_usage.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