feat(graphrag): seed the walk from what entities are, and rank with passages and vector hits

Graph retrieval tied plain vector search at best and never beat it. Measured
across five corpora, the bottleneck was seeding, not the graph: the walk
started from nodes whose embeddings were computed from bare entity names, and
a whole question shares almost nothing with a name like "Quill".

Extraction now embeds each node from "name (type): description" and each
relationship as the fact it asserts ("Alder streams_to Quill: ..."), stored on
a new nullable graph_edges.fact_embedding column that ensure_vector_schema adds
in place. Entity names are canonicalised (case, punctuation, word breaks and a
cautious plural) so "VECTOR_STORE" and "vector stores" land on one node. Extraction calls run
concurrently (GRAPHRAG_EXTRACTION_WORKERS, default 8) while embedding and graph
writes stay serial on the task thread, so ordering and idempotency are
unchanged; that measured 8.4x faster with identical output.

Retrieval gains per-source options, stored under retrieval.graph and read live
at query time:

- seed_strategy: start from matching entities (default) or matching
  relationships, which can reach an entity the question never names;
- passage_nodes (on): walk the source's passages alongside entities, with
  PageRank damping 0.5 instead of 0.85;
- blend_vector (on): fuse the graph ranking with the source's vector ranking
  by reciprocal rank.

The defaults are the measured-best configuration. Through GraphRAGRetriever,
the new seeding moved recall@4 from 0.41 to 0.68 on a multi-hop corpus and
from 0.50 to 1.00 on the docs corpus, and regressed none of the corpora
measured. Existing graphs keep name-only embeddings until rebuilt.

docsgpt/retriever/graph_rag.py

@@ -32,6 +32,7 @@ from docsgpt.graphrag.store import GraphStore
 from docsgpt.retriever.base import BaseRetriever
 from docsgpt.retriever.classic_rag import ClassicRAG
 from docsgpt.retriever.labels import labels_from_metadata
+from docsgpt.storage.db.source_config import GraphRetrievalConfig
 from docsgpt.utils import num_tokens_from_string
 from docsgpt.vectorstore.base import get_embeddings
 
@@ -39,11 +40,27 @@ SEED_NODES = 10
 SUBGRAPH_HOPS = 1
 
 
+PASSAGE_NODE_WEIGHT = 0.05
+FACT_SEED_FACTS = 5
+RRF_K = 60
+
+# PageRank damping per ranking mode — each is the value that mode was measured
+# at. Lower keeps mass nearer the seeds; with passages in the walk 0.5 measured
+# better, while entity-only ranking was measured at the conventional 0.85.
+DAMPING_WITH_PASSAGES = 0.5
+DAMPING_ENTITIES_ONLY = 0.85
+
+
 def _idf(doc_freq: Any) -> float:
     """Node-specificity weight: rarer entities (low ``doc_freq``) score higher."""
     return 1.0 / math.log(1.0 + max(int(doc_freq or 0), 0) + 1.0)
 
 
+def _damping(passage_nodes: bool) -> float:
+    """PageRank damping for a ranking mode: the value that mode was measured at."""
+    return DAMPING_WITH_PASSAGES if passage_nodes else DAMPING_ENTITIES_ONLY
+
+
 def _restart_vector(nodes: List[Any], personalization: Dict[Any, float] | None) -> Dict[Any, float]:
     """Normalized restart distribution over ``nodes``.
 
@@ -210,18 +227,9 @@ class GraphRAGRetriever(BaseRetriever):
         After PPR, each node's mass is scaled by ``1/log(2 + doc_freq)`` so a
         high-degree hub contributes less than a specific entity at equal mass.
         """
-        graph = nx.Graph()
-        for node in subgraph.get("nodes", []):
-            graph.add_node(node["id"], doc_freq=node.get("doc_freq", 0))
-        for edge in subgraph.get("edges", []):
-            src, dst = edge["src_node_id"], edge["dst_node_id"]
-            if src in graph and dst in graph:
-                raw_weight = edge.get("weight")
-                # Same rule the ranker applies: default only a missing or null
-                # weight. Coercing an explicit 0 to 1.0 here would make "these
-                # entities are not related" the strongest possible link.
-                edge_weight = 1.0 if raw_weight is None else float(raw_weight)
-                graph.add_edge(src, dst, weight=edge_weight)
+        # Through the class, not ``self``: this method reads no instance state,
+        # and callers (and tests) rely on being able to invoke it unbound.
+        graph = GraphRAGRetriever._subgraph_graph(subgraph)
         if graph.number_of_nodes() == 0:
             return {}
 
