diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index cbf683d226..5122c64ea6 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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 = [ diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 1c9ba6cb24..6090524336 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -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 diff --git a/tests/test_litellm/proxy/test_batch_expiry.py b/tests/test_litellm/proxy/test_batch_expiry.py index 25b1631792..1f54f190c6 100644 --- a/tests/test_litellm/proxy/test_batch_expiry.py +++ b/tests/test_litellm/proxy/test_batch_expiry.py @@ -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