jugalgajjar commited on
Commit
082589c
·
1 Parent(s): 911343c

feat(graph): L3 semantic entity alignment via bge embeddings

Browse files
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.graph.alignment import Candidate, align_batch
 
 
 
 
 
 
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 (aligned from %d ents, %d rels)",
 
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 unittest.mock import AsyncMock, patch
 
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