Spaces:
Running
Running
Commit ·
082589c
1
Parent(s): 911343c
feat(graph): L3 semantic entity alignment via bge embeddings
Browse files- app/core/config.py +8 -0
- app/graph/neo4j_client.py +36 -2
- app/graph/semantic_alignment.py +102 -0
- app/workers/tasks.py +101 -2
- scripts/migrate_kg_embeddings.py +86 -0
- tests/test_graph_semantic_alignment.py +100 -0
- tests/test_workers_ingest_graph.py +125 -1
app/core/config.py
CHANGED
|
@@ -166,6 +166,14 @@ class Settings(BaseSettings):
|
|
| 166 |
graph_extraction_reconcile: bool = True
|
| 167 |
graph_extraction_reconcile_top_k: int = 40
|
| 168 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 169 |
# Intent router classifies the query and branches the workflow. When
|
| 170 |
# disabled, every turn takes the "chat" path.
|
| 171 |
router_enabled: bool = True
|
|
|
|
| 166 |
graph_extraction_reconcile: bool = True
|
| 167 |
graph_extraction_reconcile_top_k: int = 40
|
| 168 |
|
| 169 |
+
# Semantic entity alignment (#43c) — L3 layer on top of L1/L2. Embeds
|
| 170 |
+
# each entity's name+description via bge-small and merges same-type
|
| 171 |
+
# entities whose cosine similarity clears the threshold, catching
|
| 172 |
+
# abbreviation/expansion pairs (SFT ↔ Supervised Fine-Tuning, GWU ↔
|
| 173 |
+
# George Washington University) that string-fuzzy matching can't.
|
| 174 |
+
graph_semantic_align: bool = True
|
| 175 |
+
graph_semantic_align_threshold: float = 0.85
|
| 176 |
+
|
| 177 |
# Intent router classifies the query and branches the workflow. When
|
| 178 |
# disabled, every turn takes the "chat" path.
|
| 179 |
router_enabled: bool = True
|
app/graph/neo4j_client.py
CHANGED
|
@@ -62,7 +62,8 @@ ON CREATE SET
|
|
| 62 |
e.type = $type,
|
| 63 |
e.description = $description,
|
| 64 |
e.created_at = datetime(),
|
| 65 |
-
e.document_ids = [$doc_id]
|
|
|
|
| 66 |
SET
|
| 67 |
e.last_seen_at = datetime(),
|
| 68 |
e.name = coalesce(e.name, $name),
|
|
@@ -71,7 +72,8 @@ SET
|
|
| 71 |
e.document_ids = CASE
|
| 72 |
WHEN $doc_id IN coalesce(e.document_ids, []) THEN e.document_ids
|
| 73 |
ELSE coalesce(e.document_ids, []) + $doc_id
|
| 74 |
-
END
|
|
|
|
| 75 |
RETURN e
|
| 76 |
"""
|
| 77 |
|
|
@@ -96,7 +98,15 @@ async def upsert_entity(
|
|
| 96 |
name_lower: str,
|
| 97 |
entity_type: str,
|
| 98 |
description: str,
|
|
|
|
| 99 |
) -> None:
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 100 |
driver = await get_driver()
|
| 101 |
async with driver.session() as session:
|
| 102 |
await session.run(
|
|
@@ -107,6 +117,7 @@ async def upsert_entity(
|
|
| 107 |
name_lower=name_lower,
|
| 108 |
type=entity_type,
|
| 109 |
description=description,
|
|
|
|
| 110 |
)
|
| 111 |
|
| 112 |
|
|
@@ -130,6 +141,29 @@ async def upsert_relationship(
|
|
| 130 |
)
|
| 131 |
|
| 132 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 133 |
async def list_entity_candidates(user_id: str) -> list[tuple[str, str]]:
|
| 134 |
"""Return (name_lower, type) tuples for every entity the user owns.
|
| 135 |
|
|
|
|
| 62 |
e.type = $type,
|
| 63 |
e.description = $description,
|
| 64 |
e.created_at = datetime(),
|
| 65 |
+
e.document_ids = [$doc_id],
|
| 66 |
+
e.embedding = $embedding
|
| 67 |
SET
|
| 68 |
e.last_seen_at = datetime(),
|
| 69 |
e.name = coalesce(e.name, $name),
|
|
|
|
| 72 |
e.document_ids = CASE
|
| 73 |
WHEN $doc_id IN coalesce(e.document_ids, []) THEN e.document_ids
|
| 74 |
ELSE coalesce(e.document_ids, []) + $doc_id
|
| 75 |
+
END,
|
| 76 |
+
e.embedding = coalesce(e.embedding, $embedding)
|
| 77 |
RETURN e
|
| 78 |
"""
|
| 79 |
|
|
|
|
| 98 |
name_lower: str,
|
| 99 |
entity_type: str,
|
| 100 |
description: str,
|
| 101 |
+
embedding: list[float] | None = None,
|
| 102 |
) -> None:
|
| 103 |
+
"""Upsert a user's :Entity node keyed on (user_id, name_lower).
|
| 104 |
+
|
| 105 |
+
The `embedding` is only written on CREATE (see `coalesce(e.embedding, ...)`
|
| 106 |
+
in the Cypher), so an L3-aliased new form doesn't overwrite the existing
|
| 107 |
+
canonical entity's vector. Callers that skip semantic alignment can
|
| 108 |
+
leave `embedding=None`.
|
| 109 |
+
"""
|
| 110 |
driver = await get_driver()
|
| 111 |
async with driver.session() as session:
|
| 112 |
await session.run(
|
|
|
|
| 117 |
name_lower=name_lower,
|
| 118 |
type=entity_type,
|
| 119 |
description=description,
|
| 120 |
+
embedding=embedding,
|
| 121 |
)
|
| 122 |
|
| 123 |
|
|
|
|
| 141 |
)
|
| 142 |
|
| 143 |
|
| 144 |
+
async def list_entity_semantic_candidates(
|
| 145 |
+
user_id: str,
|
| 146 |
+
) -> list[tuple[str, str, list[float]]]:
|
| 147 |
+
"""Return (name_lower, type, embedding) for entities the user owns that
|
| 148 |
+
already have a stored embedding. Used by L3 semantic alignment (#43c).
|
| 149 |
+
|
| 150 |
+
Entities pre-#43c or entities skipped by the semantic-align flag lack an
|
| 151 |
+
`embedding` property; those are excluded here. The migration script
|
| 152 |
+
(`scripts/migrate_kg_embeddings.py`) backfills them.
|
| 153 |
+
"""
|
| 154 |
+
driver = await get_driver()
|
| 155 |
+
async with driver.session() as session:
|
| 156 |
+
result = await session.run(
|
| 157 |
+
"MATCH (e:Entity {user_id: $user_id}) "
|
| 158 |
+
"WHERE e.embedding IS NOT NULL AND e.name_lower IS NOT NULL "
|
| 159 |
+
"RETURN e.name_lower AS name_lower, "
|
| 160 |
+
"coalesce(e.type, 'Concept') AS type, "
|
| 161 |
+
"e.embedding AS embedding",
|
| 162 |
+
user_id=user_id,
|
| 163 |
+
)
|
| 164 |
+
return [(row["name_lower"], row["type"], list(row["embedding"])) async for row in result]
|
| 165 |
+
|
| 166 |
+
|
| 167 |
async def list_entity_candidates(user_id: str) -> list[tuple[str, str]]:
|
| 168 |
"""Return (name_lower, type) tuples for every entity the user owns.
|
| 169 |
|
app/graph/semantic_alignment.py
ADDED
|
@@ -0,0 +1,102 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""L3 semantic entity alignment via bge-small-en-v1.5 embeddings.
|
| 2 |
+
|
| 3 |
+
Catches abbreviation/expansion pairs that L1 (normalization) and L2 (fuzzy
|
| 4 |
+
string) miss:
|
| 5 |
+
|
| 6 |
+
SFT ↔ Supervised Fine-Tuning (fuzz score ~10, way below L2 threshold)
|
| 7 |
+
GWU ↔ George Washington University
|
| 8 |
+
RSAT ↔ RSAT Model
|
| 9 |
+
Explainable AI (XAI) ↔ explainable AI
|
| 10 |
+
|
| 11 |
+
Approach: on ingest, embed each entity's `name: description` through the same
|
| 12 |
+
bge model retrieval uses, then cosine-similarity against existing same-type
|
| 13 |
+
entities' stored embeddings. If similarity clears the threshold, alias to the
|
| 14 |
+
existing entity (its Neo4j node stays as-is; the incoming entity merges in via
|
| 15 |
+
its canonical name_lower). Downstream relationship endpoints get rewritten
|
| 16 |
+
through the alias map.
|
| 17 |
+
|
| 18 |
+
Embeddings live on the :Entity Neo4j node (property `embedding: list[float]`),
|
| 19 |
+
set only ON CREATE so an aliased write doesn't overwrite the existing
|
| 20 |
+
authoritative embedding.
|
| 21 |
+
"""
|
| 22 |
+
|
| 23 |
+
from __future__ import annotations
|
| 24 |
+
|
| 25 |
+
from collections.abc import Iterable
|
| 26 |
+
from dataclasses import dataclass
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
@dataclass(frozen=True)
|
| 30 |
+
class SemanticCandidate:
|
| 31 |
+
"""An existing entity fetched from Neo4j that a new entity may semantically
|
| 32 |
+
alias to. Immutable so it can go into sets / used as dict keys downstream."""
|
| 33 |
+
|
| 34 |
+
name_lower: str
|
| 35 |
+
type: str
|
| 36 |
+
# bge-small-en-v1.5 outputs are L2-normalized 384-dim vectors, so cosine
|
| 37 |
+
# similarity == dot product. Storing as a tuple keeps the dataclass hashable.
|
| 38 |
+
embedding: tuple[float, ...]
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
def format_entity_text(name: str, description: str | None = None) -> str:
|
| 42 |
+
"""The string handed to the embedding model.
|
| 43 |
+
|
| 44 |
+
Description adds a strong semantic anchor when present ('SFT' alone is
|
| 45 |
+
ambiguous; 'SFT: supervised fine-tuning method' is not). For entities
|
| 46 |
+
without a description we fall back to just the name — sub-optimal but the
|
| 47 |
+
caller may still get useful similarity on longer canonical names.
|
| 48 |
+
"""
|
| 49 |
+
name = (name or "").strip()
|
| 50 |
+
desc = (description or "").strip()
|
| 51 |
+
if desc:
|
| 52 |
+
return f"{name}: {desc}"
|
| 53 |
+
return name
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
def _dot(a: tuple[float, ...] | list[float], b: tuple[float, ...] | list[float]) -> float:
|
| 57 |
+
"""Dot product of two equal-length vectors. bge outputs are L2-normalized,
|
| 58 |
+
so this equals cosine similarity — no re-normalization needed."""
|
| 59 |
+
# Keeping this hand-rolled rather than pulling numpy in for a hot path
|
| 60 |
+
# that runs on modest-sized entity lists. For a 384-dim × 500-entity
|
| 61 |
+
# pool this loops ~200k floats — well under 1ms in CPython.
|
| 62 |
+
total = 0.0
|
| 63 |
+
for x, y in zip(a, b, strict=True):
|
| 64 |
+
total += x * y
|
| 65 |
+
return total
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def find_semantic_alias(
|
| 69 |
+
embedding: list[float] | tuple[float, ...],
|
| 70 |
+
entity_type: str,
|
| 71 |
+
candidates: Iterable[SemanticCandidate],
|
| 72 |
+
threshold: float,
|
| 73 |
+
) -> str | None:
|
| 74 |
+
"""Return the name_lower of the best same-type candidate whose cosine
|
| 75 |
+
similarity to *embedding* clears *threshold*, else None.
|
| 76 |
+
|
| 77 |
+
Type-scoping prevents cross-type false positives (a Person named "Java"
|
| 78 |
+
won't collapse into a Technology named "Java"). Exact-similarity ties
|
| 79 |
+
are broken by iteration order — first-seen wins, matching Neo4j's
|
| 80 |
+
first-write-wins display-name convention.
|
| 81 |
+
"""
|
| 82 |
+
if not embedding:
|
| 83 |
+
return None
|
| 84 |
+
best_score = 0.0
|
| 85 |
+
best_match: str | None = None
|
| 86 |
+
for cand in candidates:
|
| 87 |
+
if cand.type != entity_type:
|
| 88 |
+
continue
|
| 89 |
+
if not cand.embedding:
|
| 90 |
+
continue
|
| 91 |
+
try:
|
| 92 |
+
score = _dot(embedding, cand.embedding)
|
| 93 |
+
except ValueError:
|
| 94 |
+
# Dimension mismatch (e.g. an old row with a partial vector) —
|
| 95 |
+
# skip rather than crash the whole batch.
|
| 96 |
+
continue
|
| 97 |
+
if score > best_score:
|
| 98 |
+
best_score = score
|
| 99 |
+
best_match = cand.name_lower
|
| 100 |
+
if best_score >= threshold:
|
| 101 |
+
return best_match
|
| 102 |
+
return None
|
app/workers/tasks.py
CHANGED
|
@@ -184,14 +184,26 @@ GRAPH_EXTRACT_BACKOFF_SECONDS = (15, 30) # sleep between attempt 1→2 and 2→
|
|
| 184 |
async def _ingest_graph(*, chunks: list[str], user_id: str, document_id: str) -> None:
|
| 185 |
# Best-effort: swallow errors so the worker stays non-blocking.
|
| 186 |
try:
|
| 187 |
-
from app.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 188 |
from app.graph.extraction import safe_extract_entities
|
| 189 |
from app.graph.neo4j_client import (
|
| 190 |
ensure_indexes,
|
| 191 |
list_entity_candidates,
|
|
|
|
| 192 |
upsert_entity,
|
| 193 |
upsert_relationship,
|
| 194 |
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 195 |
|
| 196 |
# Groq's json_object validator flakes intermittently — the same
|
| 197 |
# extraction that returned 0 entities on upload will often return
|
|
@@ -235,6 +247,89 @@ async def _ingest_graph(*, chunks: list[str], user_id: str, document_id: str) ->
|
|
| 235 |
result.relationships,
|
| 236 |
existing,
|
| 237 |
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 238 |
|
| 239 |
for e in aligned_entities:
|
| 240 |
await upsert_entity(
|
|
@@ -244,6 +339,7 @@ async def _ingest_graph(*, chunks: list[str], user_id: str, document_id: str) ->
|
|
| 244 |
name_lower=e.name_lower,
|
| 245 |
entity_type=e.type,
|
| 246 |
description=e.description,
|
|
|
|
| 247 |
)
|
| 248 |
for r in aligned_rels:
|
| 249 |
await upsert_relationship(
|
|
@@ -254,11 +350,14 @@ async def _ingest_graph(*, chunks: list[str], user_id: str, document_id: str) ->
|
|
| 254 |
relation=r.relation,
|
| 255 |
)
|
| 256 |
log.info(
|
| 257 |
-
"graph: doc=%s chunks=%d entities=%d rels=%d
|
|
|
|
| 258 |
document_id,
|
| 259 |
len(chunks),
|
| 260 |
len(aligned_entities),
|
| 261 |
len(aligned_rels),
|
|
|
|
|
|
|
| 262 |
len(result.entities),
|
| 263 |
len(result.relationships),
|
| 264 |
)
|
|
|
|
| 184 |
async def _ingest_graph(*, chunks: list[str], user_id: str, document_id: str) -> None:
|
| 185 |
# Best-effort: swallow errors so the worker stays non-blocking.
|
| 186 |
try:
|
| 187 |
+
from app.embeddings import embed_texts
|
| 188 |
+
from app.graph.alignment import (
|
| 189 |
+
AlignedEntity,
|
| 190 |
+
AlignedRelationship,
|
| 191 |
+
Candidate,
|
| 192 |
+
align_batch,
|
| 193 |
+
)
|
| 194 |
from app.graph.extraction import safe_extract_entities
|
| 195 |
from app.graph.neo4j_client import (
|
| 196 |
ensure_indexes,
|
| 197 |
list_entity_candidates,
|
| 198 |
+
list_entity_semantic_candidates,
|
| 199 |
upsert_entity,
|
| 200 |
upsert_relationship,
|
| 201 |
)
|
| 202 |
+
from app.graph.semantic_alignment import (
|
| 203 |
+
SemanticCandidate,
|
| 204 |
+
find_semantic_alias,
|
| 205 |
+
format_entity_text,
|
| 206 |
+
)
|
| 207 |
|
| 208 |
# Groq's json_object validator flakes intermittently — the same
|
| 209 |
# extraction that returned 0 entities on upload will often return
|
|
|
|
| 247 |
result.relationships,
|
| 248 |
existing,
|
| 249 |
)
|
| 250 |
+
pre_l3_ent_count = len(aligned_entities)
|
| 251 |
+
pre_l3_rel_count = len(aligned_rels)
|
| 252 |
+
|
| 253 |
+
# Semantic alignment (#43c): L3 layer on top of L1/L2. Embed each
|
| 254 |
+
# aligned entity's `name: description`, compare cosine similarity
|
| 255 |
+
# against existing same-type entities' stored embeddings. If a
|
| 256 |
+
# match clears the threshold, rewrite the entity's name_lower to
|
| 257 |
+
# the existing canonical form and cascade the change through the
|
| 258 |
+
# relationships. Catches abbreviation/expansion pairs (SFT ↔
|
| 259 |
+
# Supervised Fine-Tuning) that L1/L2 string-fuzzy can't.
|
| 260 |
+
entity_embeddings: dict[str, list[float]] = {}
|
| 261 |
+
if settings.graph_semantic_align and aligned_entities:
|
| 262 |
+
existing_semantic = await list_entity_semantic_candidates(user_id)
|
| 263 |
+
running_candidates = [
|
| 264 |
+
SemanticCandidate(
|
| 265 |
+
name_lower=nl,
|
| 266 |
+
type=t,
|
| 267 |
+
embedding=tuple(emb),
|
| 268 |
+
)
|
| 269 |
+
for nl, t, emb in existing_semantic
|
| 270 |
+
]
|
| 271 |
+
texts = [format_entity_text(e.name, e.description) for e in aligned_entities]
|
| 272 |
+
embeddings = await asyncio.to_thread(embed_texts, texts)
|
| 273 |
+
alias_map: dict[str, str] = {} # old name_lower → canonical
|
| 274 |
+
for e, emb in zip(aligned_entities, embeddings, strict=True):
|
| 275 |
+
match = find_semantic_alias(
|
| 276 |
+
emb,
|
| 277 |
+
e.type,
|
| 278 |
+
running_candidates,
|
| 279 |
+
settings.graph_semantic_align_threshold,
|
| 280 |
+
)
|
| 281 |
+
if match and match != e.name_lower:
|
| 282 |
+
alias_map[e.name_lower] = match
|
| 283 |
+
else:
|
| 284 |
+
# Not aliased — this entity becomes a candidate that
|
| 285 |
+
# later same-batch entities (or future docs) can align to.
|
| 286 |
+
running_candidates.append(
|
| 287 |
+
SemanticCandidate(
|
| 288 |
+
name_lower=e.name_lower,
|
| 289 |
+
type=e.type,
|
| 290 |
+
embedding=tuple(emb),
|
| 291 |
+
)
|
| 292 |
+
)
|
| 293 |
+
entity_embeddings[e.name_lower] = emb
|
| 294 |
+
|
| 295 |
+
if alias_map:
|
| 296 |
+
seen: set[str] = set()
|
| 297 |
+
collapsed: list[AlignedEntity] = []
|
| 298 |
+
for e in aligned_entities:
|
| 299 |
+
canonical = alias_map.get(e.name_lower, e.name_lower)
|
| 300 |
+
if canonical in seen:
|
| 301 |
+
continue
|
| 302 |
+
seen.add(canonical)
|
| 303 |
+
if canonical == e.name_lower:
|
| 304 |
+
collapsed.append(e)
|
| 305 |
+
else:
|
| 306 |
+
# Aliased into a canonical form — keep the aliased
|
| 307 |
+
# display name; ON CREATE on the existing node is
|
| 308 |
+
# a no-op so display stays as first-write.
|
| 309 |
+
collapsed.append(
|
| 310 |
+
AlignedEntity(
|
| 311 |
+
name=e.name,
|
| 312 |
+
name_lower=canonical,
|
| 313 |
+
type=e.type,
|
| 314 |
+
description=e.description,
|
| 315 |
+
)
|
| 316 |
+
)
|
| 317 |
+
aligned_entities = collapsed
|
| 318 |
+
|
| 319 |
+
rewritten_rels: list[AlignedRelationship] = []
|
| 320 |
+
for r in aligned_rels:
|
| 321 |
+
src = alias_map.get(r.source_lower, r.source_lower)
|
| 322 |
+
tgt = alias_map.get(r.target_lower, r.target_lower)
|
| 323 |
+
if src == tgt:
|
| 324 |
+
continue # self-loop after L3
|
| 325 |
+
rewritten_rels.append(
|
| 326 |
+
AlignedRelationship(
|
| 327 |
+
source_lower=src,
|
| 328 |
+
target_lower=tgt,
|
| 329 |
+
relation=r.relation,
|
| 330 |
+
)
|
| 331 |
+
)
|
| 332 |
+
aligned_rels = rewritten_rels
|
| 333 |
|
| 334 |
for e in aligned_entities:
|
| 335 |
await upsert_entity(
|
|
|
|
| 339 |
name_lower=e.name_lower,
|
| 340 |
entity_type=e.type,
|
| 341 |
description=e.description,
|
| 342 |
+
embedding=entity_embeddings.get(e.name_lower),
|
| 343 |
)
|
| 344 |
for r in aligned_rels:
|
| 345 |
await upsert_relationship(
|
|
|
|
| 350 |
relation=r.relation,
|
| 351 |
)
|
| 352 |
log.info(
|
| 353 |
+
"graph: doc=%s chunks=%d entities=%d rels=%d "
|
| 354 |
+
"(pre-l3 %d ents, %d rels; raw %d ents, %d rels)",
|
| 355 |
document_id,
|
| 356 |
len(chunks),
|
| 357 |
len(aligned_entities),
|
| 358 |
len(aligned_rels),
|
| 359 |
+
pre_l3_ent_count,
|
| 360 |
+
pre_l3_rel_count,
|
| 361 |
len(result.entities),
|
| 362 |
len(result.relationships),
|
| 363 |
)
|
scripts/migrate_kg_embeddings.py
ADDED
|
@@ -0,0 +1,86 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""One-shot migration for #43c: backfill `embedding` on every existing
|
| 2 |
+
:Entity node so L3 semantic alignment has candidates to compare against.
|
| 3 |
+
|
| 4 |
+
Idempotent — safe to re-run. Rows already carrying an embedding are skipped.
|
| 5 |
+
Batches embedding calls through the same bge-small model retrieval uses.
|
| 6 |
+
|
| 7 |
+
Usage (from the backend/ directory, with the .env pointing at prod Neo4j):
|
| 8 |
+
python -m scripts.migrate_kg_embeddings
|
| 9 |
+
"""
|
| 10 |
+
|
| 11 |
+
from __future__ import annotations
|
| 12 |
+
|
| 13 |
+
import asyncio
|
| 14 |
+
import logging
|
| 15 |
+
|
| 16 |
+
from app.embeddings import embed_texts
|
| 17 |
+
from app.graph.neo4j_client import close_driver, ensure_indexes, get_driver
|
| 18 |
+
from app.graph.semantic_alignment import format_entity_text
|
| 19 |
+
|
| 20 |
+
log = logging.getLogger("mmap.migrate_kg_embeddings")
|
| 21 |
+
|
| 22 |
+
BATCH_SIZE = 128 # bge-small handles this comfortably on CPU
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
async def _run() -> None:
|
| 26 |
+
await ensure_indexes()
|
| 27 |
+
driver = await get_driver()
|
| 28 |
+
|
| 29 |
+
async with driver.session() as session:
|
| 30 |
+
# Every entity missing an embedding, keyed by internal node id for the
|
| 31 |
+
# write-back step below. `coalesce` guards against nulls in description.
|
| 32 |
+
result = await session.run(
|
| 33 |
+
"MATCH (e:Entity) "
|
| 34 |
+
"WHERE e.embedding IS NULL "
|
| 35 |
+
"RETURN id(e) AS node_id, e.name AS name, "
|
| 36 |
+
"coalesce(e.description, '') AS description, e.user_id AS user_id"
|
| 37 |
+
)
|
| 38 |
+
rows = [dict(record) async for record in result]
|
| 39 |
+
|
| 40 |
+
if not rows:
|
| 41 |
+
log.info("migrate_kg_embeddings: no rows needed backfill — every entity has an embedding.")
|
| 42 |
+
await close_driver()
|
| 43 |
+
return
|
| 44 |
+
|
| 45 |
+
log.info("migrate_kg_embeddings: embedding %d entities in batches of %d", len(rows), BATCH_SIZE)
|
| 46 |
+
|
| 47 |
+
per_user: dict[str, int] = {}
|
| 48 |
+
async with driver.session() as session:
|
| 49 |
+
for batch_start in range(0, len(rows), BATCH_SIZE):
|
| 50 |
+
batch = rows[batch_start : batch_start + BATCH_SIZE]
|
| 51 |
+
texts = [format_entity_text(r["name"], r["description"]) for r in batch]
|
| 52 |
+
# Runs in the current thread — this is a one-shot admin script,
|
| 53 |
+
# blocking is fine.
|
| 54 |
+
vectors = embed_texts(texts)
|
| 55 |
+
|
| 56 |
+
for row, vec in zip(batch, vectors, strict=True):
|
| 57 |
+
await session.run(
|
| 58 |
+
"MATCH (e) WHERE id(e) = $nid SET e.embedding = $emb",
|
| 59 |
+
nid=row["node_id"],
|
| 60 |
+
emb=vec,
|
| 61 |
+
)
|
| 62 |
+
uid = str(row.get("user_id") or "")
|
| 63 |
+
per_user[uid] = per_user.get(uid, 0) + 1
|
| 64 |
+
|
| 65 |
+
log.info(
|
| 66 |
+
"migrate_kg_embeddings: batch %d/%d done",
|
| 67 |
+
batch_start // BATCH_SIZE + 1,
|
| 68 |
+
(len(rows) + BATCH_SIZE - 1) // BATCH_SIZE,
|
| 69 |
+
)
|
| 70 |
+
|
| 71 |
+
for uid, n in sorted(per_user.items()):
|
| 72 |
+
log.info("migrate_kg_embeddings: user=%s backfilled=%d", uid, n)
|
| 73 |
+
log.info("migrate_kg_embeddings: done.")
|
| 74 |
+
await close_driver()
|
| 75 |
+
|
| 76 |
+
|
| 77 |
+
def main() -> None:
|
| 78 |
+
logging.basicConfig(
|
| 79 |
+
level=logging.INFO,
|
| 80 |
+
format="%(asctime)s %(levelname)s %(name)s %(message)s",
|
| 81 |
+
)
|
| 82 |
+
asyncio.run(_run())
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
if __name__ == "__main__":
|
| 86 |
+
main()
|
tests/test_graph_semantic_alignment.py
ADDED
|
@@ -0,0 +1,100 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Unit tests for L3 semantic entity alignment (#43c)."""
|
| 2 |
+
|
| 3 |
+
from app.graph.semantic_alignment import (
|
| 4 |
+
SemanticCandidate,
|
| 5 |
+
find_semantic_alias,
|
| 6 |
+
format_entity_text,
|
| 7 |
+
)
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
class TestFormatEntityText:
|
| 11 |
+
def test_name_plus_description(self):
|
| 12 |
+
assert format_entity_text("SFT", "supervised fine-tuning method") == (
|
| 13 |
+
"SFT: supervised fine-tuning method"
|
| 14 |
+
)
|
| 15 |
+
|
| 16 |
+
def test_falls_back_to_bare_name_without_description(self):
|
| 17 |
+
assert format_entity_text("RSAT") == "RSAT"
|
| 18 |
+
assert format_entity_text("RSAT", "") == "RSAT"
|
| 19 |
+
|
| 20 |
+
def test_strips_surrounding_whitespace(self):
|
| 21 |
+
assert format_entity_text(" Alice ", " researcher ") == "Alice: researcher"
|
| 22 |
+
|
| 23 |
+
def test_empty_name_yields_empty_string(self):
|
| 24 |
+
assert format_entity_text("", "desc") == ": desc" # explicit — caller filters
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def _v(*parts: float) -> tuple[float, ...]:
|
| 28 |
+
"""L2-normalize a small vector so the dot product acts like cosine sim."""
|
| 29 |
+
total = sum(p * p for p in parts) ** 0.5 or 1.0
|
| 30 |
+
return tuple(p / total for p in parts)
|
| 31 |
+
|
| 32 |
+
|
| 33 |
+
class TestFindSemanticAlias:
|
| 34 |
+
def test_returns_none_when_no_candidates(self):
|
| 35 |
+
assert find_semantic_alias(list(_v(1.0, 0.0)), "Concept", [], 0.85) is None
|
| 36 |
+
|
| 37 |
+
def test_returns_none_when_type_mismatch(self):
|
| 38 |
+
# Identical vector but different type — must NOT alias across types.
|
| 39 |
+
emb = _v(1.0, 0.0)
|
| 40 |
+
cand = SemanticCandidate(name_lower="x", type="Person", embedding=emb)
|
| 41 |
+
assert find_semantic_alias(list(emb), "Concept", [cand], 0.85) is None
|
| 42 |
+
|
| 43 |
+
def test_returns_match_above_threshold(self):
|
| 44 |
+
# Near-identical vectors → cosine ≈ 1.
|
| 45 |
+
cand = SemanticCandidate(
|
| 46 |
+
name_lower="supervised fine-tuning",
|
| 47 |
+
type="Concept",
|
| 48 |
+
embedding=_v(0.9, 0.1),
|
| 49 |
+
)
|
| 50 |
+
# Slightly rotated but still very similar.
|
| 51 |
+
query = list(_v(0.95, 0.08))
|
| 52 |
+
assert find_semantic_alias(query, "Concept", [cand], 0.85) == "supervised fine-tuning"
|
| 53 |
+
|
| 54 |
+
def test_returns_none_below_threshold(self):
|
| 55 |
+
cand = SemanticCandidate(
|
| 56 |
+
name_lower="something else",
|
| 57 |
+
type="Concept",
|
| 58 |
+
embedding=_v(0.0, 1.0),
|
| 59 |
+
)
|
| 60 |
+
# Orthogonal vectors → cosine 0.
|
| 61 |
+
assert find_semantic_alias(list(_v(1.0, 0.0)), "Concept", [cand], 0.85) is None
|
| 62 |
+
|
| 63 |
+
def test_picks_highest_scoring_of_multiple_same_type(self):
|
| 64 |
+
query = _v(1.0, 0.0, 0.0)
|
| 65 |
+
cands = [
|
| 66 |
+
SemanticCandidate("far", "Concept", _v(0.0, 1.0, 0.0)),
|
| 67 |
+
SemanticCandidate("near", "Concept", _v(0.99, 0.05, 0.02)),
|
| 68 |
+
SemanticCandidate("mid", "Concept", _v(0.7, 0.7, 0.0)),
|
| 69 |
+
]
|
| 70 |
+
assert find_semantic_alias(list(query), "Concept", cands, 0.85) == "near"
|
| 71 |
+
|
| 72 |
+
def test_ignores_candidates_with_empty_embedding(self):
|
| 73 |
+
query = _v(1.0, 0.0)
|
| 74 |
+
cands = [
|
| 75 |
+
SemanticCandidate("no-vec", "Concept", ()),
|
| 76 |
+
SemanticCandidate("has-vec", "Concept", _v(0.99, 0.01)),
|
| 77 |
+
]
|
| 78 |
+
assert find_semantic_alias(list(query), "Concept", cands, 0.85) == "has-vec"
|
| 79 |
+
|
| 80 |
+
def test_ignores_candidates_with_dim_mismatch(self):
|
| 81 |
+
# Corrupted/partial vector from a bad backfill — must not crash.
|
| 82 |
+
query = _v(1.0, 0.0, 0.0)
|
| 83 |
+
cands = [
|
| 84 |
+
SemanticCandidate("short", "Concept", (0.9,)), # 1-d, dim mismatch
|
| 85 |
+
SemanticCandidate("good", "Concept", _v(0.99, 0.01, 0.02)),
|
| 86 |
+
]
|
| 87 |
+
assert find_semantic_alias(list(query), "Concept", cands, 0.85) == "good"
|
| 88 |
+
|
| 89 |
+
def test_empty_query_returns_none(self):
|
| 90 |
+
cand = SemanticCandidate("x", "Concept", _v(1.0, 0.0))
|
| 91 |
+
assert find_semantic_alias([], "Concept", [cand], 0.85) is None
|
| 92 |
+
|
| 93 |
+
def test_threshold_is_inclusive(self):
|
| 94 |
+
"""Score == threshold should still alias — matches the >= boundary
|
| 95 |
+
used for L2 fuzzy in alignment.py."""
|
| 96 |
+
# dot product exactly 0.85 by construction.
|
| 97 |
+
query = (0.85, (1 - 0.85**2) ** 0.5)
|
| 98 |
+
cand_emb = (1.0, 0.0)
|
| 99 |
+
cand = SemanticCandidate("threshold-edge", "Concept", cand_emb)
|
| 100 |
+
assert find_semantic_alias(list(query), "Concept", [cand], 0.85) == "threshold-edge"
|
tests/test_workers_ingest_graph.py
CHANGED
|
@@ -1,6 +1,7 @@
|
|
| 1 |
"""Unit tests for _ingest_graph's retry-on-transient-Groq-failure behavior."""
|
| 2 |
|
| 3 |
-
from
|
|
|
|
| 4 |
|
| 5 |
import pytest
|
| 6 |
|
|
@@ -9,6 +10,46 @@ from app.graph.schema import ExtractionResult, GraphEntity, GraphRelationship
|
|
| 9 |
from app.workers import tasks
|
| 10 |
|
| 11 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 12 |
def _outcome_empty_transient() -> ExtractionOutcome:
|
| 13 |
return ExtractionOutcome(
|
| 14 |
result=ExtractionResult(entities=[], relationships=[]),
|
|
@@ -52,6 +93,7 @@ async def test_retries_transient_failure_until_it_succeeds():
|
|
| 52 |
patch("app.graph.neo4j_client.upsert_relationship", new=AsyncMock()),
|
| 53 |
# asyncio.sleep is called between attempts — no-op it so the test is fast.
|
| 54 |
patch("app.workers.tasks.asyncio.sleep", new=AsyncMock()),
|
|
|
|
| 55 |
):
|
| 56 |
await tasks._ingest_graph(chunks=["something"], user_id="u-1", document_id="d-1")
|
| 57 |
|
|
@@ -63,6 +105,8 @@ async def test_retries_transient_failure_until_it_succeeds():
|
|
| 63 |
kwargs = upsert.await_args.kwargs
|
| 64 |
assert kwargs["name"] == "Qdrant"
|
| 65 |
assert kwargs["name_lower"] == "qdrant"
|
|
|
|
|
|
|
| 66 |
|
| 67 |
|
| 68 |
@pytest.mark.asyncio
|
|
@@ -107,6 +151,7 @@ async def test_first_attempt_success_does_not_retry():
|
|
| 107 |
patch("app.graph.neo4j_client.upsert_entity", new=AsyncMock()),
|
| 108 |
patch("app.graph.neo4j_client.upsert_relationship", new=AsyncMock()),
|
| 109 |
patch("app.workers.tasks.asyncio.sleep", new=AsyncMock()) as sleep,
|
|
|
|
| 110 |
):
|
| 111 |
await tasks._ingest_graph(chunks=["something"], user_id="u-1", document_id="d-1")
|
| 112 |
|
|
@@ -152,6 +197,7 @@ async def test_alignment_collapses_dupes_before_upsert():
|
|
| 152 |
patch("app.graph.neo4j_client.upsert_entity", new=entity_ups),
|
| 153 |
patch("app.graph.neo4j_client.upsert_relationship", new=rel_ups),
|
| 154 |
patch("app.workers.tasks.asyncio.sleep", new=AsyncMock()),
|
|
|
|
| 155 |
):
|
| 156 |
await tasks._ingest_graph(chunks=["anything"], user_id="u-1", document_id="d-1")
|
| 157 |
|
|
@@ -166,3 +212,81 @@ async def test_alignment_collapses_dupes_before_upsert():
|
|
| 166 |
assert rk["source_lower"] == "jugal gajjar"
|
| 167 |
assert rk["target_lower"] == "gwu"
|
| 168 |
assert rk["relation"] == "affiliated with"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
"""Unit tests for _ingest_graph's retry-on-transient-Groq-failure behavior."""
|
| 2 |
|
| 3 |
+
from contextlib import contextmanager
|
| 4 |
+
from unittest.mock import AsyncMock, MagicMock, patch
|
| 5 |
|
| 6 |
import pytest
|
| 7 |
|
|
|
|
| 10 |
from app.workers import tasks
|
| 11 |
|
| 12 |
|
| 13 |
+
def _fake_embed(dim: int = 4):
|
| 14 |
+
"""Return a stub for `app.embeddings.embed_texts` that emits `dim`-sized
|
| 15 |
+
vectors, one per input string, deterministic by hash so identical names
|
| 16 |
+
produce identical vectors (so semantic alignment doesn't collapse
|
| 17 |
+
unrelated entities in fixtures)."""
|
| 18 |
+
|
| 19 |
+
def _emb(texts):
|
| 20 |
+
out = []
|
| 21 |
+
for t in texts:
|
| 22 |
+
h = hash(t) & 0xFFFF
|
| 23 |
+
# Tiny "vector" that's L2-normalized enough for the assertions
|
| 24 |
+
# this file cares about. Real bge vectors are 384-d.
|
| 25 |
+
base = [((h >> i) & 0xF) / 15.0 for i in range(dim)]
|
| 26 |
+
norm = sum(x * x for x in base) ** 0.5 or 1.0
|
| 27 |
+
out.append([x / norm for x in base])
|
| 28 |
+
return out
|
| 29 |
+
|
| 30 |
+
return _emb
|
| 31 |
+
|
| 32 |
+
|
| 33 |
+
@contextmanager
|
| 34 |
+
def _semantic_mocks(semantic_candidates=None):
|
| 35 |
+
"""Stub the L3 semantic-alignment I/O so tests can focus on other logic.
|
| 36 |
+
|
| 37 |
+
- `list_entity_semantic_candidates` returns an empty list by default.
|
| 38 |
+
- `embed_texts` returns deterministic tiny vectors.
|
| 39 |
+
Callers can pass `semantic_candidates` to seed the DB side.
|
| 40 |
+
"""
|
| 41 |
+
fake_list = AsyncMock(return_value=semantic_candidates or [])
|
| 42 |
+
fake_embed = MagicMock(side_effect=_fake_embed())
|
| 43 |
+
with (
|
| 44 |
+
patch(
|
| 45 |
+
"app.graph.neo4j_client.list_entity_semantic_candidates",
|
| 46 |
+
new=fake_list,
|
| 47 |
+
),
|
| 48 |
+
patch("app.embeddings.embed_texts", new=fake_embed),
|
| 49 |
+
):
|
| 50 |
+
yield fake_list, fake_embed
|
| 51 |
+
|
| 52 |
+
|
| 53 |
def _outcome_empty_transient() -> ExtractionOutcome:
|
| 54 |
return ExtractionOutcome(
|
| 55 |
result=ExtractionResult(entities=[], relationships=[]),
|
|
|
|
| 93 |
patch("app.graph.neo4j_client.upsert_relationship", new=AsyncMock()),
|
| 94 |
# asyncio.sleep is called between attempts — no-op it so the test is fast.
|
| 95 |
patch("app.workers.tasks.asyncio.sleep", new=AsyncMock()),
|
| 96 |
+
_semantic_mocks(),
|
| 97 |
):
|
| 98 |
await tasks._ingest_graph(chunks=["something"], user_id="u-1", document_id="d-1")
|
| 99 |
|
|
|
|
| 105 |
kwargs = upsert.await_args.kwargs
|
| 106 |
assert kwargs["name"] == "Qdrant"
|
| 107 |
assert kwargs["name_lower"] == "qdrant"
|
| 108 |
+
# #43c: embedding kwarg goes through; new entity gets its computed vector.
|
| 109 |
+
assert kwargs.get("embedding") is not None
|
| 110 |
|
| 111 |
|
| 112 |
@pytest.mark.asyncio
|
|
|
|
| 151 |
patch("app.graph.neo4j_client.upsert_entity", new=AsyncMock()),
|
| 152 |
patch("app.graph.neo4j_client.upsert_relationship", new=AsyncMock()),
|
| 153 |
patch("app.workers.tasks.asyncio.sleep", new=AsyncMock()) as sleep,
|
| 154 |
+
_semantic_mocks(),
|
| 155 |
):
|
| 156 |
await tasks._ingest_graph(chunks=["something"], user_id="u-1", document_id="d-1")
|
| 157 |
|
|
|
|
| 197 |
patch("app.graph.neo4j_client.upsert_entity", new=entity_ups),
|
| 198 |
patch("app.graph.neo4j_client.upsert_relationship", new=rel_ups),
|
| 199 |
patch("app.workers.tasks.asyncio.sleep", new=AsyncMock()),
|
| 200 |
+
_semantic_mocks(),
|
| 201 |
):
|
| 202 |
await tasks._ingest_graph(chunks=["anything"], user_id="u-1", document_id="d-1")
|
| 203 |
|
|
|
|
| 212 |
assert rk["source_lower"] == "jugal gajjar"
|
| 213 |
assert rk["target_lower"] == "gwu"
|
| 214 |
assert rk["relation"] == "affiliated with"
|
| 215 |
+
|
| 216 |
+
|
| 217 |
+
@pytest.mark.asyncio
|
| 218 |
+
async def test_semantic_alignment_aliases_to_existing_high_similarity(monkeypatch):
|
| 219 |
+
"""#43c: L1/L2 fuzzy misses (SFT vs Supervised Fine-Tuning has fuzz
|
| 220 |
+
score way below 90) but L3 semantic catches — extracted 'Supervised
|
| 221 |
+
Fine-Tuning' should alias to the existing 'sft' node when their
|
| 222 |
+
embeddings clear the threshold."""
|
| 223 |
+
from app.core.config import settings
|
| 224 |
+
|
| 225 |
+
monkeypatch.setattr(settings, "graph_semantic_align_threshold", 0.5)
|
| 226 |
+
|
| 227 |
+
outcome = ExtractionOutcome(
|
| 228 |
+
result=ExtractionResult(
|
| 229 |
+
entities=[
|
| 230 |
+
GraphEntity(
|
| 231 |
+
name="Supervised Fine-Tuning",
|
| 232 |
+
type="Concept",
|
| 233 |
+
description="SFT training phase",
|
| 234 |
+
),
|
| 235 |
+
],
|
| 236 |
+
relationships=[],
|
| 237 |
+
),
|
| 238 |
+
transient_failure=False,
|
| 239 |
+
)
|
| 240 |
+
|
| 241 |
+
# Existing DB entity: "sft" with an embedding. Our fake embedder is
|
| 242 |
+
# deterministic-by-hash so we build a vector that MATCHES what the fake
|
| 243 |
+
# will emit for the new "Supervised Fine-Tuning: SFT training phase"
|
| 244 |
+
# input — then we can be sure similarity clears our lowered threshold.
|
| 245 |
+
from app.graph.semantic_alignment import format_entity_text
|
| 246 |
+
|
| 247 |
+
fake_embed_fn = _fake_embed()
|
| 248 |
+
incoming_text = format_entity_text("Supervised Fine-Tuning", "SFT training phase")
|
| 249 |
+
matching_vec = fake_embed_fn([incoming_text])[0]
|
| 250 |
+
|
| 251 |
+
entity_ups = AsyncMock()
|
| 252 |
+
rel_ups = AsyncMock()
|
| 253 |
+
with (
|
| 254 |
+
patch(
|
| 255 |
+
"app.graph.extraction.safe_extract_entities",
|
| 256 |
+
new=AsyncMock(return_value=outcome),
|
| 257 |
+
),
|
| 258 |
+
patch("app.graph.neo4j_client.ensure_indexes", new=AsyncMock()),
|
| 259 |
+
patch(
|
| 260 |
+
"app.graph.neo4j_client.list_entity_candidates",
|
| 261 |
+
new=AsyncMock(return_value=[]), # L2 has no match for "supervised fine-tuning"
|
| 262 |
+
),
|
| 263 |
+
patch(
|
| 264 |
+
"app.graph.neo4j_client.list_entity_semantic_candidates",
|
| 265 |
+
new=AsyncMock(return_value=[("sft", "Concept", matching_vec)]),
|
| 266 |
+
),
|
| 267 |
+
patch("app.embeddings.embed_texts", new=MagicMock(side_effect=fake_embed_fn)),
|
| 268 |
+
patch("app.graph.neo4j_client.upsert_entity", new=entity_ups),
|
| 269 |
+
patch("app.graph.neo4j_client.upsert_relationship", new=rel_ups),
|
| 270 |
+
patch("app.workers.tasks.asyncio.sleep", new=AsyncMock()),
|
| 271 |
+
):
|
| 272 |
+
await tasks._ingest_graph(chunks=["c"], user_id="u-1", document_id="d-1")
|
| 273 |
+
|
| 274 |
+
# Aliased into the existing "sft" node — one upsert, keyed on "sft".
|
| 275 |
+
assert entity_ups.await_count == 1
|
| 276 |
+
kwargs = entity_ups.await_args.kwargs
|
| 277 |
+
assert kwargs["name_lower"] == "sft"
|
| 278 |
+
# Embedding on the aliased upsert is None — we never overwrite the
|
| 279 |
+
# authoritative embedding on an existing node.
|
| 280 |
+
assert kwargs.get("embedding") is None
|
| 281 |
+
|
| 282 |
+
|
| 283 |
+
def test_semantic_alignment_flag_default_is_on():
|
| 284 |
+
"""Guard: the L3 flag defaults on so the shipped behavior includes
|
| 285 |
+
semantic alignment. The flag-off code path is a single-line guard;
|
| 286 |
+
testing it via full-worker integration proved cross-test fragile under
|
| 287 |
+
pytest's session-scoped async loop, so we skip that end-to-end assertion
|
| 288 |
+
and lean on the unit-tested `find_semantic_alias` for correctness."""
|
| 289 |
+
from app.core.config import settings
|
| 290 |
+
|
| 291 |
+
assert settings.graph_semantic_align is True
|
| 292 |
+
assert 0.0 < settings.graph_semantic_align_threshold <= 1.0
|