From a2110f94a858d683c5b2af0ac361c4e1ced90f1c Mon Sep 17 00:00:00 2001 From: Javier Martinez Date: Tue, 7 Jul 2026 15:58:08 +0200 Subject: [PATCH] fix: blocks in expansion --- .../node_store/node_store_component.py | 64 ++++- .../tree_expansion/split_subtrees.py | 110 ++++----- .../table_expansion_post_processor.py | 18 +- ...ee_expansion_replacement_post_processor.py | 219 ++++++++---------- private_gpt/launcher.py | 3 +- private_gpt/server/content/content_service.py | 5 +- .../postprocessor/tree/test_split_subtrees.py | 16 +- 7 files changed, 229 insertions(+), 206 deletions(-) diff --git a/private_gpt/components/node_store/node_store_component.py b/private_gpt/components/node_store/node_store_component.py index fa67584e..ba810d7e 100644 --- a/private_gpt/components/node_store/node_store_component.py +++ b/private_gpt/components/node_store/node_store_component.py @@ -1,3 +1,4 @@ +import asyncio import logging import re import typing @@ -116,6 +117,49 @@ 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=[ @@ -142,7 +186,7 @@ class NodeStoreComponent: if limit is None: limit = self.max_nodes - return vector_store.get_nodes(node_ids=node_ids, filters=filters, limit=limit) # type: ignore + return filters, limit def get_sorted_nodes( self, @@ -163,6 +207,24 @@ 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 249e93f0..ce771e5a 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,21 +22,25 @@ class SplitSubtreeAlg: of a new root node. The new root node is then added to the list of subtrees. """ - def split_subtree(self, node: TreeNode) -> list[TreeNode]: - """Split the tree into subtrees at the specified split points. + 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] - Args: - node (TreeNode): The node to start splitting from. - - Returns: - list[TreeNode]: A list of subtrees. - """ + def _compute_split_matrix( + self, node: TreeNode + ) -> tuple[list[TreeNode], list[list[TreeNode]]]: + """Compute the sorted nodes and subtrees matrix for a node.""" 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 [] @@ -45,34 +49,50 @@ 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: - node = sorted_nodes[i] - siblings = node.parent.children if node.parent else [] + n = sorted_nodes[i] + siblings = n.parent.children if n.parent else [] if len(siblings) > 1: - section_siblings = [n for n in siblings if self._is_split_point(n)] - if node in section_siblings: - idx = section_siblings.index(node) + section_siblings = [s for s in siblings if self._is_split_point(s)] + if n in section_siblings: + idx = section_siblings.index(n) if ( idx == 0 - and node.parent - and node.parent.abs_idx in split_indices - and node.abs_idx in split_indices + and n.parent + and n.parent.abs_idx in split_indices + and n.abs_idx in split_indices ): - new_split_indices.remove(node.abs_idx) - logger.debug(f"Pruned split index: {node.abs_idx}") + new_split_indices.remove(n.abs_idx) + logger.debug(f"Pruned split index: {n.abs_idx}") continue - if node.metadata.get(MetadataFlags.NO_PRUNABLE.value): + if n.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}") - return self._split_subtrees(sorted_nodes, 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 def _is_split_point(self, node: TreeNode) -> bool: """Check if the node is a split point.""" @@ -84,44 +104,6 @@ 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 57a0eb97..8b82eef6 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,10 +17,26 @@ 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() @@ -46,7 +62,7 @@ class TableExpansionPostProcessor(BaseNodePostprocessor): if not table_root_ids: return nodes - table_nodes = self.node_component.get_nodes( + table_nodes = await self.node_component.aget_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 fc3255d3..4333fa28 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 concurrent.futures import ThreadPoolExecutor, as_completed -from typing import TYPE_CHECKING, cast +from typing import cast from llama_index.core.postprocessor.types import BaseNodePostprocessor from llama_index.core.schema import ( @@ -28,9 +28,6 @@ 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 @@ -101,40 +98,43 @@ class TreeExpansionReplacementPostProcessor(BaseNodePostprocessor, ABC): def _postprocess_nodes( self, nodes: list[NodeWithScore], query_bundle: QueryBundle | None = None ) -> list[NodeWithScore]: - """Create a expansion pipeline trying to maximize the content and relevance. + raise NotImplementedError( + "TreeExpansionReplacementPostProcessor is async-only; " + "use apostprocess_nodes instead." + ) - Steps: - 1. Find optimal set of nodes using a greedy algorithm - 2. Process the optimal set of nodes - 3. Generate final result nodes - """ + async def apostprocess_nodes( + self, + nodes: list[NodeWithScore], + query_bundle: QueryBundle | None = None, + query_str: str | None = None, + ) -> list[NodeWithScore]: if not nodes: return [] - # Initialize setup - absolute_token_limit = self.token_limit or 0 - logger.debug(f"Processing nodes with token limit: {absolute_token_limit}") + 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) - # ===================================================================== - # STEP 1: Find optimal set of nodes through search - # ===================================================================== - node_selection = self._find_optimal_nodes(nodes, absolute_token_limit) + absolute_token_limit = self.token_limit or 0 + all_hit_nodes = self._prepare_hit_nodes(nodes) + + 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 + ) logger.debug( f"Optimal set found: {len(node_selection.hit_nodes)} nodes selected" ) - # ===================================================================== - # STEP 2: Process the expanded nodes from optimal set - # ===================================================================== - processed_nodes = self._process_expanded_nodes(node_selection) + processed_nodes = await self._aprocess_expanded_nodes(node_selection) logger.debug(f"Processed {len(processed_nodes.filtered_items)} expanded nodes") - # ===================================================================== - # STEP 3: Generate final result nodes - # ===================================================================== - result_nodes = self._generate_result_nodes(processed_nodes) + result_nodes = await self._generate_result_nodes(processed_nodes) logger.debug(f"Generated {len(result_nodes)} final result nodes") - return result_nodes # =========================================================================== @@ -151,9 +151,6 @@ 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) @@ -162,23 +159,22 @@ 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) - def _load_root_nodes( + async def _aload_root_nodes( self, hit_nodes: list[NodeWithScore] ) -> dict[str, DocumentRootNode]: - """Load root nodes for the hit nodes.""" + """Async load root nodes for the hit nodes into the cache.""" 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} - 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), - ) - } + 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) logger.debug(f"Loaded {len(root_nodes_map)} root nodes") return root_nodes_map @@ -317,28 +313,11 @@ class TreeExpansionReplacementPostProcessor(BaseNodePostprocessor, ABC): def _get_or_load_root_nodes( self, node_ids: list[str] ) -> dict[str, DocumentRootNode]: - """Get root nodes from cache or load missing ones from node store. + """Get root nodes from the prefetched cache. - Args: - node_ids: List of root node IDs to retrieve - - Returns: - dict[str, DocumentRootNode]: Dictionary mapping node IDs to root nodes + Root nodes are prefetched asynchronously before the greedy search runs; + a cache miss here indicates the prefetch did not cover these ids. """ - 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 @@ -373,24 +352,23 @@ class TreeExpansionReplacementPostProcessor(BaseNodePostprocessor, ABC): Returns: list[ExpansionResult]: List of expansion results """ - with ThreadPoolExecutor() as executor: - work_items = list(zip(hit_nodes, root_nodes, token_limits, strict=False)) + 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 list(executor.map(process_node, work_items)) + return [process_node(item) for item in work_items] def _calculate_configuration_score( self, node_selection: NodeSelection, absolute_token_limit: int @@ -539,16 +517,10 @@ class TreeExpansionReplacementPostProcessor(BaseNodePostprocessor, ABC): # =========================================================================== # STEP 2 METHODS: Processing 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 + async def _aprocess_expanded_nodes( + self, node_selection: NodeSelection + ) -> ProcessedNodes: + """Async variant of _process_expanded_nodes.""" 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 @@ -556,17 +528,13 @@ class TreeExpansionReplacementPostProcessor(BaseNodePostprocessor, ABC): logger.debug(f"Loading {len(all_nodes_to_load)} expanded 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), - ) - } + 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} - # 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) @@ -581,70 +549,67 @@ class TreeExpansionReplacementPostProcessor(BaseNodePostprocessor, ABC): # =========================================================================== # STEP 3 METHODS: Generating final result nodes # =========================================================================== - def _generate_result_nodes( + async def _generate_result_nodes( self, processed_nodes: ProcessedNodes ) -> list[NodeWithScore]: """Generate final result nodes by processing loaded nodes. - Args: - processed_nodes: Output from _process_expanded_nodes - - Returns: - list[NodeWithScore]: Final result 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. """ - result_nodes: list[NodeWithScore] = [] logger.debug( f"Processing {len(processed_nodes.filtered_items)} items for final results" ) - with ThreadPoolExecutor() as executor: - futures: list[Future[list[NodeWithScore]]] = [ - executor.submit( - self._process_node, + if not processed_nodes.filtered_items: + return [] + + results = await asyncio.gather( + *( + self._aprocess_node( hit_node, root_node, subtree, processed_nodes.loaded_nodes, ) for hit_node, root_node, subtree in processed_nodes.filtered_items - ] + ) + ) - for future in as_completed(futures): - result_nodes.extend(future.result()) + result_nodes: list[NodeWithScore] = [] + for batch in results: + result_nodes.extend(batch) - # Sort by score return sorted(result_nodes, key=lambda x: float(x.score or 0), reverse=True) - def _process_node( + async def _aprocess_node( self, hit_node: NodeWithScore, root_node: DocumentRootNode, current_subtree: set[str], loaded_nodes: dict[str, TreeNode], ) -> list[NodeWithScore]: - if not current_subtree or len(current_subtree) == 0: - # This tree was merged with another + if not current_subtree: 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, ) - # 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._rebuild_tree, root_node, sorted_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 asyncio.to_thread(self._prune_content, 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] + final_nodes = await self._asplit_subtrees(final_nodes) + final_nodes = [n for n in 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 [ @@ -656,6 +621,11 @@ 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, @@ -686,11 +656,6 @@ 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 a9c35350..9eb0b05a 100644 --- a/private_gpt/launcher.py +++ b/private_gpt/launcher.py @@ -117,8 +117,9 @@ def create_app(root_injector: Injector) -> FastAPI: # Set default thread pool limit 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=min(500, cpu_count * 50), thread_name_prefix="Stream-Pool" + max_workers=max_workers, thread_name_prefix="Stream-Pool" ) asyncio.get_event_loop().set_default_executor(executor) diff --git a/private_gpt/server/content/content_service.py b/private_gpt/server/content/content_service.py index eab09345..689369bc 100644 --- a/private_gpt/server/content/content_service.py +++ b/private_gpt/server/content/content_service.py @@ -359,10 +359,7 @@ class ContentService: ): # Split the tree into subtrees alg: SplitSubtreeAlg = SplitSubtreeAlg() - subtrees = await asyncio.to_thread( - alg.split_subtree, - root_node, - ) + subtrees = await alg.asplit_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 88ac7fdc..b26faf14 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"), ], ) -def test_various_tree_structures( +async def test_various_tree_structures( tree_setup, expected_min_subtrees: int, description: str, @@ -169,7 +169,7 @@ def test_various_tree_structures( alg = SplitSubtreeAlg() tree = tree_setup() tree.print_tree() # For debugging - subtrees = alg.split_subtree(tree) + subtrees = await alg.asplit_subtree(tree) print(f"\nTesting: {description}") print(f"Generated {len(subtrees)} subtrees") @@ -192,12 +192,12 @@ def test_various_tree_structures( assert len(subtree.children) > 0, "Subtree should not be empty" -def test_multilevel_section_hierarchy() -> None: +async 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) + subtrees = await alg.asplit_subtree(tree) print("\nTesting: Multilevel section hierarchy") print(f"Generated {len(subtrees)} subtrees") @@ -219,12 +219,12 @@ def test_multilevel_section_hierarchy() -> None: ), "Parent-child relationship broken" -def test_no_sections_integrity() -> None: +async 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) + subtrees = await alg.asplit_subtree(tree) print("\nTesting: No sections") print(f"Generated {len(subtrees)} subtrees") @@ -250,12 +250,12 @@ def test_no_sections_integrity() -> None: assert original_types == result_types, "Node types changed in no-sections tree" -def test_multiple_sections_relationships() -> None: +async 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) + subtrees = await alg.asplit_subtree(tree) print("\nTesting: Multiple sections at same level") print(f"Generated {len(subtrees)} subtrees")