spring-ai/graphrag-neo4j/src/main/java/com/ankurm/graphrag/EntityResolver.java

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));
    }
}