diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 5ed490703473..ae5905f9cdfb 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -658,7 +658,7 @@ def get_file_ids_from_messages(self, messages: List[AllMessageValues]) -> List[s if isinstance(content, str): continue for c in content: - if c["type"] == "file": + 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") diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 32ae61d7f58a..b44d21368f80 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -3,6 +3,7 @@ """ import io +import json import mimetypes import re from os import PathLike @@ -132,6 +133,39 @@ def strip_none_values_from_message(message: AllMessageValues) -> AllMessageValue return cast(AllMessageValues, {k: v for k, v in message.items() if v is not None}) +def extract_search_results_text(search_results: object) -> str: + """ + Extract model-visible text from OpenAI tool-message ``search_results``. + + Used by token estimators and TPM limiters so large search result payloads + cannot bypass preflight checks via a small ``content`` field. + + Counts every string field forwarded on Bedrock ``SearchResultBlock``: + ``source``, ``title``, ``content[].text``, and ``citations``. + """ + if not isinstance(search_results, list): + return "" + texts = "" + for result in search_results: + if not isinstance(result, dict): + continue + for key in ("source", "title"): + value = result.get(key) + if isinstance(value, str): + texts += value + content = result.get("content") + if isinstance(content, list): + for block in content: + if isinstance(block, dict): + text = block.get("text") + if isinstance(text, str): + texts += text + citations = result.get("citations") + if citations is not None: + texts += json.dumps(citations, separators=(",", ":")) + return texts + + def convert_content_list_to_str( message: Union[AllMessageValues, ChatCompletionResponseMessage], ) -> str: @@ -152,6 +186,7 @@ def convert_content_list_to_str( elif message_content is not None and isinstance(message_content, str): texts = message_content + texts += extract_search_results_text(message.get("search_results")) return texts diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 03341ce0c9b4..46e9b43a4295 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -3658,6 +3658,7 @@ def stringify_json_tool_call_content(messages: List) -> List: ToolInputSchemaBlock as BedrockToolInputSchemaBlock, ) from litellm.types.llms.bedrock import ToolJsonSchemaBlock as BedrockToolJsonSchemaBlock +from litellm.types.llms.bedrock import SearchResultBlock from litellm.types.llms.bedrock import ToolResultBlock as BedrockToolResultBlock from litellm.types.llms.bedrock import ( ToolResultContentBlock as BedrockToolResultContentBlock, @@ -4063,6 +4064,122 @@ def _convert_to_bedrock_tool_call_invoke( ) +def _append_bedrock_tool_result_media_block( + tool_result_content_blocks: List[BedrockToolResultContentBlock], + processed_block: BedrockContentBlock, + content: dict, + content_type: str, +) -> None: + if "image" in processed_block: + tool_result_content_blocks.append( + BedrockToolResultContentBlock(image=processed_block["image"]) + ) + elif "document" in processed_block: + tool_result_content_blocks.append( + BedrockToolResultContentBlock(document=processed_block["document"]) + ) + else: + verbose_logger.warning( + "Bedrock Converse: unrecognized BedrockContentBlock keys " + "%s for %s tool-result block %s; dropping.", + list(processed_block.keys()), + content_type, + content, + ) + + +def _append_bedrock_tool_result_image_url_block( + tool_result_content_blocks: List[BedrockToolResultContentBlock], + content: dict, +) -> None: + format: Optional[str] = None + if isinstance(content["image_url"], dict): + image_url = content["image_url"]["url"] + format = content["image_url"].get("format") + else: + image_url = content["image_url"] + processed_block = BedrockImageProcessor.process_image_sync( + image_url=image_url, + format=format, + ) + _append_bedrock_tool_result_media_block( + tool_result_content_blocks, processed_block, content, "image_url" + ) + + +def _append_bedrock_tool_result_file_block( + tool_result_content_blocks: List[BedrockToolResultContentBlock], + content: dict, +) -> None: + # Match the user-message path (_process_file_message): accept either + # file_data (base64 data URI) or file_id (server-side reference / URL). + file_obj = content.get("file") or {} + file_data = file_obj.get("file_data") + file_id = file_obj.get("file_id") + if file_data is None and file_id is None: + raise litellm.BadRequestError( + message="file_data and file_id cannot both be None. Got={}".format(content), + model="", + llm_provider="bedrock", + ) + processed_block = BedrockImageProcessor.process_image_sync( + image_url=cast(str, file_id or file_data), + format=file_obj.get("format"), + ) + _append_bedrock_tool_result_media_block( + tool_result_content_blocks, processed_block, content, "file" + ) + + +def _parse_bedrock_tool_result_content_list( + content_list: List, +) -> List[BedrockToolResultContentBlock]: + tool_result_content_blocks: List[BedrockToolResultContentBlock] = [] + for content in content_list: + if content["type"] == "text": + tool_result_content_blocks.append( + BedrockToolResultContentBlock(text=content["text"]) + ) + elif content["type"] == "image_url": + _append_bedrock_tool_result_image_url_block( + tool_result_content_blocks, content + ) + elif content["type"] == "file": + _append_bedrock_tool_result_file_block(tool_result_content_blocks, content) + return tool_result_content_blocks + + +def _build_bedrock_tool_result_content_blocks( + message: Union[ChatCompletionToolMessage, ChatCompletionFunctionMessage], +) -> tuple[List[BedrockToolResultContentBlock], bool]: + # Optional OpenAI tool-message extension: + # allow structured Bedrock search results on tool messages and map them + # directly to toolResult.content[].searchResult for Converse API. + # + # If `search_results` is present, we intentionally prefer it over `content` + # to avoid generating mixed text + searchResult blocks. + search_results = message.get("search_results") + if isinstance(search_results, list): + tool_result_content_blocks: List[BedrockToolResultContentBlock] = [] + for result in search_results: + if not isinstance(result, dict): + continue + tool_result_content_blocks.append( + BedrockToolResultContentBlock( + searchResult=cast(SearchResultBlock, result) + ) + ) + if tool_result_content_blocks: + return tool_result_content_blocks, True + + message_content = message["content"] + if isinstance(message_content, str): + return [BedrockToolResultContentBlock(text=message_content)], False + if isinstance(message_content, List): + return _parse_bedrock_tool_result_content_list(message_content), False + return [], False + + def _convert_to_bedrock_tool_call_result( message: Union[ChatCompletionToolMessage, ChatCompletionFunctionMessage], ) -> BedrockContentBlock: @@ -4106,90 +4223,18 @@ def _convert_to_bedrock_tool_call_result( """ - """ - tool_result_content_blocks: List[BedrockToolResultContentBlock] = [] - if isinstance(message["content"], str): - tool_result_content_blocks.append( - BedrockToolResultContentBlock(text=message["content"]) - ) - elif isinstance(message["content"], List): - content_list = message["content"] - for content in content_list: - if content["type"] == "text": - tool_result_content_blocks.append( - BedrockToolResultContentBlock(text=content["text"]) - ) - elif content["type"] == "image_url": - format: Optional[str] = None - if isinstance(content["image_url"], dict): - image_url = content["image_url"]["url"] - format = content["image_url"].get("format") - else: - image_url = content["image_url"] - _block: BedrockContentBlock = BedrockImageProcessor.process_image_sync( - image_url=image_url, - format=format, - ) - if "image" in _block: - tool_result_content_blocks.append( - BedrockToolResultContentBlock(image=_block["image"]) - ) - elif "document" in _block: - tool_result_content_blocks.append( - BedrockToolResultContentBlock(document=_block["document"]) - ) - else: - verbose_logger.warning( - "Bedrock Converse: unrecognized BedrockContentBlock keys " - "%s for image_url tool-result block %s; dropping.", - list(_block.keys()), - content, - ) - elif content["type"] == "file": - # Match the user-message path (_process_file_message): accept - # either file_data (base64 data URI) or file_id (server-side - # reference / URL) and hand off to BedrockImageProcessor. Raise - # BadRequestError on both-None rather than silently dropping. - file_obj = content.get("file") or {} - file_data = file_obj.get("file_data") - file_id = file_obj.get("file_id") - if file_data is None and file_id is None: - raise litellm.BadRequestError( - message="file_data and file_id cannot both be None. Got={}".format( - content - ), - model="", - llm_provider="bedrock", - ) - file_format = file_obj.get("format") - _file_block: BedrockContentBlock = ( - BedrockImageProcessor.process_image_sync( - image_url=cast(str, file_id or file_data), - format=file_format, - ) - ) - if "document" in _file_block: - tool_result_content_blocks.append( - BedrockToolResultContentBlock(document=_file_block["document"]) - ) - elif "image" in _file_block: - tool_result_content_blocks.append( - BedrockToolResultContentBlock(image=_file_block["image"]) - ) - else: - verbose_logger.warning( - "Bedrock Converse: unrecognized BedrockContentBlock keys " - "%s for file tool-result block %s; dropping.", - list(_file_block.keys()), - content, - ) + tool_result_content_blocks, used_search_results = ( + _build_bedrock_tool_result_content_blocks(message) + ) message.get("name", "") id = str(message.get("tool_call_id", str(uuid.uuid4()))) tool_result = BedrockToolResultBlock( - content=tool_result_content_blocks, - toolUseId=id, + content=tool_result_content_blocks, toolUseId=id ) + if used_search_results: + tool_result["status"] = cast(Literal["success"], "success") content_block = BedrockContentBlock(toolResult=tool_result) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index e6a68de07e92..7889336f416f 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -486,6 +486,14 @@ def _count_messages( use_default_image_token_count, default_token_count, ) + elif key == "search_results" and isinstance(value, list): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + extract_search_results_text, + ) + + search_results_text = extract_search_results_text(value) + if search_results_text: + num_tokens += params.count_function(search_results_text) else: # Skip unsupported keys instead of raising an error continue diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index d58d2e27595f..1a8e59ba8db4 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -41,6 +41,7 @@ from litellm.types.llms.bedrock import * from litellm.types.llms.openai import ( AllMessageValues, + ChatCompletionAnnotation, ChatCompletionAssistantMessage, ChatCompletionRedactedThinkingBlock, ChatCompletionResponseMessage, @@ -1944,6 +1945,75 @@ def _translate_message_content(self, content_blocks: List[ContentBlock]) -> Tupl return content_str, tools, reasoningContentBlocks, citationsContentBlocks + @staticmethod + def _transform_citations_to_annotations( + citations_content_blocks: Optional[List[CitationsContentBlock]], + ) -> Tuple[Optional[str], Optional[List[ChatCompletionAnnotation]]]: + """ + Convert Bedrock citationsContent blocks into OpenAI-style annotations. + + Returns: + citations_text: concatenated text from citationsContent.content + annotations: OpenAI URL citation annotations + """ + if not citations_content_blocks: + return None, None + + annotations: List[ChatCompletionAnnotation] = [] + citations_text_parts: List[str] = [] + content_offset = 0 + + for citations_block in citations_content_blocks: + block_text = "" + raw_content = citations_block.get("content") + if isinstance(raw_content, list): + for content_part in raw_content: + if isinstance(content_part, dict): + _text = content_part.get("text") + if isinstance(_text, str): + block_text += _text + + block_offset = content_offset + if block_text: + citations_text_parts.append(block_text) + content_offset += len(block_text) + + raw_citations = citations_block.get("citations") + if not isinstance(raw_citations, list): + continue + + for citation in raw_citations: + if not isinstance(citation, dict): + continue + + location = citation.get("location") + if not isinstance(location, dict): + continue + + search_location = location.get("searchResultLocation") + if not isinstance(search_location, dict): + continue + + start = search_location.get("start") + end = search_location.get("end") + if not isinstance(start, int) or not isinstance(end, int): + continue + + annotations.append( + ChatCompletionAnnotation( + type="url_citation", + url_citation={ + "start_index": block_offset + start, + "end_index": block_offset + end, + "title": str(citation.get("title") or ""), + "url": str(citation.get("source") or ""), + }, + ) + ) + + citations_text = "".join(citations_text_parts) if citations_text_parts else None + return citations_text, annotations or None + @staticmethod def _unwrap_bedrock_properties(json_str: str) -> str: """ @@ -2126,6 +2196,24 @@ def _transform_response( # noqa: PLR0915 provider_specific_fields ) + citations_text, annotations = self._transform_citations_to_annotations( + citationsContentBlocks + ) + citations_included_in_content = False + if citations_text: + stripped_content = content_str.strip() + if not stripped_content: + content_str = citations_text + citations_included_in_content = True + elif not any(char.isalnum() for char in stripped_content): + # Bedrock may emit the cited sentence in citationsContent and only + # punctuation in the text blocks; stitch citations_text in front so + # its annotation span indices stay aligned with the final content. + content_str = citations_text + content_str + citations_included_in_content = True + if annotations and citations_included_in_content: + chat_completion_message["annotations"] = annotations + if reasoningContentBlocks is not None: chat_completion_message["reasoning_content"] = ( self._transform_reasoning_content(reasoningContentBlocks) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 36a389d233dc..a9fb02a1c959 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -344,7 +344,13 @@ def get_response_headers( if litellm_call_id: return_headers["x-litellm-call-id"] = litellm_call_id if custom_headers: - return_headers.update(custom_headers) + # Ensure custom headers don't override actual upstream response headers or let framework defaults (like content-length: 0) interfere. + sanitized_custom_headers = { + key: value + for key, value in custom_headers.items() + if key.lower() not in excluded_headers + } + return_headers.update(sanitized_custom_headers) return return_headers diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 5db2a45054a1..60e640636aa1 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -49,9 +49,24 @@ class DocumentBlock(TypedDict): name: str +class SearchResultBlock(TypedDict, total=False): + """ + Search result block used in Bedrock toolResult content. + + Reference: + https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_SearchResultBlock.html + """ + + source: str + title: str + content: List[dict] + citations: dict + + class ToolResultContentBlock(TypedDict, total=False): image: ImageBlock document: DocumentBlock + searchResult: SearchResultBlock json: dict text: str @@ -106,24 +121,41 @@ class CitationWebLocationBlock(TypedDict, total=False): domain: str +class CitationSearchResultLocationBlock(TypedDict, total=False): + """ + Character span of a Nova grounding citation within the cited content, + plus the index of the search result it refers to. + """ + + start: int + end: int + searchResultIndex: int + + class CitationLocationBlock(TypedDict, total=False): """ - Location block containing the web location for a citation. + Location block describing where a citation points to. """ web: CitationWebLocationBlock + searchResultLocation: CitationSearchResultLocationBlock class CitationReferenceBlock(TypedDict, total=False): """ - Citation reference block containing a single citation with its location. - - Each citation contains: - - location.web.url: The URL of the source - - location.web.domain: The domain of the source + Citation reference block containing a single citation with its location, + source URL and title. """ location: CitationLocationBlock + source: str + title: str + + +class CitationGeneratedContentBlock(TypedDict, total=False): + """A piece of generated text associated with a citationsContent block.""" + + text: str class CitationsContentBlock(TypedDict, total=False): @@ -131,27 +163,33 @@ class CitationsContentBlock(TypedDict, total=False): Citations content block returned by Nova grounding (web search) tool. When Nova grounding is enabled via systemTool, the model may return - citationsContent blocks containing web search citation references. + citationsContent blocks containing the grounded text and its citation + references. Reference: https://docs.aws.amazon.com/nova/latest/userguide/grounding.html Example response structure: { "citationsContent": { + "content": [{"text": "The grounded answer text ..."}], "citations": [ { "location": { - "web": { - "url": "https://example.com/article", - "domain": "example.com" + "searchResultLocation": { + "start": 0, + "end": 42, + "searchResultIndex": 0 } - } + }, + "source": "https://example.com/article", + "title": "Example Article" } ] } } """ + content: List[CitationGeneratedContentBlock] citations: List[CitationReferenceBlock] diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 14114b22f39e..346909f14eb4 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -965,6 +965,7 @@ class ChatCompletionDeltaChunk(TypedDict, total=False): class ChatCompletionResponseMessage(TypedDict, total=False): content: Optional[ChatCompletionAssistantContentValue] + annotations: Optional[List[ChatCompletionAnnotation]] tool_calls: Optional[List[ChatCompletionToolCallChunk]] role: Literal["assistant"] function_call: Optional[ChatCompletionToolCallFunctionChunk] 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 b9c14739245b..7b82d1eabd93 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -38,6 +38,39 @@ def test_get_file_ids_from_messages(): ] +def test_get_file_ids_from_messages_skips_bedrock_content_blocks_without_type(): + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=MagicMock() + ) + messages = [ + { + "role": "user", + "content": [ + {"text": "What is Apptio?"}, + { + "toolResult": { + "toolUseId": "tooluse_123", + "status": "success", + "content": [ + { + "searchResult": { + "source": "source", + "title": "title", + "content": [{"text": "snippet"}], + "citations": {"enabled": True}, + } + } + ], + } + }, + {"type": "file", "file": {"file_id": "file-keep"}}, + ], + } + ] + file_ids = proxy_managed_files.get_file_ids_from_messages(messages) + assert file_ids == ["file-keep"] + + @pytest.mark.asyncio async def test_async_pre_call_hook_batch_retrieve(): from litellm.proxy._types import UserAPIKeyAuth @@ -95,9 +128,9 @@ async def test_async_pre_call_deployment_hook_resolves_model_id_from_litellm_met kwargs=kwargs, call_type=CallTypes.acreate_batch ) - assert result["input_file_id"] == provider_file_id, ( - f"Expected provider file ID '{provider_file_id}', got '{result['input_file_id']}'" - ) + assert ( + result["input_file_id"] == provider_file_id + ), f"Expected provider file ID '{provider_file_id}', got '{result['input_file_id']}'" @pytest.mark.asyncio @@ -134,9 +167,9 @@ async def test_async_pre_call_deployment_hook_prefers_top_level_model_info(): kwargs=kwargs, call_type=CallTypes.acreate_batch ) - assert result["input_file_id"] == top_level_provider_file, ( - "Should prefer top-level model_info over litellm_metadata" - ) + assert ( + result["input_file_id"] == top_level_provider_file + ), "Should prefer top-level model_info over litellm_metadata" @pytest.mark.asyncio @@ -162,9 +195,9 @@ async def test_async_pre_call_deployment_hook_no_model_info_leaves_file_id_uncha kwargs=kwargs, call_type=CallTypes.acreate_batch ) - assert result["input_file_id"] == managed_file_id, ( - "File ID should remain unchanged when model_info is not available" - ) + assert ( + result["input_file_id"] == managed_file_id + ), "File ID should remain unchanged when model_info is not available" # def test_list_managed_files(): @@ -341,7 +374,9 @@ async def test_async_pre_call_hook_for_unified_finetuning_job(): @pytest.mark.asyncio -@pytest.mark.parametrize("call_type", ["afile_content", "afile_delete", "afile_retrieve"]) +@pytest.mark.parametrize( + "call_type", ["afile_content", "afile_delete", "afile_retrieve"] +) async def test_can_user_call_unified_file_id(call_type): """ Test that on file retrieve, delete, and content we check if the user has access to the file @@ -601,7 +636,7 @@ async def test_error_file_id_for_failed_batch(): "litellm_model_name": "gpt-5.5", "unified_batch_id": "litellm_proxy;model_id:test-model-id;llm_batch_id:batch_abc123", } - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=AsyncMock() ) @@ -620,12 +655,11 @@ async def test_error_file_id_for_failed_batch(): # Mock the afile_retrieve to simulate retrieving error file metadata with patch("litellm.afile_retrieve", new_callable=AsyncMock) as mock_retrieve: mock_retrieve.return_value = error_file_object - + user_api_key_dict = UserAPIKeyAuth( - user_id="test-user-123", - parent_otel_span=MagicMock() + user_id="test-user-123", parent_otel_span=MagicMock() ) - + response = await proxy_managed_files.async_post_call_success_hook( data={}, user_api_key_dict=user_api_key_dict, @@ -636,7 +670,9 @@ async def test_error_file_id_for_failed_batch(): assert cast(LiteLLMBatch, response).error_file_id is not None assert not cast(LiteLLMBatch, response).error_file_id.startswith("error-") # Verify it's a base64 encoded managed file ID - assert _is_base64_encoded_unified_file_id(cast(LiteLLMBatch, response).error_file_id) + assert _is_base64_encoded_unified_file_id( + cast(LiteLLMBatch, response).error_file_id + ) @pytest.mark.asyncio @@ -650,7 +686,7 @@ async def test_async_post_call_success_hook_twice_assert_no_unique_violation(): # Use AsyncMock instead of real database connection prisma_client = AsyncMock() - + batch = LiteLLMBatch( id="bGl0ZWxsbV9wcm94eTttb2RlbF9pZDoxMjM0NTY3OTtsbG1fYmF0Y2hfaWQ6YmF0Y2hfNjg1YzVlNWQ2Mzk4ODE5MGI4NWJkYjIxNDdiYTEzMWQ", completion_window="24h", @@ -678,8 +714,10 @@ async def test_async_post_call_success_hook_twice_assert_no_unique_violation(): # first retrieve batch tasks = [] first_create_task = asyncio.create_task - with patch('asyncio.create_task') as mock_create_task: - mock_create_task.side_effect = lambda coro: tasks.append(first_create_task(coro)) or tasks[-1] + with patch("asyncio.create_task") as mock_create_task: + mock_create_task.side_effect = ( + lambda coro: tasks.append(first_create_task(coro)) or tasks[-1] + ) response = await proxy_managed_files.async_post_call_success_hook( data={}, @@ -700,8 +738,10 @@ async def test_async_post_call_success_hook_twice_assert_no_unique_violation(): # second retrieve batch tasks = [] second_create_task = asyncio.create_task - with patch('asyncio.create_task') as mock_create_task: - mock_create_task.side_effect = lambda coro: tasks.append(second_create_task(coro)) or tasks[-1] + with patch("asyncio.create_task") as mock_create_task: + mock_create_task.side_effect = ( + lambda coro: tasks.append(second_create_task(coro)) or tasks[-1] + ) await proxy_managed_files.async_post_call_success_hook( data={}, @@ -728,7 +768,7 @@ def test_update_responses_input_with_unified_file_id(): # Create a base64-encoded unified file ID # This decodes to: litellm_proxy:application/pdf;unified_id,6c0b5890-8914-48e0-b8f4-0ae5ed3c14a5;target_model_names,gpt-4o;llm_output_file_id,file-ECBPW7ML9g7XHdwGgUPZaM;llm_output_file_model_id,e26453f9e76e7993680d0068d98c1f4cc205bbad0967a33c664893568ca743c2 unified_file_id = "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9wZGY7dW5pZmllZF9pZCw2YzBiNTg5MC04OTE0LTQ4ZTAtYjhmNC0wYWU1ZWQzYzE0YTU7dGFyZ2V0X21vZGVsX25hbWVzLGdwdC00bztsbG1fb3V0cHV0X2ZpbGVfaWQsZmlsZS1FQ0JQVzdNTDlnN1hIZHdHZ1VQWmFNO2xsbV9vdXRwdXRfZmlsZV9tb2RlbF9pZCxlMjY0NTNmOWU3NmU3OTkzNjgwZDAwNjhkOThjMWY0Y2MyMDViYmFkMDk2N2EzM2M2NjQ4OTM1NjhjYTc0M2My" - + # Test input with unified file ID in content array input_data = [ { @@ -745,15 +785,18 @@ def test_update_responses_input_with_unified_file_id(): ], } ] - + # Update the input updated_input = update_responses_input_with_model_file_ids(input=input_data) - + # Verify the file_id was updated to the provider-specific file ID assert updated_input[0]["content"][0]["type"] == "input_file" assert updated_input[0]["content"][0]["file_id"] == "file-ECBPW7ML9g7XHdwGgUPZaM" assert updated_input[0]["content"][1]["type"] == "input_text" - assert updated_input[0]["content"][1]["text"] == "What is the first dragon in the book?" + assert ( + updated_input[0]["content"][1]["text"] + == "What is the first dragon in the book?" + ) def test_update_responses_input_with_regular_file_id(): @@ -767,7 +810,7 @@ def test_update_responses_input_with_regular_file_id(): # Regular OpenAI file ID (not a unified file ID) regular_file_id = "file-abc123xyz" - + input_data = [ { "role": "user", @@ -783,10 +826,10 @@ def test_update_responses_input_with_regular_file_id(): ], } ] - + # Update the input updated_input = update_responses_input_with_model_file_ids(input=input_data) - + # Verify the file_id was kept unchanged (regular OpenAI file ID) assert updated_input[0]["content"][0]["type"] == "input_file" assert updated_input[0]["content"][0]["file_id"] == regular_file_id @@ -800,11 +843,11 @@ def test_update_responses_input_with_string_input(): from litellm.litellm_core_utils.prompt_templates.common_utils import ( update_responses_input_with_model_file_ids, ) - + input_data = "What is AI?" - + updated_input = update_responses_input_with_model_file_ids(input=input_data) - + assert updated_input == input_data assert isinstance(updated_input, str) @@ -822,7 +865,7 @@ def test_update_responses_input_with_multiple_file_ids(): unified_file_id = "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9wZGY7dW5pZmllZF9pZCw2YzBiNTg5MC04OTE0LTQ4ZTAtYjhmNC0wYWU1ZWQzYzE0YTU7dGFyZ2V0X21vZGVsX25hbWVzLGdwdC00bztsbG1fb3V0cHV0X2ZpbGVfaWQsZmlsZS1FQ0JQVzdNTDlnN1hIZHdHZ1VQWmFNO2xsbV9vdXRwdXRfZmlsZV9tb2RlbF9pZCxlMjY0NTNmOWU3NmU3OTkzNjgwZDAwNjhkOThjMWY0Y2MyMDViYmFkMDk2N2EzM2M2NjQ4OTM1NjhjYTc0M2My" # Regular OpenAI file ID regular_file_id = "file-regular123" - + input_data = [ { "role": "user", @@ -842,9 +885,9 @@ def test_update_responses_input_with_multiple_file_ids(): ], } ] - + updated_input = update_responses_input_with_model_file_ids(input=input_data) - + # Verify unified file ID was updated assert updated_input[0]["content"][0]["file_id"] == "file-ECBPW7ML9g7XHdwGgUPZaM" # Verify regular file ID was kept unchanged @@ -864,7 +907,7 @@ def test_update_responses_input_with_model_file_id_mapping(): # Managed file ID (unified) managed_file_id = "litellm_proxy_file_123" - + # Model file ID mapping model_file_id_mapping = { managed_file_id: { @@ -872,7 +915,7 @@ def test_update_responses_input_with_model_file_id_mapping(): "model_id_2": "azure_file_xyz", } } - + input_data = [ { "role": "user", @@ -888,24 +931,24 @@ def test_update_responses_input_with_model_file_id_mapping(): ], } ] - + # Update input with model_id_1 mapping updated_input = update_responses_input_with_model_file_ids( input=input_data, model_id="model_id_1", model_file_id_mapping=model_file_id_mapping, ) - + # Verify the file_id was mapped to the correct provider-specific file ID assert updated_input[0]["content"][0]["file_id"] == "openai_file_abc" - + # Test with different model_id updated_input_2 = update_responses_input_with_model_file_ids( input=input_data, model_id="model_id_2", model_file_id_mapping=model_file_id_mapping, ) - + assert updated_input_2[0]["content"][0]["file_id"] == "azure_file_xyz" @@ -913,7 +956,7 @@ def test_update_responses_tools_with_model_file_id_mapping(): """ Test that update_responses_tools_with_model_file_ids correctly maps file IDs in code_interpreter tools with container.file_ids. - + This is a regression test for the issue where managed file IDs in tools.container.file_ids were not being replaced with provider-specific file IDs, causing "string too long" errors from OpenAI. @@ -925,7 +968,7 @@ def test_update_responses_tools_with_model_file_id_mapping(): # Managed file IDs managed_file_id_1 = "litellm_proxy_file_123" managed_file_id_2 = "litellm_proxy_file_456" - + # Model file ID mapping model_file_id_mapping = { managed_file_id_1: { @@ -935,7 +978,7 @@ def test_update_responses_tools_with_model_file_id_mapping(): "model_id_1": "openai_file_def", }, } - + tools = [ { "type": "code_interpreter", @@ -945,17 +988,20 @@ def test_update_responses_tools_with_model_file_id_mapping(): }, } ] - + # Update tools with model mapping updated_tools = update_responses_tools_with_model_file_ids( tools=tools, model_id="model_id_1", model_file_id_mapping=model_file_id_mapping, ) - + # Verify the file IDs were mapped to provider-specific file IDs assert updated_tools[0]["type"] == "code_interpreter" - assert updated_tools[0]["container"]["file_ids"] == ["openai_file_abc", "openai_file_def"] + assert updated_tools[0]["container"]["file_ids"] == [ + "openai_file_abc", + "openai_file_def", + ] def test_update_responses_tools_without_mapping(): @@ -968,7 +1014,7 @@ def test_update_responses_tools_without_mapping(): ) regular_file_id = "file-abc123" - + tools = [ { "type": "code_interpreter", @@ -978,14 +1024,14 @@ def test_update_responses_tools_without_mapping(): }, } ] - + # Update tools without mapping updated_tools = update_responses_tools_with_model_file_ids( tools=tools, model_id=None, model_file_id_mapping=None, ) - + # Verify the file ID was kept unchanged assert updated_tools[0]["container"]["file_ids"] == [regular_file_id] @@ -1001,13 +1047,13 @@ def test_update_responses_tools_with_mixed_file_ids(): managed_file_id = "litellm_proxy_file_123" regular_file_id = "file-abc123" - + model_file_id_mapping = { managed_file_id: { "model_id_1": "openai_file_abc", }, } - + tools = [ { "type": "code_interpreter", @@ -1017,16 +1063,19 @@ def test_update_responses_tools_with_mixed_file_ids(): }, } ] - + # Update tools updated_tools = update_responses_tools_with_model_file_ids( tools=tools, model_id="model_id_1", model_file_id_mapping=model_file_id_mapping, ) - + # Verify managed file ID was mapped and regular file ID was kept - assert updated_tools[0]["container"]["file_ids"] == ["openai_file_abc", regular_file_id] + assert updated_tools[0]["container"]["file_ids"] == [ + "openai_file_abc", + regular_file_id, + ] def test_get_file_ids_from_responses_tools(): @@ -1037,7 +1086,7 @@ def test_get_file_ids_from_responses_tools(): proxy_managed_files = _PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=MagicMock() ) - + tools = [ { "type": "code_interpreter", @@ -1047,9 +1096,9 @@ def test_get_file_ids_from_responses_tools(): }, } ] - + file_ids = proxy_managed_files.get_file_ids_from_responses_tools(tools) - + assert file_ids == ["file-123", "file-456"] @@ -1060,7 +1109,7 @@ def test_get_file_ids_from_responses_tools_multiple_tools(): proxy_managed_files = _PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=MagicMock() ) - + tools = [ { "type": "code_interpreter", @@ -1080,9 +1129,9 @@ def test_get_file_ids_from_responses_tools_multiple_tools(): }, }, ] - + file_ids = proxy_managed_files.get_file_ids_from_responses_tools(tools) - + # Should extract file IDs only from code_interpreter tools assert file_ids == ["file-123", "file-456", "file-789"] @@ -1094,15 +1143,15 @@ def test_get_file_ids_from_responses_tools_empty(): proxy_managed_files = _PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=MagicMock() ) - + # Test with None file_ids = proxy_managed_files.get_file_ids_from_responses_tools(None) assert file_ids == [] - + # Test with empty list file_ids = proxy_managed_files.get_file_ids_from_responses_tools([]) assert file_ids == [] - + # Test with tools without file_ids tools = [{"type": "file_search"}] file_ids = proxy_managed_files.get_file_ids_from_responses_tools(tools) @@ -1119,30 +1168,30 @@ async def test_check_file_ids_access_with_unified_file_ids(): # Create a unified file ID unified_file_id = "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9wZGY7dW5pZmllZF9pZCw2YzBiNTg5MC04OTE0LTQ4ZTAtYjhmNC0wYWU1ZWQzYzE0YTU7dGFyZ2V0X21vZGVsX25hbWVzLGdwdC00bztsbG1fb3V0cHV0X2ZpbGVfaWQsZmlsZS1FQ0JQVzdNTDlnN1hIZHdHZ1VQWmFNO2xsbV9vdXRwdXRfZmlsZV9tb2RlbF9pZCxlMjY0NTNmOWU3NmU3OTkzNjgwZDAwNjhkOThjMWY0Y2MyMDViYmFkMDk2N2EzM2M2NjQ4OTM1NjhjYTc0M2My" regular_file_id = "file-abc123" - + # Mock the access check to return True prisma_client = AsyncMock() internal_usage_cache = MagicMock() - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( internal_usage_cache=internal_usage_cache, prisma_client=prisma_client, ) - + # Mock can_user_call_unified_file_id to return True proxy_managed_files.can_user_call_unified_file_id = AsyncMock(return_value=True) - + user_api_key_dict = UserAPIKeyAuth( user_id="test_user_123", parent_otel_span=MagicMock(), ) - + # Should not raise an exception for accessible files await proxy_managed_files.check_file_ids_access( [unified_file_id, regular_file_id], user_api_key_dict, ) - + # Verify can_user_call_unified_file_id was called for the unified file ID proxy_managed_files.can_user_call_unified_file_id.assert_called_once_with( unified_file_id, user_api_key_dict @@ -1155,32 +1204,32 @@ async def test_check_file_ids_access_denied(): Test that check_file_ids_access raises HTTPException when user doesn't have access. """ from litellm.proxy._types import UserAPIKeyAuth - + unified_file_id = "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9wZGY7dW5pZmllZF9pZCw2YzBiNTg5MC04OTE0LTQ4ZTAtYjhmNC0wYWU1ZWQzYzE0YTU7dGFyZ2V0X21vZGVsX25hbWVzLGdwdC00bztsbG1fb3V0cHV0X2ZpbGVfaWQsZmlsZS1FQ0JQVzdNTDlnN1hIZHdHZ1VQWmFNO2xsbV9vdXRwdXRfZmlsZV9tb2RlbF9pZCxlMjY0NTNmOWU3NmU3OTkzNjgwZDAwNjhkOThjMWY0Y2MyMDViYmFkMDk2N2EzM2M2NjQ4OTM1NjhjYTc0M2My" - + prisma_client = AsyncMock() internal_usage_cache = MagicMock() - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( internal_usage_cache=internal_usage_cache, prisma_client=prisma_client, ) - + # Mock can_user_call_unified_file_id to return False (access denied) proxy_managed_files.can_user_call_unified_file_id = AsyncMock(return_value=False) - + user_api_key_dict = UserAPIKeyAuth( user_id="test_user_123", parent_otel_span=MagicMock(), ) - + # Should raise HTTPException with 403 status code with pytest.raises(HTTPException) as exc_info: await proxy_managed_files.check_file_ids_access( [unified_file_id], user_api_key_dict, ) - + assert exc_info.value.status_code == 403 assert "does not have access to the file" in exc_info.value.detail @@ -1191,32 +1240,32 @@ async def test_check_file_ids_access_with_regular_files_only(): Test that check_file_ids_access doesn't check access for regular (non-unified) file IDs. """ from litellm.proxy._types import UserAPIKeyAuth - + regular_file_id_1 = "file-abc123" regular_file_id_2 = "file-xyz789" - + prisma_client = AsyncMock() internal_usage_cache = MagicMock() - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( internal_usage_cache=internal_usage_cache, prisma_client=prisma_client, ) - + # Mock can_user_call_unified_file_id (should not be called for regular files) proxy_managed_files.can_user_call_unified_file_id = AsyncMock() - + user_api_key_dict = UserAPIKeyAuth( user_id="test_user_123", parent_otel_span=MagicMock(), ) - + # Should not raise exception and should not call can_user_call_unified_file_id await proxy_managed_files.check_file_ids_access( [regular_file_id_1, regular_file_id_2], user_api_key_dict, ) - + # Verify can_user_call_unified_file_id was NOT called proxy_managed_files.can_user_call_unified_file_id.assert_not_called() @@ -1227,31 +1276,31 @@ async def test_completion_with_file_access_check(): Test that completion call type checks file access before processing. """ from litellm.proxy._types import UserAPIKeyAuth - + unified_file_id = "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9wZGY7dW5pZmllZF9pZCw2YzBiNTg5MC04OTE0LTQ4ZTAtYjhmNC0wYWU1ZWQzYzE0YTU7dGFyZ2V0X21vZGVsX25hbWVzLGdwdC00bztsbG1fb3V0cHV0X2ZpbGVfaWQsZmlsZS1FQ0JQVzdNTDlnN1hIZHdHZ1VQWmFNO2xsbV9vdXRwdXRfZmlsZV9tb2RlbF9pZCxlMjY0NTNmOWU3NmU3OTkzNjgwZDAwNjhkOThjMWY0Y2MyMDViYmFkMDk2N2EzM2M2NjQ4OTM1NjhjYTc0M2My" - + prisma_client = AsyncMock() prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None) - + internal_usage_cache = MagicMock() internal_usage_cache.async_get_cache = AsyncMock(return_value=None) - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( internal_usage_cache=internal_usage_cache, prisma_client=prisma_client, ) - + # Mock the get_model_file_id_mapping to return empty dict proxy_managed_files.get_model_file_id_mapping = AsyncMock(return_value={}) - + # Mock access check to allow access proxy_managed_files.can_user_call_unified_file_id = AsyncMock(return_value=True) - + user_api_key_dict = UserAPIKeyAuth( user_id="test_user_123", parent_otel_span=MagicMock(), ) - + data = { "messages": [ { @@ -1267,7 +1316,7 @@ async def test_completion_with_file_access_check(): ], "model": "gpt-5.5", } - + # Should not raise exception result = await proxy_managed_files.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -1275,7 +1324,7 @@ async def test_completion_with_file_access_check(): data=data, call_type="acompletion", ) - + # Verify access check was called proxy_managed_files.can_user_call_unified_file_id.assert_called_once() @@ -1286,32 +1335,32 @@ async def test_responses_with_file_access_check(): Test that responses API checks file access for files in both input and tools. """ from litellm.proxy._types import UserAPIKeyAuth - + unified_file_id_1 = "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9wZGY7dW5pZmllZF9pZCw2YzBiNTg5MC04OTE0LTQ4ZTAtYjhmNC0wYWU1ZWQzYzE0YTU7dGFyZ2V0X21vZGVsX25hbWVzLGdwdC00bztsbG1fb3V0cHV0X2ZpbGVfaWQsZmlsZS1FQ0JQVzdNTDlnN1hIZHdHZ1VQWmFNO2xsbV9vdXRwdXRfZmlsZV9tb2RlbF9pZCxlMjY0NTNmOWU3NmU3OTkzNjgwZDAwNjhkOThjMWY0Y2MyMDViYmFkMDk2N2EzM2M2NjQ4OTM1NjhjYTc0M2My" unified_file_id_2 = "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9qc29uO3VuaWZpZWRfaWQsNzc3Nzc3Nzc7dGFyZ2V0X21vZGVsX25hbWVzLGdwdC00bztsbG1fb3V0cHV0X2ZpbGVfaWQsZmlsZS1YWVo7bGxtX291dHB1dF9maWxlX21vZGVsX2lkLG1vZGVsXzEyMw" - + prisma_client = AsyncMock() prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None) - + internal_usage_cache = MagicMock() internal_usage_cache.async_get_cache = AsyncMock(return_value=None) - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( internal_usage_cache=internal_usage_cache, prisma_client=prisma_client, ) - + # Mock the get_model_file_id_mapping to return empty dict proxy_managed_files.get_model_file_id_mapping = AsyncMock(return_value={}) - + # Mock access check to allow access proxy_managed_files.can_user_call_unified_file_id = AsyncMock(return_value=True) - + user_api_key_dict = UserAPIKeyAuth( user_id="test_user_123", parent_otel_span=MagicMock(), ) - + data = { "input": [ { @@ -1333,7 +1382,7 @@ async def test_responses_with_file_access_check(): ], "model": "gpt-5.5", } - + # Should not raise exception result = await proxy_managed_files.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -1341,7 +1390,7 @@ async def test_responses_with_file_access_check(): data=data, call_type="aresponses", ) - + # Verify access check was called for both file IDs assert proxy_managed_files.can_user_call_unified_file_id.call_count == 2 @@ -1353,17 +1402,19 @@ async def test_store_unified_file_id_with_none_file_object(): (e.g., for batch output files that are stored before file metadata is available). """ from litellm.proxy._types import UserAPIKeyAuth - + prisma_client = AsyncMock() - prisma_client.db.litellm_managedfiletable.create = AsyncMock(return_value=MagicMock()) + prisma_client.db.litellm_managedfiletable.create = AsyncMock( + return_value=MagicMock() + ) internal_usage_cache = MagicMock() internal_usage_cache.async_set_cache = AsyncMock() - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( internal_usage_cache=internal_usage_cache, prisma_client=prisma_client, ) - + # Store with file_object=None (simulating batch output file storage) await proxy_managed_files.store_unified_file_id( file_id="test-unified-file-id", @@ -1372,7 +1423,7 @@ async def test_store_unified_file_id_with_none_file_object(): model_mappings={"model-123": "file-provider-xyz"}, user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), ) - + # Verify DB create was called with expected data (without file_object) prisma_client.db.litellm_managedfiletable.create.assert_called_once() call_args = prisma_client.db.litellm_managedfiletable.create.call_args @@ -1387,34 +1438,38 @@ async def test_afile_delete_returns_provider_response_when_stored_file_object_no stored file_object is None (e.g., for batch output files). """ from litellm.types.llms.openai import OpenAIFileObject - + unified_file_id = "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9qc29uO3VuaWZpZWRfaWQsdGVzdC1pZDt0YXJnZXRfbW9kZWxfbmFtZXMsZ3B0LTRvO2xsbV9vdXRwdXRfZmlsZV9pZCxmaWxlLXByb3ZpZGVyLXh5ejtsbG1fb3V0cHV0X2ZpbGVfbW9kZWxfaWQsbW9kZWwtMTIz" - + prisma_client = AsyncMock() db_record = MagicMock() db_record.model_mappings = '{"model-123": "file-provider-xyz"}' - prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(return_value=db_record) + prisma_client.db.litellm_managedfiletable.find_first = AsyncMock( + return_value=db_record + ) prisma_client.db.litellm_managedfiletable.delete = AsyncMock() - + internal_usage_cache = MagicMock() - internal_usage_cache.async_get_cache = AsyncMock(return_value={ - "unified_file_id": unified_file_id, - "model_mappings": {"model-123": "file-provider-xyz"}, - "flat_model_file_ids": ["file-provider-xyz"], - "file_object": None, - "created_by": "test-user", - "updated_by": "test-user", - }) + internal_usage_cache.async_get_cache = AsyncMock( + return_value={ + "unified_file_id": unified_file_id, + "model_mappings": {"model-123": "file-provider-xyz"}, + "flat_model_file_ids": ["file-provider-xyz"], + "file_object": None, + "created_by": "test-user", + "updated_by": "test-user", + } + ) internal_usage_cache.async_set_cache = AsyncMock() - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( internal_usage_cache=internal_usage_cache, prisma_client=prisma_client, ) - + # Mock the delete_unified_file_id to return None (simulating file_object=None) proxy_managed_files.delete_unified_file_id = AsyncMock(return_value=None) - + # Mock router response provider_delete_response = OpenAIFileObject( id="file-provider-xyz", @@ -1424,16 +1479,16 @@ async def test_afile_delete_returns_provider_response_when_stored_file_object_no filename="test.jsonl", purpose="batch", ) - + mock_router = MagicMock() mock_router.afile_delete = AsyncMock(return_value=provider_delete_response) - + result = await proxy_managed_files.afile_delete( file_id=unified_file_id, litellm_parent_otel_span=None, llm_router=mock_router, ) - + # Should return the provider response with the unified file ID assert result is not None assert result.id == unified_file_id @@ -1446,21 +1501,21 @@ async def test_afile_retrieve_fetches_from_provider_when_file_object_none(): file_object is None (e.g., for batch output files). """ from litellm.types.llms.openai import OpenAIFileObject - + prisma_client = AsyncMock() internal_usage_cache = MagicMock() - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( internal_usage_cache=internal_usage_cache, prisma_client=prisma_client, ) - + # Mock get_unified_file_id to return a stored object with file_object=None stored_file = MagicMock() stored_file.file_object = None stored_file.model_mappings = {"model-123": "file-provider-xyz"} proxy_managed_files.get_unified_file_id = AsyncMock(return_value=stored_file) - + # Mock the router and provider response provider_file_response = OpenAIFileObject( id="file-provider-xyz", @@ -1470,23 +1525,25 @@ async def test_afile_retrieve_fetches_from_provider_when_file_object_none(): filename="output.jsonl", purpose="batch_output", ) - + mock_router = MagicMock() - mock_router.get_deployment_credentials_with_provider = MagicMock(return_value={ - "api_key": "test-key", - "api_base": "https://api.openai.com", - }) - + mock_router.get_deployment_credentials_with_provider = MagicMock( + return_value={ + "api_key": "test-key", + "api_base": "https://api.openai.com", + } + ) + with patch("litellm.afile_retrieve", new_callable=AsyncMock) as mock_afile_retrieve: mock_afile_retrieve.return_value = provider_file_response - + unified_file_id = "test-unified-file-id" result = await proxy_managed_files.afile_retrieve( file_id=unified_file_id, litellm_parent_otel_span=None, llm_router=mock_router, ) - + # Should return the provider response with the unified file ID assert result is not None assert result.id == unified_file_id @@ -1501,27 +1558,27 @@ async def test_afile_retrieve_raises_error_when_no_router_and_file_object_none() """ prisma_client = AsyncMock() internal_usage_cache = MagicMock() - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( internal_usage_cache=internal_usage_cache, prisma_client=prisma_client, ) - + # Mock get_unified_file_id to return a stored object with file_object=None stored_file = MagicMock() stored_file.file_object = None stored_file.model_mappings = {"model-123": "file-provider-xyz"} proxy_managed_files.get_unified_file_id = AsyncMock(return_value=stored_file) - + unified_file_id = "test-unified-file-id" - + with pytest.raises(Exception) as exc_info: await proxy_managed_files.afile_retrieve( file_id=unified_file_id, litellm_parent_otel_span=None, llm_router=None, ) - + assert "llm_router is required" in str(exc_info.value) @@ -1532,15 +1589,15 @@ async def test_afile_retrieve_returns_stored_file_object_when_exists(): (the normal case for user-uploaded files). """ from litellm.types.llms.openai import OpenAIFileObject - + prisma_client = AsyncMock() internal_usage_cache = MagicMock() - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( internal_usage_cache=internal_usage_cache, prisma_client=prisma_client, ) - + # Mock get_unified_file_id to return a stored object WITH file_object stored_file_object = OpenAIFileObject( id="test-unified-file-id", @@ -1553,13 +1610,13 @@ async def test_afile_retrieve_returns_stored_file_object_when_exists(): stored_file = MagicMock() stored_file.file_object = stored_file_object proxy_managed_files.get_unified_file_id = AsyncMock(return_value=stored_file) - + result = await proxy_managed_files.afile_retrieve( file_id="test-unified-file-id", litellm_parent_otel_span=None, llm_router=None, ) - + # Should return the stored file object directly assert result == stored_file_object @@ -1572,21 +1629,21 @@ async def test_afile_retrieve_raises_error_for_non_managed_file(): """ prisma_client = AsyncMock() internal_usage_cache = MagicMock() - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( internal_usage_cache=internal_usage_cache, prisma_client=prisma_client, ) - + # Mock get_unified_file_id to return None (file not found) proxy_managed_files.get_unified_file_id = AsyncMock(return_value=None) - + with pytest.raises(Exception) as exc_info: await proxy_managed_files.afile_retrieve( file_id="non-existent-file-id", litellm_parent_otel_span=None, ) - + assert "not found" in str(exc_info.value) @@ -1597,54 +1654,58 @@ async def test_list_batches_from_managed_objects_table(): from litellm.proxy._types import UserAPIKeyAuth prisma_client = AsyncMock() - + batch_record_1 = MagicMock() batch_record_1.unified_object_id = "unified-batch-id-1" - batch_record_1.file_object = json.dumps({ - "id": "batch_abc123", - "object": "batch", - "endpoint": "/v1/chat/completions", - "completion_window": "24h", - "status": "completed", - "created_at": 1234567890, - "input_file_id": "file-input-1", - "request_counts": {"total": 1, "completed": 1, "failed": 0}, - }) - + batch_record_1.file_object = json.dumps( + { + "id": "batch_abc123", + "object": "batch", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "status": "completed", + "created_at": 1234567890, + "input_file_id": "file-input-1", + "request_counts": {"total": 1, "completed": 1, "failed": 0}, + } + ) + batch_record_2 = MagicMock() batch_record_2.unified_object_id = "unified-batch-id-2" - batch_record_2.file_object = json.dumps({ - "id": "batch_xyz789", - "object": "batch", - "endpoint": "/v1/chat/completions", - "completion_window": "24h", - "status": "in_progress", - "created_at": 1234567891, - "input_file_id": "file-input-2", - "request_counts": {"total": 5, "completed": 2, "failed": 0}, - }) - + batch_record_2.file_object = json.dumps( + { + "id": "batch_xyz789", + "object": "batch", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "status": "in_progress", + "created_at": 1234567891, + "input_file_id": "file-input-2", + "request_counts": {"total": 5, "completed": 2, "failed": 0}, + } + ) + prisma_client.db.litellm_managedobjecttable.find_many.return_value = [ batch_record_1, batch_record_2, ] - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=prisma_client ) - + result = await proxy_managed_files.list_user_batches( user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), limit=10, ) - + assert result["object"] == "list" assert len(result["data"]) == 2 assert result["data"][0].id == "unified-batch-id-1" assert result["data"][1].id == "unified-batch-id-2" assert result["first_id"] == "unified-batch-id-1" assert result["last_id"] == "unified-batch-id-2" - + # Should filter by user_id (created_by) prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once_with( where={"file_purpose": "batch", "created_by": "test-user"}, @@ -1659,21 +1720,21 @@ async def test_list_batches_from_managed_objects_table_empty_list(): prisma_client = AsyncMock() prisma_client.db.litellm_managedobjecttable.find_many.return_value = [] - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=prisma_client ) - + result = await proxy_managed_files.list_user_batches( user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), ) - + assert result["object"] == "list" assert len(result["data"]) == 0 assert result["first_id"] is None assert result["last_id"] is None assert result["has_more"] is False - + # Verify where clause includes created_by filter # Default take is 20 when no limit is provided prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once_with( @@ -1685,6 +1746,7 @@ async def test_list_batches_from_managed_objects_table_empty_list(): def _create_unified_batch_id(model_id: str, batch_id: str) -> str: import base64 + unified_str = f"litellm_proxy;model_id:{model_id};llm_batch_id:{batch_id}" return base64.urlsafe_b64encode(unified_str.encode()).decode().rstrip("=") @@ -1694,11 +1756,11 @@ async def test_list_batches_from_managed_objects_table_provider_filter_raises_ex from litellm.proxy._types import UserAPIKeyAuth prisma_client = AsyncMock() - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=prisma_client ) - + # Filtering by provider should raise Exception with pytest.raises(Exception) as exc_info: await proxy_managed_files.list_user_batches( @@ -1706,11 +1768,11 @@ async def test_list_batches_from_managed_objects_table_provider_filter_raises_ex limit=10, provider="openai", ) - + assert str(exc_info.value) == ( "Filtering by 'provider' is not supported when using managed batches." ) - + # Verify find_many was NOT called since exception is raised before database query prisma_client.db.litellm_managedobjecttable.find_many.assert_not_called() @@ -1720,7 +1782,7 @@ async def test_list_batches_from_managed_objects_table_target_model_name_filter_ from litellm.proxy._types import UserAPIKeyAuth prisma_client = AsyncMock() - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=prisma_client ) @@ -1732,59 +1794,64 @@ async def test_list_batches_from_managed_objects_table_target_model_name_filter_ limit=10, target_model_names="gpt-5.5,gpt-3.5", ) - + assert str(exc_info.value) == ( "Filtering by 'target_model_names' is not supported when using managed batches." ) - + # Verify find_many was NOT called since exception is raised before database query prisma_client.db.litellm_managedobjecttable.find_many.assert_not_called() + @pytest.mark.asyncio async def test_list_batches_from_managed_objects_table_filters_by_created_by(): from litellm.proxy._types import UserAPIKeyAuth prisma_client = AsyncMock() - + # Create batch for user1 batch_user1 = MagicMock() batch_user1.unified_object_id = "unified-batch-user1" - batch_user1.file_object = json.dumps({ - "id": "batch_user1_abc", - "object": "batch", - "endpoint": "/v1/chat/completions", - "completion_window": "24h", - "status": "completed", - "created_at": 1234567890, - "input_file_id": "file-input-user1", - "request_counts": {"total": 1, "completed": 1, "failed": 0}, - }) - + batch_user1.file_object = json.dumps( + { + "id": "batch_user1_abc", + "object": "batch", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "status": "completed", + "created_at": 1234567890, + "input_file_id": "file-input-user1", + "request_counts": {"total": 1, "completed": 1, "failed": 0}, + } + ) + # Create batch for user2 batch_user2 = MagicMock() batch_user2.unified_object_id = "unified-batch-user2" - batch_user2.file_object = json.dumps({ - "id": "batch_user2_xyz", - "object": "batch", - "endpoint": "/v1/chat/completions", - "completion_window": "24h", - "status": "completed", - "created_at": 1234567891, - "input_file_id": "file-input-user2", - "request_counts": {"total": 2, "completed": 2, "failed": 0}, - }) - + batch_user2.file_object = json.dumps( + { + "id": "batch_user2_xyz", + "object": "batch", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "status": "completed", + "created_at": 1234567891, + "input_file_id": "file-input-user2", + "request_counts": {"total": 2, "completed": 2, "failed": 0}, + } + ) + proxy_managed_files = _PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=prisma_client ) - + # Query with user1's API key - should only return user1's batch prisma_client.db.litellm_managedobjecttable.find_many.return_value = [batch_user1] result_user1 = await proxy_managed_files.list_user_batches( user_api_key_dict=UserAPIKeyAuth(user_id="user1"), limit=10, ) - + assert len(result_user1["data"]) == 1 assert result_user1["data"][0].id == "unified-batch-user1" prisma_client.db.litellm_managedobjecttable.find_many.assert_called_with( @@ -1792,14 +1859,14 @@ async def test_list_batches_from_managed_objects_table_filters_by_created_by(): take=10, order={"created_at": "desc"}, ) - + # Query with user2's API key - should only return user2's batch prisma_client.db.litellm_managedobjecttable.find_many.return_value = [batch_user2] result_user2 = await proxy_managed_files.list_user_batches( user_api_key_dict=UserAPIKeyAuth(user_id="user2"), limit=10, ) - + assert len(result_user2["data"]) == 1 assert result_user2["data"][0].id == "unified-batch-user2" prisma_client.db.litellm_managedobjecttable.find_many.assert_called_with( @@ -1822,7 +1889,7 @@ async def test_return_unified_file_id_includes_expires_at(): filename="test.jsonl", purpose="batch", status="uploaded", - expires_at=1234657890, + expires_at=1234657890, ) file_object._hidden_params = {"model_id": "test-model-id"} @@ -1862,25 +1929,27 @@ async def test_return_unified_file_id_includes_expires_at(): async def test_user_b_cannot_retrieve_user_a_batch(): """ Test that User B cannot retrieve a batch created by User A. - + This verifies batch isolation between users at the database/hook level. """ from litellm.proxy._types import UserAPIKeyAuth - + prisma_client = AsyncMock() - + # Mock database to return User A as the creator batch_record = MagicMock() batch_record.created_by = "user_a_id" prisma_client.db.litellm_managedobjecttable.find_first.return_value = batch_record - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=prisma_client ) - + # User B tries to retrieve User A's batch - unified_batch_id = "bGl0ZWxsbV9wcm94eTttb2RlbF9pZDpteS1tb2RlbDtsbG1fYmF0Y2hfaWQ6YmF0Y2hfYWJjMTIz" - + unified_batch_id = ( + "bGl0ZWxsbV9wcm94eTttb2RlbF9pZDpteS1tb2RlbDtsbG1fYmF0Y2hfaWQ6YmF0Y2hfYWJjMTIz" + ) + with pytest.raises(HTTPException) as exc_info: await proxy_managed_files.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth( @@ -1890,7 +1959,7 @@ async def test_user_b_cannot_retrieve_user_a_batch(): data={"batch_id": unified_batch_id}, call_type="aretrieve_batch", ) - + # Should raise 403 Permission Denied assert exc_info.value.status_code == 403 @@ -1901,21 +1970,23 @@ async def test_user_b_cannot_cancel_user_a_batch(): Test that User B cannot cancel a batch created by User A. """ from litellm.proxy._types import UserAPIKeyAuth - + prisma_client = AsyncMock() - + # Mock database to return User A as the creator batch_record = MagicMock() batch_record.created_by = "user_a_id" prisma_client.db.litellm_managedobjecttable.find_first.return_value = batch_record - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=prisma_client ) - + # User B tries to cancel User A's batch - unified_batch_id = "bGl0ZWxsbV9wcm94eTttb2RlbF9pZDpteS1tb2RlbDtsbG1fYmF0Y2hfaWQ6YmF0Y2hfYWJjMTIz" - + unified_batch_id = ( + "bGl0ZWxsbV9wcm94eTttb2RlbF9pZDpteS1tb2RlbDtsbG1fYmF0Y2hfaWQ6YmF0Y2hfYWJjMTIz" + ) + with pytest.raises(HTTPException) as exc_info: await proxy_managed_files.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth( @@ -1925,7 +1996,7 @@ async def test_user_b_cannot_cancel_user_a_batch(): data={"batch_id": unified_batch_id}, call_type="acancel_batch", ) - + # Should raise 403 Permission Denied assert exc_info.value.status_code == 403 @@ -1934,26 +2005,28 @@ async def test_user_b_cannot_cancel_user_a_batch(): async def test_user_a_can_retrieve_own_batch(): """ Test that User A can successfully retrieve their own batch. - + This is a positive test case to ensure permission checks don't block legitimate access. """ from litellm.proxy._types import UserAPIKeyAuth - + prisma_client = AsyncMock() - + # Mock database to return User A as the creator batch_record = MagicMock() batch_record.created_by = "user_a_id" prisma_client.db.litellm_managedobjecttable.find_first.return_value = batch_record - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=prisma_client ) - + # User A retrieves their own batch - unified_batch_id = "bGl0ZWxsbV9wcm94eTttb2RlbF9pZDpteS1tb2RlbDtsbG1fYmF0Y2hfaWQ6YmF0Y2hfYWJjMTIz" - + unified_batch_id = ( + "bGl0ZWxsbV9wcm94eTttb2RlbF9pZDpteS1tb2RlbDtsbG1fYmF0Y2hfaWQ6YmF0Y2hfYWJjMTIz" + ) + # Should not raise an exception result = await proxy_managed_files.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth( @@ -1963,7 +2036,7 @@ async def test_user_a_can_retrieve_own_batch(): data={"batch_id": unified_batch_id}, call_type="aretrieve_batch", ) - + # Should successfully return the decoded batch_id assert "batch_id" in result assert result["model"] == "my-model" @@ -1975,21 +2048,23 @@ async def test_user_b_cannot_retrieve_user_a_file(): Test that User B cannot retrieve a file created by User A. """ from litellm.proxy._types import UserAPIKeyAuth - + prisma_client = AsyncMock() - + # Mock database to return User A as the creator file_record = MagicMock() file_record.created_by = "user_a_id" prisma_client.db.litellm_managedfiletable.find_first.return_value = file_record - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( MagicMock(), prisma_client=prisma_client ) - + # User B tries to retrieve User A's file - unified_file_id = "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9qc29uO3VuaWZpZWRfaWQsZmlsZS1hYmMxMjM" - + unified_file_id = ( + "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9qc29uO3VuaWZpZWRfaWQsZmlsZS1hYmMxMjM" + ) + with pytest.raises(HTTPException) as exc_info: await proxy_managed_files.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth( @@ -1999,7 +2074,7 @@ async def test_user_b_cannot_retrieve_user_a_file(): data={"file_id": unified_file_id}, call_type="afile_retrieve", ) - + # Should raise 403 Permission Denied assert exc_info.value.status_code == 403 @@ -2010,21 +2085,23 @@ async def test_user_b_cannot_download_user_a_file_content(): Test that User B cannot download file content for User A's file. """ from litellm.proxy._types import UserAPIKeyAuth - + prisma_client = AsyncMock() - + # Mock database to return User A as the creator file_record = MagicMock() file_record.created_by = "user_a_id" prisma_client.db.litellm_managedfiletable.find_first.return_value = file_record - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( MagicMock(), prisma_client=prisma_client ) - + # User B tries to download User A's file content - unified_file_id = "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9qc29uO3VuaWZpZWRfaWQsZmlsZS1hYmMxMjM" - + unified_file_id = ( + "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9qc29uO3VuaWZpZWRfaWQsZmlsZS1hYmMxMjM" + ) + with pytest.raises(HTTPException) as exc_info: await proxy_managed_files.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth( @@ -2034,7 +2111,7 @@ async def test_user_b_cannot_download_user_a_file_content(): data={"file_id": unified_file_id}, call_type="afile_content", ) - + # Should raise 403 Permission Denied assert exc_info.value.status_code == 403 @@ -2045,21 +2122,23 @@ async def test_user_b_cannot_delete_user_a_file(): Test that User B cannot delete a file created by User A. """ from litellm.proxy._types import UserAPIKeyAuth - + prisma_client = AsyncMock() - + # Mock database to return User A as the creator file_record = MagicMock() file_record.created_by = "user_a_id" prisma_client.db.litellm_managedfiletable.find_first.return_value = file_record - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( MagicMock(), prisma_client=prisma_client ) - + # User B tries to delete User A's file - unified_file_id = "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9qc29uO3VuaWZpZWRfaWQsZmlsZS1hYmMxMjM" - + unified_file_id = ( + "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9qc29uO3VuaWZpZWRfaWQsZmlsZS1hYmMxMjM" + ) + with pytest.raises(HTTPException) as exc_info: await proxy_managed_files.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth( @@ -2069,7 +2148,7 @@ async def test_user_b_cannot_delete_user_a_file(): data={"file_id": unified_file_id}, call_type="afile_delete", ) - + # Should raise 403 Permission Denied assert exc_info.value.status_code == 403 @@ -2078,34 +2157,38 @@ async def test_user_b_cannot_delete_user_a_file(): async def test_user_a_can_retrieve_own_file(): """ Test that User A can successfully retrieve their own file. - + Positive test case to ensure permission checks work correctly for the owner. """ from litellm.proxy._types import UserAPIKeyAuth - + prisma_client = AsyncMock() - + # Mock database to return User A as the creator file_record = MagicMock() file_record.created_by = "user_a_id" file_record.model_mappings = '{"model-123": "file-abc123"}' - file_record.file_object = json.dumps({ - "id": "file-abc123", - "object": "file", - "bytes": 1234, - "created_at": 1234567890, - "filename": "test.jsonl", - "purpose": "batch", - }) + file_record.file_object = json.dumps( + { + "id": "file-abc123", + "object": "file", + "bytes": 1234, + "created_at": 1234567890, + "filename": "test.jsonl", + "purpose": "batch", + } + ) prisma_client.db.litellm_managedfiletable.find_first.return_value = file_record - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( MagicMock(), prisma_client=prisma_client ) - + # User A retrieves their own file - unified_file_id = "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9qc29uO3VuaWZpZWRfaWQsZmlsZS1hYmMxMjM" - + unified_file_id = ( + "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9qc29uO3VuaWZpZWRfaWQsZmlsZS1hYmMxMjM" + ) + # Should not raise an exception result = await proxy_managed_files.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth( @@ -2115,7 +2198,7 @@ async def test_user_a_can_retrieve_own_file(): data={"file_id": unified_file_id}, call_type="afile_retrieve", ) - + # Should successfully return the decoded file_id assert "file_id" in result @@ -2124,44 +2207,46 @@ async def test_user_a_can_retrieve_own_file(): async def test_list_batches_only_returns_user_own_batches(): """ Test that list_user_batches only returns batches created by the requesting user. - + This ensures users cannot see other users' batches in list operations. """ from litellm.proxy._types import UserAPIKeyAuth - + prisma_client = AsyncMock() - + # Create batches for User A batch_user_a = MagicMock() batch_user_a.unified_object_id = "batch-user-a" - batch_user_a.file_object = json.dumps({ - "id": "batch_a", - "object": "batch", - "endpoint": "/v1/chat/completions", - "completion_window": "24h", - "status": "completed", - "created_at": 1234567890, - "input_file_id": "file-a", - "request_counts": {"total": 1, "completed": 1, "failed": 0}, - }) - + batch_user_a.file_object = json.dumps( + { + "id": "batch_a", + "object": "batch", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "status": "completed", + "created_at": 1234567890, + "input_file_id": "file-a", + "request_counts": {"total": 1, "completed": 1, "failed": 0}, + } + ) + # Mock database to only return User A's batches prisma_client.db.litellm_managedobjecttable.find_many.return_value = [batch_user_a] - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=prisma_client ) - + # User A requests their batches result = await proxy_managed_files.list_user_batches( user_api_key_dict=UserAPIKeyAuth(user_id="user_a_id"), limit=10, ) - + # Should only return User A's batches assert len(result["data"]) == 1 assert result["data"][0].id == "batch-user-a" - + # Verify the database query filtered by user_id prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once_with( where={"file_purpose": "batch", "created_by": "user_a_id"}, @@ -2174,51 +2259,49 @@ async def test_list_batches_only_returns_user_own_batches(): async def test_same_user_different_keys_can_access_batch(): """ Test that different API keys for the same user can access the same batch. - + This verifies that permission checks are based on user_id, not API key, allowing users to have multiple keys that can all access their resources. """ from litellm.proxy._types import UserAPIKeyAuth - + prisma_client = AsyncMock() - + # Mock database to return the user_id as creator batch_record = MagicMock() batch_record.created_by = "user_a_id" prisma_client.db.litellm_managedobjecttable.find_first.return_value = batch_record - + proxy_managed_files = _PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=prisma_client ) - - unified_batch_id = "bGl0ZWxsbV9wcm94eTttb2RlbF9pZDpteS1tb2RlbDtsbG1fYmF0Y2hfaWQ6YmF0Y2hfYWJjMTIz" - + + unified_batch_id = ( + "bGl0ZWxsbV9wcm94eTttb2RlbF9pZDpteS1tb2RlbDtsbG1fYmF0Y2hfaWQ6YmF0Y2hfYWJjMTIz" + ) + # First API key for User A retrieves the batch result1 = await proxy_managed_files.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth( - user_id="user_a_id", - api_key="key-1", - parent_otel_span=MagicMock() + user_id="user_a_id", api_key="key-1", parent_otel_span=MagicMock() ), cache=MagicMock(), data={"batch_id": unified_batch_id}, call_type="aretrieve_batch", ) - + assert "batch_id" in result1 - + # Second API key for the same User A retrieves the batch result2 = await proxy_managed_files.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth( - user_id="user_a_id", - api_key="key-2", - parent_otel_span=MagicMock() + user_id="user_a_id", api_key="key-2", parent_otel_span=MagicMock() ), cache=MagicMock(), data={"batch_id": unified_batch_id}, call_type="aretrieve_batch", ) - + assert "batch_id" in result2 # Both keys should get the same result assert result1["batch_id"] == result2["batch_id"] diff --git a/tests/pass_through_unit_tests/test_pass_through_unit_tests.py b/tests/pass_through_unit_tests/test_pass_through_unit_tests.py index 1b16177b755c..65448c6281e2 100644 --- a/tests/pass_through_unit_tests/test_pass_through_unit_tests.py +++ b/tests/pass_through_unit_tests/test_pass_through_unit_tests.py @@ -114,6 +114,40 @@ def test_update_metadata_with_tags_in_header_with_tags(mock_request): assert result == {"existing": "value", "tags": ["tag1", "tag2", "tag3"]} +def test_get_response_headers_filters_excluded_custom_headers(): + """ + Regression test: + Ensure excluded headers from FastAPI defaults (e.g. content-length: 0) + do not override passthrough response headers. + """ + upstream_headers = httpx.Headers( + { + "content-type": "application/json", + "x-amzn-requestid": "req-123", + "content-length": "999", # should be excluded + } + ) + + custom_headers = { + "x-litellm-version": "1.84.0", + "content-length": "0", # should be excluded + "server": "uvicorn", # should be excluded + } + + result = HttpPassThroughEndpointHelpers.get_response_headers( + headers=upstream_headers, + litellm_call_id="call-123", + custom_headers=custom_headers, + ) + + assert result["content-type"] == "application/json" + assert result["x-amzn-requestid"] == "req-123" + assert result["x-litellm-version"] == "1.84.0" + assert result["x-litellm-call-id"] == "call-123" + assert "content-length" not in result + assert "server" not in result + + def test_init_kwargs_for_pass_through_endpoint_basic( mock_request, mock_user_api_key_dict ): diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index c8e72b7ac5b2..a6aa35ee6d12 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -1773,6 +1773,245 @@ async def test_tool_message_string_content_cache_control(): assert tool_message_content[1]["cachePoint"]["type"] == "default" +@pytest.mark.asyncio +async def test_tool_message_search_results_maps_to_bedrock_search_result_block(): + """OpenAI tool message search_results should map to Bedrock searchResult blocks.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + BedrockConverseMessagesProcessor, + _bedrock_converse_messages_pt, + ) + + messages = [ + {"role": "user", "content": "What is Apptio?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "tooluse_a4rBqeZNRTKj2lTskvaO4H", + "type": "function", + "function": { + "name": "RAGRequest", + "arguments": '{"query":"What is Apptio?"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "tooluse_a4rBqeZNRTKj2lTskvaO4H", + "content": "Apptio is a company that makes calls to Bedrock using passthrough APIs via LiteLLM", + "search_results": [ + { + "source": "Great Source of Information About Apptio", + "title": "12adbd74-46bd-4a88-88b2-0048755f6eb5", + "content": [ + { + "text": "Apptio is a company that makes calls to Bedrock using passthrough APIs via LiteLLM" + } + ], + "citations": {"enabled": True}, + } + ], + }, + ] + + result = _bedrock_converse_messages_pt( + messages=messages, + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + llm_provider="bedrock_converse", + ) + async_result = ( + await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages=messages, + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + llm_provider="bedrock_converse", + ) + ) + assert result == async_result + + tool_result = result[2]["content"][0]["toolResult"] + assert tool_result["toolUseId"] == "tooluse_a4rBqeZNRTKj2lTskvaO4H" + assert tool_result["status"] == "success" + assert len(tool_result["content"]) == 1 + assert "searchResult" in tool_result["content"][0] + assert ( + tool_result["content"][0]["searchResult"]["title"] + == "12adbd74-46bd-4a88-88b2-0048755f6eb5" + ) + + +@pytest.mark.asyncio +async def test_tool_message_empty_search_results_falls_back_to_content(): + """Empty search_results must not skip normal tool content processing.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + BedrockConverseMessagesProcessor, + _bedrock_converse_messages_pt, + ) + + messages = [ + {"role": "user", "content": "hello"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "tooluse_empty_search", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "tooluse_empty_search", + "content": "fallback tool text", + "search_results": [], + }, + ] + + result = _bedrock_converse_messages_pt( + messages=messages, + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + llm_provider="bedrock_converse", + ) + async_result = ( + await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages=messages, + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + llm_provider="bedrock_converse", + ) + ) + assert result == async_result + + tool_result = result[2]["content"][0]["toolResult"] + assert tool_result["toolUseId"] == "tooluse_empty_search" + assert "status" not in tool_result + assert len(tool_result["content"]) == 1 + assert tool_result["content"][0]["text"] == "fallback tool text" + + +def test_transform_response_omits_annotations_when_citations_not_stitched(): + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + from litellm.types.utils import ModelResponse + + response_json = { + "metrics": {"latencyMs": 100}, + "output": { + "message": { + "role": "assistant", + "content": [ + { + "citationsContent": { + "content": [{"text": "cited sentence only in citations"}], + "citations": [ + { + "location": { + "searchResultLocation": { + "start": 0, + "end": 5, + } + }, + "source": "https://example.com", + "title": "Example", + } + ], + } + }, + {"text": "separate assistant answer"}, + ], + } + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 10, + "outputTokens": 5, + "totalTokens": 15, + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + }, + } + + class MockResponse: + def json(self): + return response_json + + @property + def text(self): + return json.dumps(response_json) + + config = AmazonConverseConfig() + result = config._transform_response( + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + response=MockResponse(), + model_response=ModelResponse(), + stream=False, + logging_obj=None, + optional_params={}, + api_key=None, + data=None, + messages=[], + encoding=None, + ) + + message = result.choices[0].message + assert message.content == "separate assistant answer" + assert message.model_dump().get("annotations") is None + + +def test_extract_search_results_text_counts_hidden_tool_payload(): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + convert_content_list_to_str, + extract_search_results_text, + ) + from litellm.litellm_core_utils.token_counter import token_counter + + hidden = "x" * 500 + message = { + "role": "tool", + "content": "small", + "search_results": [ + { + "source": "s", + "title": "t", + "content": [{"text": hidden}], + } + ], + } + + extracted = extract_search_results_text(message["search_results"]) + assert hidden in extracted + assert "st" in extracted + assert len(convert_content_list_to_str(message)) > len("small") + + tokens_with_search = token_counter( + model="gpt-3.5-turbo", + messages=[message], + ) + tokens_without_search = token_counter( + model="gpt-3.5-turbo", + messages=[{"role": "tool", "content": "small"}], + ) + assert tokens_with_search > tokens_without_search + + huge_title = "y" * 500 + title_only_message = { + "role": "tool", + "content": "small", + "search_results": [ + {"source": "s", "title": huge_title, "content": []}, + ], + } + assert len(extract_search_results_text(title_only_message["search_results"])) >= 500 + tokens_title_bypass = token_counter( + model="gpt-3.5-turbo", + messages=[title_only_message], + ) + assert tokens_title_bypass > tokens_without_search + + @pytest.mark.asyncio async def test_assistant_tool_calls_cache_control(): """Test that assistant tool_calls with cache_control generate cachePoint blocks.""" @@ -4453,6 +4692,330 @@ def text(self): assert result.choices[0].finish_reason == "stop" +def test_transform_response_citations_content_maps_to_annotations(): + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + from litellm.types.utils import ModelResponse + + response_json = { + "metrics": {"latencyMs": 100}, + "output": { + "message": { + "role": "assistant", + "content": [ + { + "citationsContent": { + "content": [ + { + "text": "Apptio is a company that makes calls to Bedrock using passthrough APIs via LiteLLM" + } + ], + "citations": [ + { + "location": { + "searchResultLocation": { + "start": 0, + "end": 42, + "searchResultIndex": 0, + } + }, + "source": "https://www.apptio.com/about", + "title": "About Apptio", + } + ], + } + }, + {"text": "."}, + ], + } + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 10, + "outputTokens": 5, + "totalTokens": 15, + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + }, + } + + class MockResponse: + def json(self): + return response_json + + @property + def text(self): + return json.dumps(response_json) + + config = AmazonConverseConfig() + model_response = ModelResponse() + + result = config._transform_response( + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + response=MockResponse(), + model_response=model_response, + stream=False, + logging_obj=None, + optional_params={}, + api_key=None, + data=None, + messages=[], + encoding=None, + ) + + message = result.choices[0].message + assert message.content.startswith("Apptio is a company") + assert message.annotations is not None + assert len(message.annotations) == 1 + annotation = message.annotations[0] + assert annotation["type"] == "url_citation" + assert annotation["url_citation"]["start_index"] == 0 + assert annotation["url_citation"]["end_index"] == 42 + assert annotation["url_citation"]["title"] == "About Apptio" + assert annotation["url_citation"]["url"] == "https://www.apptio.com/about" + + +def test_transform_response_citation_null_source_title_become_empty_strings(): + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + from litellm.types.utils import ModelResponse + + response_json = { + "metrics": {"latencyMs": 100}, + "output": { + "message": { + "role": "assistant", + "content": [ + { + "citationsContent": { + "content": [ + { + "text": "Apptio is a company that makes calls to Bedrock" + } + ], + "citations": [ + { + "location": { + "searchResultLocation": { + "start": 0, + "end": 42, + "searchResultIndex": 0, + } + }, + "source": None, + "title": None, + } + ], + } + }, + {"text": "."}, + ], + } + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 10, + "outputTokens": 5, + "totalTokens": 15, + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + }, + } + + class MockResponse: + def json(self): + return response_json + + @property + def text(self): + return json.dumps(response_json) + + config = AmazonConverseConfig() + result = config._transform_response( + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + response=MockResponse(), + model_response=ModelResponse(), + stream=False, + logging_obj=None, + optional_params={}, + api_key=None, + data=None, + messages=[], + encoding=None, + ) + + message = result.choices[0].message + annotation = message.annotations[0] + assert annotation["url_citation"]["url"] == "" + assert annotation["url_citation"]["title"] == "" + + +def test_transform_response_citations_offset_tracks_text_only_blocks(): + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + from litellm.types.utils import ModelResponse + + leading_text = "First sentence without a citation. " + cited_text = "Apptio is a company that makes calls to Bedrock" + response_json = { + "metrics": {"latencyMs": 100}, + "output": { + "message": { + "role": "assistant", + "content": [ + { + "citationsContent": { + "content": [{"text": leading_text}], + } + }, + { + "citationsContent": { + "content": [{"text": cited_text}], + "citations": [ + { + "location": { + "searchResultLocation": { + "start": 0, + "end": len(cited_text), + "searchResultIndex": 0, + } + }, + "source": "https://www.apptio.com/about", + "title": "About Apptio", + } + ], + } + }, + ], + } + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 10, + "outputTokens": 5, + "totalTokens": 15, + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + }, + } + + class MockResponse: + def json(self): + return response_json + + @property + def text(self): + return json.dumps(response_json) + + config = AmazonConverseConfig() + result = config._transform_response( + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + response=MockResponse(), + model_response=ModelResponse(), + stream=False, + logging_obj=None, + optional_params={}, + api_key=None, + data=None, + messages=[], + encoding=None, + ) + + message = result.choices[0].message + expected_start = len(leading_text) + assert message.content == leading_text + cited_text + assert ( + message.content[expected_start : expected_start + len(cited_text)] == cited_text + ) + assert message.annotations is not None + assert len(message.annotations) == 1 + assert message.annotations[0]["url_citation"]["start_index"] == expected_start + assert message.annotations[0]["url_citation"]["end_index"] == expected_start + len( + cited_text + ) + + +def test_transform_response_stitches_citations_for_whitespace_punctuation_text(): + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + from litellm.types.utils import ModelResponse + + response_json = { + "metrics": {"latencyMs": 100}, + "output": { + "message": { + "role": "assistant", + "content": [ + { + "citationsContent": { + "content": [ + { + "text": "Apptio is a company that makes calls to Bedrock using passthrough APIs via LiteLLM" + } + ], + "citations": [ + { + "location": { + "searchResultLocation": { + "start": 0, + "end": 42, + "searchResultIndex": 0, + } + }, + "source": "https://www.apptio.com/about", + "title": "About Apptio", + } + ], + } + }, + {"text": " ."}, + ], + } + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 10, + "outputTokens": 5, + "totalTokens": 15, + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + }, + } + + class MockResponse: + def json(self): + return response_json + + @property + def text(self): + return json.dumps(response_json) + + config = AmazonConverseConfig() + result = config._transform_response( + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + response=MockResponse(), + model_response=ModelResponse(), + stream=False, + logging_obj=None, + optional_params={}, + api_key=None, + data=None, + messages=[], + encoding=None, + ) + + message = result.choices[0].message + assert message.content.startswith("Apptio is a company") + assert message.annotations is not None + assert len(message.annotations) == 1 + assert message.annotations[0]["url_citation"]["start_index"] == 0 + assert message.annotations[0]["url_citation"]["end_index"] == 42 + + def test_bedrock_tool_message_openai_file_pdf_becomes_document(): """ OpenAI Chat Completions `{type: "file", file: {file_data: "data:application/pdf;...", filename}}`