From 08613b24cbbf0315249a305b743e43cc5e73411f Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Tue, 3 Mar 2026 11:03:14 -0800 Subject: [PATCH] feat(proxy): add team-level file expiry enforcement --- .../openai_files_endpoints/files_endpoints.py | 11 +- .../test_files_endpoint.py | 135 ++++++++++++++++++ 2 files changed, 145 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index ec6e9733344..a641fb1c69f 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -454,8 +454,17 @@ async def create_file( # noqa: PLR0915 model=router_model, llm_router=llm_router ) + # Apply team-level file expiry enforcement + team_metadata = user_api_key_dict.team_metadata or {} + enforced_file_expiry = team_metadata.get("enforced_file_expires_after") + if enforced_file_expiry is not None: + expires_after = FileExpiresAfter( + anchor=enforced_file_expiry["anchor"], + seconds=enforced_file_expiry["seconds"], + ) + _create_file_request = CreateFileRequest( - file=file_data, + file=file_data, purpose=cast(CREATE_FILE_REQUESTS_PURPOSE, purpose), expires_after=expires_after, **data diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index fb063ef8ee7..0239c39e67f 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -1168,3 +1168,138 @@ def test_create_file_with_deep_nested_litellm_metadata( assert captured_litellm_metadata["config"]["database"]["port"] == "5432" assert "cache" in captured_litellm_metadata["config"] assert captured_litellm_metadata["config"]["cache"]["enabled"] == "true" + + +# --------------------------------------------------------------------------- +# Team-level enforced_file_expires_after tests +# --------------------------------------------------------------------------- + + +def _make_capturing_managed_files(): + """Create a DummyManagedFiles that captures the expires_after from the request.""" + from litellm.llms.base_llm.files.transformation import BaseFileEndpoints + + captured = {} + + class CapturingManagedFiles(BaseFileEndpoints): + async def acreate_file( + self, + llm_router, + create_file_request, + target_model_names_list, + litellm_parent_otel_span, + user_api_key_dict, + ): + if isinstance(create_file_request, dict): + captured["expires_after"] = create_file_request.get("expires_after") + else: + captured["expires_after"] = getattr( + create_file_request, "expires_after", None + ) + return OpenAIFileObject( + id="file-abc123", + object="file", + bytes=100, + created_at=1234567890, + filename="mydata.jsonl", + purpose="batch", + status="uploaded", + ) + + async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router): + raise NotImplementedError + + async def afile_list(self, purpose, litellm_parent_otel_span): + raise NotImplementedError + + async def afile_delete( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): + raise NotImplementedError + + async def afile_content( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): + raise NotImplementedError + + return CapturingManagedFiles(), captured + + +def _post_file_with_team_metadata( + monkeypatch, + llm_router: Router, + team_metadata: dict, + form_data: dict, +): + """POST /v1/files with given team_metadata, return captured expires_after.""" + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + proxy_logging_obj = ProxyLogging( + user_api_key_cache=DualCache(default_in_memory_ttl=1) + ) + dummy, captured = _make_capturing_managed_files() + proxy_logging_obj.proxy_hook_mapping["managed_files"] = dummy + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj + ) + + user_key = UserAPIKeyAuth(api_key="test-key", team_metadata=team_metadata) + app.dependency_overrides[user_api_key_auth] = lambda: user_key + + test_file = ("mydata.jsonl", b'{"prompt": "Hello"}', "application/json") + try: + response = client.post( + "/v1/files", + files={"file": test_file}, + data=form_data, + headers={"Authorization": "Bearer test-key"}, + ) + assert response.status_code == 200 + finally: + app.dependency_overrides.clear() + + return captured["expires_after"] + + +def test_file_team_override_overrides_caller( + mocker: MockerFixture, monkeypatch, llm_router: Router +): + """Team enforced_file_expires_after wins over caller-provided value.""" + expires_after = _post_file_with_team_metadata( + monkeypatch, + llm_router, + team_metadata={ + "enforced_file_expires_after": { + "anchor": "created_at", + "seconds": 3600, + } + }, + form_data={ + "purpose": "batch", + "target_model_names": "gpt-3.5-turbo", + "expires_after[anchor]": "created_at", + "expires_after[seconds]": "86400", + }, + ) + assert expires_after["anchor"] == "created_at" + assert expires_after["seconds"] == 3600 + + +def test_file_no_team_setting_preserves_caller( + mocker: MockerFixture, monkeypatch, llm_router: Router +): + """No team setting = caller-provided expires_after passes through.""" + expires_after = _post_file_with_team_metadata( + monkeypatch, + llm_router, + team_metadata={}, + form_data={ + "purpose": "batch", + "target_model_names": "gpt-3.5-turbo", + "expires_after[anchor]": "created_at", + "expires_after[seconds]": "86400", + }, + ) + assert expires_after["anchor"] == "created_at" + assert expires_after["seconds"] == 86400