diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index e555d7e8ec0..8aa68f93240 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1571,7 +1571,9 @@ def infer_content_type_from_url_and_content( # Try to infer from URL extension if url: - extension: Final = url.split(".")[-1].lower().split("?")[0] # Remove query params + from urllib.parse import urlparse + + extension: Final = urlparse(url).path.split(".")[-1].lower() inferred_type: Final = extension_to_mime.get(extension) if inferred_type: return inferred_type diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 45fc93f04c1..ff87da0e7a4 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -1783,6 +1783,35 @@ class TestRequestContainsImageContent: nested = {"type": "tool_result", "content": [nested]} assert request_contains_image_content([{"role": "user", "content": [nested]}]) is False +class TestInferContentTypeQueryString: + """A dot in the query string must not be mistaken for the file extension.""" + + def _infer(self, url: str, content: bytes): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + infer_content_type_from_url_and_content, + ) + + return infer_content_type_from_url_and_content( + url=url, content=content, current_content_type="binary/octet-stream" + ) + + @pytest.mark.parametrize( + "url, content, expected", + [ + ("https://bucket.s3.amazonaws.com/report.pdf", b"%PDF-1.7", "application/pdf"), + ("https://bucket.s3.amazonaws.com/report.pdf?v=1", b"%PDF-1.7", "application/pdf"), + ("https://bucket.s3.amazonaws.com/report.pdf?v=1.0", b"%PDF-1.7", "application/pdf"), + ("https://bucket.s3.amazonaws.com/data.csv?X-Amz-Expires=3.6", b"a,b\n1,2", "text/csv"), + ("https://cdn.example.com/page.html?cb=1.2.3", b"", "text/html"), + ], + ) + def test_extension_is_read_from_the_path_not_the_query(self, url, content, expected): + assert self._infer(url, content) == expected + + def test_a_url_with_no_usable_extension_still_raises(self): + with pytest.raises(ValueError, match="Unable to determine content type from URL"): + self._infer("https://cdn.example.com/download?id=1.2", b"not-an-image") + class TestEncryptedReasoningReplay: """Regression for https://github.com/BerriAI/litellm/issues/40288."""