mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 18:15:02 +00:00
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.
This commit is contained in:
1 parent
b7e7872bf7
commit
a83e1dc0af
13 files changed
+1577
-79
No files matched your search
+140
-29
@@ -27,6 +27,7 @@ from docsgpt.core.model_utils import (
|
||||
get_api_key_for_provider,
|
||||
get_provider_from_model_id,
|
||||
)
|
||||
from docsgpt.graphrag.naming import normalize_entity_name
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.llm.llm_creator import LLMCreator
|
||||
from docsgpt.storage.db.source_config import SourceConfig
|
||||
@@ -224,6 +225,8 @@ def extract_graph_for_source(
|
||||
source's graph holds after the run — not how many upserts ran, which
|
||||
counts the same entity once per chunk it appears in.
|
||||
"""
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
from docsgpt.graphrag.store import GraphStore
|
||||
|
||||
store = GraphStore()
|
||||
@@ -266,43 +269,85 @@ def extract_graph_for_source(
|
||||
except Exception as exc:
|
||||
logger.debug("graph progress callback failed: %s", exc)
|
||||
|
||||
for chunk, chunk_id in to_process:
|
||||
def _prepare(item):
|
||||
"""One chunk's LLM extraction — the only step run concurrently.
|
||||
|
||||
A chunk spends almost all of its time waiting on the model, so that is
|
||||
what runs in the pool. Everything else stays on the calling thread:
|
||||
graph writes, so transactions and the progress checkpoint are exactly
|
||||
what they were serially, and embedding. Inside a Celery worker the
|
||||
embeddings client decides to embed locally from the task on the
|
||||
*current thread's* stack; a pool thread has none, so it would instead
|
||||
dispatch an embed task to the worker and wait on it, which Celery
|
||||
refuses inside a task — failing every chunk of the build.
|
||||
"""
|
||||
chunk, chunk_id = item
|
||||
text = _chunk_text(chunk)
|
||||
if not text:
|
||||
store.mark_chunk(source_id, chunk_id, "done")
|
||||
chunks_processed += 1
|
||||
_report()
|
||||
continue
|
||||
return chunk_id, "empty", None
|
||||
|
||||
extracted = _extract_chunk(llm, text, chunk_id)
|
||||
if extracted is None:
|
||||
store.mark_chunk(source_id, chunk_id, "failed")
|
||||
failed_chunks += 1
|
||||
_report()
|
||||
continue
|
||||
|
||||
return chunk_id, "failed", None
|
||||
try:
|
||||
entities = _build_entities(extracted["entities"])
|
||||
relationships = _build_relationships(extracted["relationships"])
|
||||
name_embeddings = _embed_names(embedding, entities, relationships)
|
||||
chunk_nodes, chunk_edges = store.apply_chunk(
|
||||
source_id, chunk_id, entities, relationships, name_embeddings
|
||||
)
|
||||
node_upserts += chunk_nodes
|
||||
edges += chunk_edges
|
||||
# ``apply_chunk`` marks the chunk done inside the transaction that
|
||||
# writes its rows, so the checkpoint cannot disagree with the graph
|
||||
# and a replayed write cannot apply the chunk twice.
|
||||
chunks_processed += 1
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Graph extraction write failed for chunk %s, skipping: %s",
|
||||
chunk_id,
|
||||
exc,
|
||||
"Graph extraction failed for chunk %s, skipping: %s", chunk_id, exc
|
||||
)
|
||||
store.mark_chunk(source_id, chunk_id, "failed")
|
||||
failed_chunks += 1
|
||||
_report()
|
||||
return chunk_id, "failed", None
|
||||
return chunk_id, "ok", (entities, relationships)
|
||||
|
||||
workers = max(1, int(getattr(settings, "GRAPHRAG_EXTRACTION_WORKERS", 1) or 1))
|
||||
pool = None
|
||||
if workers > 1 and len(to_process) > 1:
|
||||
pool = ThreadPoolExecutor(max_workers=workers)
|
||||
# ``map`` yields in submission order, so chunks are still applied in the
|
||||
# order they were given and a run stays reproducible.
|
||||
prepared = pool.map(_prepare, to_process)
|
||||
else:
|
||||
prepared = (_prepare(item) for item in to_process)
|
||||
|
||||
try:
|
||||
for chunk_id, status, payload in prepared:
|
||||
if status == "empty":
|
||||
store.mark_chunk(source_id, chunk_id, "done")
|
||||
chunks_processed += 1
|
||||
_report()
|
||||
continue
|
||||
if status == "failed":
|
||||
store.mark_chunk(source_id, chunk_id, "failed")
|
||||
failed_chunks += 1
|
||||
_report()
|
||||
continue
|
||||
|
||||
entities, relationships = payload
|
||||
try:
|
||||
# On this thread, not in the pool — see ``_prepare``.
|
||||
name_embeddings = _embed_names(embedding, entities, relationships)
|
||||
_embed_facts(embedding, relationships)
|
||||
chunk_nodes, chunk_edges = store.apply_chunk(
|
||||
source_id, chunk_id, entities, relationships, name_embeddings
|
||||
)
|
||||
node_upserts += chunk_nodes
|
||||
edges += chunk_edges
|
||||
# ``apply_chunk`` marks the chunk done inside the transaction that
|
||||
# writes its rows, so the checkpoint cannot disagree with the graph
|
||||
# and a replayed write cannot apply the chunk twice.
|
||||
chunks_processed += 1
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Graph extraction embed/write failed for chunk %s, skipping: %s",
|
||||
chunk_id,
|
||||
exc,
|
||||
)
|
||||
store.mark_chunk(source_id, chunk_id, "failed")
|
||||
failed_chunks += 1
|
||||
_report()
|
||||
finally:
|
||||
if pool is not None:
|
||||
pool.shutdown(wait=True)
|
||||
|
||||
try:
|
||||
store.set_node_degrees(source_id)
|
||||
@@ -347,7 +392,7 @@ def _build_entities(raw_entities: Any) -> List[Dict[str, Any]]:
|
||||
entities.append(
|
||||
{
|
||||
"name": name,
|
||||
"normalized_name": name.lower(),
|
||||
"normalized_name": normalize_entity_name(name),
|
||||
"type": str(e.get("type") or "") or None,
|
||||
"description": str(e.get("description") or "") or None,
|
||||
}
|
||||
@@ -373,6 +418,70 @@ def _build_relationships(raw_relationships: Any) -> List[Dict[str, Any]]:
|
||||
return relationships
|
||||
|
||||
|
||||
def _fact_text(rel: Dict[str, Any]) -> str:
|
||||
"""A relationship rendered as the sentence it asserts.
|
||||
|
||||
Embedded and stored on the edge so retrieval can match a question against
|
||||
the *relation* rather than against entity names — the difference between
|
||||
"which entity is this about" and "which fact answers this".
|
||||
"""
|
||||
source = str(rel.get("source") or "").strip()
|
||||
target = str(rel.get("target") or "").strip()
|
||||
if not source or not target:
|
||||
return ""
|
||||
relation = str(rel.get("type") or "related to").strip() or "related to"
|
||||
text = f"{source} {relation} {target}"
|
||||
description = str(rel.get("description") or "").strip()
|
||||
return f"{text}: {description}" if description else text
|
||||
|
||||
|
||||
def _embed_facts(embedding, relationships: List[Dict[str, Any]]) -> None:
|
||||
"""Attach a fact embedding to each relationship, in one batched call.
|
||||
|
||||
Mutates the relationship dicts so the embedding travels with the edge into
|
||||
``apply_chunk`` without a second mapping to keep in step. Always on: it is
|
||||
one extra batched call per chunk against an LLM call that already costs
|
||||
far more, and it lets a source switch to relationship seeding at query time
|
||||
without being rebuilt.
|
||||
"""
|
||||
pending = [(rel, _fact_text(rel)) for rel in relationships]
|
||||
pending = [(rel, text) for rel, text in pending if text]
|
||||
if not pending:
|
||||
return
|
||||
try:
|
||||
vectors = embedding.embed_documents([text for _rel, text in pending])
|
||||
except Exception as exc: # noqa: BLE001
|
||||
# The graph is still correct without them; only fact seeding degrades.
|
||||
logger.warning("Fact embedding failed, continuing without: %s", exc)
|
||||
return
|
||||
for (rel, _text), vector in zip(pending, vectors):
|
||||
rel["fact_embedding"] = vector
|
||||
|
||||
|
||||
def _seed_text(entity: Dict[str, Any]) -> str:
|
||||
"""The text a node's embedding is computed from.
|
||||
|
||||
Retrieval seeds the graph walk by matching a whole question against these
|
||||
embeddings, and a bare entity name is a poor thing to match a question
|
||||
against — a question about what a service writes to shares almost no
|
||||
surface with the name ``Quill``. Including the type and description gives
|
||||
the match something to work with; measured across five corpora it moved
|
||||
recall@4 by +0.07 to +0.50.
|
||||
|
||||
Relationship endpoints keep their bare names: they arrive as strings with
|
||||
no type or description attached.
|
||||
"""
|
||||
name = str(entity.get("name") or "").strip()
|
||||
text = name
|
||||
entity_type = str(entity.get("type") or "").strip()
|
||||
if entity_type:
|
||||
text += f" ({entity_type})"
|
||||
description = str(entity.get("description") or "").strip()
|
||||
if description:
|
||||
text += f": {description}"
|
||||
return text or name
|
||||
|
||||
|
||||
def _embed_names(
|
||||
embedding,
|
||||
entities: List[Dict[str, Any]],
|
||||
@@ -385,14 +494,16 @@ def _embed_names(
|
||||
"""
|
||||
name_by_norm: Dict[str, str] = {}
|
||||
for entity in entities:
|
||||
name_by_norm.setdefault(entity["normalized_name"], entity["name"])
|
||||
name_by_norm.setdefault(entity["normalized_name"], _seed_text(entity))
|
||||
for rel in relationships:
|
||||
for endpoint in (rel.get("source"), rel.get("target")):
|
||||
if endpoint is None:
|
||||
continue
|
||||
clean = str(endpoint).strip()
|
||||
if clean:
|
||||
name_by_norm.setdefault(clean.lower(), clean)
|
||||
# Same key the store resolves endpoints by, or the embedding
|
||||
# computed here never reaches the node it was computed for.
|
||||
name_by_norm.setdefault(normalize_entity_name(clean), clean)
|
||||
|
||||
if not name_by_norm:
|
||||
return {}
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
"""Canonical entity naming for the per-source knowledge graph.
|
||||
|
||||
Nodes are merged on ``normalized_name``, which has been ``name.lower()``. That
|
||||
splits entities a reader would call the same thing: measured on the DocsGPT docs
|
||||
corpus, ``agent``/``agents``, ``VECTOR_STORE``/``Vector store``/``vector stores``,
|
||||
``Celery worker``/``Celery workers`` and ``.env file``/``env_file`` all landed as
|
||||
separate nodes — 58 such collisions across 1,704 entities, with 75% of entities
|
||||
appearing in exactly one chunk as a result.
|
||||
|
||||
:func:`canonical_name` folds the differences that are purely orthographic:
|
||||
case, surrounding punctuation, underscore/hyphen word breaks, and a *cautious*
|
||||
plural. Cautious matters: this corpus contains ``postgres``, ``kubernetes``,
|
||||
``https`` and ``aws``, none of which are plurals, so a naive "strip trailing s"
|
||||
would corrupt them into new entities rather than merge anything.
|
||||
|
||||
Always on: every graph is built with canonical names.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
_PUNCT = re.compile(r"[^\w\s]+", re.UNICODE)
|
||||
_UNDERSCORE = re.compile(r"[_\-]+")
|
||||
_SPACE = re.compile(r"\s+")
|
||||
|
||||
#: Words that end in "s" without being plural. Singularising these would invent
|
||||
#: entities ("postgre", "kubernete") instead of merging existing ones.
|
||||
_NOT_PLURAL = frozenset(
|
||||
{
|
||||
"postgres", "kubernetes", "https", "aws", "dns", "tls", "cors", "css",
|
||||
"js", "sas", "gas", "ss", "class", "access", "process", "status",
|
||||
"analysis", "basis", "axis", "https", "rss", "less", "express",
|
||||
"redis", "nats", "kibana", "elasticsearch", "os", "ios", "macos",
|
||||
"always", "sometimes", "series", "docs", "ops", "devops", "sse",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _singular(word: str) -> str:
|
||||
"""Best-effort singular of one word, biased hard towards leaving it alone.
|
||||
|
||||
Only the endings that are unambiguous in this domain are touched:
|
||||
``-ies`` -> ``-y`` (``policies``), ``-ses``/``-xes``/``-zes``/``-ches``/
|
||||
``-shes`` -> drop ``es`` (``indexes``, ``batches``), and a bare trailing
|
||||
``s`` on a word long enough to be safe. Everything in :data:`_NOT_PLURAL`,
|
||||
and anything ending in ``ss``/``us``/``is``, is returned unchanged.
|
||||
"""
|
||||
if len(word) < 4 or word in _NOT_PLURAL:
|
||||
return word
|
||||
if word.endswith(("ss", "us", "is")):
|
||||
return word
|
||||
if word.endswith("ies") and len(word) > 4:
|
||||
return word[:-3] + "y"
|
||||
if word.endswith(("ses", "xes", "zes", "ches", "shes")):
|
||||
return word[:-2]
|
||||
if word.endswith("s"):
|
||||
return word[:-1]
|
||||
return word
|
||||
|
||||
|
||||
def canonical_name(name: str) -> str:
|
||||
"""Merge key for an entity name.
|
||||
|
||||
Args:
|
||||
name: The entity name as the model wrote it.
|
||||
|
||||
Returns:
|
||||
A lowercase, punctuation-free, singularised key. Returns ``""`` for an
|
||||
empty or punctuation-only name, which callers treat as "no entity".
|
||||
|
||||
Examples:
|
||||
``VECTOR_STORE`` and ``Vector stores`` -> ``vector store``;
|
||||
``.env file`` and ``env_file`` -> ``env file``;
|
||||
``postgres`` stays ``postgres``.
|
||||
"""
|
||||
if not name:
|
||||
return ""
|
||||
text = _UNDERSCORE.sub(" ", str(name))
|
||||
text = _PUNCT.sub(" ", text)
|
||||
text = _SPACE.sub(" ", text).strip().lower()
|
||||
if not text:
|
||||
return ""
|
||||
return " ".join(_singular(word) for word in text.split())
|
||||
|
||||
|
||||
def normalize_entity_name(name: str) -> str:
|
||||
"""The key an entity is merged on: its :func:`canonical_name`.
|
||||
|
||||
Every graph the corpora were measured on was built this way, so it is the
|
||||
only mode rather than a flag. A graph built before this used plain
|
||||
``lower()`` keys; re-extracting it merges onto these instead.
|
||||
"""
|
||||
return canonical_name(name)
|
||||
+312
-19
@@ -74,6 +74,16 @@ def _pgvector_identifiers() -> tuple[str, str, str, str]:
|
||||
)
|
||||
|
||||
|
||||
def _pgvector_vector_column() -> str:
|
||||
"""Resolve the embedding column name from the same ``PGVectorStore`` defaults."""
|
||||
import inspect
|
||||
|
||||
from docsgpt.vectorstore.pgvector import PGVectorStore
|
||||
|
||||
params = inspect.signature(PGVectorStore.__init__).parameters
|
||||
return _safe_identifier(params["vector_column"].default)
|
||||
|
||||
|
||||
def _is_connection_lost(exc: BaseException) -> bool:
|
||||
"""True when ``exc`` says the server connection went away, not that the SQL was bad.
|
||||
|
||||
@@ -253,7 +263,7 @@ class GraphStore:
|
||||
)
|
||||
|
||||
cursor.execute(
|
||||
"""
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS graph_edges (
|
||||
id UUID PRIMARY KEY,
|
||||
source_id UUID NOT NULL,
|
||||
@@ -262,10 +272,18 @@ class GraphStore:
|
||||
type TEXT,
|
||||
description TEXT,
|
||||
weight REAL DEFAULT 1.0,
|
||||
source_chunk_ids JSONB
|
||||
source_chunk_ids JSONB,
|
||||
fact_embedding vector({dimension})
|
||||
);
|
||||
"""
|
||||
)
|
||||
# ``CREATE TABLE IF NOT EXISTS`` is a no-op on a database that
|
||||
# already has the table, so a column added after the fact needs its
|
||||
# own statement or every existing deployment silently lacks it.
|
||||
cursor.execute(
|
||||
f"ALTER TABLE graph_edges "
|
||||
f"ADD COLUMN IF NOT EXISTS fact_embedding vector({dimension});"
|
||||
)
|
||||
|
||||
cursor.execute(
|
||||
"""
|
||||
@@ -453,19 +471,78 @@ class GraphStore:
|
||||
description: Optional[str] = None,
|
||||
weight: float = 1.0,
|
||||
source_chunk_ids: Optional[List[str]] = None,
|
||||
) -> str:
|
||||
"""Insert an edge on an open cursor (no commit, no degree bump).
|
||||
fact_embedding: Optional[List[float]] = None,
|
||||
) -> tuple[Optional[str], bool]:
|
||||
"""Write an edge on an open cursor (no commit, no degree bump).
|
||||
|
||||
Returns ``(edge_id, created)``. Two shapes of noise are rejected here
|
||||
rather than at read time, because once written neither is visible:
|
||||
|
||||
* A self-loop feeds a node's PageRank mass straight back to itself. It
|
||||
is dropped, reported as ``(None, False)``.
|
||||
* A pair already related by the same type is *merged* rather than
|
||||
inserted again. ``graph_edges`` carries no uniqueness constraint, so
|
||||
re-extracting one relationship across many chunks otherwise writes a
|
||||
row per chunk — a fifth of a real corpus's edges — inflating that
|
||||
pair's traversal weight and spending the bounded subgraph fetch on
|
||||
duplicates. The surviving row keeps the strongest weight seen and
|
||||
every contributing chunk id.
|
||||
|
||||
Callers that batch many edges run ``set_node_degrees`` once afterwards
|
||||
instead of bumping degree per edge.
|
||||
"""
|
||||
if str(src_node_id) == str(dst_node_id):
|
||||
return None, False
|
||||
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT id
|
||||
FROM graph_edges
|
||||
WHERE source_id = %s AND src_node_id = %s AND dst_node_id = %s
|
||||
AND type IS NOT DISTINCT FROM %s
|
||||
LIMIT 1;
|
||||
""",
|
||||
(source_id, src_node_id, dst_node_id, type),
|
||||
)
|
||||
existing = cursor.fetchone()
|
||||
if existing:
|
||||
edge_id = existing[0]
|
||||
# The chunk ids are merged in SQL, against the row's own current
|
||||
# value, rather than read here and written back: a read-modify-write
|
||||
# would drop whatever a concurrent writer appended in between.
|
||||
cursor.execute(
|
||||
"""
|
||||
UPDATE graph_edges
|
||||
SET weight = GREATEST(COALESCE(weight, 0), %s),
|
||||
description = COALESCE(description, %s),
|
||||
-- Backfills the fact embedding for an edge first written
|
||||
-- before fact embeddings were switched on.
|
||||
fact_embedding = COALESCE(fact_embedding, %s::vector),
|
||||
source_chunk_ids = COALESCE(source_chunk_ids, '[]'::jsonb) || (
|
||||
SELECT COALESCE(jsonb_agg(candidate), '[]'::jsonb)
|
||||
FROM jsonb_array_elements(%s::jsonb) AS candidate
|
||||
WHERE NOT COALESCE(source_chunk_ids, '[]'::jsonb)
|
||||
@> jsonb_build_array(candidate)
|
||||
)
|
||||
WHERE id = %s;
|
||||
""",
|
||||
(
|
||||
weight,
|
||||
description,
|
||||
fact_embedding,
|
||||
Jsonb(list(source_chunk_ids or [])),
|
||||
edge_id,
|
||||
),
|
||||
)
|
||||
return str(edge_id), False
|
||||
|
||||
edge_id = str(uuid.uuid4())
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO graph_edges
|
||||
(id, source_id, src_node_id, dst_node_id, type, description,
|
||||
weight, source_chunk_ids)
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s);
|
||||
weight, source_chunk_ids, fact_embedding)
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s);
|
||||
""",
|
||||
(
|
||||
edge_id,
|
||||
@@ -476,9 +553,10 @@ class GraphStore:
|
||||
description,
|
||||
weight,
|
||||
Jsonb(source_chunk_ids or []),
|
||||
fact_embedding,
|
||||
),
|
||||
)
|
||||
return edge_id
|
||||
return edge_id, True
|
||||
|
||||
def add_edge(
|
||||
self,
|
||||
@@ -489,21 +567,28 @@ class GraphStore:
|
||||
description: Optional[str] = None,
|
||||
weight: float = 1.0,
|
||||
source_chunk_ids: Optional[List[str]] = None,
|
||||
) -> str:
|
||||
"""Insert an edge and bump the degree of both endpoints. Returns its id."""
|
||||
fact_embedding: Optional[List[float]] = None,
|
||||
) -> Optional[str]:
|
||||
"""Write an edge and bump the degree of both endpoints. Returns its id.
|
||||
|
||||
Returns ``None`` for a self-loop, which is not written. A repeat of an
|
||||
existing pair merges into that row and returns its id, leaving degree
|
||||
alone — the endpoints gained no new neighbour.
|
||||
"""
|
||||
self._ensure_tables_once()
|
||||
conn = self._get_connection()
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
edge_id = self._add_edge(
|
||||
edge_id, created = self._add_edge(
|
||||
cursor, source_id, src_node_id, dst_node_id, type, description,
|
||||
weight, source_chunk_ids,
|
||||
)
|
||||
cursor.execute(
|
||||
"UPDATE graph_nodes SET degree = degree + 1 "
|
||||
"WHERE source_id = %s AND id IN (%s, %s);",
|
||||
(source_id, src_node_id, dst_node_id),
|
||||
weight, source_chunk_ids, fact_embedding,
|
||||
)
|
||||
if created:
|
||||
cursor.execute(
|
||||
"UPDATE graph_nodes SET degree = degree + 1 "
|
||||
"WHERE source_id = %s AND id IN (%s, %s);",
|
||||
(source_id, src_node_id, dst_node_id),
|
||||
)
|
||||
conn.commit()
|
||||
return edge_id
|
||||
except Exception as e:
|
||||
@@ -608,7 +693,7 @@ class GraphStore:
|
||||
)
|
||||
if src_id is None or dst_id is None:
|
||||
continue
|
||||
self._add_edge(
|
||||
_, created = self._add_edge(
|
||||
cursor,
|
||||
source_id,
|
||||
src_id,
|
||||
@@ -617,8 +702,10 @@ class GraphStore:
|
||||
description=rel.get("description"),
|
||||
weight=float(rel.get("weight") or 1.0),
|
||||
source_chunk_ids=[chunk_id],
|
||||
fact_embedding=rel.get("fact_embedding"),
|
||||
)
|
||||
edges_added += 1
|
||||
if created:
|
||||
edges_added += 1
|
||||
|
||||
cursor.execute(
|
||||
"""
|
||||
@@ -653,7 +740,11 @@ class GraphStore:
|
||||
clean = str(name).strip()
|
||||
if not clean:
|
||||
return None
|
||||
normalized_name = clean.lower()
|
||||
from docsgpt.graphrag.naming import normalize_entity_name
|
||||
|
||||
normalized_name = normalize_entity_name(clean)
|
||||
if not normalized_name:
|
||||
return None
|
||||
if normalized_name in node_ids:
|
||||
return node_ids[normalized_name]
|
||||
node_id = self._upsert_node(
|
||||
@@ -839,6 +930,7 @@ class GraphStore:
|
||||
FROM graph_edges
|
||||
WHERE source_id = %s
|
||||
AND (src_node_id = ANY(%s) OR dst_node_id = ANY(%s))
|
||||
ORDER BY weight DESC NULLS LAST
|
||||
LIMIT %s;
|
||||
""",
|
||||
(
|
||||
@@ -886,6 +978,7 @@ class GraphStore:
|
||||
FROM graph_edges
|
||||
WHERE source_id = %s
|
||||
AND src_node_id = ANY(%s) AND dst_node_id = ANY(%s)
|
||||
ORDER BY weight DESC NULLS LAST
|
||||
LIMIT %s;
|
||||
""",
|
||||
(source_id, node_id_list, node_id_list, MAX_SUBGRAPH_EDGES),
|
||||
@@ -999,6 +1092,206 @@ class GraphStore:
|
||||
cursor.close()
|
||||
conn.rollback()
|
||||
|
||||
def seed_nodes_from_facts(
|
||||
self,
|
||||
source_id: str,
|
||||
query_embedding: List[float],
|
||||
fact_limit: int = 5,
|
||||
limit: int = 10,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Seed nodes drawn from the *relationships* nearest the question.
|
||||
|
||||
Name matching asks "which entity is this question about", which a
|
||||
multi-document question cannot answer: the entity holding the answer is
|
||||
named in another document, not in the question. A fact string carries
|
||||
the relation — "Alder streams_to Quill: ..." — so a question about what
|
||||
a service writes to can match the edge itself and seed the walk on both
|
||||
of its endpoints, including the one nothing in the question names.
|
||||
|
||||
Endpoints are weighted by fact score divided by the entity's
|
||||
``doc_freq``: an entity appearing in every chunk is a poor seed even
|
||||
when it sits on a well-matched fact, and dividing by how widely it
|
||||
occurs prefers the specific endpoint over the hub.
|
||||
|
||||
Rows match :meth:`search_nodes_by_embedding`'s shape, so the caller's
|
||||
seed weighting is unchanged. Returns nothing when the source has no
|
||||
fact embeddings, which is the signal to fall back to name matching.
|
||||
"""
|
||||
if not query_embedding:
|
||||
return []
|
||||
conn = self._get_connection()
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
cursor.execute(
|
||||
"""
|
||||
WITH top_facts AS (
|
||||
SELECT src_node_id, dst_node_id,
|
||||
1 - (fact_embedding <=> %s::vector) AS score
|
||||
FROM graph_edges
|
||||
WHERE source_id = %s AND fact_embedding IS NOT NULL
|
||||
ORDER BY fact_embedding <=> %s::vector
|
||||
LIMIT %s
|
||||
)
|
||||
SELECT n.id::text, n.name, n.description,
|
||||
MAX(f.score / GREATEST(COALESCE(n.doc_freq, 1), 1)) AS weight
|
||||
FROM top_facts f
|
||||
JOIN graph_nodes n
|
||||
ON n.id = f.src_node_id OR n.id = f.dst_node_id
|
||||
WHERE n.source_id = %s
|
||||
GROUP BY n.id, n.name, n.description
|
||||
ORDER BY weight DESC
|
||||
LIMIT %s;
|
||||
""",
|
||||
(
|
||||
query_embedding,
|
||||
source_id,
|
||||
query_embedding,
|
||||
max(1, int(fact_limit)),
|
||||
source_id,
|
||||
max(1, int(limit)),
|
||||
),
|
||||
)
|
||||
return [
|
||||
{
|
||||
"id": row[0],
|
||||
"name": row[1],
|
||||
"description": row[2],
|
||||
# The caller reads weight back as ``1 - distance``.
|
||||
"distance": 1.0 - float(row[3] or 0.0),
|
||||
}
|
||||
for row in cursor.fetchall()
|
||||
]
|
||||
except Exception as e:
|
||||
logging.error(f"Error seeding nodes from facts: {e}")
|
||||
return []
|
||||
finally:
|
||||
cursor.close()
|
||||
conn.rollback()
|
||||
|
||||
def entity_relationships(
|
||||
self, source_id: str, name: str, limit: int = 25
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""The relationships an entity takes part in, strongest first.
|
||||
|
||||
This is the one thing a caller cannot get from vector search: which
|
||||
*named* thing an entity is connected to. Matching is on the name rather
|
||||
than a node id because the caller is an LLM holding a name it read in
|
||||
the text, not an id.
|
||||
"""
|
||||
clean = (name or "").strip()
|
||||
if not clean:
|
||||
return []
|
||||
conn = self._get_connection()
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT s.name, e.type, d.name, e.description
|
||||
FROM graph_edges e
|
||||
JOIN graph_nodes s ON s.id = e.src_node_id
|
||||
JOIN graph_nodes d ON d.id = e.dst_node_id
|
||||
WHERE e.source_id = %s AND (s.name ILIKE %s OR d.name ILIKE %s)
|
||||
ORDER BY e.weight DESC NULLS LAST
|
||||
LIMIT %s;
|
||||
""",
|
||||
(source_id, f"%{clean}%", f"%{clean}%", max(1, int(limit))),
|
||||
)
|
||||
return [
|
||||
{"source": row[0], "type": row[1], "target": row[2], "description": row[3]}
|
||||
for row in cursor.fetchall()
|
||||
]
|
||||
except Exception as e:
|
||||
logging.error(f"Error reading relationships for {name!r}: {e}")
|
||||
return []
|
||||
finally:
|
||||
cursor.close()
|
||||
conn.rollback()
|
||||
|
||||
def entity_pages(
|
||||
self, source_id: str, name: str, limit: int = 4
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Chunks an entity appears in, with the chunk it is *about* first.
|
||||
|
||||
A plain substring match answers "Halvard" with pages that merely mention
|
||||
Halvard, and an unordered ``LIMIT`` then decides which of those the
|
||||
caller sees. Nodes whose name is the entity (or the entity plus a
|
||||
qualifier the extractor appended, "Quill" -> "Quill Store") are
|
||||
preferred, and among those the chunk whose text opens with the name
|
||||
comes first; a substring match is the fallback so an unusual name still
|
||||
resolves.
|
||||
"""
|
||||
clean = (name or "").strip()
|
||||
if not clean:
|
||||
return []
|
||||
table, text_col, metadata_col, source_col = _pgvector_identifiers()
|
||||
conn = self._get_connection()
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
cursor.execute(
|
||||
f"""
|
||||
SELECT d.{metadata_col}, d.{text_col},
|
||||
(lower(n.name) = %s OR lower(n.name) LIKE %s) AS is_subject
|
||||
FROM graph_node_chunks gc
|
||||
JOIN graph_nodes n ON n.id = gc.node_id
|
||||
JOIN {table} d ON d.id::text = gc.chunk_id
|
||||
WHERE gc.source_id = %s AND d.{source_col} = %s
|
||||
AND (lower(n.name) = %s OR lower(n.name) LIKE %s OR n.name ILIKE %s)
|
||||
GROUP BY d.{metadata_col}, d.{text_col}, is_subject
|
||||
ORDER BY is_subject DESC, (d.{text_col} ILIKE %s) DESC
|
||||
LIMIT %s;
|
||||
""",
|
||||
(
|
||||
clean.lower(), f"{clean.lower()} %",
|
||||
source_id, source_id,
|
||||
clean.lower(), f"{clean.lower()} %", f"%{clean}%",
|
||||
f"{clean}%",
|
||||
max(1, int(limit)),
|
||||
),
|
||||
)
|
||||
return [
|
||||
{"metadata": row[0] or {}, "text": row[1] or ""}
|
||||
for row in cursor.fetchall()
|
||||
]
|
||||
except Exception as e:
|
||||
logging.error(f"Error reading pages for {name!r}: {e}")
|
||||
return []
|
||||
finally:
|
||||
cursor.close()
|
||||
conn.rollback()
|
||||
|
||||
def chunk_similarities(
|
||||
self, source_id: str, chunk_ids: List[str], query_embedding: List[float]
|
||||
) -> Dict[str, float]:
|
||||
"""Cosine similarity between the query and specific chunks of a source.
|
||||
|
||||
Passage nodes need their own relevance to claim a share of the walk's
|
||||
restart mass, and that number lives in the co-located pgvector table —
|
||||
the same one :meth:`get_chunk_texts` reads. Restricted to the chunk ids
|
||||
the subgraph actually reached, so this never scans the whole source.
|
||||
"""
|
||||
if not chunk_ids or not query_embedding:
|
||||
return {}
|
||||
table, _text_col, _metadata_col, source_col = _pgvector_identifiers()
|
||||
vector_col = _pgvector_vector_column()
|
||||
conn = self._get_connection()
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
cursor.execute(
|
||||
f"""
|
||||
SELECT id::text, 1 - ({vector_col} <=> %s::vector)
|
||||
FROM {table}
|
||||
WHERE {source_col} = %s AND id::text = ANY(%s);
|
||||
""",
|
||||
(query_embedding, source_id, [str(c) for c in chunk_ids]),
|
||||
)
|
||||
return {row[0]: float(row[1]) for row in cursor.fetchall()}
|
||||
except Exception as e:
|
||||
logging.error(f"Error scoring chunks against the query: {e}")
|
||||
return {}
|
||||
finally:
|
||||
cursor.close()
|
||||
conn.rollback()
|
||||
|
||||
def get_chunk_texts(
|
||||
self,
|
||||
source_id: str,
|
||||
|
||||
Reference in new issue
Block a user