diff --git a/astrbot/core/knowledge_base/parsers/pdf_parser.py b/astrbot/core/knowledge_base/parsers/pdf_parser.py index aeeea930a21..7e0c9c0d65e 100644 --- a/astrbot/core/knowledge_base/parsers/pdf_parser.py +++ b/astrbot/core/knowledge_base/parsers/pdf_parser.py @@ -4,9 +4,12 @@ """ import io +from pathlib import Path +from PIL import Image from pypdf import PdfReader +from astrbot import logger from astrbot.core.knowledge_base.parsers.base import ( BaseParser, MediaItem, @@ -47,54 +50,44 @@ async def parse(self, file_content: bytes, file_name: str) -> ParseResult: image_counter = 0 for page_num, page in enumerate(reader.pages): try: - # 安全检查 Resources - if "/Resources" not in page: - continue - - resources = page["/Resources"] - if not resources or "/XObject" not in resources: # type: ignore - continue - - xobjects = resources["/XObject"].get_object() # type: ignore - if not xobjects: - continue - - for obj_name in xobjects: + images = page.images + for image_name in images.keys(): try: - obj = xobjects[obj_name] - - if obj.get("/Subtype") != "/Image": + # pypdf encodes pixel streams and applies color spaces + # and soft masks; raw get_data() bytes are not PNG files. + image = images[image_name] + if image.image is None: continue - - # 提取图片数据 - image_data = obj.get_data() - - # 确定格式 - filter_type = obj.get("/Filter", "") - if filter_type == "/DCTDecode": - ext = "jpg" - mime_type = "image/jpeg" - elif filter_type == "/FlateDecode": - ext = "png" - mime_type = "image/png" - else: - ext = "png" - mime_type = "image/png" + ext = Path(image.name).suffix + mime_type = Image.MIME[image.image.format] image_counter += 1 media_items.append( MediaItem( media_type="image", - file_name=f"page_{page_num}_img_{image_counter}.{ext}", - content=image_data, + file_name=f"page_{page_num}_img_{image_counter}{ext}", + content=image.data, mime_type=mime_type, ), ) except Exception: - # 单个图片提取失败不影响整体 + # A broken image must not hide other images on the page. + logger.warning( + "Failed to extract image %r on PDF page %d in %r", + image_name, + page_num + 1, + file_name, + exc_info=True, + ) continue except Exception: - # 页面处理失败不影响其他页面 + # A page-level image error must not discard other pages or text. + logger.warning( + "Failed to enumerate images on PDF page %d in %r", + page_num + 1, + file_name, + exc_info=True, + ) continue full_text = "\n\n".join(text_parts) diff --git a/tests/test_pdf_parser.py b/tests/test_pdf_parser.py new file mode 100644 index 00000000000..33a69b5cd00 --- /dev/null +++ b/tests/test_pdf_parser.py @@ -0,0 +1,186 @@ +import io +from pathlib import Path +from unittest.mock import Mock + +import pytest +from PIL import Image +from pypdf import PdfWriter +from pypdf.generic import NameObject +from reportlab.lib.utils import ImageReader +from reportlab.pdfgen.canvas import Canvas + +from astrbot.core.knowledge_base.parsers.pdf_parser import PDFParser + + +def _make_pdf( + image_format: str = "PNG", + *, + transparent: bool = False, + nested: bool = False, + pages: int = 1, +) -> bytes: + """Build a PDF with text and an encoded image on each page. + + Args: + image_format: Encoding of the source image. + transparent: Whether to include a soft alpha mask. + nested: Whether to place the image inside a Form XObject. + pages: Number of pages to generate. + + Returns: + Complete PDF bytes. + """ + image_buffer = io.BytesIO() + mode = "RGBA" if transparent else "RGB" + color = (20, 80, 140, 128) if transparent else (20, 80, 140) + Image.new(mode, (8, 6), color).save(image_buffer, format=image_format) + image_buffer.seek(0) + source = ImageReader(image_buffer) + pdf_buffer = io.BytesIO() + canvas = Canvas(pdf_buffer, pagesize=(200, 200)) + for page_num in range(pages): + canvas.drawString(10, 180, f"Page {page_num + 1} text survives.") + if nested: + canvas.beginForm(f"Figure{page_num}") + canvas.drawImage(source, 10, 10, width=80, height=60, mask="auto") + if nested: + canvas.endForm() + canvas.doForm(f"Figure{page_num}") + canvas.showPage() + canvas.save() + return pdf_buffer.getvalue() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("image_format", ["PNG", "JPEG"]) +async def test_pdf_images_have_decodable_content_and_matching_metadata(image_format): + result = await PDFParser().parse(_make_pdf(image_format), "image.pdf") + + assert "Page 1 text survives." in result.text + assert len(result.media) == 1 + media = result.media[0] + assert media.media_type == "image" + with Image.open(io.BytesIO(media.content)) as image: + image.load() + assert image.size == (8, 6) + assert image.format == image_format + assert media.mime_type == Image.MIME[image.format] + assert Path(media.file_name).suffix == ( + ".jpg" if image_format == "JPEG" else ".png" + ) + if image_format == "PNG": + assert image.getpixel((0, 0)) == (20, 80, 140) + + +@pytest.mark.asyncio +async def test_pdf_images_preserve_soft_mask_transparency(): + result = await PDFParser().parse(_make_pdf(transparent=True), "alpha.pdf") + + assert len(result.media) == 1 + with Image.open(io.BytesIO(result.media[0].content)) as image: + image.load() + assert image.mode == "RGBA" + assert image.getpixel((0, 0)) == (20, 80, 140, 128) + + +@pytest.mark.asyncio +async def test_pdf_images_inside_forms_are_extracted(): + result = await PDFParser().parse(_make_pdf(nested=True), "nested.pdf") + + assert "Page 1 text survives." in result.text + assert len(result.media) == 1 + with Image.open(io.BytesIO(result.media[0].content)) as image: + image.load() + assert image.getpixel((0, 0)) == (20, 80, 140) + + +@pytest.mark.asyncio +async def test_pdf_images_have_unique_names_across_pages(): + result = await PDFParser().parse(_make_pdf(pages=2), "pages.pdf") + + assert "Page 1 text survives." in result.text + assert "Page 2 text survives." in result.text + assert [media.file_name for media in result.media] == [ + "page_0_img_1.png", + "page_1_img_2.png", + ] + + +@pytest.mark.asyncio +async def test_pdf_bad_image_does_not_discard_text_or_other_images(monkeypatch): + warning = Mock() + monkeypatch.setattr( + "astrbot.core.knowledge_base.parsers.pdf_parser.logger.warning", warning + ) + writer = PdfWriter(clone_from=io.BytesIO(_make_pdf())) + xobjects = writer.pages[0]["/Resources"]["/XObject"] + image_name = next(iter(xobjects)) + # Put a malformed image before the valid one on the same page. + invalid = xobjects[image_name].clone(writer, force_duplicate=True) + del invalid["/Width"] + valid = xobjects.raw_get(image_name) + xobjects.clear() + xobjects[NameObject("/Broken")] = invalid.indirect_reference + xobjects[image_name] = valid + buffer = io.BytesIO() + writer.write(buffer) + + result = await PDFParser().parse(buffer.getvalue(), "broken.pdf") + + assert "Page 1 text survives." in result.text + assert len(result.media) == 1 + warning.assert_called_once_with( + "Failed to extract image %r on PDF page %d in %r", + "/Broken", + 1, + "broken.pdf", + exc_info=True, + ) + with Image.open(io.BytesIO(result.media[0].content)) as image: + image.load() + assert image.size == (8, 6) + + +@pytest.mark.asyncio +async def test_pdf_page_image_error_is_logged_and_later_pages_survive(monkeypatch): + from pypdf._page import PageObject + + warning = Mock() + monkeypatch.setattr( + "astrbot.core.knowledge_base.parsers.pdf_parser.logger.warning", warning + ) + images_property = PageObject.images + calls = 0 + + def images(page): + nonlocal calls + calls += 1 + if calls == 1: + raise ValueError("broken image resources") + return images_property.fget(page) + + monkeypatch.setattr(PageObject, "images", property(images)) + result = await PDFParser().parse(_make_pdf(pages=2), "broken-page.pdf") + + assert "Page 1 text survives." in result.text + assert "Page 2 text survives." in result.text + assert len(result.media) == 1 + warning.assert_called_once_with( + "Failed to enumerate images on PDF page %d in %r", + 1, + "broken-page.pdf", + exc_info=True, + ) + + +@pytest.mark.asyncio +async def test_pdf_without_images_still_returns_text(): + buffer = io.BytesIO() + canvas = Canvas(buffer) + canvas.drawString(10, 180, "Page 1 text survives.") + canvas.save() + + result = await PDFParser().parse(buffer.getvalue(), "text.pdf") + + assert "Page 1 text survives." in result.text + assert result.media == []