add option to add metadata into text chunks (microsoft#1681)

* add new options

* add metadata json into input document

* remove doc change

* add metadata column into text loader

* prepend_metadata

* run fix

* fix tests and patch

* fix test

* add watrning for metadata tokens > config size

* fix typo and run fix

* fix test_integration

* fix test

* run check

* rename and fix chunking

* fix

* fix

* fiz test verbs

* fix

* fix tests

* fix chunking

* fix index

* fix cosmos test

* fix vars

* fix after PR

* fix

Author: dayesouza

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

@@ -0,0 +1,4 @@
+{
+  "type": "patch",
+  "description": "add option to prepend metadata into chunks"
+}

graphrag/config/defaults.py

@@ -67,6 +67,8 @@
 CHUNK_OVERLAP = 100
 CHUNK_GROUP_BY_COLUMNS = ["id"]
 CHUNK_STRATEGY = ChunkStrategyType.tokens
+CHUNK_PREPEND_METADATA = False
+CHUNK_SIZE_INCLUDES_METADATA = False
 
 # Claim extraction
 DESCRIPTION = "Any claims or facts that could be relevant to information discovery."

graphrag/config/models/chunking_config.py

@@ -26,3 +26,11 @@ class ChunkingConfig(BaseModel):
     encoding_model: str = Field(
         description="The encoding model to use.", default=defs.ENCODING_MODEL
     )
+    prepend_metadata: bool = Field(
+        description="Prepend metadata into each chunk.",
+        default=defs.CHUNK_PREPEND_METADATA,
+    )
+    chunk_size_includes_metadata: bool = Field(
+        description="Count metadata in max tokens.",
+        default=defs.CHUNK_SIZE_INCLUDES_METADATA,
+    )

graphrag/index/flows/create_base_text_units.py

@@ -3,13 +3,15 @@
 
 """All the steps to transform base text_units."""
 
-from typing import cast
+import json
+from typing import Any, cast
 
 import pandas as pd
 
 from graphrag.callbacks.workflow_callbacks import WorkflowCallbacks
 from graphrag.config.models.chunking_config import ChunkStrategyType
 from graphrag.index.operations.chunk_text.chunk_text import chunk_text
+from graphrag.index.operations.chunk_text.strategies import get_encoding_fn
 from graphrag.index.utils.hashing import gen_sha512_hash
 from graphrag.logger.progress import Progress
 
@@ -22,6 +24,8 @@ def create_base_text_units(
     overlap: int,
     encoding_model: str,
     strategy: ChunkStrategyType,
+    prepend_metadata: bool = False,
+    chunk_size_includes_metadata: bool = False,
 ) -> pd.DataFrame:
     """All the steps to transform base text_units."""
     sort = documents.sort_values(by=["id"], ascending=[True])
@@ -32,25 +36,66 @@ def create_base_text_units(
 
     callbacks.progress(Progress(percent=0))
 
+    agg_dict = {"text_with_ids": list}
+    if "metadata" in documents:
+        agg_dict["metadata"] = "first"  # type: ignore
+
     aggregated = (
         (
             sort.groupby(group_by_columns, sort=False)
             if len(group_by_columns) > 0
             else sort.groupby(lambda _x: True)
         )
-        .agg(texts=("text_with_ids", list))
+        .agg(agg_dict)
         .reset_index()
     )
+    aggregated.rename(columns={"text_with_ids": "texts"}, inplace=True)
 
-    aggregated["chunks"] = chunk_text(
-        aggregated,
-        column="texts",
-        size=size,
-        overlap=overlap,
-        encoding_model=encoding_model,
-        strategy=strategy,
-        callbacks=callbacks,
-    )
+    def chunker(row: dict[str, Any]) -> Any:
+        line_delimiter = ".\n"
+        metadata_str = ""
+        metadata_tokens = 0
+
+        if prepend_metadata and "metadata" in row:
+            metadata = row["metadata"]
+            if isinstance(metadata, str):
+                metadata = json.loads(metadata)
+            if isinstance(metadata, dict):
+                metadata_str = (
+                    line_delimiter.join(f"{k}: {v}" for k, v in metadata.items())
+                    + line_delimiter
+                )
+
+            if chunk_size_includes_metadata:
+                encode, _ = get_encoding_fn(encoding_model)
+                metadata_tokens = len(encode(metadata_str))
+                if metadata_tokens >= size:
+                    message = "Metadata tokens exceeds the maximum tokens per chunk. Please increase the tokens per chunk."
+                    raise ValueError(message)
+
+        chunked = chunk_text(
+            pd.DataFrame([row]).reset_index(drop=True),
+            column="texts",
+            size=size - metadata_tokens,
+            overlap=overlap,
+            encoding_model=encoding_model,
+            strategy=strategy,
+            callbacks=callbacks,
+        )[0]
+
+        if prepend_metadata:
+            for index, chunk in enumerate(chunked):
+                if isinstance(chunk, str):
+                    chunked[index] = metadata_str + chunk
+                else:
+                    chunked[index] = (
+                        (chunk[0], metadata_str + chunk[1], chunk[2]) if chunk else None
+                    )
+
+        row["chunks"] = chunked
+        return row
+
+    aggregated = aggregated.apply(lambda row: chunker(row), axis=1)
 
     aggregated = cast("pd.DataFrame", aggregated[[*group_by_columns, "chunks"]])
     aggregated = aggregated.explode("chunks")

graphrag/index/flows/create_final_documents.py

@@ -7,9 +7,7 @@
 
 
 def create_final_documents(
-    documents: pd.DataFrame,
-    text_units: pd.DataFrame,
-    metadata: list[str] | None = None,
+    documents: pd.DataFrame, text_units: pd.DataFrame
 ) -> pd.DataFrame:
     """All the steps to transform final documents."""
     exploded = (
@@ -46,27 +44,17 @@ def create_final_documents(
     rejoined["id"] = rejoined["id"].astype(str)
     rejoined["human_readable_id"] = rejoined.index + 1
 
-    # Convert metadata columns to strings and collapse them into a JSON object
-    if metadata:
-        # Convert all specified columns to string at once
-        rejoined[metadata] = rejoined[metadata].astype(str)
-
-        # Collapse the metadata columns into a single JSON object column
-        rejoined["metadata"] = rejoined[metadata].to_dict(orient="records")
-
-        # Drop the original metadata columns after collapsing them
-        rejoined.drop(columns=metadata, inplace=True)
-
     # set the final column order, but adjust for metadata
     core_columns = [
         "id",
         "human_readable_id",
         "title",
         "text",
         "text_unit_ids",
+        "creation_date",
     ]
     final_columns = [column for column in core_columns if column in rejoined.columns]
-    if metadata:
+    if "metadata" in rejoined.columns:
         final_columns.append("metadata")
 
     return rejoined.loc[:, final_columns]

graphrag/index/input/csv.py

@@ -50,7 +50,7 @@ async def load_file(path: str, group: dict | None) -> pd.DataFrame:
                 )
             else:
                 data["text"] = data.apply(lambda x: x[config.text_column], axis=1)
-        if config.title_column is not None and "title" not in data.columns:
+        if config.title_column is not None:
             if config.title_column not in data.columns:
                 log.warning(
                     "title_column %s not found in csv file %s",
@@ -59,6 +59,11 @@ async def load_file(path: str, group: dict | None) -> pd.DataFrame:
                 )
             else:
                 data["title"] = data.apply(lambda x: x[config.title_column], axis=1)
+        else:
+            data["title"] = data.apply(lambda _: path, axis=1)
+
+        creation_date = await storage.get_creation_date(path)
+        data["creation_date"] = data.apply(lambda _: creation_date, axis=1)
 
         return data
 

graphrag/index/input/factory.py

@@ -72,8 +72,23 @@ async def create_input(
             f"Loading Input ({config.file_type})", transient=False
         )
         loader = loaders[config.file_type]
-        results = await loader(config, progress, storage)
-        return cast("pd.DataFrame", results)
+        result = await loader(config, progress, storage)
+        # Convert metadata columns to strings and collapse them into a JSON object
+        if config.metadata:
+            if all(col in result.columns for col in config.metadata):
+                # Collapse the metadata columns into a single JSON object column
+                result["metadata"] = result[config.metadata].apply(
+                    lambda row: row.to_dict(), axis=1
+                )
+            else:
+                value_error_msg = (
+                    "One or more metadata columns not found in the DataFrame."
+                )
+                raise ValueError(value_error_msg)
+
+            result[config.metadata] = result[config.metadata].astype(str)
+
+        return cast("pd.DataFrame", result)
 
     msg = f"Unknown input type {config.file_type}"
     raise ValueError(msg)

graphrag/index/input/text.py

@@ -38,6 +38,7 @@ async def load_file(
         new_item = {**group, "text": text}
         new_item["id"] = gen_sha512_hash(new_item, new_item.keys())
         new_item["title"] = str(Path(path).name)
+        new_item["creation_date"] = await storage.get_creation_date(path)
         return new_item
 
     files = list(

graphrag/index/operations/chunk_text/chunk_text.py

@@ -58,14 +58,21 @@ def chunk_text(
 
     num_total = _get_num_total(input, column)
     tick = progress_ticker(callbacks.progress, num_total)
+
     # collapse the config back to a single object to support "polymorphic" function call
     config = ChunkingConfig(size=size, overlap=overlap, encoding_model=encoding_model)
+
     return cast(
         "pd.Series",
         input.apply(
             cast(
                 "Any",
-                lambda x: run_strategy(strategy_exec, x[column], config, tick),
+                lambda x: run_strategy(
+                    strategy_exec,
+                    x[column],
+                    config,
+                    tick,
+                ),
             ),
             axis=1,
         ),
@@ -85,12 +92,7 @@ def run_strategy(
     # We can work with both just a list of text content
     # or a list of tuples of (document_id, text content)
     # text_to_chunk = '''
-    texts = []
-    for item in input:
-        if isinstance(item, str):
-            texts.append(item)
-        else:
-            texts.append(item[1])
+    texts = [item if isinstance(item, str) else item[1] for item in input]
 
     strategy_results = strategy_exec(texts, config, tick)
 

graphrag/index/operations/chunk_text/strategies.py

@@ -17,13 +17,8 @@
 from graphrag.logger.progress import ProgressTicker
 
 
-def run_tokens(
-    input: list[str], config: ChunkingConfig, tick: ProgressTicker
-) -> Iterable[TextChunk]:
-    """Chunks text into chunks based on encoding tokens."""
-    tokens_per_chunk = config.size
-    chunk_overlap = config.overlap
-    encoding_name = config.encoding_model
+def get_encoding_fn(encoding_name):
+    """Get the encoding model."""
     enc = tiktoken.get_encoding(encoding_name)
 
     def encode(text: str) -> list[int]:
@@ -34,6 +29,20 @@ def encode(text: str) -> list[int]:
     def decode(tokens: list[int]) -> str:
         return enc.decode(tokens)
 
+    return encode, decode
+
+
+def run_tokens(
+    input: list[str],
+    config: ChunkingConfig,
+    tick: ProgressTicker,
+) -> Iterable[TextChunk]:
+    """Chunks text into chunks based on encoding tokens."""
+    tokens_per_chunk = config.size
+    chunk_overlap = config.overlap
+    encoding_name = config.encoding_model
+
+    encode, decode = get_encoding_fn(encoding_name)
     return split_multiple_texts_on_tokens(
         input,
         Tokenizer(