spring-ai/llm-gateway/src/main/java/com/ankurm/gateway/SemanticCache.java

package com.ankurm.gateway;

import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;

import org.springframework.ai.embedding.EmbeddingModel;

/**
 * Serves a stored answer when a new question embeds close enough to an old one. The key includes
 * the tenant: a cache shared across tenants answers tenant B with tenant A's data.
 */
public final class SemanticCache {

    private record Entry(float[] vector, String answer) {
    }

    private final EmbeddingModel embeddings;

    private final double threshold;

    private final Map<String, List<Entry>> byKey = new ConcurrentHashMap<>();

    public SemanticCache(EmbeddingModel embeddings, double threshold) {
        this.embeddings = embeddings;
        this.threshold = threshold;
    }

    public Optional<String> lookup(String tenant, String feature, String question) {
        float[] q = embeddings.embed(question);
        Entry best = null;
        double bestScore = -1;
        for (Entry e : byKey.getOrDefault(tenant + "|" + feature, List.of())) {
            double s = cosine(q, e.vector());
            if (s >= threshold && s > bestScore) {
                best = e;
                bestScore = s;
            }
        }
        return Optional.ofNullable(best).map(Entry::answer);
    }

    public void put(String tenant, String feature, String question, String answer) {
        byKey.computeIfAbsent(tenant + "|" + feature, k -> new CopyOnWriteArrayList<>())
                .add(new Entry(embeddings.embed(question), answer));
    }

    static double cosine(float[] a, float[] b) {
        double dot = 0, na = 0, 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) + 1e-10);
    }
}