mirror of
https://github.com/imartinez/privateGPT.git
synced 2026-07-17 01:48:03 +00:00
293 lines
9.4 KiB
Python
293 lines
9.4 KiB
Python
import pandas as pd
|
|
import pytest
|
|
|
|
from private_gpt.components.postprocessor.tree_expansion.split_subtrees import (
|
|
SplitSubtreeAlg,
|
|
)
|
|
from private_gpt.components.readers.nodes import (
|
|
DocumentRootNode,
|
|
ListNode,
|
|
SectionNode,
|
|
TableNode,
|
|
TextNode,
|
|
)
|
|
|
|
|
|
def create_multilevel_sections() -> DocumentRootNode:
|
|
"""Creates a tree with multiple nested section levels."""
|
|
root = DocumentRootNode()
|
|
|
|
# First main section with subsections
|
|
section1 = SectionNode(text="Main Section 1")
|
|
root.add_child(section1)
|
|
text1 = TextNode(text="Main section 1 content")
|
|
section1.add_child(text1)
|
|
|
|
# Subsections under section1
|
|
subsection1_1 = SectionNode(text="Subsection 1.1")
|
|
section1.add_child(subsection1_1)
|
|
text1_1 = TextNode(text="Subsection 1.1 content")
|
|
subsection1_1.add_child(text1_1)
|
|
|
|
subsection1_2 = SectionNode(text="Subsection 1.2")
|
|
section1.add_child(subsection1_2)
|
|
text1_2 = TextNode(text="Subsection 1.2 content")
|
|
subsection1_2.add_child(text1_2)
|
|
|
|
# Second main section with deeper nesting
|
|
section2 = SectionNode(text="Main Section 2")
|
|
root.add_child(section2)
|
|
|
|
subsection2_1 = SectionNode(text="Subsection 2.1")
|
|
section2.add_child(subsection2_1)
|
|
|
|
# Even deeper nesting
|
|
subsubsection2_1_1 = SectionNode(text="SubSubsection 2.1.1")
|
|
subsection2_1.add_child(subsubsection2_1_1)
|
|
text2_1_1 = TextNode(text="Deep nested content")
|
|
subsubsection2_1_1.add_child(text2_1_1)
|
|
|
|
# Update references
|
|
root.update_references()
|
|
|
|
return root
|
|
|
|
|
|
def create_no_sections() -> DocumentRootNode:
|
|
"""Creates a tree with no sections, just different types of content."""
|
|
root = DocumentRootNode()
|
|
|
|
# Add some text nodes
|
|
text1 = TextNode(text="First paragraph")
|
|
root.add_child(text1)
|
|
|
|
text2 = TextNode(text="Second paragraph")
|
|
root.add_child(text2)
|
|
|
|
# Add a list
|
|
list_node = ListNode()
|
|
root.add_child(list_node)
|
|
list_item1 = TextNode(text="List item 1")
|
|
list_item2 = TextNode(text="List item 2")
|
|
list_node.add_child(list_item1)
|
|
list_node.add_child(list_item2)
|
|
|
|
# Add a table
|
|
df = pd.DataFrame({"Column1": ["A", "B"], "Column2": [1, 2]})
|
|
table = TableNode(df=df, description="Sample table")
|
|
root.add_child(table)
|
|
|
|
# Update references
|
|
root.update_references()
|
|
|
|
return root
|
|
|
|
|
|
def create_multiple_sections_same_level() -> DocumentRootNode:
|
|
"""Creates a tree with multiple sections at the same level."""
|
|
root = DocumentRootNode()
|
|
|
|
# Section 1 with content
|
|
section1 = SectionNode(text="Section 1")
|
|
root.add_child(section1)
|
|
text1 = TextNode(text="Content for section 1")
|
|
section1.add_child(text1)
|
|
|
|
# Section 2 with list
|
|
section2 = SectionNode(text="Section 2")
|
|
root.add_child(section2)
|
|
list_node = ListNode()
|
|
section2.add_child(list_node)
|
|
list_item = TextNode(text="List item in section 2")
|
|
list_node.add_child(list_item)
|
|
|
|
# Section 3 with table
|
|
section3 = SectionNode(text="Section 3")
|
|
root.add_child(section3)
|
|
df = pd.DataFrame({"A": [1, 2], "B": [3, 4]})
|
|
table = TableNode(df=df, description="Table in section 3")
|
|
section3.add_child(table)
|
|
|
|
# Update references
|
|
root.update_references()
|
|
|
|
return root
|
|
|
|
|
|
def create_mixed_depth_sections() -> DocumentRootNode:
|
|
"""Creates a tree with sections at different depths mixed together."""
|
|
root = DocumentRootNode()
|
|
|
|
# First section with deep nesting
|
|
section1 = SectionNode(text="Section 1")
|
|
root.add_child(section1)
|
|
|
|
subsection1_1 = SectionNode(text="Subsection 1.1")
|
|
section1.add_child(subsection1_1)
|
|
|
|
subsubsection1_1_1 = SectionNode(text="SubSubsection 1.1.1")
|
|
subsection1_1.add_child(subsubsection1_1_1)
|
|
text1 = TextNode(text="Deep nested content")
|
|
subsubsection1_1_1.add_child(text1)
|
|
|
|
# Second section at root level
|
|
section2 = SectionNode(text="Section 2")
|
|
root.add_child(section2)
|
|
text2 = TextNode(text="Root level content")
|
|
section2.add_child(text2)
|
|
|
|
# Third section with medium nesting
|
|
section3 = SectionNode(text="Section 3")
|
|
root.add_child(section3)
|
|
|
|
subsection3_1 = SectionNode(text="Subsection 3.1")
|
|
section3.add_child(subsection3_1)
|
|
text3 = TextNode(text="Medium nested content")
|
|
subsection3_1.add_child(text3)
|
|
|
|
# Update references
|
|
root.update_references()
|
|
|
|
return root
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("tree_setup", "expected_min_subtrees", "description"),
|
|
[
|
|
(create_no_sections, 1, "No sections"),
|
|
(create_multiple_sections_same_level, 3, "Multiple sections at same level"),
|
|
(create_mixed_depth_sections, 3, "Mixed depth sections"),
|
|
(create_multilevel_sections, 3, "Multilevel nested sections"),
|
|
],
|
|
)
|
|
def test_various_tree_structures(
|
|
tree_setup,
|
|
expected_min_subtrees: int,
|
|
description: str,
|
|
) -> None:
|
|
"""Test the splitter with various tree structures."""
|
|
alg = SplitSubtreeAlg()
|
|
tree = tree_setup()
|
|
tree.print_tree() # For debugging
|
|
subtrees = alg.split_subtree(tree)
|
|
|
|
print(f"\nTesting: {description}")
|
|
print(f"Generated {len(subtrees)} subtrees")
|
|
for i, st in enumerate(subtrees):
|
|
print(f"\nSubtree {i}:")
|
|
st.print_tree()
|
|
|
|
# Basic quantity check
|
|
assert len(subtrees) >= expected_min_subtrees, (
|
|
f"Expected at least {expected_min_subtrees} subtrees for {description}, "
|
|
f"but got {len(subtrees)}"
|
|
)
|
|
|
|
# Verify each subtree is valid
|
|
for subtree in subtrees:
|
|
assert isinstance(
|
|
subtree, DocumentRootNode
|
|
), "Subtree root should be DocumentRoot"
|
|
# Check that each subtree has at least one child
|
|
assert len(subtree.children) > 0, "Subtree should not be empty"
|
|
|
|
|
|
def test_multilevel_section_hierarchy() -> None:
|
|
"""Detailed test for multilevel section hierarchy preservation."""
|
|
alg = SplitSubtreeAlg()
|
|
tree = create_multilevel_sections()
|
|
tree.print_tree() # For debugging
|
|
subtrees = alg.split_subtree(tree)
|
|
|
|
print("\nTesting: Multilevel section hierarchy")
|
|
print(f"Generated {len(subtrees)} subtrees")
|
|
for i, st in enumerate(subtrees):
|
|
print(f"\nSubtree {i}:")
|
|
st.print_tree()
|
|
|
|
# Verify that deep nested sections stay together
|
|
for subtree in subtrees:
|
|
if len(subtree.children) > 0 and isinstance(subtree.children[0], SectionNode):
|
|
section = subtree.children[0]
|
|
# If this section has subsections, verify the parent-child relationship
|
|
if any(isinstance(child, SectionNode) for child in section.children):
|
|
# Verify parent references are correct
|
|
for child in section.children:
|
|
if isinstance(child, SectionNode):
|
|
assert (
|
|
child.parent.id_ == section.id_
|
|
), "Parent-child relationship broken"
|
|
|
|
|
|
def test_no_sections_integrity() -> None:
|
|
"""Detailed test for trees without sections."""
|
|
alg = SplitSubtreeAlg()
|
|
tree = create_no_sections()
|
|
tree.print_tree() # For debugging
|
|
subtrees = alg.split_subtree(tree)
|
|
|
|
print("\nTesting: No sections")
|
|
print(f"Generated {len(subtrees)} subtrees")
|
|
for i, st in enumerate(subtrees):
|
|
print(f"\nSubtree {i}:")
|
|
st.print_tree()
|
|
|
|
# Should return exactly one subtree as there are no section splits
|
|
assert len(subtrees) == 1, "Tree without sections should not be split"
|
|
|
|
# Verify all content is preserved
|
|
original_nodes = list(tree.flatten())
|
|
result_nodes = list(subtrees[0].flatten())
|
|
assert len(result_nodes) >= len(
|
|
original_nodes
|
|
), "Content was lost in no-sections tree"
|
|
|
|
# Verify the types of nodes are preserved
|
|
original_types = [type(n) for n in original_nodes]
|
|
result_types = [type(n) for n in result_nodes]
|
|
result_types = result_types[len(result_types) - len(original_types) :]
|
|
|
|
assert original_types == result_types, "Node types changed in no-sections tree"
|
|
|
|
|
|
def test_multiple_sections_relationships() -> None:
|
|
"""Detailed test for multiple sections at the same level."""
|
|
alg = SplitSubtreeAlg()
|
|
tree = create_multiple_sections_same_level()
|
|
tree.print_tree() # For debugging
|
|
subtrees = alg.split_subtree(tree)
|
|
|
|
print("\nTesting: Multiple sections at same level")
|
|
print(f"Generated {len(subtrees)} subtrees")
|
|
for i, st in enumerate(subtrees):
|
|
print(f"\nSubtree {i}:")
|
|
st.print_tree()
|
|
|
|
# Verify each section is in its own subtree
|
|
section_counts = [
|
|
sum(1 for child in subtree.flatten() if isinstance(child, SectionNode))
|
|
for subtree in subtrees
|
|
]
|
|
|
|
# Each subtree should have exactly one section
|
|
assert all(
|
|
count == 1 for count in section_counts
|
|
), "Each subtree should contain exactly one section"
|
|
|
|
# Verify content type preservation
|
|
for subtree in subtrees:
|
|
section = next(
|
|
(child for child in subtree.children if isinstance(child, SectionNode)),
|
|
None,
|
|
)
|
|
if section:
|
|
# Verify that the section's content is preserved
|
|
if "Section 2" in section.text:
|
|
assert any(
|
|
isinstance(node, ListNode) for node in section.flatten()
|
|
), "List content not preserved in Section 2"
|
|
elif "Section 3" in section.text:
|
|
assert any(
|
|
isinstance(node, TableNode) for node in section.flatten()
|
|
), "Table content not preserved in Section 3"
|