diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 28c4cf5e66f..7671e1c5425 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -2,6 +2,7 @@ ## This hook is used to check for LiteLLM managed files in the request body, and replace them with model-specific file id import base64 +import itertools import json from collections.abc import Iterator, Mapping, Sequence from types import MappingProxyType @@ -255,6 +256,22 @@ def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, s _MANAGED_FILES_TARGET: Final = "managed_files" +def _user_content_parts(message: AllMessageValues) -> tuple[object, ...]: + match message: + case {"role": "user", "content": list() as parts}: + return tuple(parts) + case _: + return () + + +def _file_id_of_content_part(part: object) -> str | None: + match part: + case {"type": "file", "file": {"file_id": str() as file_id}} if file_id: + return file_id + case _: + return None + + class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Class variables or attributes def __init__(self, internal_usage_cache: InternalUsageCache, prisma_client: PrismaClient): @@ -1056,21 +1073,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): """ Gets file ids from messages """ - file_ids = [] - for message in messages: - if message.get("role") == "user": - content = message.get("content") - if content: - if isinstance(content, str): - continue - for c in content: - if c.get("type") == "file": - file_object = cast(ChatCompletionFileObject, c) - file_object_file_field = file_object["file"] - file_id = file_object_file_field.get("file_id") - if file_id: - file_ids.append(file_id) - return file_ids + parts: Final = itertools.chain.from_iterable(_user_content_parts(message) for message in messages) + return [file_id for file_id in map(_file_id_of_content_part, parts) if file_id] def get_file_ids_from_responses_input(self, input: Union[str, List[Dict[str, object]]]) -> List[str]: """ diff --git a/tests/integration/providers/test_ollama_prompt_tools_wire.py b/tests/integration/providers/test_ollama_prompt_tools_wire.py index f3aa18cb366..2e02835d32f 100644 --- a/tests/integration/providers/test_ollama_prompt_tools_wire.py +++ b/tests/integration/providers/test_ollama_prompt_tools_wire.py @@ -1198,14 +1198,20 @@ def test_malformed_content_is_a_400_naming_the_message_and_field( @pytest.mark.parametrize( - "content", + ("content", "detail"), [ - pytest.param(22, id="user-int-content"), - pytest.param({"text": _QUESTION}, id="user-dict-content"), - pytest.param(["just a string"], id="user-string-part"), + pytest.param(22, "has int content; content must be a string or a list of content parts", id="user-int-content"), + pytest.param( + {"text": _QUESTION}, + "has dict content; content must be a string or a list of content parts", + id="user-dict-content", + ), + pytest.param(["just a string"], "has a str content part; content parts must be objects", id="user-string-part"), ], ) -def test_non_list_user_content_is_rejected_before_it_reaches_ollama(gateway: Gateway, content: JsonValue) -> None: +def test_non_list_user_content_is_rejected_before_it_reaches_ollama( + gateway: Gateway, content: JsonValue, detail: str +) -> None: with _ollama_server(lambda _: _generate_reply(_ANSWER)) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) code, text = _post( @@ -1213,8 +1219,8 @@ def test_non_list_user_content_is_rejected_before_it_reaches_ollama(gateway: Gat "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": content}], "tools": [_WEATHER_TOOL]}, ) - assert code >= 400, text - assert _error_message(text), text + assert code == 400, text + assert f"the user message at index 0 {detail}" in _error_message(text), text assert _generate_calls(wire) == () payload: Final = _post_chat(gateway, model, _second_turn(), tools=[_WEATHER_TOOL]) assert json.dumps(payload["choices"]).count(_ANSWER) == 1, payload diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py index 11d9dbec49c..5d2d1661efe 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_files.py +++ b/tests/unit/enterprise/proxy/hooks/test_managed_files.py @@ -73,6 +73,43 @@ def test_get_file_ids_from_messages_skips_bedrock_content_blocks_without_type(): assert file_ids == ["file-keep"] +_UNSCANNABLE_USER_CONTENT = [ + pytest.param(22, id="int-content"), + pytest.param({"text": "hi"}, id="dict-content"), + pytest.param(["hi"], id="string-part"), +] + + +@pytest.mark.parametrize("content", _UNSCANNABLE_USER_CONTENT) +def test_get_file_ids_from_messages_skips_content_it_cannot_scan(content): + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=MagicMock()) + messages = [ + {"role": "user", "content": content}, + {"role": "user", "content": [{"type": "file", "file": {"file_id": "file-keep"}}]}, + ] + assert proxy_managed_files.get_file_ids_from_messages(messages) == ["file-keep"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("content", _UNSCANNABLE_USER_CONTENT) +async def test_async_pre_call_hook_passes_unscannable_user_content_through_untouched(content): + from litellm.proxy._types import UserAPIKeyAuth + + prisma_client = AsyncMock() + proxy_managed_files = PROXY_LiteLLMManagedFiles(DualCache(), prisma_client=prisma_client) + data = {"messages": [{"role": "user", "content": content}], "model": "gpt-5.5"} + + result = await proxy_managed_files.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(user_id="123", parent_otel_span=MagicMock()), + cache=DualCache(), + data=data, + call_type="acompletion", + ) + + assert result == {"messages": [{"role": "user", "content": content}], "model": "gpt-5.5"} + assert prisma_client.mock_calls == [] + + @pytest.mark.asyncio async def test_async_pre_call_hook_batch_retrieve(): from litellm.proxy._types import UserAPIKeyAuth