graphrag/tests/verbs/test_create_base_text_units.py
Nathan Evans 8fd7730067
Some checks failed
Python Build and Type Check / python-ci (ubuntu-latest, 3.11) (push) Has been cancelled
Python Build and Type Check / python-ci (ubuntu-latest, 3.13) (push) Has been cancelled
Python Build and Type Check / python-ci (windows-latest, 3.11) (push) Has been cancelled
Python Build and Type Check / python-ci (windows-latest, 3.13) (push) Has been cancelled
Python Integration Tests / python-ci (ubuntu-latest, 3.13) (push) Has been cancelled
Python Integration Tests / python-ci (windows-latest, 3.13) (push) Has been cancelled
Python Notebook Tests / python-ci (ubuntu-latest, 3.13) (push) Has been cancelled
Python Notebook Tests / python-ci (windows-latest, 3.13) (push) Has been cancelled
Python Smoke Tests / python-ci (ubuntu-latest, 3.13) (push) Has been cancelled
Python Smoke Tests / python-ci (windows-latest, 3.13) (push) Has been cancelled
Python Unit Tests / python-ci (ubuntu-latest, 3.13) (push) Has been cancelled
Python Unit Tests / python-ci (windows-latest, 3.13) (push) Has been cancelled
Chunker factory (#2156)
* Delete NoopTextSplitter

* Delete unused check_token_limit

* Add base chunking factory and migrate workflow to use it

* Split apart chunker module

* Co-locate chunking/splitting

* Collapse token splitting functionality into one class/function

* Restore create_base_text_units parameterization

* Move Tokenizer base class to common package

* Move pre-pending into chunkers

* Streamline config

* Fix defaults construction

* Add prepending tests

* Remove chunk_size_includes_metadata config

* Revert ChunkingDocument interface

* Move metadata prepending to a util

* Move Tokenizer back to GR core

* Fix tokenizer removal from chunker

* Set defaults for chunking config

* Move chunking to monorepo package

* Format

* Typo

* Add ChunkResult model

* Streamline chunking config

* Add missing version updates for graphrag_chunking
2026-01-06 15:39:44 -08:00

46 lines
1.4 KiB
Python

# Copyright (c) 2024 Microsoft Corporation.
# Licensed under the MIT License
from graphrag.config.models.graph_rag_config import GraphRagConfig
from graphrag.index.workflows.create_base_text_units import run_workflow
from graphrag.utils.storage import load_table_from_storage
from .util import (
DEFAULT_MODEL_CONFIG,
compare_outputs,
create_test_context,
load_test_table,
update_document_metadata,
)
async def test_create_base_text_units():
expected = load_test_table("text_units")
context = await create_test_context()
config = GraphRagConfig(models=DEFAULT_MODEL_CONFIG) # type: ignore
await run_workflow(config, context)
actual = await load_table_from_storage("text_units", context.output_storage)
compare_outputs(actual, expected, columns=["text", "document_id", "n_tokens"])
async def test_create_base_text_units_metadata():
expected = load_test_table("text_units_metadata")
context = await create_test_context()
config = GraphRagConfig(models=DEFAULT_MODEL_CONFIG) # type: ignore
config.input.metadata = ["title"]
config.chunking.prepend_metadata = True
await update_document_metadata(config.input.metadata, context)
await run_workflow(config, context)
actual = await load_table_from_storage("text_units", context.output_storage)
compare_outputs(actual, expected, ["text", "document_id", "n_tokens"])