chore: import upstream snapshot with attribution
CodeQL / Analyze (python) (push) Has been cancelled
Update Platform Components Table / update (push) Has been cancelled
Docker image release / Build base image (push) Has been cancelled
Sync docs with Docusaurus / sync (push) Has been cancelled
Tests / Check if changed (push) Has been cancelled
Tests / format (push) Has been cancelled
Tests / check-imports (push) Has been cancelled
Tests / Unit / macos-latest (push) Has been cancelled
Tests / Unit / ubuntu-latest (push) Has been cancelled
Tests / Unit / windows-latest (push) Has been cancelled
Tests / mypy (push) Has been cancelled
Tests / Integration / ubuntu-latest (push) Has been cancelled
Tests / Integration / macos-latest (push) Has been cancelled
Tests / Integration / windows-latest (push) Has been cancelled
Tests / notify-slack-on-failure (push) Has been cancelled
Tests / Mark tests as completed (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
Update Platform Components Table / update (push) Has been cancelled
Docker image release / Build base image (push) Has been cancelled
Sync docs with Docusaurus / sync (push) Has been cancelled
Tests / Check if changed (push) Has been cancelled
Tests / format (push) Has been cancelled
Tests / check-imports (push) Has been cancelled
Tests / Unit / macos-latest (push) Has been cancelled
Tests / Unit / ubuntu-latest (push) Has been cancelled
Tests / Unit / windows-latest (push) Has been cancelled
Tests / mypy (push) Has been cancelled
Tests / Integration / ubuntu-latest (push) Has been cancelled
Tests / Integration / macos-latest (push) Has been cancelled
Tests / Integration / windows-latest (push) Has been cancelled
Tests / notify-slack-on-failure (push) Has been cancelled
Tests / Mark tests as completed (push) Has been cancelled
This commit is contained in:
@@ -0,0 +1,3 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
@@ -0,0 +1,143 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
from haystack import Document
|
||||
from haystack.components.converters.image.document_to_image import DocumentToImageContent
|
||||
from haystack.core.serialization import component_from_dict, component_to_dict
|
||||
from haystack.dataclasses import ByteStream
|
||||
|
||||
|
||||
class TestDocumentToImageContent:
|
||||
def test_to_dict(self) -> None:
|
||||
converter = DocumentToImageContent()
|
||||
assert component_to_dict(converter, "converter") == {
|
||||
"init_parameters": {"file_path_meta_field": "file_path", "root_path": "", "detail": None, "size": None},
|
||||
"type": "haystack.components.converters.image.document_to_image.DocumentToImageContent",
|
||||
}
|
||||
|
||||
def test_to_dict_not_defaults(self) -> None:
|
||||
converter = DocumentToImageContent(
|
||||
file_path_meta_field="image_path", root_path="/data", detail="high", size=(800, 600)
|
||||
)
|
||||
assert component_to_dict(converter, "converter") == {
|
||||
"init_parameters": {
|
||||
"file_path_meta_field": "image_path",
|
||||
"root_path": "/data",
|
||||
"detail": "high",
|
||||
"size": (800, 600),
|
||||
},
|
||||
"type": "haystack.components.converters.image.document_to_image.DocumentToImageContent",
|
||||
}
|
||||
|
||||
def test_from_dict(self) -> None:
|
||||
data = {
|
||||
"init_parameters": {
|
||||
"file_path_meta_field": "image_path",
|
||||
"root_path": "/test",
|
||||
"detail": "auto",
|
||||
"size": (512, 512),
|
||||
},
|
||||
"type": "haystack.components.converters.image.document_to_image.DocumentToImageContent",
|
||||
}
|
||||
converter = component_from_dict(DocumentToImageContent, data, "name")
|
||||
assert component_to_dict(converter, "converter") == data
|
||||
|
||||
def test_run_with_empty_documents_list(self) -> None:
|
||||
converter = DocumentToImageContent()
|
||||
results = converter.run(documents=[])
|
||||
assert results == {"image_contents": []}
|
||||
|
||||
def test_run_with_missing_file_path_metadata(self) -> None:
|
||||
converter = DocumentToImageContent()
|
||||
# Document without file_path in metadata
|
||||
doc_no_path = Document(content="test", meta={})
|
||||
# Document with file_path but file doesn't exist
|
||||
doc_no_file = Document(content="test", meta={"file_path": "nonexistent.jpg"})
|
||||
with pytest.raises(ValueError, match="is missing the 'file_path' key"):
|
||||
_ = converter.run(documents=[doc_no_path, doc_no_file])
|
||||
|
||||
def test_run_with_non_image_documents(self) -> None:
|
||||
converter = DocumentToImageContent()
|
||||
docx_doc = Document(content="test", meta={"file_path": "test/test_files/docx/sample_docx.docx"})
|
||||
with pytest.raises(ValueError, match="has an unsupported MIME type"):
|
||||
_ = converter.run(documents=[docx_doc])
|
||||
|
||||
def test_run_with_invalid_file_path(self, caplog) -> None:
|
||||
converter = DocumentToImageContent()
|
||||
pdf_doc = Document(content="test", meta={"file_path": "wrong_name.jpg"})
|
||||
with pytest.raises(ValueError, match="has an invalid file path 'wrong_name.jpg'"):
|
||||
_ = converter.run(documents=[pdf_doc])
|
||||
|
||||
def test_run_with_pdf_missing_page_number(self, caplog) -> None:
|
||||
converter = DocumentToImageContent()
|
||||
pdf_doc = Document(content="test", meta={"file_path": "test/test_files/pdf/sample_pdf_1.pdf"})
|
||||
with pytest.raises(ValueError, match="is missing the 'page_number' key"):
|
||||
_ = converter.run(documents=[pdf_doc])
|
||||
|
||||
def test_run_with_image_documents(self) -> None:
|
||||
converter = DocumentToImageContent(root_path="test/test_files/images")
|
||||
image_doc = Document(content="test", meta={"file_path": "apple.jpg"})
|
||||
results = converter.run(documents=[image_doc])
|
||||
assert len(results["image_contents"]) == 1
|
||||
assert results["image_contents"][0].meta == {"file_path": "apple.jpg"}
|
||||
|
||||
def test_run_with_pdf_documents(self) -> None:
|
||||
converter = DocumentToImageContent()
|
||||
pdf_doc = Document(content="test", meta={"file_path": "test/test_files/pdf/sample_pdf_1.pdf", "page_number": 1})
|
||||
results = converter.run(documents=[pdf_doc])
|
||||
assert len(results["image_contents"]) == 1
|
||||
assert results["image_contents"][0].meta == {
|
||||
"file_path": "test/test_files/pdf/sample_pdf_1.pdf",
|
||||
"page_number": 1,
|
||||
}
|
||||
|
||||
def test_run_with_mixed_document_types(self) -> None:
|
||||
converter = DocumentToImageContent(root_path="test/test_files")
|
||||
documents = [
|
||||
Document(content="", meta={"file_path": "images/apple.jpg"}),
|
||||
Document(content="", meta={"file_path": "pdf/sample_pdf_1.pdf", "page_number": 1}),
|
||||
Document(content="text", meta={"file_path": "docx/sample_docx.docx"}),
|
||||
]
|
||||
with pytest.raises(ValueError, match="has an unsupported MIME type"):
|
||||
_ = converter.run(documents=documents)
|
||||
|
||||
@patch("haystack.components.converters.image.document_to_image._extract_image_sources_info")
|
||||
@patch("haystack.components.converters.image.document_to_image._batch_convert_pdf_pages_to_images")
|
||||
@patch("PIL.Image.open")
|
||||
@patch("haystack.components.converters.image.document_to_image.ByteStream")
|
||||
def test_run_none_images(
|
||||
self,
|
||||
mocked_byte_stream,
|
||||
mocked_pil_open,
|
||||
mocked_batch_convert_pdf_pages_to_images,
|
||||
mocked_extract_image_sources_info,
|
||||
caplog,
|
||||
):
|
||||
converter = DocumentToImageContent()
|
||||
|
||||
mocked_extract_image_sources_info.return_value = [
|
||||
{"path": "doc1.pdf", "mime_type": "application/pdf", "page_number": 999}, # Page 999 doesn't exist
|
||||
{"path": "image1.jpg", "mime_type": "image/jpeg"},
|
||||
]
|
||||
mocked_batch_convert_pdf_pages_to_images.return_value = {} # Empty dict because page was skipped
|
||||
mocked_pil_open.return_value = Image.new("RGB", (100, 100))
|
||||
mocked_byte_stream.from_file_path.return_value = ByteStream(b"")
|
||||
|
||||
documents = [
|
||||
Document(content="PDF 1", meta={"file_path": "doc1.pdf", "page_number": 999}),
|
||||
Document(content="Image 1", meta={"file_path": "image1.jpg"}),
|
||||
]
|
||||
|
||||
image_contents = converter.run(documents=documents)["image_contents"]
|
||||
|
||||
assert caplog.records[-1].levelname == "WARNING"
|
||||
assert "Conversion failed for some documents." in caplog.records[-1].message
|
||||
|
||||
assert image_contents[0] is None
|
||||
assert image_contents[1] is not None
|
||||
@@ -0,0 +1,82 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack.components.converters.image.file_to_document import ImageFileToDocument
|
||||
from haystack.core.serialization import component_from_dict, component_to_dict
|
||||
from haystack.dataclasses import ByteStream
|
||||
|
||||
|
||||
class TestImageFileToDocument:
|
||||
def test_to_dict(self) -> None:
|
||||
converter = ImageFileToDocument()
|
||||
assert component_to_dict(converter, "converter") == {
|
||||
"init_parameters": {"store_full_path": False},
|
||||
"type": "haystack.components.converters.image.file_to_document.ImageFileToDocument",
|
||||
}
|
||||
|
||||
def test_to_dict_not_defaults(self) -> None:
|
||||
converter = ImageFileToDocument(store_full_path=True)
|
||||
assert component_to_dict(converter, "converter") == {
|
||||
"init_parameters": {"store_full_path": True},
|
||||
"type": "haystack.components.converters.image.file_to_document.ImageFileToDocument",
|
||||
}
|
||||
|
||||
def test_from_dict(self) -> None:
|
||||
data = {
|
||||
"init_parameters": {"store_full_path": False},
|
||||
"type": "haystack.components.converters.image.file_to_document.ImageFileToDocument",
|
||||
}
|
||||
converter = component_from_dict(ImageFileToDocument, data, "name")
|
||||
assert component_to_dict(converter, "converter") == data
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("image_path", "mime_type"),
|
||||
[
|
||||
("./test/test_files/images/haystack-logo.png", "image/png"),
|
||||
("./test/test_files/images/apple.jpg", "image/jpeg"),
|
||||
],
|
||||
)
|
||||
def test_run_with_valid_sources(self, image_path: str, mime_type: str) -> None:
|
||||
converter = ImageFileToDocument(store_full_path=True)
|
||||
results = converter.run(sources=[image_path], meta={"source": "test_source"})
|
||||
|
||||
assert len(results["documents"]) == 1
|
||||
assert results["documents"][0].content is None
|
||||
assert results["documents"][0].meta == {"source": "test_source", "file_path": image_path}
|
||||
|
||||
def test_run_with_no_sources(self) -> None:
|
||||
converter = ImageFileToDocument()
|
||||
results = converter.run(sources=[])
|
||||
assert len(results["documents"]) == 0
|
||||
assert results == {"documents": []}
|
||||
|
||||
def test_run_with_invalid_source_type(self, caplog) -> None:
|
||||
converter = ImageFileToDocument()
|
||||
converter.run(sources=[123]) # Invalid source type
|
||||
assert "Could not read" in caplog.text
|
||||
|
||||
def test_run_with_non_existent_file(self, caplog) -> None:
|
||||
converter = ImageFileToDocument()
|
||||
converter.run(sources=["./non_existent_file.png"])
|
||||
assert "Could not read" in caplog.text
|
||||
assert "No such file or directory:" in caplog.text
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("image_path", "mime_type"),
|
||||
[
|
||||
("./test/test_files/images/haystack-logo.png", "image/png"),
|
||||
("./test/test_files/images/apple.jpg", "image/jpeg"),
|
||||
],
|
||||
)
|
||||
def test_run_with_bytestream_sources(self, image_path: str, mime_type: str) -> None:
|
||||
byte_stream = ByteStream.from_file_path(Path(image_path), mime_type=mime_type, meta={"file_path": image_path})
|
||||
converter = ImageFileToDocument(store_full_path=True)
|
||||
results = converter.run(sources=[byte_stream])
|
||||
assert len(results["documents"]) == 1
|
||||
assert results["documents"][0].content is None
|
||||
assert results["documents"][0].meta == {"file_path": image_path}
|
||||
@@ -0,0 +1,116 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack.components.converters.image.file_to_image import ImageFileToImageContent
|
||||
from haystack.components.converters.image.image_utils import _encode_image_to_base64
|
||||
from haystack.core.serialization import component_from_dict, component_to_dict
|
||||
from haystack.dataclasses import ByteStream
|
||||
|
||||
|
||||
class TestImageFileToImageContent:
|
||||
def test_to_dict(self) -> None:
|
||||
converter = ImageFileToImageContent()
|
||||
assert component_to_dict(converter, "converter") == {
|
||||
"init_parameters": {"detail": None, "size": None},
|
||||
"type": "haystack.components.converters.image.file_to_image.ImageFileToImageContent",
|
||||
}
|
||||
|
||||
def test_to_dict_not_defaults(self) -> None:
|
||||
converter = ImageFileToImageContent(detail="low", size=(128, 128))
|
||||
assert component_to_dict(converter, "converter") == {
|
||||
"init_parameters": {"detail": "low", "size": (128, 128)},
|
||||
"type": "haystack.components.converters.image.file_to_image.ImageFileToImageContent",
|
||||
}
|
||||
|
||||
def test_from_dict(self) -> None:
|
||||
data = {
|
||||
"init_parameters": {"detail": "auto", "size": None},
|
||||
"type": "haystack.components.converters.image.file_to_image.ImageFileToImageContent",
|
||||
}
|
||||
converter = component_from_dict(ImageFileToImageContent, data, "name")
|
||||
assert component_to_dict(converter, "converter") == data
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("image_path", "mime_type"),
|
||||
[
|
||||
("./test/test_files/images/haystack-logo.png", "image/png"),
|
||||
("./test/test_files/images/apple.jpg", "image/jpeg"),
|
||||
],
|
||||
)
|
||||
def test_run_with_valid_sources(self, image_path: str, mime_type: str) -> None:
|
||||
converter = ImageFileToImageContent()
|
||||
results = converter.run(sources=[image_path], size=(128, 128))
|
||||
|
||||
byte_stream = ByteStream.from_file_path(
|
||||
Path(image_path), mime_type=mime_type, meta={"file_name": image_path.rsplit("/", maxsplit=1)[-1]}
|
||||
)
|
||||
assert len(results["image_contents"]) == 1
|
||||
assert results["image_contents"][0].base64_image is not None
|
||||
assert (
|
||||
results["image_contents"][0].base64_image
|
||||
== _encode_image_to_base64(bytestream=byte_stream, size=(128, 128))[1]
|
||||
)
|
||||
assert results["image_contents"][0].mime_type == mime_type
|
||||
assert results["image_contents"][0].detail is None
|
||||
assert results["image_contents"][0].meta["file_path"] == str(Path(image_path))
|
||||
|
||||
def test_run_with_no_sources(self) -> None:
|
||||
converter = ImageFileToImageContent()
|
||||
results = converter.run(sources=[])
|
||||
assert len(results["image_contents"]) == 0
|
||||
assert results == {"image_contents": []}
|
||||
|
||||
def test_run_with_invalid_source_type(self, caplog) -> None:
|
||||
converter = ImageFileToImageContent()
|
||||
converter.run(sources=[123]) # Invalid source type
|
||||
assert "Could not read" in caplog.text
|
||||
|
||||
def test_run_with_non_existent_file(self, caplog) -> None:
|
||||
converter = ImageFileToImageContent()
|
||||
converter.run(sources=["./non_existent_file.png"])
|
||||
assert "Could not read" in caplog.text
|
||||
assert "No such file or directory:" in caplog.text
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("image_path", "mime_type"),
|
||||
[
|
||||
("./test/test_files/images/haystack-logo.png", "image/png"),
|
||||
("./test/test_files/images/apple.jpg", "image/jpeg"),
|
||||
],
|
||||
)
|
||||
def test_run_with_bytestream_sources(self, image_path: str, mime_type: str) -> None:
|
||||
byte_stream = ByteStream.from_file_path(Path(image_path), mime_type=mime_type, meta={"file_path": image_path})
|
||||
|
||||
# Initialize the converter
|
||||
converter = ImageFileToImageContent(size=(128, 128))
|
||||
|
||||
# Run the converter with the ByteStream
|
||||
results = converter.run(sources=[byte_stream])
|
||||
|
||||
# Assertions
|
||||
assert len(results["image_contents"]) == 1
|
||||
assert results["image_contents"][0].base64_image is not None
|
||||
assert (
|
||||
results["image_contents"][0].base64_image
|
||||
== _encode_image_to_base64(bytestream=byte_stream, size=(128, 128))[1]
|
||||
)
|
||||
assert results["image_contents"][0].mime_type == mime_type
|
||||
assert results["image_contents"][0].detail is None
|
||||
assert results["image_contents"][0].meta["file_path"] == image_path
|
||||
|
||||
def test_run_with_empty_bytestream(self) -> None:
|
||||
# Create an empty ByteStream object
|
||||
byte_stream = ByteStream(data=b"", meta={"file_path": "empty_file.png"})
|
||||
|
||||
# Initialize the converter
|
||||
converter = ImageFileToImageContent()
|
||||
|
||||
# Run the converter with the empty ByteStream
|
||||
results = converter.run(sources=[byte_stream])
|
||||
|
||||
assert results["image_contents"] == []
|
||||
@@ -0,0 +1,227 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import glob
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from PIL import Image
|
||||
from pytest import LogCaptureFixture
|
||||
|
||||
from haystack.components.converters.image.image_utils import (
|
||||
_batch_convert_pdf_pages_to_images,
|
||||
_convert_pdf_to_images,
|
||||
_encode_image_to_base64,
|
||||
_encode_pil_image_to_base64,
|
||||
_extract_image_sources_info,
|
||||
_PDFPageInfo,
|
||||
)
|
||||
from haystack.components.converters.utils import get_bytestream_from_source
|
||||
from haystack.dataclasses import ByteStream, Document
|
||||
|
||||
|
||||
class TestToBase64Jpeg:
|
||||
def test_to_base64_jpeg(self) -> None:
|
||||
image_array = np.array(
|
||||
[
|
||||
[[34.215402, 132.78745697, 71.04739979], [24.23156181, 35.26147199, 124.95610316]],
|
||||
[[155.47443501, 196.98050276, 154.74734292], [253.24590033, 84.62392497, 157.34396641]],
|
||||
]
|
||||
)
|
||||
image = Image.fromarray(image_array.astype("uint8"))
|
||||
b64_str = _encode_pil_image_to_base64(image=image, mime_type="image/jpeg")
|
||||
assert (
|
||||
b64_str
|
||||
== "/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAAgGBgcGBQgHBwcJCQgKDBQNDAsLDBkSEw8UHRofHh0aHBwgJC4nICIsIxwcKDcpLDAxNDQ0Hyc5PTgyPC4zNDL/2wBDAQkJCQwLDBgNDRgyIRwhMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjL/wAARCAACAAIDASIAAhEBAxEB/8QAHwAAAQUBAQEBAQEAAAAAAAAAAAECAwQFBgcICQoL/8QAtRAAAgEDAwIEAwUFBAQAAAF9AQIDAAQRBRIhMUEGE1FhByJxFDKBkaEII0KxwRVS0fAkM2JyggkKFhcYGRolJicoKSo0NTY3ODk6Q0RFRkdISUpTVFVWV1hZWmNkZWZnaGlqc3R1dnd4eXqDhIWGh4iJipKTlJWWl5iZmqKjpKWmp6ipqrKztLW2t7i5usLDxMXGx8jJytLT1NXW19jZ2uHi4+Tl5ufo6erx8vP09fb3+Pn6/8QAHwEAAwEBAQEBAQEBAQAAAAAAAAECAwQFBgcICQoL/8QAtREAAgECBAQDBAcFBAQAAQJ3AAECAxEEBSExBhJBUQdhcRMiMoEIFEKRobHBCSMzUvAVYnLRChYkNOEl8RcYGRomJygpKjU2Nzg5OkNERUZHSElKU1RVVldYWVpjZGVmZ2hpanN0dXZ3eHl6goOEhYaHiImKkpOUlZaXmJmaoqOkpaanqKmqsrO0tba3uLm6wsPExcbHyMnK0tPU1dbX2Nna4uPk5ebn6Onq8vP09fb3+Pn6/9oADAMBAAIRAxEAPwCe50jTTdTE6daffb/livr9KKKK9aHwo9ql8EfRH//Z" # noqa: E501
|
||||
)
|
||||
|
||||
|
||||
class TestConvertPdfToImages:
|
||||
def test_convert_pdf_to_images(self) -> None:
|
||||
bytestream = get_bytestream_from_source(Path("test/test_files/pdf/sample_pdf_1.pdf"))
|
||||
output = _convert_pdf_to_images(bytestream=bytestream, page_range=[1])
|
||||
assert output is not None
|
||||
|
||||
def test_convert_pdf_to_images_with_size(self) -> None:
|
||||
bytestream = get_bytestream_from_source(Path("test/test_files/pdf/sample_pdf_1.pdf"))
|
||||
pages_images = _convert_pdf_to_images(bytestream=bytestream, page_range=[1], size=(100, 100))
|
||||
|
||||
assert len(pages_images) == 1
|
||||
assert pages_images[0][0] == 1
|
||||
assert pages_images[0][1].width <= 100
|
||||
assert pages_images[0][1].height <= 100
|
||||
|
||||
def test_convert_pdf_to_images_invalid_page(self, caplog: LogCaptureFixture) -> None:
|
||||
bytestream = get_bytestream_from_source(Path("test/test_files/pdf/sample_pdf_1.pdf"))
|
||||
out = _convert_pdf_to_images(bytestream=bytestream, page_range=[5])
|
||||
assert out == []
|
||||
assert "Page 5 is out of range for the PDF file. Skipping it." in caplog.text
|
||||
|
||||
def test_convert_pdf_to_images_error_reading_file(self, caplog: LogCaptureFixture) -> None:
|
||||
bytestream = ByteStream(data=b"", mime_type="application/pdf")
|
||||
out = _convert_pdf_to_images(bytestream=bytestream, page_range=[1])
|
||||
assert out == []
|
||||
assert "Could not read PDF file" in caplog.text
|
||||
|
||||
def test_convert_pdf_to_images_empty_file(self, caplog: LogCaptureFixture) -> None:
|
||||
bytestream = get_bytestream_from_source(Path("test/test_files/pdf/sample_pdf_1.pdf"))
|
||||
|
||||
with patch("haystack.components.converters.image.image_utils.PdfDocument") as mock_pdf_document:
|
||||
mock_pdf_document.__len__.return_value = 0
|
||||
out = _convert_pdf_to_images(bytestream=bytestream, page_range=[1])
|
||||
|
||||
assert out == []
|
||||
assert "PDF file is empty" in caplog.text
|
||||
|
||||
def test_convert_pdf_to_images_scale_if_large_pdf(self, caplog: LogCaptureFixture) -> None:
|
||||
bytestream = get_bytestream_from_source(Path("test/test_files/pdf/sample_pdf_1.pdf"))
|
||||
|
||||
caplog.set_level(logging.INFO)
|
||||
|
||||
mock_pdf_document = MagicMock()
|
||||
mock_pdf_document.__len__.return_value = 1
|
||||
mock_page = MagicMock()
|
||||
mock_page.get_mediabox.return_value = (0, 0, 1e6, 1e6)
|
||||
mock_pdf_document.__getitem__.return_value = mock_page
|
||||
|
||||
with patch("haystack.components.converters.image.image_utils.PdfDocument", return_value=mock_pdf_document):
|
||||
_convert_pdf_to_images(bytestream=bytestream, page_range=[1])
|
||||
|
||||
assert "Large PDF detected" in caplog.text
|
||||
|
||||
|
||||
class TestEncodeImageToBase64:
|
||||
def test_encode_image_to_base64(self) -> None:
|
||||
bytestream = get_bytestream_from_source(Path("test/test_files/images/haystack-logo.png"))
|
||||
base64_str = _encode_image_to_base64(bytestream=bytestream)
|
||||
assert base64_str is not None
|
||||
|
||||
def test_encode_image_to_base64_downsize(self) -> None:
|
||||
bytestream = get_bytestream_from_source(Path("test/test_files/images/haystack-logo.png"))
|
||||
base64_str = _encode_image_to_base64(bytestream=bytestream, size=(128, 128))
|
||||
assert base64_str is not None
|
||||
|
||||
|
||||
class TestExtractImageSourcesInfo:
|
||||
def test_extract_image_source_info(self, test_files_path):
|
||||
image_paths = glob.glob(str(test_files_path / "images" / "*.*")) + glob.glob(
|
||||
str(test_files_path / "pdf" / "*.pdf")
|
||||
)
|
||||
|
||||
documents = []
|
||||
for i, path in enumerate(image_paths):
|
||||
document = Document(content=f"document number {i}", meta={"file_path": path})
|
||||
if path.endswith(".pdf"):
|
||||
document.meta["page_number"] = 1
|
||||
documents.append(document)
|
||||
|
||||
images_source_info = _extract_image_sources_info(
|
||||
documents=documents, file_path_meta_field="file_path", root_path=""
|
||||
)
|
||||
assert len(images_source_info) == len(documents)
|
||||
|
||||
for image_source_info in images_source_info:
|
||||
assert str(image_source_info["path"]) in image_paths
|
||||
assert image_source_info["mime_type"] in ["image/jpeg", "image/png", "application/pdf"]
|
||||
if image_source_info["mime_type"] == "application/pdf":
|
||||
assert image_source_info.get("page_number") == 1
|
||||
else:
|
||||
assert "page_number" not in image_source_info
|
||||
|
||||
def test_extract_image_source_info_errors(self, test_files_path):
|
||||
document = Document(content="test")
|
||||
with pytest.raises(ValueError, match="missing the 'file_path' key"):
|
||||
_extract_image_sources_info(documents=[document], file_path_meta_field="file_path", root_path="")
|
||||
|
||||
document = Document(content="test", meta={"file_path": "invalid_path"})
|
||||
with pytest.raises(ValueError, match="has an invalid file path"):
|
||||
_extract_image_sources_info(documents=[document], file_path_meta_field="file_path", root_path="")
|
||||
|
||||
document = Document(content="test", meta={"file_path": str(test_files_path / "docx" / "sample_docx.docx")})
|
||||
with pytest.raises(ValueError, match="has an unsupported MIME type"):
|
||||
_extract_image_sources_info(documents=[document], file_path_meta_field="file_path", root_path="")
|
||||
|
||||
document = Document(content="test", meta={"file_path": str(test_files_path / "pdf" / "sample_pdf_1.pdf")})
|
||||
with pytest.raises(ValueError, match="missing the 'page_number' key"):
|
||||
_extract_image_sources_info(documents=[document], file_path_meta_field="file_path", root_path="")
|
||||
|
||||
def test_extract_image_source_info_rejects_path_traversal(self, test_files_path):
|
||||
# Attacker-controlled document metadata attempts to escape the configured root.
|
||||
document = Document(content="test", meta={"file_path": "../../../../../../etc/passwd"})
|
||||
with pytest.raises(ValueError, match="escapes the configured root"):
|
||||
_extract_image_sources_info(
|
||||
documents=[document], file_path_meta_field="file_path", root_path=str(test_files_path / "images")
|
||||
)
|
||||
|
||||
def test_extract_image_source_info_rejects_absolute_outside_root(self, test_files_path):
|
||||
# Absolute path that lies outside the configured root must be rejected before any IO.
|
||||
document = Document(content="test", meta={"file_path": "/etc/passwd"})
|
||||
with pytest.raises(ValueError, match="escapes the configured root"):
|
||||
_extract_image_sources_info(
|
||||
documents=[document], file_path_meta_field="file_path", root_path=str(test_files_path / "images")
|
||||
)
|
||||
|
||||
def test_extract_image_source_info_accepts_path_inside_root(self, test_files_path):
|
||||
# When the resolved path is inside the configured root, processing must succeed.
|
||||
document = Document(content="test", meta={"file_path": "haystack-logo.png"})
|
||||
images_source_info = _extract_image_sources_info(
|
||||
documents=[document], file_path_meta_field="file_path", root_path=str(test_files_path / "images")
|
||||
)
|
||||
assert len(images_source_info) == 1
|
||||
assert images_source_info[0]["mime_type"] == "image/png"
|
||||
|
||||
|
||||
class TestBatchConvertPdfPagesToImages:
|
||||
@patch("haystack.components.converters.image.image_utils._convert_pdf_to_images")
|
||||
def test_batch_convert_pdf_pages_to_images(self, mocked_convert_pdf_to_images, test_files_path):
|
||||
mocked_convert_pdf_to_images.return_value = [
|
||||
(1, Image.new("RGB", (100, 100))),
|
||||
(2, Image.new("RGB", (100, 100))),
|
||||
]
|
||||
|
||||
pdf_path = test_files_path / "pdf" / "sample_pdf_1.pdf"
|
||||
pdf_doc_1: _PDFPageInfo = {"doc_idx": 0, "path": pdf_path, "page_number": 1}
|
||||
pdf_doc_2: _PDFPageInfo = {"doc_idx": 1, "path": pdf_path, "page_number": 2}
|
||||
pdf_documents = [pdf_doc_1, pdf_doc_2]
|
||||
|
||||
result = _batch_convert_pdf_pages_to_images(pdf_page_infos=pdf_documents, return_base64=False)
|
||||
|
||||
pdf_bytestream = ByteStream.from_file_path(pdf_path)
|
||||
|
||||
mocked_convert_pdf_to_images.assert_called_once_with(
|
||||
bytestream=pdf_bytestream, page_range=[1, 2], size=None, return_base64=False
|
||||
)
|
||||
|
||||
assert len(result) == len(pdf_documents)
|
||||
assert 0 in result and 1 in result
|
||||
assert isinstance(result[0], Image.Image)
|
||||
assert isinstance(result[1], Image.Image)
|
||||
|
||||
@patch("haystack.components.converters.image.image_utils._convert_pdf_to_images")
|
||||
def test_batch_convert_pdf_pages_to_images_base64(self, mocked_convert_pdf_to_images, test_files_path):
|
||||
mocked_convert_pdf_to_images.return_value = [(1, "base64_image_1"), (2, "base64_image_2")]
|
||||
|
||||
pdf_path = test_files_path / "pdf" / "sample_pdf_1.pdf"
|
||||
pdf_doc_1: _PDFPageInfo = {"doc_idx": 0, "path": pdf_path, "page_number": 1}
|
||||
pdf_doc_2: _PDFPageInfo = {"doc_idx": 1, "path": pdf_path, "page_number": 2}
|
||||
pdf_documents = [pdf_doc_1, pdf_doc_2]
|
||||
|
||||
result = _batch_convert_pdf_pages_to_images(pdf_page_infos=pdf_documents, return_base64=True)
|
||||
|
||||
pdf_bytestream = ByteStream.from_file_path(pdf_path)
|
||||
|
||||
mocked_convert_pdf_to_images.assert_called_once_with(
|
||||
bytestream=pdf_bytestream, page_range=[1, 2], size=None, return_base64=True
|
||||
)
|
||||
|
||||
assert result == {0: "base64_image_1", 1: "base64_image_2"}
|
||||
|
||||
def test_batch_convert_pdf_pages_to_images_no_pages_info(self):
|
||||
result = _batch_convert_pdf_pages_to_images(pdf_page_infos=[])
|
||||
|
||||
assert isinstance(result, dict)
|
||||
assert len(result) == 0
|
||||
@@ -0,0 +1,107 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from haystack.components.converters.image.image_utils import _convert_pdf_to_images
|
||||
from haystack.components.converters.image.pdf_to_image import PDFToImageContent
|
||||
from haystack.core.serialization import component_from_dict, component_to_dict
|
||||
from haystack.dataclasses import ByteStream
|
||||
|
||||
|
||||
class TestPDFToImageContent:
|
||||
def test_to_dict(self) -> None:
|
||||
converter = PDFToImageContent()
|
||||
assert component_to_dict(converter, "converter") == {
|
||||
"init_parameters": {"detail": None, "size": None, "page_range": None},
|
||||
"type": "haystack.components.converters.image.pdf_to_image.PDFToImageContent",
|
||||
}
|
||||
|
||||
def test_to_dict_not_defaults(self) -> None:
|
||||
converter = PDFToImageContent(detail="low", size=(128, 128), page_range=[1])
|
||||
assert component_to_dict(converter, "converter") == {
|
||||
"init_parameters": {"detail": "low", "size": (128, 128), "page_range": [1]},
|
||||
"type": "haystack.components.converters.image.pdf_to_image.PDFToImageContent",
|
||||
}
|
||||
|
||||
def test_from_dict(self) -> None:
|
||||
data = {
|
||||
"init_parameters": {"detail": "auto", "size": None, "page_range": [1]},
|
||||
"type": "haystack.components.converters.image.pdf_to_image.PDFToImageContent",
|
||||
}
|
||||
converter = component_from_dict(PDFToImageContent, data, "name")
|
||||
assert component_to_dict(converter, "converter") == data
|
||||
|
||||
def test_run_with_valid_source(self) -> None:
|
||||
file_path = "./test/test_files/pdf/sample_pdf_1.pdf"
|
||||
mime_type = "application/pdf"
|
||||
converter = PDFToImageContent()
|
||||
results = converter.run(sources=[file_path])
|
||||
|
||||
byte_stream = ByteStream.from_file_path(Path(file_path), mime_type=mime_type, meta={"file_path": file_path})
|
||||
assert len(results["image_contents"]) == 4
|
||||
assert results["image_contents"][0].base64_image is not None
|
||||
assert (
|
||||
results["image_contents"][0].base64_image
|
||||
== _convert_pdf_to_images(bytestream=byte_stream, size=None, page_range=[1], return_base64=True)[0][1]
|
||||
)
|
||||
assert results["image_contents"][0].mime_type == "image/jpeg"
|
||||
assert results["image_contents"][0].detail is None
|
||||
assert results["image_contents"][0].meta["file_path"] == str(Path(file_path))
|
||||
assert results["image_contents"][0].meta["page_number"] == 1
|
||||
assert results["image_contents"][1].meta["page_number"] == 2
|
||||
assert results["image_contents"][2].meta["page_number"] == 3
|
||||
assert results["image_contents"][3].meta["page_number"] == 4
|
||||
|
||||
def test_run_with_no_sources(self) -> None:
|
||||
converter = PDFToImageContent()
|
||||
results = converter.run(sources=[])
|
||||
assert len(results["image_contents"]) == 0
|
||||
assert results == {"image_contents": []}
|
||||
|
||||
def test_run_with_invalid_source_type(self, caplog) -> None:
|
||||
converter = PDFToImageContent()
|
||||
converter.run(sources=[123]) # Invalid source type
|
||||
assert "Could not read" in caplog.text
|
||||
|
||||
def test_run_with_non_existent_file(self, caplog) -> None:
|
||||
converter = PDFToImageContent()
|
||||
converter.run(sources=["./non_existent_file.png"])
|
||||
assert "Could not read" in caplog.text
|
||||
assert "No such file or directory:" in caplog.text
|
||||
|
||||
def test_run_with_bytestream_sources(self) -> None:
|
||||
file_path = "./test/test_files/pdf/sample_pdf_1.pdf"
|
||||
mime_type = "application/pdf"
|
||||
byte_stream = ByteStream.from_file_path(Path(file_path), mime_type=mime_type, meta={"file_path": file_path})
|
||||
|
||||
# Initialize the converter
|
||||
converter = PDFToImageContent()
|
||||
|
||||
# Run the converter with the ByteStream
|
||||
results = converter.run(sources=[byte_stream])
|
||||
|
||||
# Assertions
|
||||
assert len(results["image_contents"]) == 4
|
||||
assert results["image_contents"][0].base64_image is not None
|
||||
assert (
|
||||
results["image_contents"][0].base64_image
|
||||
== _convert_pdf_to_images(bytestream=byte_stream, size=None, page_range=[1], return_base64=True)[0][1]
|
||||
)
|
||||
assert results["image_contents"][0].mime_type == "image/jpeg"
|
||||
assert results["image_contents"][0].detail is None
|
||||
assert results["image_contents"][0].meta["file_path"] == file_path
|
||||
assert results["image_contents"][0].meta["page_number"] == 1
|
||||
|
||||
def test_run_with_empty_bytestream(self) -> None:
|
||||
# Create an empty ByteStream object
|
||||
byte_stream = ByteStream(data=b"", meta={"file_path": "empty_file.pdf"})
|
||||
|
||||
# Initialize the converter
|
||||
converter = PDFToImageContent()
|
||||
|
||||
# Run the converter with the empty ByteStream
|
||||
results = converter.run(sources=[byte_stream])
|
||||
|
||||
assert results["image_contents"] == []
|
||||
Reference in New Issue
Block a user