From 88e4d6bddb4f49f36d40bf5b0c0026ac745f659f Mon Sep 17 00:00:00 2001 From: cleanjunc Date: Tue, 9 Jun 2026 15:50:45 +0300 Subject: [PATCH] Fix: restore GraphRAG entity ranking by indexing pagerank and n-hop paths (#15797) ### Summary Closes #15795 Knowledge-graph queries rank entities by `pagerank * sim` in `KGSearch`, but the entity chunks written at index time stopped carrying the values that ranking depends on. `graph_node_to_chunk` only stored `entity_type`, `description`, and `source_id`, dropping the node `pagerank` and the n-hop neighbour paths, while `search.py` still read them back as `rank_flt` and `n_hop_with_weight`. The producer of these fields, `update_nodes_pagerank_nhop_neighbour`, was removed in #6513, but the read side in `KGSearch` was never updated. The result is that on every knowledge-graph query: - `pagerank` resolves to `0`, so the `pagerank * sim` sort key is `0` for every entity and selection falls back to arbitrary order. - Every displayed entity score is `0.00`. - The n-hop relation-enrichment block is dead code because `n_hop_ents` is always empty, leaving `merge_tuples` and `is_continuous_subsequence` orphaned. This PR restores the missing index-time fields so the documented `P(E|Q) = pagerank * sim` ranking and the n-hop enrichment work again. What changed: - `graph_node_to_chunk` now writes `rank_flt` from the node pagerank and `n_hop_with_weight` from the recomputed n-hop neighbour paths. - Reintroduced the n-hop path computation (`n_neighbor`) in `rag/graphrag/utils.py`, reusing the previously orphaned `merge_tuples` / `is_continuous_subsequence` helpers, with a direction-agnostic edge-weight lookup for undirected graphs. `set_graph` computes the paths per added or updated node and passes them through. - `KGSearch` now selects `n_hop_with_weight` in the entity keyword search so Infinity and OceanBase return it (Elasticsearch and OpenSearch already read it from `_source`), and the read is hardened against missing keys or empty strings before `json.loads`. - Added the `n_hop_with_weight` column to OceanBase, including the `EXTRA_COLUMNS` migration entry so existing tables get it. The other engines already map both fields via dynamic templates or the Infinity mapping. Scope note: pagerank and n-hop are re-indexed for the added or updated nodes in each pass, consistent with the existing incremental indexing design. ### Type of change - [x] Bug Fix (non-breaking change which fixes an issue) ### Testing Added unit tests in `test/unit_test/rag/graphrag/test_graphrag_utils.py`: - `n_neighbor`: path and weight shape, one-hop vs two-hop, isolated nodes, missing weights, and direction-agnostic lookup. - `graph_node_to_chunk`: `rank_flt` populated from pagerank and defaulting to `0`, `n_hop_with_weight` serialized and defaulting to an empty list. ``` uv run pytest test/unit_test/rag/graphrag/ # 106 passed uv run ruff check rag/graphrag/ rag/utils/ob_conn.py ``` --- rag/graphrag/search.py | 12 +- rag/graphrag/utils.py | 46 +++++++- rag/utils/ob_conn.py | 4 + .../rag/graphrag/test_graphrag_utils.py | 108 ++++++++++++++++++ 4 files changed, 166 insertions(+), 4 deletions(-) diff --git a/rag/graphrag/search.py b/rag/graphrag/search.py index 0c8a37a56..0cbb884f2 100644 --- a/rag/graphrag/search.py +++ b/rag/graphrag/search.py @@ -78,10 +78,18 @@ class KGSearch(Dealer): continue if isinstance(ent["entity_kwd"], list): ent["entity_kwd"] = ent["entity_kwd"][0] + # n_hop_with_weight may be absent (older chunks) or an empty string + # (the Infinity column default), neither of which json.loads handles. + n_hop_raw = ent.get("n_hop_with_weight") or "[]" + try: + n_hop_ents = json.loads(n_hop_raw) + except (json.JSONDecodeError, TypeError): + logging.warning(f"Failed to parse n_hop_with_weight for entity {ent.get('entity_kwd')}: {n_hop_raw}") + n_hop_ents = [] res[ent["entity_kwd"]] = { "sim": get_float(ent.get("_score", 0)), "pagerank": get_float(ent.get("rank_flt", 0)), - "n_hop_ents": json.loads(ent.get("n_hop_with_weight", "[]")), + "n_hop_ents": n_hop_ents, "description": ent.get("content_with_weight", "{}") } return res @@ -111,7 +119,7 @@ class KGSearch(Dealer): filters = deepcopy(filters) filters["knowledge_graph_kwd"] = "entity" matchDense = self.get_vector(", ".join(keywords), emb_mdl, 1024, sim_thr) - es_res = self.dataStore.search(["content_with_weight", "entity_kwd", "rank_flt"], [], filters, [matchDense], + es_res = self.dataStore.search(["content_with_weight", "entity_kwd", "rank_flt", "n_hop_with_weight"], [], filters, [matchDense], OrderByExpr(), 0, N, idxnms, kb_ids) return self._ent_info_from_(es_res, sim_thr) diff --git a/rag/graphrag/utils.py b/rag/graphrag/utils.py index c92fe0c6f..038646199 100644 --- a/rag/graphrag/utils.py +++ b/rag/graphrag/utils.py @@ -18,6 +18,7 @@ import os import re import time from collections import defaultdict +from copy import deepcopy from hashlib import md5 from typing import Any, Callable, Set, Tuple @@ -370,7 +371,7 @@ def chunk_id(chunk): return xxhash.xxh64((chunk["content_with_weight"] + chunk["kb_id"]).encode("utf-8")).hexdigest() -async def graph_node_to_chunk(kb_id, embd_mdl, ent_name, meta, chunks): +async def graph_node_to_chunk(kb_id, embd_mdl, ent_name, meta, chunks, nhop_neighbors=None): global chat_limiter enable_timeout_assertion = os.environ.get("ENABLE_TIMEOUT_ASSERTION") chunk = { @@ -383,6 +384,11 @@ async def graph_node_to_chunk(kb_id, embd_mdl, ent_name, meta, chunks): "content_with_weight": json.dumps(meta, ensure_ascii=False), "content_ltks": rag_tokenizer.tokenize(meta["description"]), "source_id": meta["source_id"], + # pagerank drives the P(E|Q) = pagerank * sim ranking in KGSearch; the + # n-hop neighbour paths feed its relation-enrichment step. Both are read + # back as `rank_flt` / `n_hop_with_weight` in rag/graphrag/search.py. + "rank_flt": float(meta.get("pagerank", 0) or 0), + "n_hop_with_weight": json.dumps(nhop_neighbors or [], ensure_ascii=False), "kb_id": kb_id, "available_int": 0, } @@ -549,8 +555,9 @@ async def set_graph(tenant_id: str, kb_id: str, embd_mdl, graph: nx.Graph, chang tasks = [] for ii, node in enumerate(change.added_updated_nodes): node_attrs = graph.nodes[node] + nhop_neighbors = n_neighbor(graph, node) tasks.append(asyncio.create_task( - graph_node_to_chunk(kb_id, embd_mdl, node, node_attrs, chunks) + graph_node_to_chunk(kb_id, embd_mdl, node, node_attrs, chunks, nhop_neighbors) )) if ii % 100 == 9 and callback: callback(msg=f"Get embedding of nodes: {ii}/{len(change.added_updated_nodes)}") @@ -697,6 +704,41 @@ def merge_tuples(list1, list2): return result +def n_neighbor(graph: nx.Graph, node, n_hop: int = 2): + """Enumerate paths of up to ``n_hop`` edges starting at ``node`` together + with the edge weight along each step. + + Returns a list of ``{"path": (n0, n1, ...), "weights": [w0, w1, ...]}`` + dicts (``len(weights) == len(path) - 1``). This is the structure consumed + by :class:`rag.graphrag.search.KGSearch` for n-hop relation enrichment and + is stored per entity chunk as ``n_hop_with_weight``. + """ + source_edge = list(graph.edges(node)) + if not source_edge: + return [] + count = 1 + while count < n_hop: + count += 1 + sc_edge = deepcopy(source_edge) + source_edge = [] + for pair in sc_edge: + append_edge = list(graph.edges(pair[-1])) + for tuples in merge_tuples([pair], append_edge): + source_edge.append(tuples) + wts = nx.get_edge_attributes(graph, "weight") + nbrs = [] + for path in source_edge: + nbr = {"path": path, "weights": []} + for i in range(len(path) - 1): + f, t = path[i], path[i + 1] + w = wts.get((f, t)) + if w is None: + w = wts.get((t, f), 0) + nbr["weights"].append(w) + nbrs.append(nbr) + return nbrs + + async def get_entity_type2samples(idxnms, kb_ids: list): es_res = await settings.retriever.search({"knowledge_graph_kwd": "ty2ents", "kb_id": kb_ids, "size": 10000, "fields": ["content_with_weight"]},idxnms,kb_ids) diff --git a/rag/utils/ob_conn.py b/rag/utils/ob_conn.py index fde2138f0..a51c923b6 100644 --- a/rag/utils/ob_conn.py +++ b/rag/utils/ob_conn.py @@ -48,6 +48,8 @@ column_mom_id = Column("mom_id", String(256), nullable=True, comment="parent chu column_chunk_data = Column("chunk_data", JSON, nullable=True, comment="table parser row data") column_raptor_kwd = Column("raptor_kwd", String(256), nullable=True, comment="RAPTOR summary marker") column_raptor_layer_int = Column("raptor_layer_int", Integer, nullable=True, comment="RAPTOR summary layer") +column_n_hop_with_weight = Column("n_hop_with_weight", LONGTEXT, nullable=True, + comment="JSON-encoded n-hop neighbour paths and weights for a graph entity") column_definitions: list[Column] = [ Column("id", String(256), primary_key=True, comment="chunk id"), @@ -86,6 +88,7 @@ column_definitions: list[Column] = [ Column("weight_flt", Double, nullable=True, comment="the weight of community report"), Column("entities_kwd", ARRAY(String(256)), nullable=True, comment="node ids of entities"), Column("rank_flt", Double, nullable=True, comment="rank of this entity"), + column_n_hop_with_weight, Column("removed_kwd", String(256), nullable=True, index=True, server_default="'N'", comment="whether it has been deleted"), column_raptor_kwd, @@ -138,6 +141,7 @@ EXTRA_COLUMNS: list[Column] = [ column_chunk_data, column_raptor_kwd, column_raptor_layer_int, + column_n_hop_with_weight, ] diff --git a/test/unit_test/rag/graphrag/test_graphrag_utils.py b/test/unit_test/rag/graphrag/test_graphrag_utils.py index 8a15d30d2..06ee172b8 100644 --- a/test/unit_test/rag/graphrag/test_graphrag_utils.py +++ b/test/unit_test/rag/graphrag/test_graphrag_utils.py @@ -14,9 +14,14 @@ # limitations under the License. # +import json +from types import SimpleNamespace + import networkx as nx +import numpy as np import pytest +import rag.graphrag.utils as graphrag_utils from rag.graphrag.utils import ( GRAPH_FIELD_SEP, GraphChange, @@ -31,6 +36,7 @@ from rag.graphrag.utils import ( is_continuous_subsequence, is_float_regex, merge_tuples, + n_neighbor, pack_user_ass_to_openai_messages, perform_variable_replacements, split_string_by_multi_markers, @@ -510,6 +516,108 @@ class TestMergeTuples: assert merge_tuples([], []) == [] +@pytest.mark.p1 +class TestNNeighbor: + """Tests for n_neighbor function (n-hop neighbour path enumeration). + + Regression coverage for the GraphRAG entity-ranking pipeline: the result + is serialized into each entity chunk as ``n_hop_with_weight`` and consumed + by KGSearch for n-hop relation enrichment. + """ + + def _line_graph(self): + # A -1- B -2- C -3- D + g = nx.Graph() + g.add_edge("A", "B", weight=1) + g.add_edge("B", "C", weight=2) + g.add_edge("C", "D", weight=3) + return g + + def test_isolated_node_returns_empty(self): + g = nx.Graph() + g.add_node("A") + assert n_neighbor(g, "A") == [] + + def test_result_shape(self): + nbrs = n_neighbor(self._line_graph(), "A") + assert isinstance(nbrs, list) + for nbr in nbrs: + assert set(nbr.keys()) == {"path", "weights"} + assert len(nbr["weights"]) == len(nbr["path"]) - 1 + + def test_two_hop_paths_and_weights(self): + # From A, 2-hop reaches the path A -> B -> C with weights [1, 2]. + nbrs = n_neighbor(self._line_graph(), "A", n_hop=2) + paths = {tuple(n["path"]): n["weights"] for n in nbrs} + assert ("A", "B", "C") in paths + assert paths[("A", "B", "C")] == [1, 2] + + def test_one_hop_only(self): + nbrs = n_neighbor(self._line_graph(), "A", n_hop=1) + paths = {tuple(n["path"]) for n in nbrs} + assert paths == {("A", "B")} + + def test_missing_weight_defaults_to_zero(self): + g = nx.Graph() + g.add_edge("A", "B") # no weight attribute + nbrs = n_neighbor(g, "A", n_hop=1) + assert nbrs[0]["weights"] == [0] + + def test_weight_lookup_is_direction_agnostic(self): + # Undirected graph: edge attributes may be keyed either way; the + # weight must still be recovered regardless of traversal direction. + nbrs = n_neighbor(self._line_graph(), "D", n_hop=1) + assert nbrs[0]["path"][0] == "D" + assert nbrs[0]["weights"] == [3] + + +@pytest.mark.p1 +class TestGraphNodeToChunk: + """Tests for graph_node_to_chunk field population. + + Regression coverage for the dropped ranking fields: the entity chunk must + carry ``rank_flt`` (pagerank) and ``n_hop_with_weight`` so KGSearch's + ``pagerank * sim`` ranking and n-hop enrichment are not permanently dead. + """ + + @pytest.fixture + def fake_embd(self, monkeypatch): + # Skip the real encode/Redis path by returning a cached embedding. + monkeypatch.setattr(graphrag_utils, "get_embed_cache", lambda *_a, **_k: np.array([0.1, 0.2, 0.3])) + return graphrag_utils + + @pytest.mark.asyncio + async def test_writes_rank_flt_from_pagerank(self, fake_embd): + chunks = [] + meta = {"entity_type": "PERSON", "description": "desc", "source_id": ["s1"], "pagerank": 0.42} + await fake_embd.graph_node_to_chunk("kb1", SimpleNamespace(llm_name="m"), "ALICE", meta, chunks) + assert len(chunks) == 1 + assert chunks[0]["rank_flt"] == pytest.approx(0.42) + + @pytest.mark.asyncio + async def test_rank_flt_defaults_to_zero_without_pagerank(self, fake_embd): + chunks = [] + meta = {"entity_type": "PERSON", "description": "desc", "source_id": ["s1"]} + await fake_embd.graph_node_to_chunk("kb1", SimpleNamespace(llm_name="m"), "ALICE", meta, chunks) + assert chunks[0]["rank_flt"] == 0.0 + + @pytest.mark.asyncio + async def test_writes_n_hop_with_weight(self, fake_embd): + chunks = [] + meta = {"entity_type": "PERSON", "description": "desc", "source_id": ["s1"], "pagerank": 0.1} + nhop = [{"path": ("ALICE", "BOB"), "weights": [3]}] + await fake_embd.graph_node_to_chunk("kb1", SimpleNamespace(llm_name="m"), "ALICE", meta, chunks, nhop) + stored = json.loads(chunks[0]["n_hop_with_weight"]) + assert stored == [{"path": ["ALICE", "BOB"], "weights": [3]}] + + @pytest.mark.asyncio + async def test_n_hop_defaults_to_empty_list(self, fake_embd): + chunks = [] + meta = {"entity_type": "PERSON", "description": "desc", "source_id": ["s1"]} + await fake_embd.graph_node_to_chunk("kb1", SimpleNamespace(llm_name="m"), "ALICE", meta, chunks) + assert json.loads(chunks[0]["n_hop_with_weight"]) == [] + + class TestFlatUniqList: """Tests for flat_uniq_list function."""