tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-20 00:23:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(proxy): add team-level batch output expiry enforcement

Browse Source
This commit is contained in:
Ryan Crabbe
2026-03-03 10:58:31 -08:00
parent 43cec8c980
commit 3d15bcdb11
3 changed files with 153 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/_types.py
+6
View File
@@ -1551,6 +1551,8 @@ class NewTeamRequest(TeamBase):
] = None # allow user to set TPM limit for all team members
team_member_key_duration: Optional[str] = None # e.g. "1d", "1w", "1m"
allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None
enforced_batch_output_expires_after: Optional[dict] = None
enforced_file_expires_after: Optional[dict] = None
model_config = ConfigDict(protected_namespaces=())
@@ -1606,6 +1608,8 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
model_rpm_limit: Optional[Dict[str, int]] = None
model_tpm_limit: Optional[Dict[str, int]] = None
allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None
enforced_batch_output_expires_after: Optional[dict] = None
enforced_file_expires_after: Optional[dict] = None
router_settings: Optional[dict] = None
access_group_ids: Optional[List[str]] = None
@@ -3783,6 +3787,8 @@ LiteLLM_ManagementEndpoint_MetadataFields = [
"temp_budget_increase",
"temp_budget_expiry",
"allowed_vector_store_indexes",
"enforced_batch_output_expires_after",
"enforced_file_expires_after",
]
LiteLLM_ManagementEndpoint_MetadataFields_Premium = [
litellm/proxy/batches_endpoints/endpoints.py
+9
View File
@@ -118,6 +118,15 @@ async def create_batch( # noqa: PLR0915
or "openai"
)
_create_batch_data = LiteLLMBatchCreateRequest(**data)
# Apply team-level batch output expiry enforcement
team_metadata = user_api_key_dict.team_metadata or {}
enforced_batch_expiry = team_metadata.get(
"enforced_batch_output_expires_after"
)
if enforced_batch_expiry is not None:
_create_batch_data["output_expires_after"] = enforced_batch_expiry
input_file_id = _create_batch_data.get("input_file_id", None)
unified_file_id: Union[str, Literal[False]] = False
tests/test_litellm/proxy/test_batch_expiry.py
+138 -50
View File
@@ -4,7 +4,7 @@ Tests for batch output_expires_after passthrough and team-level expiry enforceme
import os
import sys
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import MagicMock, patch
import pytest
@@ -13,62 +13,150 @@ sys.path.insert(
) # Adds the parent directory to the system path
import litellm
from litellm.types.llms.openai import CreateBatchRequest
from litellm.caching.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.proxy_server import app
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
from litellm.types.utils import LiteLLMBatch
from fastapi.testclient import TestClient
client = TestClient(app)
TEAM_EXPIRY = {"anchor": "created_at", "seconds": 3600}
CALLER_EXPIRY = {"anchor": "created_at", "seconds": 86400}
class TestCreateBatchOutputExpiresAfterPassthrough:
"""Verify output_expires_after flows through create_batch to the provider."""
@pytest.fixture
def llm_router() -> Router:
return Router(
model_list=[
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "openai/gpt-3.5-turbo",
"api_key": "test-key",
},
"model_info": {"id": "gpt-3.5-turbo-id"},
},
]
)
def test_output_expires_after_included_in_request(self):
"""When output_expires_after is provided, it reaches the openai batches instance."""
captured = {}
original_create = None
def _setup_proxy(monkeypatch, llm_router: Router):
proxy_logging_obj = ProxyLogging(
user_api_key_cache=DualCache(default_in_memory_ttl=1)
)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router)
monkeypatch.setattr(
"litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj
)
def capturing_create(**kwargs):
captured.update(kwargs)
mock_response = MagicMock()
mock_response.id = "batch_123"
return mock_response
with patch(
"litellm.batches.main.openai_batches_instance"
) as mock_instance:
mock_instance.create_batch.side_effect = capturing_create
litellm.create_batch(
completion_window="24h",
endpoint="/v1/chat/completions",
input_file_id="file-abc123",
output_expires_after={"anchor": "created_at", "seconds": 86400},
custom_llm_provider="openai",
def _make_batch_response() -> LiteLLMBatch:
return LiteLLMBatch(
id="batch_abc123",
completion_window="24h",
created_at=1234567890,
endpoint="/v1/chat/completions",
input_file_id="file-abc123",
object="batch",
status="validating",
)
def test_output_expires_after_passthrough():
"""output_expires_after flows through create_batch to the provider."""
captured = {}
def capturing_create(**kwargs):
captured.update(kwargs)
mock_response = MagicMock()
mock_response.id = "batch_123"
return mock_response
with patch("litellm.batches.main.openai_batches_instance") as mock_instance:
mock_instance.create_batch.side_effect = capturing_create
litellm.create_batch(
completion_window="24h",
endpoint="/v1/chat/completions",
input_file_id="file-abc123",
output_expires_after=CALLER_EXPIRY,
custom_llm_provider="openai",
)
assert captured["create_batch_data"]["output_expires_after"] == CALLER_EXPIRY
class TestBatchEndpointTeamOverride:
"""Verify team-level enforced_batch_output_expires_after in the proxy endpoint."""
def _post_batch(
self,
monkeypatch,
llm_router: Router,
team_metadata: dict,
request_body: dict,
) -> dict:
"""POST /v1/batches with given team_metadata and body, return captured kwargs."""
_setup_proxy(monkeypatch, llm_router)
user_key = UserAPIKeyAuth(
api_key="test-key",
team_metadata=team_metadata,
)
app.dependency_overrides[user_api_key_auth] = lambda: user_key
captured_kwargs = {}
async def mock_acreate_batch(**kwargs):
captured_kwargs.update(kwargs)
return _make_batch_response()
monkeypatch.setattr(litellm, "acreate_batch", mock_acreate_batch)
try:
response = client.post(
"/v1/batches",
json=request_body,
headers={"Authorization": "Bearer test-key"},
)
assert response.status_code == 200
finally:
app.dependency_overrides.clear()
create_batch_data = captured["create_batch_data"]
assert create_batch_data["output_expires_after"] == {
"anchor": "created_at",
"seconds": 86400,
}
return captured_kwargs
def test_output_expires_after_absent_when_not_provided(self):
"""Backward compat: output_expires_after not in request when omitted."""
captured = {}
def test_team_override_overrides_caller(self, monkeypatch, llm_router):
"""Team enforcement wins over caller-provided value."""
kwargs = self._post_batch(
monkeypatch,
llm_router,
team_metadata={
"enforced_batch_output_expires_after": TEAM_EXPIRY,
},
request_body={
"input_file_id": "file-abc123",
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
"output_expires_after": CALLER_EXPIRY,
},
)
assert kwargs["output_expires_after"] == TEAM_EXPIRY
def capturing_create(**kwargs):
captured.update(kwargs)
mock_response = MagicMock()
mock_response.id = "batch_123"
return mock_response
with patch(
"litellm.batches.main.openai_batches_instance"
) as mock_instance:
mock_instance.create_batch.side_effect = capturing_create
litellm.create_batch(
completion_window="24h",
endpoint="/v1/chat/completions",
input_file_id="file-abc123",
custom_llm_provider="openai",
)
create_batch_data = captured["create_batch_data"]
assert "output_expires_after" not in create_batch_data
def test_no_team_setting_preserves_caller(self, monkeypatch, llm_router):
"""No team setting = caller value passes through."""
kwargs = self._post_batch(
monkeypatch,
llm_router,
team_metadata={},
request_body={
"input_file_id": "file-abc123",
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
"output_expires_after": CALLER_EXPIRY,
},
)
assert kwargs["output_expires_after"] == CALLER_EXPIRY