@@ -230,13 +238,33 @@ class GraphRAGRetriever(BaseRetriever):
             personalization = None
 
         ranks = _personalized_pagerank(
-            graph, personalization=personalization, weight="weight"
+            graph,
+            personalization=personalization,
+            weight="weight",
+            alpha=_damping(passage_nodes=False),
         )
         return {
             node: rank * _idf(graph.nodes[node].get("doc_freq", 0))
             for node, rank in ranks.items()
         }
 
+    @staticmethod
+    def _subgraph_graph(subgraph) -> "nx.Graph":
+        """The fetched subgraph as a weighted undirected graph."""
+        graph = nx.Graph()
+        for node in subgraph.get("nodes", []):
+            graph.add_node(node["id"], doc_freq=node.get("doc_freq", 0))
+        for edge in subgraph.get("edges", []):
+            src, dst = edge["src_node_id"], edge["dst_node_id"]
+            if src in graph and dst in graph:
+                raw_weight = edge.get("weight")
+                # Default only a missing or null weight. Coercing an explicit 0
+                # to 1.0 would make "these entities are not related" the
+                # strongest possible link.
+                edge_weight = 1.0 if raw_weight is None else float(raw_weight)
+                graph.add_edge(src, dst, weight=edge_weight)
+        return graph
+
     def _rank_chunks(self, store, source_id, node_scores) -> List[str]:
         """Score chunks by summed (PPR mass x IDF) of their linked nodes; top candidates.
 
@@ -254,6 +282,74 @@ class GraphRAGRetriever(BaseRetriever):
         candidates = max(self.chunks * 2, self.chunks + 5)
         return ranked[: max(1, candidates)]
 
+    def _rank_chunks_with_passages(
+        self, store, source_id, subgraph, seeds, query_embedding
+    ) -> List[str]:
+        """Rank chunks by walking a graph that contains the chunks themselves.
+
+        :meth:`_rank_chunks` reads a chunk's score *off* its entities, summing
+        their PPR mass — so a chunk touching many mid-scoring generic entities
+        outranks one touching the few entities the question is about. Putting
+        the chunks in the walk instead, each joined to its own entities and
+        carrying a small share of the restart mass proportional to its own
+        vector similarity, makes a chunk reachable both ways: by being about the
+        question, and by being connected to what is. Graph retrieval then
+        contains vector retrieval rather than competing with it.
+
+        Only an improvement when the seeds are good: measured across five
+        corpora it helped alongside richer seed embeddings and *hurt* with
+        bare-name seeds (0.73 -> 0.57 on one corpus). Graphs are now always
+        built with the richer seed text; one built before that change should be
+        rebuilt before this is relied on.
+        """
+        node_ids = [node["id"] for node in subgraph.get("nodes", [])]
+        chunk_links = store.get_chunk_ids_for_nodes(source_id, node_ids)
+        candidate_ids = sorted({c for chunks in chunk_links.values() for c in chunks})
+        if not candidate_ids:
+            return []
+
+        graph = self._subgraph_graph(subgraph)
+        similarities = store.chunk_similarities(
+            source_id, candidate_ids, query_embedding
+        )
+        # Normalised so the passage share is a fixed fraction of the restart
+        # mass rather than whatever absolute cosine this embedding model emits.
+        scores = [similarities.get(c, 0.0) for c in candidate_ids]
+        low, high = (min(scores), max(scores)) if scores else (0.0, 0.0)
+        spread = high - low
+
+        personalization = dict(seeds)
+        passage_of: Dict[str, str] = {}
+        for chunk_id in candidate_ids:
+            linked = [n for n, chunks in chunk_links.items() if chunk_id in chunks]
+            linked = [n for n in linked if n in graph]
+            if not linked:
+                continue
+            passage_node = f"chunk::{chunk_id}"
+            passage_of[passage_node] = chunk_id
+            for node in linked:
+                graph.add_edge(passage_node, node, weight=1.0)
+            similarity = similarities.get(chunk_id, 0.0)
+            normalized = (similarity - low) / spread if spread > 0 else 0.0
+            personalization[passage_node] = normalized * PASSAGE_NODE_WEIGHT
+
+        if graph.number_of_nodes() == 0 or not any(personalization.values()):
+            return []
+
+        ranks = _personalized_pagerank(
+            graph,
+            personalization=personalization,
+            weight="weight",
+            alpha=_damping(passage_nodes=True),
+        )
+        chunk_scores = {
+            chunk_id: ranks.get(passage_node, 0.0)
+            for passage_node, chunk_id in passage_of.items()
+        }
+        ranked = sorted(chunk_scores, key=lambda c: chunk_scores[c], reverse=True)
+        candidates = max(self.chunks * 2, self.chunks + 5)
+        return ranked[: max(1, candidates)]
+
     def _source_top_k(self, source_id) -> int:
         """How many chunks this source may contribute — its own top-k.
 
