Fix summarization and relationship grouping on Inc Indexing (#1768)

* Finx sumarization for large descriptions on incremental indexing

* Semver

* Ruff

Author: AlonsoGuevara

.semversioner/next-release/patch-20250225224555418932.json

@@ -0,0 +1,4 @@
+{
+  "type": "patch",
+  "description": "Fix summarization over large datasets for inc indexing. Fix relationship summarization"
+}

graphrag/index/update/entities.py

@@ -3,19 +3,12 @@
 
 """Entity related operations and utils for Incremental Indexing."""
 
-import asyncio
 import itertools
 
 import numpy as np
 import pandas as pd
 
-from graphrag.cache.pipeline_cache import PipelineCache
-from graphrag.callbacks.workflow_callbacks import WorkflowCallbacks
-from graphrag.config.models.graph_rag_config import GraphRagConfig
 from graphrag.data_model.schemas import ENTITIES_FINAL_COLUMNS
-from graphrag.index.operations.summarize_descriptions.graph_intelligence_strategy import (
-    run_graph_intelligence as run_entity_summarization,
-)
 
 
 def _group_and_resolve_entities(
@@ -83,61 +76,3 @@ def _group_and_resolve_entities(
     resolved = resolved.loc[:, ENTITIES_FINAL_COLUMNS]
 
     return resolved, id_mapping
-
-
-async def _run_entity_summarization(
-    entities_df: pd.DataFrame,
-    config: GraphRagConfig,
-    cache: PipelineCache,
-    callbacks: WorkflowCallbacks,
-) -> pd.DataFrame:
-    """Run entity summarization.
-
-    Parameters
-    ----------
-    entities_df : pd.DataFrame
-        The entities dataframe.
-    config : GraphRagConfig
-        The pipeline configuration.
-    cache : PipelineCache
-        Pipeline cache used during the summarization process.
-
-    Returns
-    -------
-    pd.DataFrame
-        The updated entities dataframe with summarized descriptions.
-    """
-    summarization_llm_settings = config.get_language_model_config(
-        config.summarize_descriptions.model_id
-    )
-    summarization_strategy = config.summarize_descriptions.resolved_strategy(
-        config.root_dir, summarization_llm_settings
-    )
-
-    # Prepare tasks for async summarization where needed
-    async def process_row(row):
-        # Accessing attributes directly from the named tuple.
-        description = row.description
-        if isinstance(description, list) and len(description) > 1:
-            # Run entity summarization asynchronously
-            result = await run_entity_summarization(
-                row.title,
-                description,
-                callbacks,
-                cache,
-                summarization_strategy,
-            )
-            return result.description
-        # Handle case where description is a single-item list or not a list
-        return description[0] if isinstance(description, list) else description
-
-    # Create a list of async tasks for summarization
-    tasks = [
-        process_row(row) for row in entities_df.itertuples(index=False, name="Entity")
-    ]
-    results = await asyncio.gather(*tasks)
-
-    # Update the 'description' column in the DataFrame
-    entities_df["description"] = results
-
-    return entities_df

graphrag/index/update/incremental_index.py

@@ -18,9 +18,9 @@
 )
 from graphrag.index.update.entities import (
     _group_and_resolve_entities,
-    _run_entity_summarization,
 )
 from graphrag.index.update.relationships import _update_and_merge_relationships
+from graphrag.index.workflows.extract_graph import get_summarized_entities_relationships
 from graphrag.index.workflows.generate_text_embeddings import generate_text_embeddings
 from graphrag.logger.print_progress import ProgressLogger
 from graphrag.storage.pipeline_storage import PipelineStorage
@@ -104,18 +104,16 @@ async def update_dataframe_outputs(
         "documents", previous_storage, delta_storage, output_storage
     )
 
-    # Update entities and merge them
-    progress_logger.info("Updating Entities")
-    merged_entities_df, entity_id_mapping = await _update_entities(
+    # Update entities, relationships and merge them
+    progress_logger.info("Updating Entities and Relationships")
+    (
+        merged_entities_df,
+        merged_relationships_df,
+        entity_id_mapping,
+    ) = await _update_entities_and_relationships(
         previous_storage, delta_storage, output_storage, config, cache, callbacks
     )
 
-    # Update relationships with the entities id mapping
-    progress_logger.info("Updating Relationships")
-    merged_relationships_df = await _update_relationships(
-        previous_storage, delta_storage, output_storage
-    )
-
     # Update and merge final text units
     progress_logger.info("Updating Text Units")
     merged_text_units = await _update_text_units(
@@ -166,8 +164,11 @@ async def update_dataframe_outputs(
 
 
 async def _update_community_reports(
-    previous_storage, delta_storage, output_storage, community_id_mapping
-):
+    previous_storage: PipelineStorage,
+    delta_storage: PipelineStorage,
+    output_storage: PipelineStorage,
+    community_id_mapping: dict,
+) -> pd.DataFrame:
     """Update the community reports output."""
     old_community_reports = await load_table_from_storage(
         "community_reports", previous_storage
@@ -186,7 +187,11 @@ async def _update_community_reports(
     return merged_community_reports
 
 
-async def _update_communities(previous_storage, delta_storage, output_storage):
+async def _update_communities(
+    previous_storage: PipelineStorage,
+    delta_storage: PipelineStorage,
+    output_storage: PipelineStorage,
+) -> dict:
     """Update the communities output."""
     old_communities = await load_table_from_storage("communities", previous_storage)
     delta_communities = await load_table_from_storage("communities", delta_storage)
@@ -199,7 +204,11 @@ async def _update_communities(previous_storage, delta_storage, output_storage):
     return community_id_mapping
 
 
-async def _update_covariates(previous_storage, delta_storage, output_storage):
+async def _update_covariates(
+    previous_storage: PipelineStorage,
+    delta_storage: PipelineStorage,
+    output_storage: PipelineStorage,
+) -> None:
     """Update the covariates output."""
     old_covariates = await load_table_from_storage("covariates", previous_storage)
     delta_covariates = await load_table_from_storage("covariates", delta_storage)
@@ -209,8 +218,11 @@ async def _update_covariates(previous_storage, delta_storage, output_storage):
 
 
 async def _update_text_units(
-    previous_storage, delta_storage, output_storage, entity_id_mapping
-):
+    previous_storage: PipelineStorage,
+    delta_storage: PipelineStorage,
+    output_storage: PipelineStorage,
+    entity_id_mapping: dict,
+) -> pd.DataFrame:
     """Update the text units output."""
     old_text_units = await load_table_from_storage("text_units", previous_storage)
     delta_text_units = await load_table_from_storage("text_units", delta_storage)
@@ -223,48 +235,65 @@ async def _update_text_units(
     return merged_text_units
 
 
-async def _update_relationships(previous_storage, delta_storage, output_storage):
-    """Update the relationships output."""
+async def _update_entities_and_relationships(
+    previous_storage: PipelineStorage,
+    delta_storage: PipelineStorage,
+    output_storage: PipelineStorage,
+    config: GraphRagConfig,
+    cache: PipelineCache,
+    callbacks: WorkflowCallbacks,
+) -> tuple[pd.DataFrame, pd.DataFrame, dict]:
+    """Update Final Entities  and Relationships output."""
+    old_entities = await load_table_from_storage("entities", previous_storage)
+    delta_entities = await load_table_from_storage("entities", delta_storage)
+
+    merged_entities_df, entity_id_mapping = _group_and_resolve_entities(
+        old_entities, delta_entities
+    )
+
+    # Update Relationships
     old_relationships = await load_table_from_storage("relationships", previous_storage)
     delta_relationships = await load_table_from_storage("relationships", delta_storage)
     merged_relationships_df = _update_and_merge_relationships(
         old_relationships,
         delta_relationships,
     )
 
-    await write_table_to_storage(
-        merged_relationships_df, "relationships", output_storage
+    summarization_llm_settings = config.get_language_model_config(
+        config.summarize_descriptions.model_id
     )
-
-    return merged_relationships_df
-
-
-async def _update_entities(
-    previous_storage, delta_storage, output_storage, config, cache, callbacks
-):
-    """Update Final Entities output."""
-    old_entities = await load_table_from_storage("entities", previous_storage)
-    delta_entities = await load_table_from_storage("entities", delta_storage)
-
-    merged_entities_df, entity_id_mapping = _group_and_resolve_entities(
-        old_entities, delta_entities
+    summarization_strategy = config.summarize_descriptions.resolved_strategy(
+        config.root_dir, summarization_llm_settings
     )
 
-    # Re-run description summarization
-    merged_entities_df = await _run_entity_summarization(
+    (
         merged_entities_df,
-        config,
-        cache,
-        callbacks,
+        merged_relationships_df,
+    ) = await get_summarized_entities_relationships(
+        extracted_entities=merged_entities_df,
+        extracted_relationships=merged_relationships_df,
+        callbacks=callbacks,
+        cache=cache,
+        summarization_strategy=summarization_strategy,
+        summarization_num_threads=summarization_llm_settings.concurrent_requests,
     )
 
     # Save the updated entities back to storage
     await write_table_to_storage(merged_entities_df, "entities", output_storage)
 
-    return merged_entities_df, entity_id_mapping
+    await write_table_to_storage(
+        merged_relationships_df, "relationships", output_storage
+    )
+
+    return merged_entities_df, merged_relationships_df, entity_id_mapping
 
 
-async def _concat_dataframes(name, previous_storage, delta_storage, output_storage):
+async def _concat_dataframes(
+    name: str,
+    previous_storage: PipelineStorage,
+    delta_storage: PipelineStorage,
+    output_storage: PipelineStorage,
+) -> pd.DataFrame:
     """Concatenate dataframes."""
     old_df = await load_table_from_storage(name, previous_storage)
     delta_df = await load_table_from_storage(name, delta_storage)

graphrag/index/update/relationships.py

@@ -3,6 +3,8 @@
 
 """Relationship related operations and utils for Incremental Indexing."""
 
+import itertools
+
 import numpy as np
 import pandas as pd
 
@@ -42,10 +44,28 @@ def _update_and_merge_relationships(
     )
 
     # Merge the DataFrames without copying if possible
-    final_relationships = pd.concat(
+    merged_relationships = pd.concat(
         [old_relationships, delta_relationships], ignore_index=True, copy=False
     )
 
+    # Group by title and resolve conflicts
+    aggregated = (
+        merged_relationships.groupby(["source", "target"])
+        .agg({
+            "id": "first",
+            "human_readable_id": "first",
+            "description": lambda x: list(x.astype(str)),  # Ensure str
+            # Concatenate nd.array into a single list
+            "text_unit_ids": lambda x: list(itertools.chain(*x.tolist())),
+            "weight": "mean",
+            "combined_degree": "sum",
+        })
+        .reset_index()
+    )
+
+    # Force the result into a DataFrame
+    final_relationships: pd.DataFrame = pd.DataFrame(aggregated)
+
     # Recalculate target and source degrees
     final_relationships["source_degree"] = final_relationships.groupby("source")[
         "target"

graphrag/index/validate_config.py

@@ -40,6 +40,8 @@ def validate_config_names(logger: ProgressLogger, parameters: GraphRagConfig) ->
     embedding_llm_settings = parameters.get_language_model_config(
         parameters.embed_text.model_id
     )
+    if embedding_llm_settings.max_retries == -1:
+        embedding_llm_settings.max_retries = language_model_defaults.max_retries
     embed_llm = ModelManager().register_embedding(
         name="test-embed-llm",
         model_type=embedding_llm_settings.type,

graphrag/index/workflows/extract_graph.py

@@ -105,6 +105,27 @@ async def extract_graph(
         callbacks.error(error_msg)
         raise ValueError(error_msg)
 
+    entities, relationships = await get_summarized_entities_relationships(
+        extracted_entities=extracted_entities,
+        extracted_relationships=extracted_relationships,
+        callbacks=callbacks,
+        cache=cache,
+        summarization_strategy=summarization_strategy,
+        summarization_num_threads=summarization_num_threads,
+    )
+
+    return (entities, relationships)
+
+
+async def get_summarized_entities_relationships(
+    extracted_entities: pd.DataFrame,
+    extracted_relationships: pd.DataFrame,
+    callbacks: WorkflowCallbacks,
+    cache: PipelineCache,
+    summarization_strategy: dict[str, Any] | None = None,
+    summarization_num_threads: int = 4,
+) -> tuple[pd.DataFrame, pd.DataFrame]:
+    """Summarize the entities and relationships."""
     entity_summaries, relationship_summaries = await summarize_descriptions(
         entities_df=extracted_entities,
         relationships_df=extracted_relationships,
@@ -120,8 +141,7 @@ async def extract_graph(
 
     extracted_entities.drop(columns=["description"], inplace=True)
     entities = extracted_entities.merge(entity_summaries, on="title", how="left")
-
-    return (entities, relationships)
+    return entities, relationships
 
 
 def _validate_data(df: pd.DataFrame) -> bool: