Feature: speed up graph index (#968)

* feat: refine query context

* feat: refine query context

* feat: refine query context

* feat: refine query context

* feat: speed up graph index

* feat: speed up graph index

Author: earayu

aperag/graph/lightrag/constants.py

@@ -41,11 +41,5 @@
 
 # Default values for environment variables
 DEFAULT_MAX_TOKEN_SUMMARY = 500
-DEFAULT_FORCE_LLM_SUMMARY_ON_MERGE = 6
-DEFAULT_WOKERS = 2
+DEFAULT_FORCE_LLM_SUMMARY_ON_MERGE = 10
 DEFAULT_TIMEOUT = 150
-
-# Logging configuration defaults
-DEFAULT_LOG_MAX_BYTES = 10485760  # Default 10MB
-DEFAULT_LOG_BACKUP_COUNT = 5  # Default 5 backups
-DEFAULT_LOG_FILENAME = "lightrag.log"  # Default log filename

aperag/graph/lightrag/lightrag.py

@@ -47,7 +47,6 @@
     final,
 )
 
-from aperag.concurrent_control import MultiLock, get_or_create_lock
 from aperag.graph.lightrag.constants import (
     DEFAULT_FORCE_LLM_SUMMARY_ON_MERGE,
     DEFAULT_MAX_TOKEN_SUMMARY,
@@ -214,7 +213,7 @@ class LightRAG:
     llm_model_max_token_size: int = field(default=32768)
     """Maximum number of tokens allowed per LLM response."""
 
-    llm_model_max_async: int = field(default=4)
+    llm_model_max_async: int = field(default=8)
     """Maximum number of concurrent LLM calls."""
 
     llm_model_kwargs: dict[str, Any] = field(default_factory=dict)
@@ -440,29 +439,28 @@ def _find_connected_components(self, chunk_results: List[tuple[dict, dict]]) ->
         logger.info(f"Found {len(components)} connected components from {len(adjacency)} entities")
         return components
 
-    async def _process_entity_groups(
+    async def _grouping_process_chunk_results(
         self,
         chunk_results: List[tuple[dict, dict]],
-        components: List[List[str]],
         collection_id: str | None = None,
     ) -> dict[str, Any]:
         """
         Process entities and relationships in groups based on connected components.
 
         Args:
             chunk_results: List of (nodes_dict, edges_dict) from entity extraction
-            components: List of entity groups (connected components)
             collection_id: Optional collection ID for logging
 
         Returns:
             Dict with processing results
         """
+        components = self._find_connected_components(chunk_results)
+
         workspace = self.workspace if self.workspace else "default"
         lightrag_logger = create_lightrag_logger(prefix="LightRAG-GraphIndex", workspace=workspace)
 
-        total_entities = 0
-        total_relations = 0
-        processed_groups = 0
+        # Prepare component data for parallel processing
+        component_tasks = []
 
         for i, component in enumerate(components):
             # Create a set for quick lookup
@@ -492,19 +490,31 @@ async def _process_entity_groups(
             if not component_chunk_results:
                 continue
 
-            # Create locks for all entities in this component
-            entity_locks = []
-            for entity_name in sorted(component):  # Sort to prevent deadlock
-                lock = get_or_create_lock(f"entity:{entity_name}:{self.workspace}")
-                entity_locks.append(lock)
+            # Add task data for this component
+            component_tasks.append(
+                {
+                    "index": i,
+                    "component": component,
+                    "component_chunk_results": component_chunk_results,
+                    "total_components": len(components),
+                }
+            )
+
+        # Process components concurrently with semaphore
+        semaphore = asyncio.Semaphore(self.llm_model_max_async)
 
-            lightrag_logger.info(f"Processing component {i + 1}/{len(components)} with {len(component)} entities")
+        async def _process_component_with_semaphore(task_data):
+            async with semaphore:
+                lightrag_logger.info(
+                    f"Processing component {task_data['index'] + 1}/{task_data['total_components']} "
+                    f"with {len(task_data['component'])} entities"
+                )
 
-            # Use MultiLock to acquire all locks for this component
-            async with MultiLock(entity_locks):
-                # Merge nodes and edges for this component
-                await merge_nodes_and_edges(
-                    chunk_results=component_chunk_results,
+                # Call merge_nodes_and_edges with component information
+                result = await merge_nodes_and_edges(
+                    chunk_results=task_data["component_chunk_results"],
+                    component=task_data["component"],
+                    workspace=self.workspace,
                     knowledge_graph_inst=self.chunk_entity_relation_graph,
                     entity_vdb=self.entities_vdb,
                     relationships_vdb=self.relationships_vdb,
@@ -517,17 +527,43 @@ async def _process_entity_groups(
                     lightrag_logger=lightrag_logger,
                 )
 
-            # Count entities and relations in this component
-            component_entity_count = sum(len(nodes) for nodes, _ in component_chunk_results)
-            component_relation_count = sum(len(edges) for _, edges in component_chunk_results)
+                lightrag_logger.info(
+                    f"Completed component {task_data['index'] + 1}: "
+                    f"{result['entity_count']} entities, {result['relation_count']} relations"
+                )
 
-            total_entities += component_entity_count
-            total_relations += component_relation_count
-            processed_groups += 1
+                return result
 
-            lightrag_logger.info(
-                f"Completed component {i + 1}: {component_entity_count} entities, {component_relation_count} relations"
-            )
+        # Create tasks for concurrent processing
+        tasks = []
+        for task_data in component_tasks:
+            task = asyncio.create_task(_process_component_with_semaphore(task_data))
+            tasks.append(task)
+
+        # Wait for all tasks to complete or for the first exception
+        done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION)
+
+        # Check if any task raised an exception
+        for task in done:
+            if task.exception():
+                # Cancel all pending tasks
+                for pending_task in pending:
+                    pending_task.cancel()
+
+                # Wait for cancellation to complete
+                if pending:
+                    await asyncio.wait(pending)
+
+                # Re-raise the exception
+                raise task.exception()
+
+        # Collect results from all tasks
+        results = [task.result() for task in tasks]
+
+        # Calculate totals
+        total_entities = sum(r["entity_count"] for r in results)
+        total_relations = sum(r["relation_count"] for r in results)
+        processed_groups = len(results)
 
         return {
             "groups_processed": processed_groups,
@@ -699,11 +735,8 @@ async def aprocess_graph_indexing(
                 lightrag_logger=lightrag_logger,
             )
 
-            # 2. Find connected components in the extracted entities and relationships
-            components = self._find_connected_components(chunk_results)
-
-            # 3. Process each component group with its own lock scope
-            result = await self._process_entity_groups(chunk_results, components, collection_id)
+            # 2. Process each component group with its own lock scope
+            result = await self._grouping_process_chunk_results(chunk_results, collection_id)
 
             # Count total results
             entity_count = sum(len(nodes) for nodes, _ in chunk_results)

aperag/graph/lightrag/operate.py

@@ -41,6 +41,8 @@
 from collections import Counter, defaultdict
 from typing import Any, AsyncIterator
 
+from aperag.concurrent_control import MultiLock, get_or_create_lock
+
 from .base import (
     BaseGraphStorage,
     BaseKVStorage,
@@ -242,38 +244,86 @@ async def _merge_nodes_then_upsert(
     force_llm_summary_on_merge: int,
     lightrag_logger: LightRAGLogger | None = None,
 ):
-    """Get existing nodes from knowledge graph use name,if exists, merge data, else create, then upsert."""
+    """
+    Merge multiple entity nodes with the same name and upsert the result to knowledge graph.
+
+    This function handles entity deduplication by:
+    1. Retrieving existing entity data from knowledge graph
+    2. Merging existing data with new entity data
+    3. Determining the final entity properties through aggregation
+    4. Optionally using LLM to summarize lengthy descriptions
+    5. Upserting the merged entity back to knowledge graph
+
+    Args:
+        entity_name: The name of the entity to merge
+        nodes_data: List of new entity data dictionaries to merge
+        knowledge_graph_inst: Knowledge graph storage instance
+        llm_model_func: LLM function for description summarization
+        tokenizer: Tokenizer for text processing
+        llm_model_max_token_size: Maximum token size for LLM input
+        summary_to_max_tokens: Maximum tokens for summary output
+        language: Language for LLM summarization
+        force_llm_summary_on_merge: Threshold for triggering LLM summarization
+        lightrag_logger: Optional logger instance
+
+    Returns:
+        dict: The merged node data that was upserted
+    """
+
+    # 1. Initialize containers for collecting existing entity data
     already_entity_types = []
     already_source_ids = []
     already_description = []
     already_file_paths = []
 
+    # 2. Retrieve existing entity from knowledge graph if it exists
     already_node = await knowledge_graph_inst.get_node(entity_name)
     if already_node:
+        # 2.1. Collect existing entity type
         already_entity_types.append(already_node["entity_type"])
+
+        # 2.2. Split and collect existing source IDs (multiple IDs separated by GRAPH_FIELD_SEP)
         already_source_ids.extend(split_string_by_multi_markers(already_node["source_id"], [GRAPH_FIELD_SEP]))
+
+        # 2.3. Split and collect existing file paths (multiple paths separated by GRAPH_FIELD_SEP)
         already_file_paths.extend(split_string_by_multi_markers(already_node["file_path"], [GRAPH_FIELD_SEP]))
+
+        # 2.4. Collect existing description
         already_description.append(already_node["description"])
 
+    # 3. Merge and determine final entity properties
+
+    # 3.1. Determine entity type by frequency count (most common type wins)
     entity_type = sorted(
         Counter([dp["entity_type"] for dp in nodes_data] + already_entity_types).items(),
         key=lambda x: x[1],
         reverse=True,
     )[0][0]
+
+    # 3.2. Merge descriptions with field separator, sorted and deduplicated
     description = GRAPH_FIELD_SEP.join(sorted(set([dp["description"] for dp in nodes_data] + already_description)))
+
+    # 3.3. Merge source IDs, deduplicated
     source_id = GRAPH_FIELD_SEP.join(set([dp["source_id"] for dp in nodes_data] + already_source_ids))
+
+    # 3.4. Merge file paths, deduplicated
     file_path = GRAPH_FIELD_SEP.join(set([dp["file_path"] for dp in nodes_data] + already_file_paths))
 
-    num_fragment = description.count(GRAPH_FIELD_SEP) + 1
-    num_new_fragment = len(set([dp["description"] for dp in nodes_data]))
+    # 4. Calculate description fragment counts for summarization decision
+    num_fragment = description.count(GRAPH_FIELD_SEP) + 1  # Total description fragments
+    num_new_fragment = len(set([dp["description"] for dp in nodes_data]))  # New unique descriptions
 
+    # 5. Handle description summarization if there are multiple fragments
     if num_fragment > 1:
+        # 5.1. Check if LLM summarization is needed based on fragment count threshold
         if num_fragment >= force_llm_summary_on_merge:
+            # 5.1.1. Log LLM summarization decision
             if lightrag_logger:
                 lightrag_logger.log_entity_merge(entity_name, num_fragment, num_new_fragment, is_llm_summary=True)
             else:
                 logger.info(f"LLM merge N: {entity_name} | {num_new_fragment}+{num_fragment - num_new_fragment}")
 
+            # 5.1.2. Use LLM to summarize lengthy descriptions
             description = await _handle_entity_relation_summary(
                 entity_name,
                 description,
@@ -285,11 +335,13 @@ async def _merge_nodes_then_upsert(
                 lightrag_logger,
             )
         else:
+            # 5.2. Simple merge without LLM summarization (fragment count below threshold)
             if lightrag_logger:
                 lightrag_logger.log_entity_merge(entity_name, num_fragment, num_new_fragment, is_llm_summary=False)
             else:
                 logger.info(f"Merge N: {entity_name} | {num_new_fragment}+{num_fragment - num_new_fragment}")
 
+    # 6. Create final node data structure
     node_data = dict(
         entity_id=entity_name,
         entity_type=entity_type,
@@ -298,10 +350,14 @@ async def _merge_nodes_then_upsert(
         file_path=file_path,
         created_at=int(time.time()),
     )
+
+    # 7. Upsert the merged entity to knowledge graph
     await knowledge_graph_inst.upsert_node(
         entity_name,
         node_data=node_data,
     )
+
+    # 8. Add entity_name to returned data and return the final merged entity
     node_data["entity_name"] = entity_name
     return node_data
 
@@ -445,6 +501,65 @@ async def _merge_edges_then_upsert(
 
 @timing_wrapper("Merge & Update")
 async def merge_nodes_and_edges(
+    chunk_results: list,
+    component: list[str],
+    workspace: str,
+    knowledge_graph_inst: BaseGraphStorage,
+    entity_vdb: BaseVectorStorage,
+    relationships_vdb: BaseVectorStorage,
+    llm_model_func,
+    tokenizer,
+    llm_model_max_token_size,
+    summary_to_max_tokens,
+    addon_params,
+    force_llm_summary_on_merge,
+    lightrag_logger: LightRAGLogger | None = None,
+) -> dict[str, int]:
+    """Merge nodes and edges from extraction results
+
+    Args:
+        chunk_results: List of (nodes_dict, edges_dict) tuples
+        component: List of entity names in this component (for locking)
+        workspace: Workspace identifier for lock creation
+        knowledge_graph_inst: Knowledge graph storage
+        entity_vdb: Entity vector database
+        relationships_vdb: Relationships vector database
+        llm_model_func: LLM function
+        tokenizer: Tokenizer
+        llm_model_max_token_size: Max token size for LLM
+        summary_to_max_tokens: Max tokens for summary
+        addon_params: Additional parameters
+        force_llm_summary_on_merge: Force LLM summary threshold
+        lightrag_logger: Optional logger
+
+    Returns:
+        Dict with entity_count and relation_count
+    """
+
+    # Create locks for all entities in this component
+    entity_locks = []
+    for entity_name in sorted(component):  # Sort to prevent deadlock
+        lock = get_or_create_lock(f"entity:{entity_name}:{workspace}")
+        entity_locks.append(lock)
+
+    # Use MultiLock to acquire all locks for this component
+    async with MultiLock(entity_locks):
+        return await _merge_nodes_and_edges_impl(
+            chunk_results,
+            knowledge_graph_inst,
+            entity_vdb,
+            relationships_vdb,
+            llm_model_func,
+            tokenizer,
+            llm_model_max_token_size,
+            summary_to_max_tokens,
+            addon_params,
+            force_llm_summary_on_merge,
+            lightrag_logger,
+        )
+
+
+async def _merge_nodes_and_edges_impl(
     chunk_results: list,
     knowledge_graph_inst: BaseGraphStorage,
     entity_vdb: BaseVectorStorage,
@@ -456,8 +571,8 @@ async def merge_nodes_and_edges(
     addon_params,
     force_llm_summary_on_merge,
     lightrag_logger: LightRAGLogger | None = None,
-) -> None:
-    """Merge nodes and edges from extraction results"""
+) -> dict[str, int]:
+    """Internal implementation of merge_nodes_and_edges"""
 
     # Collect all nodes and edges from all chunks
     all_nodes = defaultdict(list)
@@ -539,6 +654,10 @@ async def merge_nodes_and_edges(
         }
         await relationships_vdb.upsert(data_for_vdb)
 
+    entity_count = len(entities_data)
+    relation_count = len(relationships_data)
+    return {"entity_count": entity_count, "relation_count": relation_count}
+
 
 @timing_wrapper("Entity Extraction")
 async def extract_entities(