diff --git a/litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs b/litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs index 0df7cab4acf..42fabe1bd1e 100644 --- a/litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs @@ -290,13 +290,13 @@ fn resolve_document_mime( if !header.eq_ignore_ascii_case("application/octet-stream") && !header.eq_ignore_ascii_case("binary/octet-stream") { - return header.clone(); + return header.to_ascii_lowercase(); } } sniff_document_mime(bytes) .or_else(|| mime_from_file_name(url_path)) .map(str::to_string) - .or(header_content_type) + .or_else(|| header_content_type.map(|header| header.to_ascii_lowercase())) .unwrap_or_else(|| "application/octet-stream".to_string()) } @@ -623,6 +623,26 @@ mod tests { ); } + #[test] + fn resolve_document_mime_normalizes_specific_header_casing() { + assert_eq!( + resolve_document_mime(Some("Application/PDF".to_string()), b"%PDF-1.4", "/x"), + "application/pdf" + ); + } + + #[test] + fn resolve_document_mime_sniffs_when_header_is_binary_octet_stream() { + assert_eq!( + resolve_document_mime( + Some("Binary/Octet-Stream".to_string()), + b"%PDF-1.4 payload", + "/x" + ), + "application/pdf" + ); + } + #[test] fn resolve_document_mime_sniffs_when_header_is_octet_stream() { assert_eq!( diff --git a/litellm-rust/crates/core/src/ocr/mime.rs b/litellm-rust/crates/core/src/ocr/mime.rs index 34b5d7d3df1..abff428b39a 100644 --- a/litellm-rust/crates/core/src/ocr/mime.rs +++ b/litellm-rust/crates/core/src/ocr/mime.rs @@ -18,9 +18,6 @@ pub fn sniff_document_mime(bytes: &[u8]) -> Option<&'static str> { { return Some("image/tiff"); } - if bytes.starts_with(b"BM") { - return Some("image/bmp"); - } None } @@ -71,7 +68,12 @@ mod tests { sniff_document_mime(&[0x4d, 0x4d, 0x00, 0x2a]), Some("image/tiff") ); - assert_eq!(sniff_document_mime(b"BMxxxx"), Some("image/bmp")); + } + + #[test] + fn sniff_document_mime_does_not_detect_bmp_by_magic_bytes() { + assert_eq!(sniff_document_mime(b"BMxxxx"), None); + assert_eq!(sniff_document_mime(b"BM\x36\x00\x00\x00random"), None); } #[test] diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index cab7e7d7164..2aa9fd3ed61 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -372,6 +372,8 @@ async def aocr( _MIME_PATTERN = re.compile(r"^[\w.+-]+/[\w.+-]+$") +_GENERIC_MIME_TYPES = frozenset({"application/octet-stream", "binary/octet-stream"}) + _RUST_BRIDGE_INTERNAL_PARAMS = {"original_generic_function"} _MIME_TYPE_MAP = { @@ -414,11 +416,27 @@ def _sniff_mime_type_from_bytes(data: bytes) -> str | None: return "image/webp" if data.startswith((b"II*\x00", b"MM\x00*")): return "image/tiff" - if data.startswith(b"BM"): - return "image/bmp" return None +def _normalize_mime_type(value: str) -> str | None: + normalized = value.split(";", 1)[0].strip().lower() + if not _MIME_PATTERN.match(normalized): + return None + return normalized + + +def _resolve_declared_mime_type(raw: object) -> str | None: + if raw is None: + return None + if not isinstance(raw, str): + raise ValueError("OCR document 'mime_type' must be a string when provided") + normalized = _normalize_mime_type(raw) + if normalized is None: + raise ValueError("OCR document 'mime_type' is not a valid MIME type") + return normalized + + def convert_file_document_to_url_document(document: dict[str, Any]) -> dict[str, str]: """ Convert a file-type document dict to a document_url-type document dict @@ -488,25 +506,26 @@ def convert_file_document_to_url_document(document: dict[str, Any]) -> dict[str, if not file_bytes: raise ValueError("File is empty or could not be read") - raw_mime = document.get("mime_type") - explicit_mime = raw_mime if isinstance(raw_mime, str) else None - resolved_mime = next( + declared_mime = ( + _resolve_declared_mime_type(cast(object, document["mime_type"])) + if "mime_type" in document + else None + ) + filename_mime = _normalize_mime_type(mime_type) + specific_mime = next( ( candidate - for candidate in (explicit_mime, mime_type) - if candidate and candidate != "application/octet-stream" + for candidate in (declared_mime, filename_mime) + if candidate is not None and candidate not in _GENERIC_MIME_TYPES ), None, ) mime_type = ( - resolved_mime + specific_mime or _sniff_mime_type_from_bytes(file_bytes) or "application/octet-stream" ) - if not _MIME_PATTERN.match(mime_type): - raise ValueError(f"Invalid MIME type: {mime_type}") - base64_data = base64.b64encode(file_bytes).decode("utf-8") data_uri = f"data:{mime_type};base64,{base64_data}" diff --git a/litellm/proxy/ocr_endpoints/endpoints.py b/litellm/proxy/ocr_endpoints/endpoints.py index e32fee6afc5..552bbddadda 100644 --- a/litellm/proxy/ocr_endpoints/endpoints.py +++ b/litellm/proxy/ocr_endpoints/endpoints.py @@ -27,10 +27,10 @@ def _build_document_from_upload( Delegates to convert_file_document_to_url_document after resolving MIME type from the upload's content_type header or filename. """ - mime_type = content_type.split(";")[0].strip() if content_type else None - if not mime_type or mime_type == "application/octet-stream": - if filename: - mime_type = get_mime_type(filename) + mime_type = content_type.split(";")[0].strip().lower() if content_type else None + generic_mime_types = ("application/octet-stream", "binary/octet-stream") + if (not mime_type or mime_type in generic_mime_types) and filename: + mime_type = get_mime_type(filename) return convert_file_document_to_url_document( { diff --git a/tests/e2e/gateway/test_ocr_rust_e2e.py b/tests/e2e/gateway/test_ocr_rust_e2e.py index 7559c18f9d4..2c47b31b19b 100644 --- a/tests/e2e/gateway/test_ocr_rust_e2e.py +++ b/tests/e2e/gateway/test_ocr_rust_e2e.py @@ -11,13 +11,12 @@ from __future__ import annotations import os from dataclasses import dataclass from pathlib import Path -from typing import Any import httpx import pytest import yaml -from litellm.proxy.ocr_endpoints.endpoints import _build_document_from_upload +from litellm.llms.base_llm.ocr.transformation import OCRResponse TEST_PDF_URL = ( "https://cdn.jsdelivr.net/gh/BerriAI/litellm" @@ -34,12 +33,12 @@ REPO_ROOT = Path(__file__).resolve().parents[3] FIXTURE_PDF = REPO_ROOT / "tests" / "llm_translation" / "fixtures" / "dummy.pdf" FIXTURE_IMAGE = REPO_ROOT / "tests" / "image_gen_tests" / "test_image.png" -RUST_OCR_UPLOAD_CASES = [ - pytest.param(FIXTURE_PDF, "application/pdf", "document_url", id="pdf_octet_stream"), - pytest.param(FIXTURE_IMAGE, "image/png", "image_url", id="image_octet_stream"), -] +RUST_OCR_UPLOAD_CASES = ( + pytest.param(FIXTURE_PDF, id="pdf_octet_stream"), + pytest.param(FIXTURE_IMAGE, id="image_octet_stream"), +) -RUST_OCR_GATEWAY_CASES = [ +RUST_OCR_GATEWAY_CASES = ( pytest.param( "rust-ocr-mistral", {"type": "document_url", "document_url": TEST_PDF_URL}, @@ -68,7 +67,7 @@ RUST_OCR_GATEWAY_CASES = [ }, id="vertex_deepseek", ), -] +) CONFIG_PATH = Path(__file__).with_name("litellm-config.yml") @@ -103,7 +102,9 @@ class OcrGateway: json={"model": model, "document": document}, ) - def ocr_upload(self, model: str, content: bytes, upload_name: str) -> httpx.Response: + def ocr_upload( + self, model: str, content: bytes, upload_name: str + ) -> httpx.Response: with httpx.Client( timeout=float(os.getenv("E2E_REQUEST_TIMEOUT", "120")) ) as client: @@ -135,13 +136,13 @@ def resources() -> OcrResources: ) -def _assert_ocr_response_shape(response_json: dict[str, Any]) -> None: - assert response_json["object"] == "ocr" - assert response_json["model"] - assert isinstance(response_json["pages"], list) - assert len(response_json["pages"]) > 0 - assert "index" in response_json["pages"][0] - assert "markdown" in response_json["pages"][0] +def _assert_ocr_response(response: httpx.Response) -> None: + assert response.status_code == 200, response.text + parsed = OCRResponse.model_validate(response.json()) + assert parsed.object == "ocr" + assert parsed.model + assert len(parsed.pages) > 0 + assert any(page.markdown.strip() for page in parsed.pages) class TestRustOcrGateway: @@ -166,29 +167,17 @@ class TestRustOcrGateway: ) -> None: response = resources.gateway.ocr(model, document) - assert response.status_code == 200, response.text - _assert_ocr_response_shape(response.json()) + _assert_ocr_response(response) - @pytest.mark.parametrize( - ("fixture_path", "expected_mime", "expected_document_type"), - RUST_OCR_UPLOAD_CASES, - ) + @pytest.mark.parametrize("fixture_path", RUST_OCR_UPLOAD_CASES) def test_rust_ocr_octet_stream_upload_response( self, resources: OcrResources, fixture_path: Path, - expected_mime: str, - expected_document_type: str, ) -> None: - if not fixture_path.is_file(): - pytest.skip(f"Missing OCR fixture: {fixture_path}") + assert fixture_path.is_file(), f"Missing committed OCR fixture: {fixture_path}" content = fixture_path.read_bytes() - document = _build_document_from_upload(content, "document", "application/octet-stream") - assert document["type"] == expected_document_type - assert document[expected_document_type].startswith(f"data:{expected_mime};base64,") - response = resources.gateway.ocr_upload("rust-ocr-mistral", content, "document") - assert response.status_code == 200, response.text - _assert_ocr_response_shape(response.json()) + _assert_ocr_response(response) diff --git a/tests/test_litellm/ocr/test_ocr_file_input.py b/tests/test_litellm/ocr/test_ocr_file_input.py index 16d10d2a520..5e9c166fcf4 100644 --- a/tests/test_litellm/ocr/test_ocr_file_input.py +++ b/tests/test_litellm/ocr/test_ocr_file_input.py @@ -85,8 +85,9 @@ class TestSniffMimeTypeFromBytes: assert _sniff_mime_type_from_bytes(b"II*\x00 rest") == "image/tiff" assert _sniff_mime_type_from_bytes(b"MM\x00* rest") == "image/tiff" - def test_should_detect_bmp_from_magic_bytes(self): - assert _sniff_mime_type_from_bytes(b"BM rest") == "image/bmp" + def test_should_not_sniff_bmp_from_magic_bytes(self) -> None: + assert _sniff_mime_type_from_bytes(b"BM rest") is None + assert _sniff_mime_type_from_bytes(b"BM\x00\x00\x00\x00random") is None def test_should_return_none_for_unknown_prefix(self): assert _sniff_mime_type_from_bytes(b"not a known file") is None @@ -333,18 +334,119 @@ class TestConvertFileDocumentToUrlDocument: with pytest.raises(ValueError, match="Unsupported file input type"): convert_file_document_to_url_document({"type": "file", "file": 12345}) - def test_should_raise_error_for_invalid_mime_type(self): - """MIME types with special characters should be rejected.""" + def test_should_reject_malformed_mime_type(self) -> None: content = b"some content" - with pytest.raises(ValueError, match="Invalid MIME type"): + with pytest.raises(ValueError, match="not a valid MIME type"): convert_file_document_to_url_document( { "type": "file", "file": content, - "mime_type": "text/html; charset=utf-8\nX-Injected: true", + "mime_type": "text/html\r\nX-Injected: true", } ) + def test_should_reject_blank_mime_type(self) -> None: + content = b"%PDF-1.4 content" + with pytest.raises(ValueError, match="not a valid MIME type"): + convert_file_document_to_url_document( + {"type": "file", "file": content, "mime_type": " "} + ) + + def test_should_reject_non_string_mime_type(self) -> None: + content = b"%PDF-1.4 content" + with pytest.raises(ValueError, match="must be a string"): + convert_file_document_to_url_document( + {"type": "file", "file": content, "mime_type": 123} + ) + + def test_should_not_echo_raw_mime_value_in_error(self) -> None: + content = b"some content" + secret_marker = "X-Injected: super-secret" + with pytest.raises(ValueError) as exc_info: + convert_file_document_to_url_document( + { + "type": "file", + "file": content, + "mime_type": f"text/html\r\n{secret_marker}", + } + ) + assert secret_marker not in str(exc_info.value) + + def test_should_strip_parameters_from_explicit_mime_type(self) -> None: + content = b"raw bytes" + result = convert_file_document_to_url_document( + { + "type": "file", + "file": content, + "mime_type": "application/pdf; charset=utf-8", + } + ) + assert result["type"] == "document_url" + assert result["document_url"].startswith("data:application/pdf;base64,") + + def test_should_normalize_explicit_mime_casing(self) -> None: + content = b"raw bytes" + result = convert_file_document_to_url_document( + {"type": "file", "file": content, "mime_type": "Application/PDF"} + ) + assert result["type"] == "document_url" + assert result["document_url"].startswith("data:application/pdf;base64,") + + def test_should_sniff_when_explicit_mime_is_uppercase_octet_stream(self) -> None: + content = b"%PDF-1.4\n1 0 obj\n<< >>\nendobj\n" + result = convert_file_document_to_url_document( + {"type": "file", "file": content, "mime_type": "Application/Octet-Stream"} + ) + assert result["type"] == "document_url" + assert result["document_url"].startswith("data:application/pdf;base64,") + + def test_should_sniff_when_explicit_mime_is_binary_octet_stream(self) -> None: + content = b"%PDF-1.4\n1 0 obj\n<< >>\nendobj\n" + result = convert_file_document_to_url_document( + {"type": "file", "file": content, "mime_type": "Binary/Octet-Stream"} + ) + assert result["type"] == "document_url" + assert result["document_url"].startswith("data:application/pdf;base64,") + + def test_should_treat_generic_octet_stream_with_parameters_as_ambiguous( + self, + ) -> None: + content = b"%PDF-1.4\n1 0 obj\n<< >>\nendobj\n" + result = convert_file_document_to_url_document( + { + "type": "file", + "file": content, + "mime_type": "application/octet-stream; charset=binary", + } + ) + assert result["type"] == "document_url" + assert result["document_url"].startswith("data:application/pdf;base64,") + + def test_should_not_sniff_bmp_from_raw_bytes(self) -> None: + content = b"BM\x36\x00\x00\x00 fake bitmap header" + result = convert_file_document_to_url_document( + {"type": "file", "file": content} + ) + assert result["type"] == "document_url" + assert result["document_url"].startswith( + "data:application/octet-stream;base64," + ) + + def test_should_resolve_bmp_from_extension(self) -> None: + content = b"BM\x36\x00\x00\x00 fake bitmap header" + with tempfile.NamedTemporaryFile(suffix=".bmp", delete=False) as f: + f.write(content) + f.flush() + tmp_path = Path(f.name) + try: + result = convert_file_document_to_url_document( + {"type": "file", "file": tmp_path} + ) + assert result["type"] == "image_url" + assert result["image_url"].startswith("data:image/bmp;base64,") + finally: + os.unlink(str(tmp_path)) + def test_should_override_mime_type_for_pathlib_path(self): """Explicit mime_type should override auto-detection from extension.""" content = b"some content" @@ -493,6 +595,30 @@ class TestBuildDocumentFromUpload: assert result["type"] == "image_url" assert result["image_url"].startswith("data:image/png;base64,") + def test_should_detect_mime_from_filename_for_uppercase_octet_stream(self) -> None: + content = b"pdf content" + + result = self._build( + file_content=content, + filename="report.pdf", + content_type="APPLICATION/OCTET-STREAM", + ) + + assert result["type"] == "document_url" + assert result["document_url"].startswith("data:application/pdf;base64,") + + def test_should_detect_mime_from_filename_for_binary_octet_stream(self) -> None: + content = b"png content" + + result = self._build( + file_content=content, + filename="image.png", + content_type="binary/octet-stream", + ) + + assert result["type"] == "image_url" + assert result["image_url"].startswith("data:image/png;base64,") + class TestProxySecurityGuard: """Test that the proxy rejects type='file' documents in JSON requests