asmhatre/spring-ai
1
0
Fork 0
Code Activity

Add graphrag-neo4j module: structured-output entity and relation extraction into a typed Neo4j graph with a vector index, multi-hop retrieval by span linking and variable-length Cypher paths, compared with plain vector RAG over the same 35 chunks and 9 questions (complete evidence 7/9 vs 5/9 at k=8); Testcontainers on neo4j:2026.10.1 and 5.26. The extraction step is a rule-based stand-in, not a language model.

Browse Source 
Co-Authored-By: Claude Sonnet 5.5 <[email protected]>
Claude-Session: https://claude.ai/code/session_01CifR2BE7PBvRy5hT8TS9oz
This commit is contained in:
asmhatre
2026-10-26 02:03:40 +00:00
co-authored by Claude Sonnet 5.5
parent 279be7eb9f
commit 108fc44542
66 changed files with 1997 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
graphrag-neo4j/src/main/java/com/ankurm/graphrag/EntityResolver.java
+126
View File
@@ -0,0 +1,126 @@
package com.ankurm.graphrag;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import org.springframework.ai.embedding.EmbeddingModel;
/**
* Step 2: decide which names are the same thing. A model that reads each chunk on its own writes "Card
* Gateway" in one post-mortem and "Payments Gateway" in a service page; unless something merges them they
* become two nodes and every path through one is invisible from the other.
*
* <p>Three layers, applied in order: names that differ only in case, punctuation or spacing are the same key;
* a hand-kept alias table; and, optionally, cosine similarity of the name embeddings between entities of the
* same type above a threshold.
*/
public final class EntityResolver {
/** Two distinct names of the same type and how close their embeddings are. */
public record Pair(String a, String b, EntityType type, double cosine) {
}
/** What every mention maps to, and what was merged on the way. */
public record Resolution(Map<String, String> canonicalByName, Map<String, EntityType> typeByCanonical,
List<Pair> merged) {
}
private final EmbeddingModel embeddings;
public EntityResolver(EmbeddingModel embeddings) {
this.embeddings = embeddings;
}
public static String key(String name) {
return name.toLowerCase(Locale.ROOT).replaceAll("[^a-z0-9]", "");
}
/**
* @param mentions every entity mention from every chunk, repeats included (the most frequent spelling wins)
* @param aliases lower-case-insensitive alias to canonical name
* @param threshold merge same-type names whose cosine is at least this; pass a value above 1 to switch it off
*/
public Resolution resolve(List<Extraction.Entity> mentions, Map<String, String> aliases, double threshold) {
Map<String, Integer> counts = new HashMap<>();
Map<String, EntityType> types = new LinkedHashMap<>();
Map<String, String> spelling = new LinkedHashMap<>();
for (Extraction.Entity e : mentions) {
counts.merge(key(e.name()), 1, Integer::sum);
types.putIfAbsent(key(e.name()), e.type());
spelling.putIfAbsent(key(e.name()), e.name());
}
Map<String, String> aliasByKey = new HashMap<>();
aliases.forEach((alias, canonical) -> aliasByKey.put(key(alias), canonical));
// key -> canonical key after the alias table
Map<String, String> parent = new LinkedHashMap<>();
for (String k : types.keySet()) {
String canonical = aliasByKey.get(k);
parent.put(k, canonical == null ? k : key(canonical));
if (canonical != null) {
types.putIfAbsent(key(canonical), types.get(k));
spelling.putIfAbsent(key(canonical), canonical);
}
}
List<Pair> merged = new ArrayList<>();
List<String> roots = parent.values().stream().distinct().toList();
if (threshold <= 1.0) {
for (Pair p : similarPairs(roots, types, spelling)) {
if (p.cosine() >= threshold) {
String ka = key(p.a());
String kb = key(p.b());
String winner = counts.getOrDefault(ka, 0) >= counts.getOrDefault(kb, 0) ? ka : kb;
String loser = winner.equals(ka) ? kb : ka;
final String w = winner;
final String l = loser;
parent.replaceAll((k, v) -> v.equals(l) ? w : v);
merged.add(p);
}
}
}
Map<String, String> canonicalByName = new LinkedHashMap<>();
Map<String, EntityType> typeByCanonical = new LinkedHashMap<>();
for (Extraction.Entity e : mentions) {
String root = parent.get(key(e.name()));
String canonical = spelling.get(root);
canonicalByName.put(e.name(), canonical);
typeByCanonical.putIfAbsent(canonical, types.get(root));
}
return new Resolution(canonicalByName, typeByCanonical, merged);
}
/** All pairs of distinct same-type names, most similar first. Used to merge and to print the evidence. */
public List<Pair> similarPairs(List<String> keys, Map<String, EntityType> types, Map<String, String> spelling) {
List<float[]> vectors = embeddings.embed(keys.stream().map(spelling::get).toList());
List<Pair> out = new ArrayList<>();
for (int i = 0; i < keys.size(); i++) {
for (int j = i + 1; j < keys.size(); j++) {
if (types.get(keys.get(i)) == types.get(keys.get(j))) {
out.add(new Pair(spelling.get(keys.get(i)), spelling.get(keys.get(j)), types.get(keys.get(i)),
cosine(vectors.get(i), vectors.get(j))));
}
}
}
out.sort(Comparator.comparingDouble(Pair::cosine).reversed());
return out;
}
public static double cosine(float[] a, float[] b) {
double dot = 0;
double na = 0;
double nb = 0;
for (int i = 0; i < a.length; i++) {
dot += a[i] * b[i];
na += a[i] * a[i];
nb += b[i] * b[i];
}
return dot / (Math.sqrt(na) * Math.sqrt(nb));
}
}