mirror of
https://github.com/imartinez/privateGPT.git
synced 2026-07-17 01:48:03 +00:00
317 lines
13 KiB
Python
317 lines
13 KiB
Python
import unittest
|
|
|
|
from private_gpt.components.ingest.transformations.markdown_to_tree_transform import (
|
|
MarkdownTreeNodeParser,
|
|
)
|
|
from private_gpt.components.readers.nodes.tree_node import TreeMetadataMode, TreeNode
|
|
|
|
SAMPLE_MARKDOWN = """\
|
|
### Document Title
|
|
#### Subtitle
|
|
This is a sample document generated randomly with various elements such as tables, lists, nested lists, code blocks, and headers.
|
|
|
|
|
|
### Table Example
|
|
| Column A | Column B |
|
|
|----------|----------|
|
|
| Cell 1 | Cell 2 |
|
|
|----------|----------|
|
|
| Cell 3 | Cell 4 |
|
|
|
|
### Nested List Example
|
|
* Item 1
|
|
+ Subitem 1
|
|
+ Subitem 2
|
|
* Item 2
|
|
+ Subsubitem 1
|
|
- Further subsubitem 1
|
|
- Further subsubitem 2
|
|
"""
|
|
|
|
|
|
class TestMarkdownParser(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
"""Set up the parser instance."""
|
|
self.parser = MarkdownTreeNodeParser(include_metadata=True)
|
|
|
|
def _serialize_tree(self, node: TreeNode) -> str:
|
|
"""Recursively serialize a tree node to markdown."""
|
|
return node.get_content(TreeMetadataMode.USER)
|
|
|
|
def _markdown_to_html(self, markdown: str) -> str:
|
|
"""Convert markdown to html."""
|
|
from mistune import HTMLRenderer, create_markdown
|
|
|
|
md = create_markdown(
|
|
renderer=HTMLRenderer(), # type: ignore
|
|
plugins=["strikethrough", "table"],
|
|
) # type: ignore
|
|
return md(markdown) # type: ignore
|
|
|
|
def _round_trip_test(
|
|
self, markdown: str, check: bool = True, exact: bool = True
|
|
) -> str:
|
|
"""Perform a round-trip parse/serialize/parse/serialize test."""
|
|
# Parse input markdown
|
|
root = self.parser.parse(markdown)
|
|
# Serialize the tree back to markdown
|
|
serialized = self._serialize_tree(root)
|
|
# Parse serialized markdown
|
|
reparsed_root = self.parser.parse(serialized)
|
|
# Serialize reparsed tree
|
|
reserialized = self._serialize_tree(reparsed_root)
|
|
# Assert equality between input and re-serialized markdown
|
|
|
|
if check:
|
|
if exact:
|
|
assert reserialized.strip() == markdown.strip()
|
|
else:
|
|
# Compare the input markdown and the re-serialized markdown
|
|
# as html strings, since the Markdown parser may add extra
|
|
input_html = self._markdown_to_html(markdown)
|
|
reserialized_html = self._markdown_to_html(reserialized)
|
|
assert reserialized_html.strip() == input_html.strip()
|
|
return reserialized
|
|
|
|
def test_headers(self) -> None:
|
|
"""Test header parsing."""
|
|
markdown = "# Header 1\n\n## Header 2\n\n### Header 3"
|
|
self._round_trip_test(markdown)
|
|
|
|
markdown = "# Header 1\n## Header 2\n### Header 3"
|
|
self._round_trip_test(markdown, exact=False)
|
|
|
|
def test_lists(self) -> None:
|
|
"""Test ordered and unordered list parsing."""
|
|
markdown = "1. First item\n2. Second item"
|
|
self._round_trip_test(markdown)
|
|
|
|
markdown = "- First item\n- Second item"
|
|
self._round_trip_test(markdown, exact=False)
|
|
|
|
markdown = "+ First item\n\t+ Second item"
|
|
self._round_trip_test(markdown, exact=False)
|
|
|
|
def test_code_blocks(self) -> None:
|
|
"""Test inline, indented, and fenced code blocks."""
|
|
markdown = "Inline `code`"
|
|
self._round_trip_test(markdown)
|
|
|
|
markdown = "```\nSample text\n```"
|
|
self._round_trip_test(markdown)
|
|
|
|
markdown = "```js \nvar foo = 'bar';\n```"
|
|
self._round_trip_test(markdown, exact=False)
|
|
|
|
def test_tables(self) -> None:
|
|
"""Test table parsing."""
|
|
# Case 1: Single row table
|
|
markdown = (
|
|
"| Header 1 | Header 2 |\n"
|
|
"|----------|----------|\n"
|
|
"| Cell 1 | Cell 2 |\n"
|
|
"| Cell 3 | Cell 4 |"
|
|
)
|
|
self._round_trip_test(markdown, exact=False)
|
|
|
|
# Case 2: Duplicate header name
|
|
markdown = (
|
|
"| Header 1 | Header 1 |\n"
|
|
"|----------|----------|\n"
|
|
"| Cell 1 | Cell 2 |\n"
|
|
"| Cell 3 | Cell 4 |"
|
|
)
|
|
root = self.parser.parse(markdown)
|
|
serialized = self._serialize_tree(root)
|
|
expected = (
|
|
"| Header 1_1 | Header 1_2 |\n"
|
|
"| - | - |\n"
|
|
"| Cell 1 | Cell 2 |\n"
|
|
"| Cell 3 | Cell 4 |\n\n"
|
|
)
|
|
assert serialized.strip() == expected.strip()
|
|
|
|
# Case 3: Has default header without header row
|
|
markdown = "| 0 | 1 |\n| - | - |\n| Alice | 25 |\n| Bob | 30 |"
|
|
self._round_trip_test(markdown)
|
|
|
|
# Case 4: Has default header with header row
|
|
markdown = "| 0 | 1 |\n| - | - |\n| Name | Age |\n| Alice | 25 |\n| Bob | 30 |"
|
|
root = self.parser.parse(markdown)
|
|
serialized = self._serialize_tree(root)
|
|
expected = "| Name | Age |\n| - | - |\n| Alice | 25 |\n| Bob | 30 |\n\n"
|
|
assert serialized.strip() == expected.strip()
|
|
|
|
# Case 5. Empty headers
|
|
markdown = "| | |\n| - | - |\n| Alice | 25 |\n| Bob | 30 |"
|
|
root = self.parser.parse(markdown)
|
|
serialized = self._serialize_tree(root)
|
|
expected = (
|
|
"| Unknown_1 | Unknown_2 |\n| - | - |\n| Alice | 25 |\n| Bob | 30 |\n\n"
|
|
)
|
|
assert serialized.strip() == expected.strip()
|
|
|
|
# Case 6. Nan header in some position
|
|
markdown = (
|
|
"| 0 | 1 | 2 |\n"
|
|
"| - | - | - |\n"
|
|
"| nan | Name | Age |\n"
|
|
"| 1 | Alice | 25 |\n"
|
|
"| 2 | Bob | 30 |"
|
|
)
|
|
root = self.parser.parse(markdown)
|
|
serialized = self._serialize_tree(root)
|
|
expected = (
|
|
"| Unknown | Name | Age |\n"
|
|
"| - | - | - |\n"
|
|
"| 1 | Alice | 25 |\n"
|
|
"| 2 | Bob | 30 |\n\n"
|
|
)
|
|
assert serialized.strip() == expected.strip()
|
|
|
|
def test_blockquotes(self) -> None:
|
|
"""Test blockquote parsing."""
|
|
markdown = "> Blockquote level 1\n> Blockquote level 2\n> Blockquote level 3"
|
|
self._round_trip_test(markdown)
|
|
markdown = "> Blockquote level 1\n>> Blockquote level 2\n>>> Blockquote level 3"
|
|
self._round_trip_test(markdown)
|
|
|
|
def test_typographic_replacements(self) -> None:
|
|
"""Test typographic replacements."""
|
|
markdown = "(c) (C) (r) (R) (tm) (TM)\nSmart quotes: 'single' and \"double\""
|
|
self._round_trip_test(markdown)
|
|
|
|
def test_images_and_links(self) -> None:
|
|
"""Test images and links."""
|
|
markdown = (
|
|
"[Example link](http://example.com)\n"
|
|
""
|
|
)
|
|
self._round_trip_test(markdown)
|
|
|
|
def test_sanity_heading_format(self) -> None:
|
|
"""Test the sanitization of headings for proper trimming and formatting."""
|
|
markdown = " ## Improper Heading"
|
|
new_markdown = self._round_trip_test(markdown, check=False)
|
|
expected = "## Improper Heading"
|
|
assert expected.strip() == new_markdown.strip(), (
|
|
"Sanitized headings should remove excessive spaces and normalize formatting."
|
|
)
|
|
|
|
def test_sanity_paragraph_format(self) -> None:
|
|
"""Test the sanitization of paragraphs with inconsistent line breaks."""
|
|
markdown = "Paragraph one.\n\n Paragraph two with extra indentation.\n\n"
|
|
new_markdown = self._round_trip_test(markdown, check=False)
|
|
expected = "Paragraph one.\n\nParagraph two with extra indentation.\n\n"
|
|
assert expected.strip() == new_markdown.strip(), (
|
|
"Sanitized paragraphs should have consistent formatting and spacing."
|
|
)
|
|
|
|
def test_sanity_image_format(self) -> None:
|
|
"""Test the sanitization of image syntax."""
|
|
markdown = ""
|
|
new_markdown = self._round_trip_test(markdown, check=False)
|
|
expected = ""
|
|
assert expected.strip() == new_markdown.strip(), (
|
|
"Sanitized images should have no unnecessary spaces and correct syntax."
|
|
)
|
|
|
|
def test_sanity_link_format(self) -> None:
|
|
"""Test the sanitization of link syntax."""
|
|
markdown = "[ Link Text ]( http://example.com )"
|
|
new_markdown = self._round_trip_test(markdown, check=False)
|
|
expected = "[Link Text](http://example.com)\n\n"
|
|
assert expected.strip() == new_markdown.strip(), (
|
|
"Sanitized links should remove extra spaces and preserve valid syntax."
|
|
)
|
|
|
|
def test_sanity_bold_format(self) -> None:
|
|
"""Test the sanitization of bold syntax."""
|
|
markdown = "**Bold Text** and __More Bold__"
|
|
new_markdown = self._round_trip_test(markdown, check=False)
|
|
expected = "Bold Text and More Bold\n\n"
|
|
assert expected.strip() == new_markdown.strip(), (
|
|
"Sanitized bold should remove asterisks/underscores and preserve text content."
|
|
)
|
|
|
|
def test_sanity_italic_format(self) -> None:
|
|
"""Test the sanitization of italic syntax."""
|
|
markdown = "*Italic Text* and _More Italic_"
|
|
new_markdown = self._round_trip_test(markdown, check=False)
|
|
expected = "Italic Text and More Italic\n\n"
|
|
assert expected.strip() == new_markdown.strip(), (
|
|
"Sanitized italic should remove single asterisks/underscores and preserve text content."
|
|
)
|
|
|
|
def test_sanity_bold_italic_format(self) -> None:
|
|
"""Test the sanitization of combined bold and italic syntax."""
|
|
markdown = "***Bold Italic*** and ___More Bold Italic___"
|
|
new_markdown = self._round_trip_test(markdown, check=False)
|
|
expected = "Bold Italic and More Bold Italic\n\n"
|
|
assert expected.strip() == new_markdown.strip(), (
|
|
"Sanitized bold-italic should remove triple asterisks/underscores and preserve text content."
|
|
)
|
|
|
|
def test_sanity_mixed_format(self) -> None:
|
|
"""Test the sanitization of mixed markdown formatting."""
|
|
markdown = "**Bold** with *italic* and ***bold-italic*** mixed"
|
|
new_markdown = self._round_trip_test(markdown, check=False)
|
|
expected = "Bold with italic and bold-italic mixed\n\n"
|
|
assert expected.strip() == new_markdown.strip(), (
|
|
"Sanitized mixed formatting should remove all markdown syntax and preserve text content."
|
|
)
|
|
|
|
def test_sanity_item_format(self) -> None:
|
|
markdown = "* ● Hypertension: Diagnosed at age 32, currently managed with losartan 50 mg once daily."
|
|
new_markdown = self._round_trip_test(markdown, check=False)
|
|
expected = "* Hypertension: Diagnosed at age 32, currently managed with losartan 50 mg once daily.\n"
|
|
assert expected == new_markdown, (
|
|
"Sanitized items should remove extra icons and preserve valid syntax."
|
|
)
|
|
|
|
markdown = "* ◦ Hypertension: Diagnosed at age 32, currently managed with losartan 50 mg once daily."
|
|
new_markdown = self._round_trip_test(markdown, check=False)
|
|
expected = "* Hypertension: Diagnosed at age 32, currently managed with losartan 50 mg once daily.\n"
|
|
assert expected == new_markdown, (
|
|
"Sanitized items should remove extra spaces and preserve valid syntax."
|
|
)
|
|
|
|
markdown = "* ◦ Hypertension: Diagnosed at age 32, currently managed with losartan 50 mg once daily."
|
|
new_markdown = self._round_trip_test(markdown, check=False)
|
|
expected = "* Hypertension: Diagnosed at age 32, currently managed with losartan 50 mg once daily.\n"
|
|
assert expected == new_markdown, (
|
|
"Sanitized items should remove extra spaces and preserve valid syntax."
|
|
)
|
|
|
|
markdown = "* + Hypertension: Diagnosed at age 32, currently managed with losartan 50 mg once daily."
|
|
new_markdown = self._round_trip_test(markdown, check=False)
|
|
expected = "* Hypertension: Diagnosed at age 32, currently managed with losartan 50 mg once daily.\n"
|
|
assert expected == new_markdown, (
|
|
"Sanitized items should remove extra spaces and preserve valid syntax."
|
|
)
|
|
|
|
markdown = "* - Hypertension: Diagnosed at age 32, currently managed with losartan 50 mg once daily."
|
|
new_markdown = self._round_trip_test(markdown, check=False)
|
|
expected = "* Hypertension: Diagnosed at age 32, currently managed with losartan 50 mg once daily.\n"
|
|
assert expected == new_markdown, (
|
|
"Sanitized items should remove extra spaces and preserve valid syntax."
|
|
)
|
|
|
|
markdown = " - 1. Hypertension: Diagnosed at age 32, currently managed with losartan 50 mg once daily."
|
|
new_markdown = self._round_trip_test(markdown, check=False)
|
|
expected = "1. Hypertension: Diagnosed at age 32, currently managed with losartan 50 mg once daily.\n"
|
|
assert expected == new_markdown, (
|
|
"Sanitized items should remove extra spaces and preserve valid syntax."
|
|
)
|
|
|
|
markdown = " 5. + Hypertension: Diagnosed at age 32, currently managed with losartan 50 mg once daily."
|
|
new_markdown = self._round_trip_test(markdown, check=False)
|
|
expected = "5. Hypertension: Diagnosed at age 32, currently managed with losartan 50 mg once daily.\n"
|
|
assert expected == new_markdown, (
|
|
"Sanitized items should remove extra spaces and preserve valid syntax."
|
|
)
|
|
|
|
def test_sample_markdown(self) -> None:
|
|
"""Test the parser on the sample markdown."""
|
|
self._round_trip_test(SAMPLE_MARKDOWN, exact=False)
|