fix: blocks in expansion

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,

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],

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),

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] = []

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)
 

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] = []

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")