Add more verb tests (#1773)

* Add NLP verb test

* Add finalize_graph tests

* Add more thorough final column assertions

Author: natoverse

tests/verbs/test_create_communities.py

@@ -3,6 +3,7 @@
 
 from graphrag.callbacks.noop_workflow_callbacks import NoopWorkflowCallbacks
 from graphrag.config.create_graphrag_config import create_graphrag_config
+from graphrag.data_model.schemas import COMMUNITIES_FINAL_COLUMNS
 from graphrag.index.workflows.create_communities import (
     run_workflow,
 )
@@ -36,11 +37,14 @@ async def test_create_communities():
 
     actual = await load_table_from_storage("communities", context.storage)
 
-    assert "period" in expected.columns
     columns = list(expected.columns.values)
+    # don't compare period since it is created with the current date each time
     columns.remove("period")
     compare_outputs(
         actual,
         expected,
         columns=columns,
     )
+
+    for column in COMMUNITIES_FINAL_COLUMNS:
+        assert column in actual.columns

tests/verbs/test_create_community_reports.py

@@ -5,6 +5,7 @@
 from graphrag.callbacks.noop_workflow_callbacks import NoopWorkflowCallbacks
 from graphrag.config.create_graphrag_config import create_graphrag_config
 from graphrag.config.enums import ModelType
+from graphrag.data_model.schemas import COMMUNITY_REPORTS_FINAL_COLUMNS
 from graphrag.index.operations.summarize_communities.community_reports_extractor import (
     CommunityReportResponse,
     FindingModel,
@@ -80,3 +81,6 @@ async def test_create_community_reports():
     # assert a handful of mock data items to confirm they get put in the right spot
     assert actual["rank"][:1][0] == 2
     assert actual["rating_explanation"][:1][0] == "<rating_explanation>"
+
+    for column in COMMUNITY_REPORTS_FINAL_COLUMNS:
+        assert column in actual.columns

tests/verbs/test_create_final_documents.py

@@ -3,6 +3,7 @@
 
 from graphrag.callbacks.noop_workflow_callbacks import NoopWorkflowCallbacks
 from graphrag.config.create_graphrag_config import create_graphrag_config
+from graphrag.data_model.schemas import DOCUMENTS_FINAL_COLUMNS
 from graphrag.index.workflows.create_final_documents import (
     run_workflow,
 )
@@ -36,6 +37,9 @@ async def test_create_final_documents():
 
     compare_outputs(actual, expected)
 
+    for column in DOCUMENTS_FINAL_COLUMNS:
+        assert column in actual.columns
+
 
 async def test_create_final_documents_with_metadata_column():
     context = await create_test_context(
@@ -58,12 +62,7 @@ async def test_create_final_documents_with_metadata_column():
 
     actual = await load_table_from_storage("documents", context.storage)
 
-    # our test dataframe does not have metadata, so we'll assert without it
-    # and separately confirm it is in the output
-    compare_outputs(
-        actual, expected, columns=["id", "human_readable_id", "text", "metadata"]
-    )
-    assert len(actual.columns) == 7
-    assert "title" in actual.columns
-    assert "text_unit_ids" in actual.columns
-    assert "metadata" in actual.columns
+    compare_outputs(actual, expected)
+
+    for column in DOCUMENTS_FINAL_COLUMNS:
+        assert column in actual.columns

tests/verbs/test_create_final_text_units.py

@@ -3,6 +3,7 @@
 
 from graphrag.callbacks.noop_workflow_callbacks import NoopWorkflowCallbacks
 from graphrag.config.create_graphrag_config import create_graphrag_config
+from graphrag.data_model.schemas import TEXT_UNITS_FINAL_COLUMNS
 from graphrag.index.workflows.create_final_text_units import (
     run_workflow,
 )
@@ -39,37 +40,7 @@ async def test_create_final_text_units():
 
     actual = await load_table_from_storage("text_units", context.storage)
 
-    compare_outputs(actual, expected)
-
-
-async def test_create_final_text_units_no_covariates():
-    expected = load_test_table("text_units")
-
-    context = await create_test_context(
-        storage=[
-            "text_units",
-            "entities",
-            "relationships",
-            "covariates",
-        ],
-    )
+    for column in TEXT_UNITS_FINAL_COLUMNS:
+        assert column in actual.columns
 
-    config = create_graphrag_config({"models": DEFAULT_MODEL_CONFIG})
-    config.extract_claims.enabled = False
-
-    await run_workflow(
-        config,
-        context,
-        NoopWorkflowCallbacks(),
-    )
-
-    actual = await load_table_from_storage("text_units", context.storage)
-
-    # we're short a covariate_ids column
-    columns = list(expected.columns.values)
-    columns.remove("covariate_ids")
-    compare_outputs(
-        actual,
-        expected,
-        columns=columns,
-    )
+    compare_outputs(actual, expected)

tests/verbs/test_extract_covariates.py

@@ -6,6 +6,7 @@
 from graphrag.callbacks.noop_workflow_callbacks import NoopWorkflowCallbacks
 from graphrag.config.create_graphrag_config import create_graphrag_config
 from graphrag.config.enums import ModelType
+from graphrag.data_model.schemas import COVARIATES_FINAL_COLUMNS
 from graphrag.index.workflows.extract_covariates import (
     run_workflow,
 )
@@ -26,7 +27,6 @@
 
 async def test_extract_covariates():
     input = load_test_table("text_units")
-    expected = load_test_table("covariates")
 
     context = await create_test_context(
         storage=["text_units"],
@@ -52,7 +52,9 @@ async def test_extract_covariates():
 
     actual = await load_table_from_storage("covariates", context.storage)
 
-    assert len(actual.columns) == len(expected.columns)
+    for column in COVARIATES_FINAL_COLUMNS:
+        assert column in actual.columns
+
     # our mock only returns one covariate per text unit, so that's a 1:1 mapping versus the LLM-extracted content in the test data
     assert len(actual) == len(input)
 

tests/verbs/test_extract_graph_nlp.py

@@ -0,0 +1,38 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+from graphrag.callbacks.noop_workflow_callbacks import NoopWorkflowCallbacks
+from graphrag.config.create_graphrag_config import create_graphrag_config
+from graphrag.index.workflows.extract_graph_nlp import (
+    run_workflow,
+)
+from graphrag.utils.storage import load_table_from_storage
+
+from .util import (
+    DEFAULT_MODEL_CONFIG,
+    create_test_context,
+)
+
+
+async def test_extract_graph_nlp():
+    context = await create_test_context(
+        storage=["text_units"],
+    )
+
+    config = create_graphrag_config({"models": DEFAULT_MODEL_CONFIG})
+
+    await run_workflow(
+        config,
+        context,
+        NoopWorkflowCallbacks(),
+    )
+
+    nodes_actual = await load_table_from_storage("entities", context.storage)
+    edges_actual = await load_table_from_storage("relationships", context.storage)
+
+    # this will be the raw count of entities and edges with no pruning
+    # with NLP it is deterministic, so we can assert exact row counts
+    assert len(nodes_actual) == 1148
+    assert len(nodes_actual.columns) == 5
+    assert len(edges_actual) == 29445
+    assert len(edges_actual.columns) == 5

tests/verbs/test_finalize_graph.py

@@ -0,0 +1,91 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+from graphrag.callbacks.noop_workflow_callbacks import NoopWorkflowCallbacks
+from graphrag.config.create_graphrag_config import create_graphrag_config
+from graphrag.data_model.schemas import (
+    ENTITIES_FINAL_COLUMNS,
+    RELATIONSHIPS_FINAL_COLUMNS,
+)
+from graphrag.index.workflows.finalize_graph import (
+    run_workflow,
+)
+from graphrag.utils.storage import load_table_from_storage, write_table_to_storage
+
+from .util import (
+    DEFAULT_MODEL_CONFIG,
+    create_test_context,
+    load_test_table,
+)
+
+
+async def test_finalize_graph():
+    context = await _prep_tables()
+
+    config = create_graphrag_config({"models": DEFAULT_MODEL_CONFIG})
+
+    await run_workflow(
+        config,
+        context,
+        NoopWorkflowCallbacks(),
+    )
+
+    nodes_actual = await load_table_from_storage("entities", context.storage)
+    edges_actual = await load_table_from_storage("relationships", context.storage)
+
+    assert len(nodes_actual) == 251
+    assert len(edges_actual) == 372
+
+    # x and y will be zero with the default configuration, because we do not embed/umap
+    assert nodes_actual["x"].sum() == 0
+    assert nodes_actual["y"].sum() == 0
+
+    for column in ENTITIES_FINAL_COLUMNS:
+        assert column in nodes_actual.columns
+    for column in RELATIONSHIPS_FINAL_COLUMNS:
+        assert column in edges_actual.columns
+
+
+async def test_finalize_graph_umap():
+    context = await _prep_tables()
+
+    config = create_graphrag_config({"models": DEFAULT_MODEL_CONFIG})
+
+    config.embed_graph.enabled = True
+    config.umap.enabled = True
+
+    await run_workflow(
+        config,
+        context,
+        NoopWorkflowCallbacks(),
+    )
+
+    nodes_actual = await load_table_from_storage("entities", context.storage)
+    edges_actual = await load_table_from_storage("relationships", context.storage)
+
+    assert len(nodes_actual) == 251
+    assert len(edges_actual) == 372
+
+    # x and y should have some value other than zero due to umap
+    assert nodes_actual["x"].sum() != 0
+    assert nodes_actual["y"].sum() != 0
+
+    for column in ENTITIES_FINAL_COLUMNS:
+        assert column in nodes_actual.columns
+    for column in RELATIONSHIPS_FINAL_COLUMNS:
+        assert column in edges_actual.columns
+
+
+async def _prep_tables():
+    context = await create_test_context(
+        storage=["entities", "relationships"],
+    )
+
+    # edit the tables to eliminate final fields that wouldn't be on the inputs
+    entities = load_test_table("entities")
+    entities.drop(columns=["x", "y", "degree"], inplace=True)
+    await write_table_to_storage(entities, "entities", context.storage)
+    relationships = load_test_table("relationships")
+    relationships.drop(columns=["combined_degree"], inplace=True)
+    await write_table_to_storage(relationships, "relationships", context.storage)
+    return context