@@ -273,6 +369,111 @@ class GraphRAGRetriever(BaseRetriever):
         base = self.base_chunks if self.base_chunks is not None else self.chunks
         return max(1, base // max(1, len(self.vectorstores)))
 
+    def _vector_ranking(self, source_id, query_embedding: List[float]) -> List[tuple]:
+        """The source's own vector ranking, as ``(text, metadata)`` in score order.
+
+        Used only by the hybrid path. Vector hits carry no row id, so the fused
+        ranking is keyed on the chunk text itself — the one identifier both
+        rankings share — and the metadata travels with it so a hit the graph
+        never surfaced can still be emitted as a document.
+        """
+        from docsgpt.vectorstore.vector_creator import VectorCreator
+
+        store = None
+        try:
+            store = VectorCreator.create_vectorstore(
+                settings.VECTOR_STORE, source_id, settings.EMBEDDINGS_KEY
+            )
+            hits = store.search(
+                self._classic._get_rephrased_question(),
+                k=max(self.chunks * 4, 20),
+                query_vector=query_embedding,
+            )
+        except Exception as e:
+            logging.error(
+                "GraphRAG hybrid: vector ranking failed for %s: %s", source_id, e
+            )
+            return []
+        finally:
+            close = getattr(store, "close", None)
+            if close is not None:
+                try:
+                    close()
+                except Exception as e:
+                    logging.debug("Error closing hybrid vector store: %s", e)
+        ranked = []
+        for hit in hits:
+            text = getattr(hit, "page_content", None)
+            metadata = getattr(hit, "metadata", None)
+            if text is None and isinstance(hit, dict):
+                text = hit.get("text") or hit.get("page_content")
+                metadata = hit.get("metadata")
+            if text:
+                ranked.append((text, metadata or {}))
+        return ranked
+
+    @staticmethod
+    def _rrf_order(rankings: List[List[str]], k: int) -> Dict[str, float]:
+        """Reciprocal rank fusion over ranked lists of the same key type.
+
+        Rank-based on purpose: PPR mass and cosine similarity are not on
+        comparable scales, and normalising either one invents a calibration
+        that does not exist.
+        """
+        scores: Dict[str, float] = {}
+        for ranking in rankings:
+            for position, key in enumerate(ranking):
+                scores[key] = scores.get(key, 0.0) + 1.0 / (k + position + 1)
+        return scores
+
+    def _graph_options(self, source_id) -> GraphRetrievalConfig:
+        """This source's graph retrieval options, or the recommended defaults.
+
+        Options travel on the per-source retrieval config the Dispatcher hands
+        over. A request that carries no per-source detail gets the defaults,
+        which are the measured-best configuration rather than a neutral one.
+        """
+        cfg = (getattr(self, "per_source_retrieval", None) or {}).get(source_id)
+        options = cfg.get("graph") if isinstance(cfg, dict) else getattr(cfg, "graph", None)
+        if isinstance(options, GraphRetrievalConfig):
+            return options
+        try:
+            return GraphRetrievalConfig.model_validate(options or {})
+        except Exception:
+            return GraphRetrievalConfig()
+
+    def _seed_rows(
+        self, store, source_id, query_embedding: List[float]
+    ) -> List[Dict[str, Any]]:
+        """The nodes the walk restarts from, per the source's ``seed_strategy``.
+
+        Seeding decides more than ranking does: a walk that starts on the wrong
+        nodes cannot be rescued downstream.
+
+        ``entities``
+            Cosine NN over entity embeddings, built from each entity's name,
+            type and description so a whole question has something to match.
+            The default: best or tied-best on every corpus measured.
+        ``relationships``
+            Cosine NN over relationship sentences ("A streams_to B: ..."),
+            seeding both endpoints of the best-matching facts. The only way to
+            start on an entity the question never names; strongest on
+            chain-structured content, weaker on ordinary prose.
+
+        Relationship seeding falls back to entity matching for a source with no
+        fact embeddings (one built before they were recorded), so it still
+        retrieves rather than returning nothing.
+        """
+        if self._graph_options(source_id).seed_strategy == "relationships":
+            by_fact = store.seed_nodes_from_facts(
+                source_id, query_embedding, fact_limit=FACT_SEED_FACTS, limit=SEED_NODES
+            )
+            if by_fact:
+                return by_fact
+        return store.search_nodes_by_embedding(
+            source_id, query_embedding, k=SEED_NODES
+        )
+
     def _graph_docs_for_source(
         self, store, source_id, query_embedding: List[float]
     ) -> List[Dict[str, Any]]:
@@ -284,9 +485,7 @@ class GraphRAGRetriever(BaseRetriever):
             query_embedding: Embedding of the rephrased question, computed once
                 by the caller for the whole retrieval.
         """
-        seed_rows = store.search_nodes_by_embedding(
-            source_id, query_embedding, k=SEED_NODES
-        )
+        seed_rows = self._seed_rows(store, source_id, query_embedding)
         if not seed_rows:
             return []
 
@@ -300,26 +499,63 @@ class GraphRAGRetriever(BaseRetriever):
             for row in seed_rows
         }
 
+        options = self._graph_options(source_id)
         subgraph = store.get_subgraph(source_id, seed_ids, hops=SUBGRAPH_HOPS)
-        node_scores = self._ppr_scores(subgraph, seeds)
-        if not node_scores:
+        if options.passage_nodes:
+            chunk_ids = self._rank_chunks_with_passages(
+                store, source_id, subgraph, seeds, query_embedding
+            )
+        else:
+            node_scores = self._ppr_scores(subgraph, seeds)
+            if not node_scores:
+                return []
+            chunk_ids = self._rank_chunks(store, source_id, node_scores)
+        if not chunk_ids:
             return []
 
-        chunk_ids = self._rank_chunks(store, source_id, node_scores)
         chunk_data = store.get_chunk_texts(source_id, chunk_ids)
 
+        # ``(text, metadata)`` in rank order. Chunk ids stop being the currency
+        # here: a hit contributed by the vector ranking has no graph chunk id,
+        # and keying on ids is what made an earlier version of this fusion able
+        # only to reorder the graph's own candidates.
+        candidates: List[tuple] = []
+        for chunk_id in chunk_ids:
+            chunk = chunk_data.get(chunk_id)
+            text = chunk.get("text") if chunk else None
+            if text:
+                candidates.append((text, chunk.get("metadata")))
+
+        if options.blend_vector:
+            # The graph ranks by how much PPR mass landed on a chunk's
+            # entities, which says nothing about whether the chunk is about the
+            # question. Fusing with the source's own vector ranking keeps the
+            # graph's reach while letting plain relevance back in — including
+            # chunks the graph never surfaced, which is where most of the value
+            # is: no reordering can rescue a question whose answer the graph
+            # missed entirely.
+            vector_hits = self._vector_ranking(source_id, query_embedding)
+            if vector_hits:
+                metadata_by_text = {text: meta for text, meta in candidates}
+                for text, meta in vector_hits:
+                    metadata_by_text.setdefault(text, meta)
+                fused = self._rrf_order(
+                    [[t for t, _ in candidates], [t for t, _ in vector_hits]],
+                    RRF_K,
+                )
+                candidates = [
+                    (text, metadata_by_text.get(text))
+                    for text in sorted(fused, key=lambda t: fused[t], reverse=True)
+                ]
+
         docs: List[Dict[str, Any]] = []
         token_budget = max(int(self.doc_token_limit * 0.9), 100)
         cumulative_tokens = 0
         source_top_k = self._source_top_k(source_id)
-        for chunk_id in chunk_ids:
+        for text, metadata in candidates:
             if len(docs) >= source_top_k:
                 break
-            chunk = chunk_data.get(chunk_id)
-            text = chunk.get("text") if chunk else None
-            if not text:
-                continue
-            labels = labels_from_metadata(chunk.get("metadata"), text, source_id)
+            labels = labels_from_metadata(metadata, text, source_id)
             doc_tokens = num_tokens_from_string(f"{labels['filename']}\n{text}")
             if cumulative_tokens + doc_tokens >= token_budget:
                 break