diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index e11c6e70540..bc1eb6cebc2 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -46,6 +46,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.openai_files_endpoints.common_utils import ( + BATCH_CREATE_HIDDEN_PARAM, FILE_LIST_CONTINUATION_CHUNK_SIZE, MAX_FILE_LIST_LIMIT, _is_base64_encoded_unified_file_id, @@ -1321,7 +1322,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ## Check if unified_file_id is in the response unified_file_id = response._hidden_params.get("unified_file_id") # managed file id unified_batch_id = response._hidden_params.get("unified_batch_id") # managed batch id - is_batch_create: Final = unified_file_id is not None + is_batch_create: Final = response._hidden_params.get(BATCH_CREATE_HIDDEN_PARAM) is True model_id = cast(Optional[str], response._hidden_params.get("model_id")) model_name = cast(Optional[str], response._hidden_params.get("model_name")) @@ -1410,10 +1411,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): litellm_parent_otel_span=user_api_key_dict.parent_otel_span, ) - # Only record batch creation metric on actual create (not retrieve/cancel). - # unified_file_id in _hidden_params is only set by the create_batch endpoint. - original_unified_file_id = response._hidden_params.get("unified_file_id") - if original_unified_file_id: + if is_batch_create: prom_logger = self._get_prometheus_logger() if prom_logger: batch_provider = "" diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index be889a22cae..5c4bacd757c 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -25,6 +25,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( get_custom_llm_provider_from_request_query, ) from litellm.proxy.openai_files_endpoints.common_utils import ( + BATCH_CREATE_HIDDEN_PARAM, _is_base64_encoded_unified_file_id, add_internal_model_credentials, apply_team_provider_credentials, @@ -347,6 +348,8 @@ async def create_batch( **_create_batch_data, ) + response._hidden_params[BATCH_CREATE_HIDDEN_PARAM] = True + ### CALL HOOKS ### - modify outgoing data response = await proxy_logging_obj.post_call_success_hook( data=data, user_api_key_dict=user_api_key_dict, response=response diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 992ed0d814d..15eeddbc489 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -37,6 +37,8 @@ MAX_FILE_LIST_LIMIT: Final = 10000 FILE_LIST_CONTINUATION_CHUNK_SIZE: Final = 500 +BATCH_CREATE_HIDDEN_PARAM: Final = "batch_create" + def validate_file_list_limit(limit: int | None) -> None: """Reject a ``limit`` outside the range OpenAI documents for GET /v1/files.""" diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index 57394f1cebe..7e96c956664 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -10,6 +10,7 @@ from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFi from litellm.caching import DualCache from litellm.proxy._types import CallTypes from litellm.proxy.openai_files_endpoints.common_utils import ( + BATCH_CREATE_HIDDEN_PARAM, _is_base64_encoded_unified_file_id, encode_file_id_with_model, ) @@ -3185,7 +3186,7 @@ def _batch_response(batch_id, output_file_id=None, is_create=False): output_file_id=output_file_id, ) if is_create: - batch._hidden_params["unified_file_id"] = "unified-input-file-id" + batch._hidden_params[BATCH_CREATE_HIDDEN_PARAM] = True return batch @@ -3411,11 +3412,8 @@ async def test_provider_format_file_without_ownership_row_stays_accessible(): @pytest.mark.asyncio -async def test_post_call_batch_create_stores_ownership_row(): - """ - Batch creation (response hidden params carry the unified input file id) - must write an ownership row attributed to the creating key. - """ +@pytest.mark.parametrize("batch_id", [MODEL_ENCODED_BATCH_ID, RAW_PROVIDER_BATCH_ID]) +async def test_post_call_batch_create_stores_ownership_row(batch_id): from litellm.proxy._types import UserAPIKeyAuth prisma_client = AsyncMock() @@ -3432,13 +3430,11 @@ async def test_post_call_batch_create_stores_ownership_row(): user_api_key_dict=UserAPIKeyAuth( user_id="user_a", team_id="team_a", parent_otel_span=MagicMock() ), - response=_batch_response(MODEL_ENCODED_BATCH_ID, is_create=True), + response=_batch_response(batch_id, is_create=True), ) upsert_call = prisma_client.db.litellm_managedobjecttable.upsert.await_args - assert upsert_call.kwargs["where"] == { - "unified_object_id": MODEL_ENCODED_BATCH_ID - } + assert upsert_call.kwargs["where"] == {"unified_object_id": batch_id} create_data = upsert_call.kwargs["data"]["create"] assert create_data["created_by"] == "user_a" assert create_data["team_id"] == "team_a" diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index f3ad8a8592e..091b958d7c3 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -15,6 +15,7 @@ from typing import Optional from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth +from litellm.proxy.openai_files_endpoints.common_utils import BATCH_CREATE_HIDDEN_PARAM from litellm.types.llms.openai import FileListPage, OpenAIFileObject from litellm.types.utils import LiteLLMBatch @@ -1540,6 +1541,11 @@ async def test_batch_create_hook_persists_creating_key_and_tags(): managed_files = _make_managed_files_instance() creator = UserAPIKeyAuth(api_key="sk-the-creator", user_id="alice", parent_otel_span=None) create_response = _make_batch_response(status="validating", output_file_id=None) + create_response._hidden_params = { + BATCH_CREATE_HIDDEN_PARAM: True, + "model_id": "model-deploy-xyz", + "model_name": "azure/gpt-4", + } await managed_files.async_post_call_success_hook( data={"litellm_metadata": {"tags": ["env:prod", "team:ml"], "user_api_key": creator.api_key}}, @@ -1554,6 +1560,52 @@ async def test_batch_create_hook_persists_creating_key_and_tags(): assert stored["user_api_key_dict"] is creator +@pytest.mark.asyncio +async def test_batch_create_hook_records_created_metric_once(): + managed_files = _make_managed_files_instance() + prometheus_logger = MagicMock() + managed_files._get_prometheus_logger = MagicMock(return_value=prometheus_logger) + create_response = _make_batch_response(status="validating", output_file_id=None) + create_response._hidden_params = { + BATCH_CREATE_HIDDEN_PARAM: True, + "model_id": "model-deploy-xyz", + "model_name": "azure/gpt-4", + } + + await managed_files.async_post_call_success_hook( + data={}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-the-creator", user_id="alice", parent_otel_span=None), + response=create_response, + ) + + prometheus_logger.record_managed_batch_created.assert_called_once() + recorded = prometheus_logger.record_managed_batch_created.call_args.kwargs + assert recorded["model"] == "azure/gpt-4" + assert recorded["api_provider"] == "azure" + assert recorded["user"] == "alice" + + +@pytest.mark.asyncio +async def test_batch_retrieve_hook_does_not_record_created_metric(): + managed_files = _make_managed_files_instance() + prometheus_logger = MagicMock() + managed_files._get_prometheus_logger = MagicMock(return_value=prometheus_logger) + retrieve_response = _make_batch_response(status="in_progress", output_file_id=None) + retrieve_response._hidden_params = { + "unified_batch_id": "some-unified-batch-id", + "model_id": "model-deploy-xyz", + "model_name": "azure/gpt-4", + } + + await managed_files.async_post_call_success_hook( + data={}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-the-poller", user_id="bob", parent_otel_span=None), + response=retrieve_response, + ) + + prometheus_logger.record_managed_batch_created.assert_not_called() + + @pytest.mark.asyncio async def test_batch_retrieve_hook_does_not_claim_attribution(): """A retrieve carries unified_batch_id but no unified_file_id, so it must not rewrite diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index 2f6a5a3b0e0..a37c8ff2bb4 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -29,6 +29,7 @@ added to this layer raises instead of silently passing - the inventory of seams cannot drift without a test failure. """ +import base64 import json from contextlib import ExitStack from dataclasses import dataclass @@ -36,7 +37,7 @@ from typing import Any, Dict, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest - +from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles import litellm import litellm.proxy.batches_endpoints.endpoints as endpoints @@ -989,6 +990,67 @@ async def test_create__uses_acreate_batch_route_type(harness, openai_env_creds): assert harness.pre_call.call_args.kwargs["route_type"] == "acreate_batch" +def install_managed_files_hook(harness: Harness) -> AsyncMock: + prisma_client = AsyncMock() + managed_files = _PROXY_LiteLLMManagedFiles(MagicMock(async_set_cache=AsyncMock()), prisma_client=prisma_client) + harness.logging.post_call_success_hook = AsyncMock(side_effect=managed_files.async_post_call_success_hook) + harness.router.model_list = [] + return prisma_client + + +TEAM_A_KEY = UserAPIKeyAuth(api_key="sk-team-a", user_id="user_a", team_id="team_a") + + +def assert_ownership_registered_for_team_a(prisma_client: AsyncMock, batch_id: str) -> None: + upsert = prisma_client.db.litellm_managedobjecttable.upsert + upsert.assert_awaited_once() + assert upsert.await_args.kwargs["where"] == {"unified_object_id": batch_id} + created = upsert.await_args.kwargs["data"]["create"] + assert created["created_by"] == "user_a" + assert created["team_id"] == "team_a" + prisma_client.db.litellm_managedobjecttable.update_many.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body", + [ + {"input_file_id": AZURE_FILE_ID}, + {"input_file_id": "file-plain", "model": "vertex-model"}, + {"input_file_id": "file-plain"}, + ], + ids=["model_encoded_file_id", "model_param", "provider_fallback"], +) +async def test_create__registers_ownership_for_creator(harness, openai_env_creds, body): + set_body(harness, {**body, "endpoint": "/v1/chat/completions", "completion_window": "24h"}) + prisma_client = install_managed_files_hook(harness) + + resp = await call_create(harness, user=TEAM_A_KEY) + + assert_ownership_registered_for_team_a(prisma_client, resp.id) + + +@pytest.mark.asyncio +async def test_create__unified_file_id_registers_ownership_for_creator(harness): + unified_input_file_id = base64.urlsafe_b64encode( + b"litellm_proxy:application/octet-stream;unified_id,input-uuid;target_model_names,gpt-4o-mini" + ).decode() + set_body( + harness, + { + "input_file_id": unified_input_file_id, + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + prisma_client = install_managed_files_hook(harness) + + resp = await call_create(harness, user=TEAM_A_KEY) + + assert harness.router_acreate.call_count == 1 + assert_ownership_registered_for_team_a(prisma_client, resp.id) + + @pytest.mark.asyncio async def test_create__metadata_sanitized_before_forwarding(harness, openai_env_creds): set_body(