Skip to content

[ISSUE-844] fix(ai): invalidate an embedding vector when its text changes #858

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
E2ern1ty wants to merge 2 commits into
base: master
Choose a base branch
from
Open

[ISSUE-844] fix(ai): invalidate an embedding vector when its text changes #858

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
61bd968
[ISSUE-844] fix(ai): invalidate an embedding vector when its text cha…
E2ern1ty Aug 26, 2026
923347f
[ISSUE-844] fix(ai): key the embedding index by the entity key, and p…
E2ern1ty Aug 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions geaflow-ai/src/main/java/org/apache/geaflow/ai/common/model/EmbeddingService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,13 @@ public String embedding(String... inputs) {
public static class EmbeddingResult {
public String input;
public double[] embedding;
/**
* Fingerprint of the text this vector was produced from, see
* {@link ModelUtils#getEmbeddedTextFingerprint}. It is set by whoever persists the result,
* so it is null on a result fresh from the model, and null on a record read from an index
* file written before the field existed.
*/
public String contentHash;

public EmbeddingResult(String input, double[] embedding) {
this.input = input;
Expand Down
36 changes: 36 additions & 0 deletions geaflow-ai/src/main/java/org/apache/geaflow/ai/common/model/ModelUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,9 @@

package org.apache.geaflow.ai.common.model;

import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.ArrayList;
import java.util.List;
import org.apache.geaflow.ai.common.config.Constants;
Expand All @@ -28,6 +31,8 @@

public class ModelUtils {

private static final char[] HEX = "0123456789abcdef".toCharArray();

public static List<String> splitLongText(int maxChunkSize, String... textList) {
List<String> chunks = new ArrayList<>();
for (String text : textList) {
Expand All @@ -39,6 +44,37 @@ public static List<String> splitLongText(int maxChunkSize, String... textList) {
return chunks;
}

/**
* A fingerprint of the text that was handed to the embedding model, so that a stored vector can
* later be told apart from one produced by a different text. The entity key alone cannot do
* this: it is derived from the id and the label, and stays the same when the value changes.
*
* <p>The chunk count and a separator are folded in, so that the same characters split
* differently do not fingerprint alike.
*
* @param chunks the texts embedded for one entity, in the order they were sent
* @return the fingerprint as lower case hexadecimal
*/
public static String getEmbeddedTextFingerprint(List<String> chunks) {
MessageDigest digest;
try {
digest = MessageDigest.getInstance("SHA-256");
} catch (NoSuchAlgorithmException e) {
throw new IllegalStateException("SHA-256 is required to fingerprint embedded text", e);
}
digest.update(Integer.toString(chunks.size()).getBytes(StandardCharsets.UTF_8));
for (String chunk : chunks) {
digest.update((byte) 0);
digest.update(chunk.getBytes(StandardCharsets.UTF_8));
}
byte[] bytes = digest.digest();
StringBuilder hex = new StringBuilder(bytes.length * 2);
for (byte b : bytes) {
hex.append(HEX[(b >> 4) & 0xF]).append(HEX[b & 0xF]);
}
return hex.toString();
}

public static String getGraphEntityKey(GraphEntity entity) {
if (entity instanceof GraphVertex) {
return Constants.PREFIX_V + ((GraphVertex) entity).getVertex().getId() + entity.getLabel();
Expand Down
141 changes: 106 additions & 35 deletions geaflow-ai/src/main/java/org/apache/geaflow/ai/index/EmbeddingIndexStore.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,15 +47,25 @@ public class EmbeddingIndexStore implements IndexStore {
private VerbalizationFunction verbFunc;
private String indexFilePath;
private ModelConfig modelConfig;
private Map<GraphEntity, List<EmbeddingService.EmbeddingResult>> indexStoreMap;
/**
* Keyed by {@link ModelUtils#getGraphEntityKey}, the same string the index file is keyed by, so
* that one notion of which entity a vector belongs to serves both. Keying it by the entity would
* hand that decision to {@code GraphVertex.equals}, which answers on the id and the label and
* would take the index with it if it ever came to answer on the values as well.
*
* <p>Built aside and assigned once, so a reader during initStore sees the index it had before
* rather than one half way through being built, and empty until there has been one.
*/
private volatile Map<String, List<EmbeddingService.EmbeddingResult>> indexStoreMap =
Collections.emptyMap();

public void initStore(GraphAccessor graphAccessor, VerbalizationFunction func,
String indexFilePath, ModelConfig modelInfo) {
this.graphAccessor = graphAccessor;
this.verbFunc = func;
this.indexFilePath = indexFilePath;
this.modelConfig = modelInfo;
this.indexStoreMap = new HashMap<>();
Map<String, List<EmbeddingService.EmbeddingResult>> loading = new HashMap<>();

//Read index items from indexFilePath
Map<String, GraphEntity> key2EntityMap = new HashMap<>();
Expand Down Expand Up @@ -85,7 +95,14 @@ public void initStore(GraphAccessor graphAccessor, VerbalizationFunction func,
}


// A record is accepted only when its fingerprint matches the text the entity would be
// embedded from now. The key tells which entity a vector belongs to but nothing about the
// value it came from, so without this an entity that kept its id and changed its value
// would count as indexed and keep the vector of the value it no longer has.
long count = 0;
long staleRecords = 0;
long unversionedRecords = 0;
Map<String, String> key2Fingerprint = new HashMap<>();
try (BufferedReader reader = new BufferedReader(
new InputStreamReader(
new FileInputStream(this.indexFilePath),
Expand All @@ -96,25 +113,46 @@ public void initStore(GraphAccessor graphAccessor, VerbalizationFunction func,
if (line.isEmpty()) {
continue;
}
EmbeddingService.EmbeddingResult embedding;
try {
EmbeddingService.EmbeddingResult embedding =
new Gson().fromJson(line, EmbeddingService.EmbeddingResult.class);
String key = embedding.input;
GraphEntity entity = key2EntityMap.get(key);
if (entity != null) {
this.indexStoreMap.computeIfAbsent(entity, k -> new ArrayList<>()).add(embedding);

@kitalkuyo-gita kitalkuyo-gita Aug 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the indexStoreMap is read and written by different threads during runtime (e.g., index initialization on the startup thread, querying on the worker thread), concurrent read/write safety must be guaranteed. Replace the map with a ConcurrentHashMap.

@E2ern1ty E2ern1ty Aug 27, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in 923347f, but not with a ConcurrentHashMap, because I do not think that fixes the case you describe.

Take it as stated: initialization on the startup thread, a query on a worker thread. initStore assigned the map to the field up front and then mutated it for the whole of its run — the file read, then batch after batch of embedding. A worker reading in that window sees whatever has been filled in so far. With a ConcurrentHashMap it still does: every individual get is safe, no exception is thrown, and the answer is a partial index. For a recall path that is worse than an error, because it comes back as a plausible result rather than a failure. The same goes for a second initStore on a live store, which under either map empties the index for readers before filling it again.

So the map is built aside and assigned to the field in one go, at the end of initStore. The field is volatile, and starts as Collections.emptyMap(). A reader now sees either the index that was there before or the finished one, and a reader that arrives before any build has happened gets an empty list rather than the NullPointerException it would have got on master. getEntityIndex reads the field once per call, so it works against one map throughout.

Two notes on the trade-off, since it is not free:

  • During a rebuild both maps are held, so peak memory is roughly doubled for that window. I think consistency is worth it here, and there is no caller in tree that rebuilds a live store — nothing in geaflow-ai/src/main constructs this store at all today.
  • This does not make initStore safe to call concurrently with itself. Two builds at once still race on the index file, which is appended to. That was true before and I have not tried to fix it here; if you would like it enforced, the smallest thing would be to reject a concurrent call outright rather than to make it work.

There is a test: it holds the embeddings endpoint open with a latch, runs a second initStore on the same store on another thread, and asserts that a reader in that window gets the vector from before the build. It fails if the field is assigned before the build instead of after — with the assignment moved back, the reader gets zero vectors, which is the partial index a ConcurrentHashMap would also have shown.

}
count++;
embedding = new Gson().fromJson(line, EmbeddingService.EmbeddingResult.class);
} catch (Throwable e) {
// Only the parse is guarded. What follows verbalises the entity, and a
// verbaliser that throws must not be reported as a malformed file.
LOGGER.info("Cannot parse embedding item: " + line);
continue;
}
String key = embedding == null ? null : embedding.input;
GraphEntity entity = key == null ? null : key2EntityMap.get(key);
if (entity != null) {
if (embedding.contentHash == null) {
unversionedRecords++;
} else if (embedding.contentHash.equals(
currentFingerprint(key, entity, key2Fingerprint))) {
loading.computeIfAbsent(key, k -> new ArrayList<>()).add(embedding);
} else {
staleRecords++;
}
}
count++;
}
} catch (Throwable e) {
throw new RuntimeException(e);
}

LOGGER.info("Success to read index store file. items num: " + count);
LOGGER.info("Success to rebuild index with file. index num: " + this.indexStoreMap.size());
if (staleRecords > 0) {
LOGGER.info("{} records were produced from text their entity no longer has. Not "
+ "loading them, so those entities are embedded again.", staleRecords);
}
if (unversionedRecords > 0) {
// Written before the fingerprint existed, so there is no way to tell whether they match
// the current value. Trusting them would keep exactly the staleness this check is for.
LOGGER.warn("{} records carry no fingerprint, so the text they were produced from "
+ "cannot be established. Not loading them, which costs one round of embedding, "
+ "after which they carry one.", unversionedRecords);
}
LOGGER.info("Success to rebuild index with file. index num: " + loading.size());


//Scan entities in the graph, make new index items
Expand All @@ -123,21 +161,22 @@ public void initStore(GraphAccessor graphAccessor, VerbalizationFunction func,

final int BATCH_SIZE = Constants.EMBEDDING_INDEX_STORE_BATCH_SIZE;
List<GraphEntity> pendingEntities = new ArrayList<>(BATCH_SIZE);
Set<GraphEntity> batchEntitiesBuffer = new HashSet<>(BATCH_SIZE);
Set<String> batchEntitiesBuffer = new HashSet<>(BATCH_SIZE);
List<String> result = new ArrayList<>();
final int REPORT_SIZE = Constants.EMBEDDING_INDEX_STORE_REPORT_SIZE;
long reportedCount = this.indexStoreMap.size();
long addedCount = this.indexStoreMap.size();
long reportedCount = loading.size();
long addedCount = loading.size();
for (Iterator<GraphVertex> itV = graphAccessor.scanVertex(); itV.hasNext(); ) {
GraphVertex vertex = itV.next();

// Scan vertices or edges, skip already indexed data,
// add un-indexed data to batch processing collection
if (!indexStoreMap.containsKey(vertex) && !batchEntitiesBuffer.contains(vertex)) {
batchEntitiesBuffer.add(vertex);
String vertexKey = ModelUtils.getGraphEntityKey(vertex);
if (!loading.containsKey(vertexKey) && !batchEntitiesBuffer.contains(vertexKey)) {
batchEntitiesBuffer.add(vertexKey);
pendingEntities.add(vertex);
if (pendingEntities.size() >= BATCH_SIZE) {
result.addAll(indexBatch(embeddingService, pendingEntities));
result.addAll(indexBatch(embeddingService, pendingEntities, loading));
flushBatchIndex(result, false);
pendingEntities.clear();
batchEntitiesBuffer.clear();
Expand All @@ -147,11 +186,12 @@ public void initStore(GraphAccessor graphAccessor, VerbalizationFunction func,

for (Iterator<GraphEdge> itE = graphAccessor.scanEdge(vertex); itE.hasNext(); ) {
GraphEdge edge = itE.next();
if (!indexStoreMap.containsKey(edge) && !batchEntitiesBuffer.contains(edge)) {
batchEntitiesBuffer.add(edge);
String edgeKey = ModelUtils.getGraphEntityKey(edge);
if (!loading.containsKey(edgeKey) && !batchEntitiesBuffer.contains(edgeKey)) {
batchEntitiesBuffer.add(edgeKey);
pendingEntities.add(edge);
if (pendingEntities.size() >= BATCH_SIZE) {
result.addAll(indexBatch(embeddingService, pendingEntities));
result.addAll(indexBatch(embeddingService, pendingEntities, loading));
flushBatchIndex(result, false);
pendingEntities.clear();
batchEntitiesBuffer.clear();
Expand All @@ -165,30 +205,55 @@ public void initStore(GraphAccessor graphAccessor, VerbalizationFunction func,
}
}
if (pendingEntities.size() > 0) {
result.addAll(indexBatch(embeddingService, pendingEntities));
result.addAll(indexBatch(embeddingService, pendingEntities, loading));
flushBatchIndex(result, true);
addedCount += pendingEntities.size();
pendingEntities.clear();
batchEntitiesBuffer.clear();
}

this.indexStoreMap = loading;
LOGGER.info("Successfully added {} new index items. Total indexed: {}",
addedCount, indexStoreMap.size());
addedCount, loading.size());
}

private List<String> indexBatch(EmbeddingService service, List<GraphEntity> pendingEntities) {
/**
* The fingerprint of what the entity would be embedded from now, worked out once per entity
* since one file holds every chunk of it as a record of its own. This takes
* {@link VerbalizationFunction#verbalize(GraphEntity)} to be a function of the entity alone: one
* that answers differently for an unchanged entity would have its vectors produced again on
* every run.
*/
private String currentFingerprint(String key, GraphEntity entity, Map<String, String> memo) {
String fingerprint = memo.get(key);
if (fingerprint == null) {
fingerprint = ModelUtils.getEmbeddedTextFingerprint(embeddedChunks(this.verbFunc, entity));
memo.put(key, fingerprint);
}
return fingerprint;
}

/** The texts an entity is embedded from, which is also what its fingerprint covers. */
private static List<String> embeddedChunks(VerbalizationFunction func, GraphEntity entity) {
return ModelUtils.splitLongText(Constants.EMBEDDING_INDEX_STORE_SPLIT_TEXT_CHUNK_SIZE,
func.verbalize(entity).toArray(new String[0]));
}

private List<String> indexBatch(EmbeddingService service, List<GraphEntity> pendingEntities,
Map<String, List<EmbeddingService.EmbeddingResult>> loading) {
if (pendingEntities == null || service == null || pendingEntities.isEmpty()) {
return new ArrayList<>();
}
List<String> pendingTexts = new ArrayList<>(pendingEntities.size());
Map<GraphEntity, Pair<Integer, Integer>> entity2StartEndPair = new HashMap<>();
Map<GraphEntity, String> entity2Fingerprint = new HashMap<>();
for (GraphEntity e : pendingEntities) {
Integer start = pendingTexts.size();
pendingTexts.addAll(ModelUtils.splitLongText(
Constants.EMBEDDING_INDEX_STORE_SPLIT_TEXT_CHUNK_SIZE,
verbFunc.verbalize(e).toArray(new String[0])));
List<String> chunks = embeddedChunks(verbFunc, e);
pendingTexts.addAll(chunks);
Integer end = pendingTexts.size();
entity2StartEndPair.put(e, Pair.of(start, end));
entity2Fingerprint.put(e, ModelUtils.getEmbeddedTextFingerprint(chunks));
}

Gson gson = new Gson();
Expand All @@ -209,17 +274,19 @@ private List<String> indexBatch(EmbeddingService service, List<GraphEntity> pend
List<String> formatResult = new ArrayList<>();
for (Map.Entry<GraphEntity, Pair<Integer, Integer>> entry : entity2StartEndPair.entrySet()) {
GraphEntity e = entry.getKey();
String key = ModelUtils.getGraphEntityKey(e);
List<EmbeddingService.EmbeddingResult> embeddings = new ArrayList<>();
for (int i = entry.getValue().getLeft(); i < entry.getValue().getRight(); i++) {
if (StringUtils.isNotBlank(result.get(i))) {
EmbeddingService.EmbeddingResult res = gson.fromJson(result.get(i),
EmbeddingService.EmbeddingResult.class);
res.input = ModelUtils.getGraphEntityKey(e);
res.input = key;
res.contentHash = entity2Fingerprint.get(e);
formatResult.add(gson.toJson(res));
embeddings.add(res);
}
}
indexStoreMap.put(e, embeddings);
loading.put(key, embeddings);
}
return formatResult;
}
Expand All @@ -243,14 +310,18 @@ private void flushBatchIndex(List<String> newItemStrings, boolean force) {

@Override
public List<IVector> getEntityIndex(GraphEntity entity) {
if (entity != null && indexStoreMap.get(entity) != null) {
List<EmbeddingService.EmbeddingResult> resultList = indexStoreMap.get(entity);
List<IVector> result = new ArrayList<>();
for (EmbeddingService.EmbeddingResult res : resultList) {
double[] embedding = res.embedding;
result.add(new EmbeddingVector(embedding));
if (entity != null) {
// Read once: the field is replaced wholesale when an index is built.
List<EmbeddingService.EmbeddingResult> resultList =
indexStoreMap.get(ModelUtils.getGraphEntityKey(entity));
if (resultList != null) {
List<IVector> result = new ArrayList<>();
for (EmbeddingService.EmbeddingResult res : resultList) {
double[] embedding = res.embedding;
result.add(new EmbeddingVector(embedding));
}
return result;
}
return result;
}
return Collections.emptyList();
}
Expand Down
16 changes: 7 additions & 9 deletions geaflow-ai/src/test/java/org/apache/geaflow/ai/GraphMemoryTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@
import org.apache.geaflow.ai.graph.EmptyGraphAccessor;
import org.apache.geaflow.ai.graph.GraphAccessor;
import org.apache.geaflow.ai.graph.LocalMemoryGraphAccessor;
import org.apache.geaflow.ai.index.EmbeddingIndexStore;
import org.apache.geaflow.ai.index.EntityAttributeIndexStore;
import org.apache.geaflow.ai.index.IndexStore;
import org.apache.geaflow.ai.index.vector.EmbeddingVector;
Expand Down Expand Up @@ -97,18 +96,17 @@ public void testLdbcMainPipeline() {
indexStore.initStore(new SubgraphSemanticPromptFunction(graphAccessor));
LOGGER.info("Success to init EntityAttributeIndexStore.");

// No embedding store here. This test used to load one from a checked in index file, but
// every query below embeds to double[0] through MockChatRobot, and a vector of a different
// length scores 0.0, so no stored vector could ever pass the threshold: what is asserted
// comes from the keyword path alone. Most of that file also no longer lined up with the text
// the current code would embed, so it could not be carried over to a format that records
// what a vector was produced from; see ISSUE-844 for the figures. The embedding store is
// covered against a local endpoint by EmbeddingIndexInvalidationTest instead.
ModelConfig modelInfo = new ModelConfig(null, null, null, null);
EmbeddingIndexStore embeddingStore = new EmbeddingIndexStore();
embeddingStore.initStore(graphAccessor,
new SubgraphSemanticPromptFunction(graphAccessor),
"src/test/resources/index/LDBCEmbeddingIndexStore",
modelInfo);
LOGGER.info("Success to init EmbeddingIndexStore.");

GraphMemoryServer server = new GraphMemoryServer();
server.addGraphAccessor(graphAccessor);
server.addIndexStore(indexStore);
server.addIndexStore(embeddingStore);
MockChatRobot robot = new MockChatRobot();
robot.setModelInfo(modelInfo);

Expand Down
Loading