diff --git a/private_gpt/components/node_store/node_store_component.py b/private_gpt/components/node_store/node_store_component.py index ba810d7e..fa67584e 100644 --- a/private_gpt/components/node_store/node_store_component.py +++ b/private_gpt/components/node_store/node_store_component.py @@ -1,4 +1,3 @@ -import asyncio import logging import re import typing @@ -117,49 +116,6 @@ class NodeStoreComponent: f"Vector store for collection {collection} does not support get_nodes" ) - filters, limit = self._resolve_get_args(artifacts, node_ids, filters, limit) - - return vector_store.get_nodes(node_ids=node_ids, filters=filters, limit=limit) # type: ignore - - async def aget_nodes( - self, - collection: str, - artifacts: list[str] | None = None, - node_ids: list[str] | None = None, - filters: MetadataFilters | None = None, - limit: int | None = None, - ) -> list[BaseNode]: - """Async variant of get_nodes. - - Falls back to the sync implementation via asyncio.to_thread when the - vector store does not expose an async aget_nodes method. - """ - vector_store = self._vector_store_component.vector_store(collection) - if vector_store is None: - raise ValueError(f"Vector store for collection {collection} not found") - if not hasattr(vector_store, "aget_nodes"): - return await asyncio.to_thread( - self.get_nodes, - collection, - artifacts, - node_ids, - filters, - limit, - ) - - filters, limit = self._resolve_get_args(artifacts, node_ids, filters, limit) - - return await vector_store.aget_nodes( # type: ignore - node_ids=node_ids, filters=filters, limit=limit - ) - - def _resolve_get_args( - self, - artifacts: list[str] | None, - node_ids: list[str] | None, - filters: MetadataFilters | None, - limit: int | None, - ) -> tuple[MetadataFilters | None, int | None]: if artifacts: artifact_filters = MetadataFilters( filters=[ @@ -186,7 +142,7 @@ class NodeStoreComponent: if limit is None: limit = self.max_nodes - return filters, limit + return vector_store.get_nodes(node_ids=node_ids, filters=filters, limit=limit) # type: ignore def get_sorted_nodes( self, @@ -207,24 +163,6 @@ class NodeStoreComponent: ] return sorted_nodes - async def aget_sorted_nodes( - self, - collection: str, - artifacts: list[str] | None = None, - node_ids: list[str] | None = None, - filters: MetadataFilters | None = None, - limit: int | None = None, - ) -> list[BaseNode]: - """Async variant of get_sorted_nodes.""" - unsorted_nodes = await self.aget_nodes( - collection, artifacts, node_ids, filters, limit - ) - index = {node.id_: node for node in unsorted_nodes} - sorted_nodes = [ - index[node_id] for node_id in node_ids if node_id in index # type: ignore - ] - return sorted_nodes - def get_node( self, collection: str, diff --git a/private_gpt/components/postprocessor/tree_expansion/split_subtrees.py b/private_gpt/components/postprocessor/tree_expansion/split_subtrees.py index ce771e5a..249e93f0 100644 --- a/private_gpt/components/postprocessor/tree_expansion/split_subtrees.py +++ b/private_gpt/components/postprocessor/tree_expansion/split_subtrees.py @@ -1,6 +1,6 @@ -import asyncio import logging import uuid +from concurrent.futures import ThreadPoolExecutor from private_gpt.components.ingest.metadata_helper import MetadataFlags, MetadataNode from private_gpt.components.readers.nodes import SectionNode, TreeNode @@ -22,25 +22,21 @@ class SplitSubtreeAlg: of a new root node. The new root node is then added to the list of subtrees. """ - async def asplit_subtree(self, node: TreeNode) -> list[TreeNode]: - """Async split: parallelizes _create_subtree via asyncio.gather.""" - _, subtrees_matrix = self._compute_split_matrix(node) - if not subtrees_matrix: - return [] - results = await asyncio.gather( - *(asyncio.to_thread(self._create_subtree, sn) for sn in subtrees_matrix) - ) - return [s for s in results if s] + def split_subtree(self, node: TreeNode) -> list[TreeNode]: + """Split the tree into subtrees at the specified split points. - def _compute_split_matrix( - self, node: TreeNode - ) -> tuple[list[TreeNode], list[list[TreeNode]]]: - """Compute the sorted nodes and subtrees matrix for a node.""" + Args: + node (TreeNode): The node to start splitting from. + + Returns: + list[TreeNode]: A list of subtrees. + """ sorted_nodes = list(node.flatten()) split_nodes: list[TreeNode] = [ n for n in sorted_nodes if self._is_split_point(n) ] + # Find split points where sections have siblings split_indices: list[int] = [] for n in split_nodes: siblings = n.parent.children if n.parent else [] @@ -49,50 +45,34 @@ class SplitSubtreeAlg: split_indices = sorted(set(split_indices)) logger.debug(f"Split indices: {split_indices}") + # Prune first element in a nested object + # since below content of the section has been joined + # with some content, and it is the nearest split point new_split_indices: list[int] = split_indices.copy() for i in split_indices: - n = sorted_nodes[i] - siblings = n.parent.children if n.parent else [] + node = sorted_nodes[i] + siblings = node.parent.children if node.parent else [] if len(siblings) > 1: - section_siblings = [s for s in siblings if self._is_split_point(s)] - if n in section_siblings: - idx = section_siblings.index(n) + section_siblings = [n for n in siblings if self._is_split_point(n)] + if node in section_siblings: + idx = section_siblings.index(node) if ( idx == 0 - and n.parent - and n.parent.abs_idx in split_indices - and n.abs_idx in split_indices + and node.parent + and node.parent.abs_idx in split_indices + and node.abs_idx in split_indices ): - new_split_indices.remove(n.abs_idx) - logger.debug(f"Pruned split index: {n.abs_idx}") + new_split_indices.remove(node.abs_idx) + logger.debug(f"Pruned split index: {node.abs_idx}") continue - if n.metadata.get(MetadataFlags.NO_PRUNABLE.value): + if node.metadata.get(MetadataFlags.NO_PRUNABLE.value): new_split_indices.remove(i) logger.debug(f"Pruned no-prunable split index: {i}") + # Split the tree into subtrees at the specified indices logger.debug(f"Processed Split indices: {new_split_indices}") - subtrees_matrix = self._build_subtrees_matrix(sorted_nodes, new_split_indices) - return sorted_nodes, subtrees_matrix - - def _build_subtrees_matrix( - self, nodes: list[TreeNode], split_indices: list[int] - ) -> list[list[TreeNode]]: - """Build the matrix of node lists to be turned into subtrees.""" - if not nodes: - return [] - if not split_indices: - split_indices = [0] - - subtrees_matrix: list[list[TreeNode]] = [] - start = 0 - current_node: list[TreeNode] = [] - for idx in split_indices: - subtrees_matrix.append(current_node + nodes[start:idx]) - start = idx + 1 - current_node = [nodes[idx]] - subtrees_matrix.append(current_node + nodes[start:]) - return subtrees_matrix + return self._split_subtrees(sorted_nodes, new_split_indices) def _is_split_point(self, node: TreeNode) -> bool: """Check if the node is a split point.""" @@ -104,6 +84,44 @@ class SplitSubtreeAlg: node = node.parent return node + def _split_subtrees( + self, nodes: list[TreeNode], split_indices: list[int] + ) -> list[TreeNode]: + """Split the tree into subtrees at the specified indices.""" + if not nodes: + return [] + if not split_indices: + # We don't have any split points, so return the tree as is + # processing the result + split_indices = [0] + + # Split the nodes into subtrees based on the split indices + subtrees_matrix: list[list[TreeNode]] = [] + start = 0 + current_node: list[TreeNode] = [] + for idx in split_indices: + subtrees_matrix.append(current_node + nodes[start:idx]) + start = idx + 1 + current_node = [nodes[idx]] + subtrees_matrix.append(current_node + nodes[start:]) + + # Rebuild the subtrees from the split nodes + subtrees = [] + with ThreadPoolExecutor() as executor: + results = executor.map( + lambda subtree_nodes: self._create_subtree(subtree_nodes), + subtrees_matrix, + ) + for new_subtree in results: + if new_subtree: + logger.debug( + f"Generated a new subtree with {len(new_subtree.children)} nodes" + ) + subtrees.append(new_subtree) + else: + logger.debug("Failed to create subtree. Skipping.") + return subtrees + def _create_subtree( self, nodes: list[TreeNode], diff --git a/private_gpt/components/postprocessor/tree_expansion/table_expansion_post_processor.py b/private_gpt/components/postprocessor/tree_expansion/table_expansion_post_processor.py index 8b82eef6..57a0eb97 100644 --- a/private_gpt/components/postprocessor/tree_expansion/table_expansion_post_processor.py +++ b/private_gpt/components/postprocessor/tree_expansion/table_expansion_post_processor.py @@ -17,26 +17,10 @@ class TableExpansionPostProcessor(BaseNodePostprocessor): def _postprocess_nodes( self, nodes: list[NodeWithScore], query_bundle: QueryBundle | None = None - ) -> list[NodeWithScore]: - raise RuntimeError( - "TableExpansionPostProcessor is async-only; " - "use apostprocess_nodes instead." - ) - - async def apostprocess_nodes( - self, - nodes: list[NodeWithScore], - query_bundle: QueryBundle | None = None, - query_str: str | None = None, ) -> list[NodeWithScore]: if not nodes: return [] - if query_str is not None and query_bundle is not None: - raise ValueError("Cannot specify both query_str and query_bundle") - elif query_str is not None: - query_bundle = QueryBundle(query_str) - expanded_nodes = [] table_root_ids = set() @@ -62,7 +46,7 @@ class TableExpansionPostProcessor(BaseNodePostprocessor): if not table_root_ids: return nodes - table_nodes = await self.node_component.aget_nodes( + table_nodes = self.node_component.get_nodes( collection=self.collection, node_ids=list(table_root_ids), limit=len(table_root_ids), diff --git a/private_gpt/components/postprocessor/tree_expansion/tree_expansion_replacement_post_processor.py b/private_gpt/components/postprocessor/tree_expansion/tree_expansion_replacement_post_processor.py index 4333fa28..fc3255d3 100644 --- a/private_gpt/components/postprocessor/tree_expansion/tree_expansion_replacement_post_processor.py +++ b/private_gpt/components/postprocessor/tree_expansion/tree_expansion_replacement_post_processor.py @@ -1,7 +1,7 @@ -import asyncio import logging from abc import ABC -from typing import cast +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import TYPE_CHECKING, cast from llama_index.core.postprocessor.types import BaseNodePostprocessor from llama_index.core.schema import ( @@ -28,6 +28,9 @@ from private_gpt.components.readers.nodes.tree_node import TreeNode from private_gpt.settings.settings import settings from private_gpt.utils.random import generate_deterministic_uuid_from_seed +if TYPE_CHECKING: + from concurrent.futures import Future + config = settings() debug_mode = config.server.debug_mode @@ -98,43 +101,40 @@ class TreeExpansionReplacementPostProcessor(BaseNodePostprocessor, ABC): def _postprocess_nodes( self, nodes: list[NodeWithScore], query_bundle: QueryBundle | None = None ) -> list[NodeWithScore]: - raise NotImplementedError( - "TreeExpansionReplacementPostProcessor is async-only; " - "use apostprocess_nodes instead." - ) + """Create a expansion pipeline trying to maximize the content and relevance. - async def apostprocess_nodes( - self, - nodes: list[NodeWithScore], - query_bundle: QueryBundle | None = None, - query_str: str | None = None, - ) -> list[NodeWithScore]: + Steps: + 1. Find optimal set of nodes using a greedy algorithm + 2. Process the optimal set of nodes + 3. Generate final result nodes + """ if not nodes: return [] - if query_str is not None and query_bundle is not None: - raise ValueError("Cannot specify both query_str and query_bundle") - elif query_str is not None: - query_bundle = QueryBundle(query_str) - + # Initialize setup absolute_token_limit = self.token_limit or 0 - all_hit_nodes = self._prepare_hit_nodes(nodes) + logger.debug(f"Processing nodes with token limit: {absolute_token_limit}") - if all_hit_nodes: - await self._aload_root_nodes(all_hit_nodes) - - node_selection = await asyncio.to_thread( - self._find_optimal_nodes, nodes, absolute_token_limit - ) + # ===================================================================== + # STEP 1: Find optimal set of nodes through search + # ===================================================================== + node_selection = self._find_optimal_nodes(nodes, absolute_token_limit) logger.debug( f"Optimal set found: {len(node_selection.hit_nodes)} nodes selected" ) - processed_nodes = await self._aprocess_expanded_nodes(node_selection) + # ===================================================================== + # STEP 2: Process the expanded nodes from optimal set + # ===================================================================== + processed_nodes = self._process_expanded_nodes(node_selection) logger.debug(f"Processed {len(processed_nodes.filtered_items)} expanded nodes") - result_nodes = await self._generate_result_nodes(processed_nodes) + # ===================================================================== + # STEP 3: Generate final result nodes + # ===================================================================== + result_nodes = self._generate_result_nodes(processed_nodes) logger.debug(f"Generated {len(result_nodes)} final result nodes") + return result_nodes # =========================================================================== @@ -151,6 +151,9 @@ class TreeExpansionReplacementPostProcessor(BaseNodePostprocessor, ABC): # Filter and sort nodes all_hit_nodes = self._prepare_hit_nodes(nodes) + # Load root nodes once for efficiency + self._root_nodes = self._root_nodes or self._load_root_nodes(all_hit_nodes) + # Run greedy search to find optimal nodes return self._greedy_search(all_hit_nodes, absolute_token_limit) @@ -159,22 +162,23 @@ class TreeExpansionReplacementPostProcessor(BaseNodePostprocessor, ABC): hit_nodes = [node for node in nodes if isinstance(node.node, TreeNode)] return sorted(hit_nodes, key=lambda x: float(x.score or 0), reverse=True) - async def _aload_root_nodes( + def _load_root_nodes( self, hit_nodes: list[NodeWithScore] ) -> dict[str, DocumentRootNode]: - """Async load root nodes for the hit nodes into the cache.""" + """Load root nodes for the hit nodes.""" root_ids = [ node.node.root_id for node in hit_nodes if hasattr(node.node, "root_id") ] unique_root_ids = {root_id for root_id in root_ids if root_id} - nodes = await self.node_component.aget_sorted_nodes( - collection=self.collection, - node_ids=list(unique_root_ids), - limit=len(unique_root_ids), - ) - root_nodes_map = {node.node_id: cast(DocumentRootNode, node) for node in nodes} - self._root_nodes.update(root_nodes_map) + root_nodes_map = { + node.node_id: cast(DocumentRootNode, node) + for node in self.node_component.get_sorted_nodes( + collection=self.collection, + node_ids=list(unique_root_ids), + limit=len(unique_root_ids), + ) + } logger.debug(f"Loaded {len(root_nodes_map)} root nodes") return root_nodes_map @@ -313,11 +317,28 @@ class TreeExpansionReplacementPostProcessor(BaseNodePostprocessor, ABC): def _get_or_load_root_nodes( self, node_ids: list[str] ) -> dict[str, DocumentRootNode]: - """Get root nodes from the prefetched cache. + """Get root nodes from cache or load missing ones from node store. - Root nodes are prefetched asynchronously before the greedy search runs; - a cache miss here indicates the prefetch did not cover these ids. + Args: + node_ids: List of root node IDs to retrieve + + Returns: + dict[str, DocumentRootNode]: Dictionary mapping node IDs to root nodes """ + missing_ids = [ + node_id for node_id in node_ids if node_id not in self._root_nodes + ] + if missing_ids: + logger.debug(f"Loading {len(missing_ids)} missing root nodes") + loaded_nodes = { + node.node_id: cast(DocumentRootNode, node) + for node in self.node_component.get_sorted_nodes( + collection=self.collection, + node_ids=missing_ids, + limit=len(missing_ids), + ) + } + self._root_nodes.update(loaded_nodes) return { node_id: self._root_nodes[node_id] for node_id in node_ids @@ -352,23 +373,24 @@ class TreeExpansionReplacementPostProcessor(BaseNodePostprocessor, ABC): Returns: list[ExpansionResult]: List of expansion results """ - work_items = list(zip(hit_nodes, root_nodes, token_limits, strict=False)) + with ThreadPoolExecutor() as executor: + work_items = list(zip(hit_nodes, root_nodes, token_limits, strict=False)) - def process_node( - work_item: tuple[NodeWithScore, TreeNode, int] - ) -> ExpansionResult: - hit_node, root_node, token_limit = work_item - partial_hit_node = root_node.find_self_or_child_by_id(hit_node.node.id_) + def process_node( + work_item: tuple[NodeWithScore, TreeNode, int] + ) -> ExpansionResult: + hit_node, root_node, token_limit = work_item + partial_hit_node = root_node.find_self_or_child_by_id(hit_node.node.id_) - if partial_hit_node: - expansion_result = self._expand( - hit_node=partial_hit_node, token_limit=token_limit - ) - return expansion_result + if partial_hit_node: + expansion_result = self._expand( + hit_node=partial_hit_node, token_limit=token_limit + ) + return expansion_result - return ExpansionResult(node_ids=set(), token_count=0) + return ExpansionResult(node_ids=set(), token_count=0) - return [process_node(item) for item in work_items] + return list(executor.map(process_node, work_items)) def _calculate_configuration_score( self, node_selection: NodeSelection, absolute_token_limit: int @@ -517,10 +539,16 @@ class TreeExpansionReplacementPostProcessor(BaseNodePostprocessor, ABC): # =========================================================================== # STEP 2 METHODS: Processing expanded nodes # =========================================================================== - async def _aprocess_expanded_nodes( - self, node_selection: NodeSelection - ) -> ProcessedNodes: - """Async variant of _process_expanded_nodes.""" + def _process_expanded_nodes(self, node_selection: NodeSelection) -> ProcessedNodes: + """Process expanded nodes by loading them from storage. + + Args: + node_selection: Selected nodes and their expansions + + Returns: + ProcessedNodes: Processed node data + """ + # Prepare the list of all nodes to load active_subtrees = [s for s in node_selection.subtrees if s.node_ids] all_nodes_to_load = { id_ for subtree in active_subtrees for id_ in subtree.node_ids @@ -528,13 +556,17 @@ class TreeExpansionReplacementPostProcessor(BaseNodePostprocessor, ABC): logger.debug(f"Loading {len(all_nodes_to_load)} expanded nodes") - nodes = await self.node_component.aget_nodes( - collection=self.collection, - node_ids=list(all_nodes_to_load), - limit=len(all_nodes_to_load), - ) - loaded_nodes = {n.id_: cast(TreeNode, n) for n in nodes} + # Load all nodes + loaded_nodes = { + n.id_: cast(TreeNode, n) + for n in self.node_component.get_nodes( + collection=self.collection, + node_ids=list(all_nodes_to_load), + limit=len(all_nodes_to_load), + ) + } + # Filter out hit nodes/root nodes with empty subtrees filtered_items = [] for i, (hit_node, root_node) in enumerate( zip(node_selection.hit_nodes, node_selection.root_nodes, strict=False) @@ -549,67 +581,70 @@ class TreeExpansionReplacementPostProcessor(BaseNodePostprocessor, ABC): # =========================================================================== # STEP 3 METHODS: Generating final result nodes # =========================================================================== - async def _generate_result_nodes( + def _generate_result_nodes( self, processed_nodes: ProcessedNodes ) -> list[NodeWithScore]: """Generate final result nodes by processing loaded nodes. - Parallelizes across filtered items via asyncio.gather; each item's - rebuild/prune/split runs in worker threads so GIL-releasing - clone/serialize work (pydantic-core) gets real parallelism. + Args: + processed_nodes: Output from _process_expanded_nodes + + Returns: + list[NodeWithScore]: Final result nodes """ + result_nodes: list[NodeWithScore] = [] logger.debug( f"Processing {len(processed_nodes.filtered_items)} items for final results" ) - if not processed_nodes.filtered_items: - return [] - - results = await asyncio.gather( - *( - self._aprocess_node( + with ThreadPoolExecutor() as executor: + futures: list[Future[list[NodeWithScore]]] = [ + executor.submit( + self._process_node, hit_node, root_node, subtree, processed_nodes.loaded_nodes, ) for hit_node, root_node, subtree in processed_nodes.filtered_items - ) - ) + ] - result_nodes: list[NodeWithScore] = [] - for batch in results: - result_nodes.extend(batch) + for future in as_completed(futures): + result_nodes.extend(future.result()) + # Sort by score return sorted(result_nodes, key=lambda x: float(x.score or 0), reverse=True) - async def _aprocess_node( + def _process_node( self, hit_node: NodeWithScore, root_node: DocumentRootNode, current_subtree: set[str], loaded_nodes: dict[str, TreeNode], ) -> list[NodeWithScore]: - if not current_subtree: + if not current_subtree or len(current_subtree) == 0: + # This tree was merged with another return [] + # Sort the nodes by their absolute index sorted_nodes = sorted( [loaded_nodes[node_id] for node_id in current_subtree], key=lambda x: x.abs_idx, ) - final_nodes = await asyncio.to_thread( - self._rebuild_tree, root_node, sorted_nodes - ) - final_nodes = [n for n in final_nodes if n] + # Post-process the nodes to remove irrelevant content and add diffs + final_nodes = [n for n in self._rebuild_tree(root_node, sorted_nodes) if n] - final_nodes = await asyncio.to_thread(self._prune_content, final_nodes) - final_nodes = [n for n in final_nodes if n] + # Prune the content of the final node + final_nodes = [n for n in self._prune_content(final_nodes) if n] - final_nodes = await self._asplit_subtrees(final_nodes) - final_nodes = [n for n in final_nodes if n] + # Split the final node into smaller subtrees + final_nodes = [n for n in self._split_subtrees(final_nodes) if n] + # Update metadata for the final node final_nodes = [self._update_node(hit_node.node, n) for n in final_nodes if n] + + # Frozen and deduplicate the final nodes final_nodes = self._deduplicate_subtrees(final_nodes) return [ @@ -621,11 +656,6 @@ class TreeExpansionReplacementPostProcessor(BaseNodePostprocessor, ABC): if final_node ] - async def _asplit_subtrees(self, nodes: list[TreeNode]) -> list[TreeNode]: - alg = SplitSubtreeAlg() - results = await asyncio.gather(*(alg.asplit_subtree(n) for n in nodes if n)) - return [s for subtrees in results for s in subtrees] - def _expand( self, hit_node: TreeNode, @@ -656,6 +686,11 @@ class TreeExpansionReplacementPostProcessor(BaseNodePostprocessor, ABC): result = [node.prune() for node in nodes] return [node for node in result if node] + def _split_subtrees(self, nodes: list[TreeNode]) -> list[TreeNode]: + alg: SplitSubtreeAlg = SplitSubtreeAlg() + subtrees = [alg.split_subtree(node) for node in nodes if node] + return [subtree for subtrees in subtrees for subtree in subtrees] + def _deduplicate_subtrees(self, nodes: list[TreeNode]) -> list[TreeNode]: # Convert into frozen elements frozen_nodes: list[TreeNode] = [] diff --git a/private_gpt/launcher.py b/private_gpt/launcher.py index 83ce0470..0af5bae2 100644 --- a/private_gpt/launcher.py +++ b/private_gpt/launcher.py @@ -82,9 +82,8 @@ def create_app(root_injector: Injector) -> FastAPI: # ``scheduler.chat.mode=celery`` is enabled, so a small I/O-only pool is enough # and stops the GIL from being contended with the event loop. cpu_count = os.cpu_count() or 1 - max_workers = settings.server.max_workers or min(64, cpu_count * 5) executor = concurrent.futures.ThreadPoolExecutor( - max_workers=max_workers, thread_name_prefix="Stream-Pool" + max_workers=min(500, cpu_count * 50), thread_name_prefix="Stream-Pool" ) asyncio.get_running_loop().set_default_executor(executor) diff --git a/private_gpt/server/content/content_service.py b/private_gpt/server/content/content_service.py index 689369bc..eab09345 100644 --- a/private_gpt/server/content/content_service.py +++ b/private_gpt/server/content/content_service.py @@ -359,7 +359,10 @@ class ContentService: ): # Split the tree into subtrees alg: SplitSubtreeAlg = SplitSubtreeAlg() - subtrees = await alg.asplit_subtree(root_node) + subtrees = await asyncio.to_thread( + alg.split_subtree, + root_node, + ) # Concat the content of each subtree until the maximum chunk size is reached current_chunk: list[BaseNode] = [] diff --git a/tests/components/postprocessor/tree/test_split_subtrees.py b/tests/components/postprocessor/tree/test_split_subtrees.py index b26faf14..88ac7fdc 100644 --- a/tests/components/postprocessor/tree/test_split_subtrees.py +++ b/tests/components/postprocessor/tree/test_split_subtrees.py @@ -160,7 +160,7 @@ def create_mixed_depth_sections() -> DocumentRootNode: (create_multilevel_sections, 3, "Multilevel nested sections"), ], ) -async def test_various_tree_structures( +def test_various_tree_structures( tree_setup, expected_min_subtrees: int, description: str, @@ -169,7 +169,7 @@ async def test_various_tree_structures( alg = SplitSubtreeAlg() tree = tree_setup() tree.print_tree() # For debugging - subtrees = await alg.asplit_subtree(tree) + subtrees = alg.split_subtree(tree) print(f"\nTesting: {description}") print(f"Generated {len(subtrees)} subtrees") @@ -192,12 +192,12 @@ async def test_various_tree_structures( assert len(subtree.children) > 0, "Subtree should not be empty" -async def test_multilevel_section_hierarchy() -> None: +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 = await alg.asplit_subtree(tree) + subtrees = alg.split_subtree(tree) print("\nTesting: Multilevel section hierarchy") print(f"Generated {len(subtrees)} subtrees") @@ -219,12 +219,12 @@ async def test_multilevel_section_hierarchy() -> None: ), "Parent-child relationship broken" -async def test_no_sections_integrity() -> None: +def test_no_sections_integrity() -> None: """Detailed test for trees without sections.""" alg = SplitSubtreeAlg() tree = create_no_sections() tree.print_tree() # For debugging - subtrees = await alg.asplit_subtree(tree) + subtrees = alg.split_subtree(tree) print("\nTesting: No sections") print(f"Generated {len(subtrees)} subtrees") @@ -250,12 +250,12 @@ async def test_no_sections_integrity() -> None: assert original_types == result_types, "Node types changed in no-sections tree" -async def test_multiple_sections_relationships() -> None: +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 = await alg.asplit_subtree(tree) + subtrees = alg.split_subtree(tree) print("\nTesting: Multiple sections at same level") print(f"Generated {len(subtrees)} subtrees")