diff --git a/api/core/rag/extractor/markdown_extractor.py b/api/core/rag/extractor/markdown_extractor.py index 195fe56ff1cfe0..4463680a85f2d8 100644 --- a/api/core/rag/extractor/markdown_extractor.py +++ b/api/core/rag/extractor/markdown_extractor.py @@ -98,16 +98,15 @@ def markdown_to_tups(self, markdown_text: str) -> list[tuple[str | None, str]]: return markdown_tups def remove_images(self, content: str) -> str: - """Get a dictionary of a markdown file from its path.""" - pattern = r"!{1}\[\[(.*)\]\]" - content = re.sub(pattern, "", content) - return content + """Remove wiki and inline image references, keeping the surrounding text.""" + content = re.sub(r"!\[\[[^\]]*\]\]", "", content) + # A destination may contain one level of balanced parentheses, e.g. + # ![plot](figures/plot(v2).png), so it cannot stop at the first ")". + return re.sub(r"!\[[^\]]*\]\((?:[^()\\]|\\.|\([^()]*\))*\)", "", content) def remove_hyperlinks(self, content: str) -> str: - """Get a dictionary of a markdown file from its path.""" - pattern = r"\[(.*?)\]\((.*?)\)" - content = re.sub(pattern, r"\1", content) - return content + """Unwrap ``[text](url)`` links and leave images untouched.""" + return re.sub(r"(? list[tuple[str | None, str]]: """Parse file into tuples.""" diff --git a/api/tests/unit_tests/core/rag/extractor/test_markdown_extractor.py b/api/tests/unit_tests/core/rag/extractor/test_markdown_extractor.py index 8f2a5e2cf75fcd..cd9ce50a36b6db 100644 --- a/api/tests/unit_tests/core/rag/extractor/test_markdown_extractor.py +++ b/api/tests/unit_tests/core/rag/extractor/test_markdown_extractor.py @@ -159,3 +159,52 @@ def test_extract_builds_documents_for_header_and_non_header(self, monkeypatch: p docs = extractor.extract() assert [doc.page_content for doc in docs] == ["plain", "\n\nHeader\nvalue"] + + def test_remove_images_removes_inline_and_wiki_images(self): + extractor = MarkdownExtractor(file_path="dummy_path", remove_images=True) + + assert extractor.remove_images("before ![diagram](diagram.png) after") == "before after" + assert extractor.remove_images("before ![[diagram.png]] after") == "before after" + + def test_remove_images_keeps_text_between_images(self): + extractor = MarkdownExtractor(file_path="dummy_path", remove_images=True) + + assert extractor.remove_images("![[one.png]] useful text ![[two.png]]") == " useful text " + + def test_remove_hyperlinks_keeps_images_intact(self): + extractor = MarkdownExtractor(file_path="dummy_path", remove_hyperlinks=True) + + assert extractor.remove_hyperlinks("![diagram](diagram.png) and [guide](guide.html)") == ( + "![diagram](diagram.png) and guide" + ) + + def test_extract_removes_inline_images_but_keeps_link_text(self, tmp_path: Path): + path = tmp_path / "images.md" + path.write_text("before ![diagram](diagram.png) after", encoding="utf-8") + + docs = MarkdownExtractor(str(path), remove_images=True, encoding="utf-8").extract() + + assert [doc.page_content for doc in docs] == ["before after"] + + def test_remove_images_handles_parenthesized_destinations(self): + extractor = MarkdownExtractor(file_path="dummy_path", remove_images=True) + + assert extractor.remove_images("before ![plot](figures/plot(v2).png) after") == "before after" + assert extractor.remove_images("before ![chart](https://example.com/charts/(a)/render.png) after") == "before after" + assert extractor.remove_images(r"before ![esc](a\(b\).png) after") == "before after" + + def test_extract_removes_parenthesized_image_with_surrounding_prose(self, tmp_path: Path): + path = tmp_path / "plots.md" + path.write_text( + "Quarterly revenue rose. ![plot](figures/plot(v2).png) Costs fell.", + encoding="utf-8", + ) + + docs = MarkdownExtractor(str(path), remove_images=True, encoding="utf-8").extract() + + assert [doc.page_content for doc in docs] == ["Quarterly revenue rose. Costs fell."] + + def test_remove_hyperlinks_handles_parenthesized_destinations(self): + extractor = MarkdownExtractor(file_path="dummy_path", remove_hyperlinks=True) + + assert extractor.remove_hyperlinks("see [plot](figures/plot(v2).png) here") == "see plot here"