diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 9553d2c5246..e73f73d2f33 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -112,6 +112,7 @@ async def acreate_batch( metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, + output_expires_after: Optional[Dict[str, Any]] = None, **kwargs, ) -> LiteLLMBatch: """ @@ -133,6 +134,7 @@ async def acreate_batch( metadata, extra_headers, extra_body, + output_expires_after, **kwargs, ) @@ -160,6 +162,7 @@ def create_batch( metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, + output_expires_after: Optional[Dict[str, Any]] = None, **kwargs, ) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]: """ @@ -215,6 +218,8 @@ def create_batch( extra_headers=extra_headers, extra_body=extra_body, ) + if output_expires_after is not None: + _create_batch_request["output_expires_after"] = output_expires_after if model is not None: provider_config = ProviderConfigManager.get_provider_batches_config( model=model, diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index d06d879dad1..c5d610e639b 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -424,6 +424,7 @@ class CreateBatchRequest(TypedDict, total=False): endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"] input_file_id: str metadata: Optional[Dict[str, str]] + output_expires_after: Optional[FileExpiresAfter] extra_headers: Optional[Dict[str, str]] extra_body: Optional[Dict[str, str]] timeout: Optional[float] diff --git a/tests/test_litellm/proxy/test_batch_expiry.py b/tests/test_litellm/proxy/test_batch_expiry.py new file mode 100644 index 00000000000..25b1631792e --- /dev/null +++ b/tests/test_litellm/proxy/test_batch_expiry.py @@ -0,0 +1,74 @@ +""" +Tests for batch output_expires_after passthrough and team-level expiry enforcement. +""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +import litellm +from litellm.types.llms.openai import CreateBatchRequest + + +class TestCreateBatchOutputExpiresAfterPassthrough: + """Verify output_expires_after flows through create_batch to the provider.""" + + 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 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", + ) + + create_batch_data = captured["create_batch_data"] + assert create_batch_data["output_expires_after"] == { + "anchor": "created_at", + "seconds": 86400, + } + + def test_output_expires_after_absent_when_not_provided(self): + """Backward compat: output_expires_after not in request when omitted.""" + 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", + custom_llm_provider="openai", + ) + + create_batch_data = captured["create_batch_data"] + assert "output_expires_after" not in create_batch_data