wehub-resource-sync
6b7e6b44f1
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
gh-pages / build (push) Has been cancelled
Python Publish (pypi) / Upload release to PyPI (push) Has been cancelled
Spellcheck / spellcheck (push) Has been cancelled
302 lines
10 KiB
Python
302 lines
10 KiB
Python
# Copyright (C) 2026 Microsoft Corporation.
|
|
# Licensed under the MIT License
|
|
|
|
"""Tests for extract_graph merge and orphan-filtering operations.
|
|
|
|
Validates that _merge_entities, _merge_relationships, and
|
|
filter_orphan_relationships correctly aggregate per-text-unit
|
|
extraction results and remove relationships whose source or
|
|
target has no corresponding entity.
|
|
"""
|
|
|
|
import pandas as pd
|
|
from graphrag.index.operations.extract_graph.extract_graph import (
|
|
_merge_entities,
|
|
_merge_relationships,
|
|
)
|
|
from graphrag.index.operations.extract_graph.utils import (
|
|
filter_orphan_relationships,
|
|
)
|
|
|
|
|
|
def _entity_row(
|
|
title: str,
|
|
entity_type: str = "THING",
|
|
description: str = "desc",
|
|
source_id: str = "tu1",
|
|
) -> dict:
|
|
"""Build a single raw entity row as produced by the graph extractor."""
|
|
return {
|
|
"title": title,
|
|
"type": entity_type,
|
|
"description": description,
|
|
"source_id": source_id,
|
|
}
|
|
|
|
|
|
def _relationship_row(
|
|
source: str,
|
|
target: str,
|
|
weight: float = 1.0,
|
|
description: str = "desc",
|
|
source_id: str = "tu1",
|
|
) -> dict:
|
|
"""Build a single raw relationship row as produced by the graph extractor."""
|
|
return {
|
|
"source": source,
|
|
"target": target,
|
|
"weight": weight,
|
|
"description": description,
|
|
"source_id": source_id,
|
|
}
|
|
|
|
|
|
class TestMergeEntities:
|
|
"""Tests for the _merge_entities aggregation helper."""
|
|
|
|
def test_groups_by_title_and_type(self):
|
|
"""Entities with the same title+type merge into one row."""
|
|
df1 = pd.DataFrame([_entity_row("A", "PERSON")])
|
|
df2 = pd.DataFrame([_entity_row("A", "PERSON", source_id="tu2")])
|
|
merged = _merge_entities([df1, df2])
|
|
|
|
assert len(merged) == 1
|
|
assert merged.iloc[0]["title"] == "A"
|
|
assert merged.iloc[0]["frequency"] == 2
|
|
|
|
def test_different_types_stay_separate(self):
|
|
"""Same title but different type should not merge."""
|
|
df = pd.DataFrame([
|
|
_entity_row("A", "PERSON"),
|
|
_entity_row("A", "ORG"),
|
|
])
|
|
merged = _merge_entities([df])
|
|
|
|
assert len(merged) == 2
|
|
|
|
def test_empty_input(self):
|
|
"""Empty entity list should produce an empty DataFrame."""
|
|
df = pd.DataFrame(columns=["title", "type", "description", "source_id"])
|
|
merged = _merge_entities([df])
|
|
|
|
assert len(merged) == 0
|
|
|
|
|
|
class TestMergeRelationships:
|
|
"""Tests for the _merge_relationships aggregation helper."""
|
|
|
|
def test_groups_by_source_target(self):
|
|
"""Relationships with same source+target merge and sum weight."""
|
|
df1 = pd.DataFrame([_relationship_row("A", "B", weight=2.0)])
|
|
df2 = pd.DataFrame([_relationship_row("A", "B", weight=3.0)])
|
|
merged = _merge_relationships([df1, df2])
|
|
|
|
assert len(merged) == 1
|
|
assert merged.iloc[0]["weight"] == 5.0
|
|
|
|
def test_distinct_pairs_stay_separate(self):
|
|
"""Different source-target pairs remain separate rows."""
|
|
df = pd.DataFrame([
|
|
_relationship_row("A", "B"),
|
|
_relationship_row("B", "C"),
|
|
])
|
|
merged = _merge_relationships([df])
|
|
|
|
assert len(merged) == 2
|
|
|
|
def test_empty_input(self):
|
|
"""Empty relationship list should produce an empty DataFrame."""
|
|
df = pd.DataFrame(
|
|
columns=["source", "target", "weight", "description", "source_id"]
|
|
)
|
|
merged = _merge_relationships([df])
|
|
|
|
assert len(merged) == 0
|
|
|
|
|
|
class TestFilterOrphanRelationships:
|
|
"""Tests for orphan relationship filtering.
|
|
|
|
After LLM graph extraction, relationships may reference entity
|
|
names that have no corresponding entity row. These must be
|
|
removed before downstream processing.
|
|
"""
|
|
|
|
def test_all_valid_relationships_kept(self):
|
|
"""Relationships whose endpoints all exist should be retained."""
|
|
entities = pd.DataFrame([
|
|
_entity_row("A"),
|
|
_entity_row("B"),
|
|
_entity_row("C"),
|
|
])
|
|
merged_entities = _merge_entities([entities])
|
|
|
|
relationships = pd.DataFrame([
|
|
_relationship_row("A", "B"),
|
|
_relationship_row("B", "C"),
|
|
])
|
|
merged_rels = _merge_relationships([relationships])
|
|
|
|
filtered = filter_orphan_relationships(merged_rels, merged_entities)
|
|
|
|
assert len(filtered) == 2
|
|
|
|
def test_removes_relationship_with_missing_source(self):
|
|
"""Relationship whose source has no entity entry is dropped."""
|
|
entities = pd.DataFrame([_entity_row("B")])
|
|
merged_entities = _merge_entities([entities])
|
|
|
|
relationships = pd.DataFrame([
|
|
_relationship_row("PHANTOM", "B"),
|
|
])
|
|
merged_rels = _merge_relationships([relationships])
|
|
|
|
filtered = filter_orphan_relationships(merged_rels, merged_entities)
|
|
|
|
assert len(filtered) == 0
|
|
|
|
def test_removes_relationship_with_missing_target(self):
|
|
"""Relationship whose target has no entity entry is dropped."""
|
|
entities = pd.DataFrame([_entity_row("A")])
|
|
merged_entities = _merge_entities([entities])
|
|
|
|
relationships = pd.DataFrame([
|
|
_relationship_row("A", "PHANTOM"),
|
|
])
|
|
merged_rels = _merge_relationships([relationships])
|
|
|
|
filtered = filter_orphan_relationships(merged_rels, merged_entities)
|
|
|
|
assert len(filtered) == 0
|
|
|
|
def test_removes_relationship_with_both_missing(self):
|
|
"""Relationship where both endpoints are missing is dropped."""
|
|
entities = pd.DataFrame([_entity_row("A")])
|
|
merged_entities = _merge_entities([entities])
|
|
|
|
relationships = pd.DataFrame([
|
|
_relationship_row("GHOST_1", "GHOST_2"),
|
|
])
|
|
merged_rels = _merge_relationships([relationships])
|
|
|
|
filtered = filter_orphan_relationships(merged_rels, merged_entities)
|
|
|
|
assert len(filtered) == 0
|
|
|
|
def test_keeps_valid_drops_orphan_mixed(self):
|
|
"""Valid and orphaned relationships coexist; only valid survive."""
|
|
entities = pd.DataFrame([
|
|
_entity_row("A"),
|
|
_entity_row("B"),
|
|
])
|
|
merged_entities = _merge_entities([entities])
|
|
|
|
relationships = pd.DataFrame([
|
|
_relationship_row("A", "B"),
|
|
_relationship_row("A", "PHANTOM"),
|
|
_relationship_row("PHANTOM", "B"),
|
|
_relationship_row("GHOST_1", "GHOST_2"),
|
|
])
|
|
merged_rels = _merge_relationships([relationships])
|
|
|
|
filtered = filter_orphan_relationships(merged_rels, merged_entities)
|
|
|
|
assert len(filtered) == 1
|
|
assert filtered.iloc[0]["source"] == "A"
|
|
assert filtered.iloc[0]["target"] == "B"
|
|
|
|
def test_empty_entities_drops_all_relationships(self):
|
|
"""If there are no entities, all relationships are orphaned."""
|
|
entities = pd.DataFrame(columns=["title", "type", "description", "source_id"])
|
|
merged_entities = _merge_entities([entities])
|
|
|
|
relationships = pd.DataFrame([
|
|
_relationship_row("A", "B"),
|
|
])
|
|
merged_rels = _merge_relationships([relationships])
|
|
|
|
filtered = filter_orphan_relationships(merged_rels, merged_entities)
|
|
|
|
assert len(filtered) == 0
|
|
|
|
def test_empty_relationships_returns_empty(self):
|
|
"""If there are no relationships, result is empty DataFrame."""
|
|
entities = pd.DataFrame([_entity_row("A")])
|
|
merged_entities = _merge_entities([entities])
|
|
|
|
relationships = pd.DataFrame(
|
|
columns=["source", "target", "weight", "description", "source_id"]
|
|
)
|
|
merged_rels = _merge_relationships([relationships])
|
|
|
|
filtered = filter_orphan_relationships(merged_rels, merged_entities)
|
|
|
|
assert len(filtered) == 0
|
|
|
|
def test_preserves_all_columns(self):
|
|
"""Filtered DataFrame retains all original columns."""
|
|
entities = pd.DataFrame([
|
|
_entity_row("A"),
|
|
_entity_row("B"),
|
|
])
|
|
merged_entities = _merge_entities([entities])
|
|
|
|
relationships = pd.DataFrame([
|
|
_relationship_row("A", "B", weight=5.0, description="linked"),
|
|
])
|
|
merged_rels = _merge_relationships([relationships])
|
|
|
|
filtered = filter_orphan_relationships(merged_rels, merged_entities)
|
|
|
|
assert set(filtered.columns) == set(merged_rels.columns)
|
|
assert filtered.iloc[0]["weight"] == 5.0
|
|
assert filtered.iloc[0]["description"] == ["linked"]
|
|
|
|
def test_multi_text_unit_orphan(self):
|
|
"""Orphan detected across multiple text units after merge."""
|
|
df1 = pd.DataFrame([
|
|
_entity_row("A", source_id="tu1"),
|
|
_relationship_row("A", "HALLUCINATED", source_id="tu1"),
|
|
])
|
|
df2 = pd.DataFrame([
|
|
_entity_row("A", source_id="tu2"),
|
|
_relationship_row("A", "HALLUCINATED", source_id="tu2"),
|
|
])
|
|
|
|
entity_dfs = [
|
|
df1[["title", "type", "description", "source_id"]],
|
|
df2[["title", "type", "description", "source_id"]],
|
|
]
|
|
rel_dfs = [
|
|
df1[["source", "target", "weight", "description", "source_id"]],
|
|
df2[["source", "target", "weight", "description", "source_id"]],
|
|
]
|
|
|
|
merged_entities = _merge_entities(entity_dfs)
|
|
merged_rels = _merge_relationships(rel_dfs)
|
|
|
|
filtered = filter_orphan_relationships(merged_rels, merged_entities)
|
|
|
|
assert len(filtered) == 0
|
|
|
|
def test_resets_index_after_filter(self):
|
|
"""Filtered DataFrame should have a clean 0-based index."""
|
|
entities = pd.DataFrame([
|
|
_entity_row("A"),
|
|
_entity_row("B"),
|
|
_entity_row("C"),
|
|
])
|
|
merged_entities = _merge_entities([entities])
|
|
|
|
relationships = pd.DataFrame([
|
|
_relationship_row("PHANTOM", "B"),
|
|
_relationship_row("A", "B"),
|
|
_relationship_row("A", "PHANTOM"),
|
|
_relationship_row("B", "C"),
|
|
])
|
|
merged_rels = _merge_relationships([relationships])
|
|
|
|
filtered = filter_orphan_relationships(merged_rels, merged_entities)
|
|
|
|
assert list(filtered.index) == list(range(len(filtered)))
|