Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
Original file line number Diff line number Diff line change
Expand Up @@ -172,8 +172,16 @@ public void initStore(GraphAccessor graphAccessor, VerbalizationFunction func,
batchEntitiesBuffer.clear();
}

LOGGER.info("Successfully added {} new index items. Total indexed: {}",
addedCount, indexStoreMap.size());
// Count entries that actually carry vectors, not entities that were queued: an entity with
// no embeddable text is registered with an empty list and must not be reported as indexed.
long withVectors = 0;
for (List<EmbeddingService.EmbeddingResult> vectors : indexStoreMap.values()) {

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.

Insufficient Exception Handling in EmbeddingIndexStore

Risk Description: In the indexBatch method, when an entity lacks embeddable text, a warning log is generated, but the entity is still registered in indexStoreMap (with an empty list of warnings). This can lead to the following issues:

During a second initialization, the containsKey check skips these entities (assuming they are already indexed).
Queries return empty results instead of a clear signal indicating "not indexed."

// Inside the indexBatch method - lines 210-223
for (Map.Entry<GraphEntity, Pair<Integer, Integer>> entry : entity2StartEndPair.entrySet()) {
GraphEntity e = entry.getKey();
List<EmbeddingService.EmbeddingResult> embeddings = new ArrayList<>();
for (int i = entry.getValue().getLeft(); i < entry.getValue().getRight(); i++) {
// If start == end (entities in withoutText), this loop does not execute
if (StringUtils.isNotBlank(result.get(i))) {
// ...
embeddings.add(res);
}
}
// ⚠️ Registered regardless of whether embeddings is empty
indexStoreMap.put(e, embeddings);  // embeddings might be an empty list
}

Problem scenario: Assume the graph contains two entities:

1st run: Entity A has text (embedding successful); Entity B has no text (empty list registered).
2nd run: Entity B's value is updated to meaningful text.
containsKey(B) returns true → skipped.
Entity B is never re-embedded.
Recommendation:

  1. Do not register entities that yield no embeddings, or use a separate set for "checked but empty" entities:
// New field
private Set<GraphEntity> checkedButEmpty = new HashSet<>();

// Modified indexBatch
for (Map.Entry<GraphEntity, Pair<Integer, Integer>> entry : entity2StartEndPair.entrySet()) {
GraphEntity e = entry.getKey();
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);
formatResult.add(gson.toJson(res)); 
embeddings.add(res);
}
}
if (!embeddings.isEmpty()) {
indexStoreMap.put(e, embeddings);
} else {
checkedButEmpty.add(e);  // Explicitly mark as checked but containing no content
}
}

// Modify the check logic in initStore
if (!indexStoreMap.containsKey(vertex) && !checkedButEmpty.contains(vertex)
&& !batchEntitiesBuffer.contains(vertex)) {
// Only entities that have truly not been processed enter the batch
batchEntitiesBuffer.add(vertex);
pendingEntities.add(vertex);
}
  1. Enhance the distinction of return values ​​for getEntityIndex:
@Override
public List<IVector> getEntityIndex(GraphEntity entity) {
if (entity != null) {
List<EmbeddingService.EmbeddingResult> resultList = indexStoreMap.get(entity);
if (resultList != null) {
if (resultList.isEmpty()) {
LOGGER.debug("Entity {} was checked but contains no embeddable text", entity);
} else {
List<IVector> result = new ArrayList<>();
for (EmbeddingService.EmbeddingResult res : resultList) {
result.add(new EmbeddingVector(res.embedding));
}
return result;
}
}
}
return Collections.emptyList();
}

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.

Good thing to check, and I went looking for it rather than reasoning about it. The skip you describe does not happen, and I have added a test that would fail if it did: EmbeddingIndexLifecycleTest, in b007b1a4.

The reason is in initStore rather than in indexBatch. Every call begins with

this.indexStoreMap = new HashMap<>();

and then repopulates the map from the index file only. Nothing is ever written to that file for an entity with no embeddable text, because indexBatch returns records for embedded texts alone. So on the next run the empty entry from the previous run does not exist, containsKey is false, and the entity is queued again. Your scenario ends the other way round: entity B does get re-embedded as soon as its value carries meaning.

The test runs three initStore calls against a local /v1/embeddings endpoint, so the request and response serialisation, the file and the decision about what to embed are all real, with no model service and no key. Run 1: a meaningful vertex and a meaningful edge are embedded and persisted, a date-only vertex is not, 2 lines in the file. Run 2, unchanged graph: 0 requests, vectors restored from the file, still 2 lines. Run 3, the date replaced by the master replied in the temple: exactly 1 request, its body contains the new text and does not contain the already-indexed text, the file grows to 3 lines, and getEntityIndex for that entity is no longer empty. The log from run 2 shows the mechanism directly, 1 of 1 entities have no embeddable text, the 1 being that entity queued again rather than skipped.

On not registering the empty list at all: I looked at that and it would cost something. MemoryGraph.scanEdge(vertex) returns both out-edges and in-edges, so initStore reaches the same edge from each of its endpoints. For an edge with no embeddable text, the empty entry in indexStoreMap is what makes the second visit a no-op. Dropping it without a replacement means verbalizing that edge twice per run, and the checkedButEmpty set restores the skip but also adds a second piece of state that has to be kept in step with the map and reset on every initStore, for a saving of one local verbalize call. Given that the cross-run recovery above depends on those entities not being remembered, I would rather keep one map whose meaning is "these are the vectors I have" than two structures that have to agree.

Your second suggestion I have taken, in a smaller form. getEntityIndex returned an empty list for two different situations, an entity that was checked and held nothing, and an entity nobody has looked at. It still returns the same value, since a caller wants vectors either way and has nothing useful to do with the distinction, but the first case now logs at debug level instead of being silent.

One nearby gap that is real and that I have deliberately not touched: an entity that is already in the index file and whose value then changes is not re-embedded, because the index key is ModelUtils.getGraphEntityKey, which is id and label with no content in it, so the stale vector keeps matching. That is independent of this fix, it predates it, and fixing it means a content hash in the record and an invalidation path. Happy to open a separate issue for it if you agree it is worth one.

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.

There is indeed a vulnerability nearby that I deliberately avoided: entities already existing in the index file using ModelUtils.getGraphEntityKey will not be re-embedded after their values ​​change because the index key is an empty id and tag, so outdated vectors still match. This is unrelated to this fix; it predates it, and fixing it would require adding a content hash and an invalidation path to the record. If you think it's worth creating a separate issue for this, I'd be happy to do so.

Regarding this, which could potentially cause fact-checking issues when old and new facts are switched, you're welcome to open a new issue so we can discuss this.

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.

Opened as #844. Before writing it up I reproduced it, and it is worse than a stale vector: the values get swapped in both directions.

One vertex, value changed between two runs against the same index file, on master (3f73eb55), with a local endpoint that records what it is asked to embed:

requests index lines vectors held
run 1, the master said learning without thought is labour lost 1 1 1
run 2, same id, the stable burned down and nobody asked about the horse 0 1 1

The vector the entity holds after run 2 belongs to the value that is no longer there, cosine(stored, old value) = 1.0000 against cosine(stored, current value) = 0.2357. Recall then goes the wrong way on both queries: a query for the old value, which is not in the graph any more, returns the entity, while a query for the current value does not. And because the subgraph is verbalised from the current graph, the first case hands the model a context containing the new value while having matched on the old one. That is your fact-switching concern, and I think it is the part that matters most: the retrieval reason and the returned content disagree, so nothing downstream can tell that it happened.

Small correction to the paraphrase, in case it matters for the issue: the key is not empty, it is PREFIX_V + id + label, for example Vv1chunk. The problem is not that it is empty but that nothing in it derives from the embedded text, so a value change leaves the key identical and initStore reads that as already indexed.

One more thing the same run turned up, which I put in #844 as a secondary observation: a record whose entity has been deleted from the graph is never pruned. It is dropped during the rebuild, since key2EntityMap has no such key, but the line stays in the file, so the file only grows, and a later entity reusing that id would silently adopt the old vector.

Directions are sketched in the issue rather than settled: a hash of the embedded text in each record, compared during the rebuild, with a mismatch counting as not indexed. That changes the file format, so it needs a decision on how to treat existing index files, and pruning wants a rewrite rather than an append, which is probably separate compaction work. Happy to pick it up once you have a view on the record format. Keeping it out of #832, which only changed which values count as embeddable.

if (vectors != null && !vectors.isEmpty()) {
withVectors++;
}
}
LOGGER.info("Successfully added {} new index items. Entities holding vectors: {} of {}",
addedCount, withVectors, indexStoreMap.size());
}

private List<String> indexBatch(EmbeddingService service, List<GraphEntity> pendingEntities) {
Expand All @@ -182,14 +190,25 @@ private List<String> indexBatch(EmbeddingService service, List<GraphEntity> pend
}
List<String> pendingTexts = new ArrayList<>(pendingEntities.size());
Map<GraphEntity, Pair<Integer, Integer>> entity2StartEndPair = new HashMap<>();
List<GraphEntity> withoutText = new ArrayList<>();
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])));
Integer end = pendingTexts.size();
if (start.equals(end)) {
withoutText.add(e);
}
entity2StartEndPair.put(e, Pair.of(start, end));
}
if (!withoutText.isEmpty()) {
// Say so rather than reporting these as indexed. An entity whose values are all
// ignorable yields no text, so it gets an empty vector list and can never be recalled.
LOGGER.warn("{} of {} entities have no embeddable text and will hold no vectors, "
+ "for example {}", withoutText.size(), pendingEntities.size(),
ModelUtils.getGraphEntityKey(withoutText.get(0)));
}

Gson gson = new Gson();
int batchSize = pendingEntities.size();
Expand Down Expand Up @@ -243,15 +262,24 @@ 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));
}
return result;
if (entity == null) {
return Collections.emptyList();
}
List<EmbeddingService.EmbeddingResult> resultList = indexStoreMap.get(entity);
if (resultList == null) {
return Collections.emptyList();
}
if (resultList.isEmpty()) {
// An entity that was looked at but held no embeddable text is indistinguishable from an
// entity nobody has looked at yet, since both give an empty result. Say which it was.
LOGGER.debug("Entity {} was checked and holds no embeddable text",
ModelUtils.getGraphEntityKey(entity));
return Collections.emptyList();
}
List<IVector> result = new ArrayList<>(resultList.size());
for (EmbeddingService.EmbeddingResult res : resultList) {
result.add(new EmbeddingVector(res.embedding));
}
return Collections.emptyList();
return result;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -31,17 +31,17 @@ public class SearchUtils {
'*', '#', '-', '?', '`', '{', '}', '[', ']', '(', ')', '>', '<', ':', '/', '.'
));

// Set of allowed characters for validation in isAllAllowedChars
// Includes: digits (0-9), and some common safe symbols
private static final Set<Character> IGNORE_CHARS = buildIgnoredChars();
// Characters that carry no meaning on their own: digits and a few common symbols.
// A value made up entirely of these is not worth indexing or embedding.
private static final Set<Character> IGNORABLE_CHARS = buildIgnorableChars();

/**
* Builds the set of allowed characters for input validation.
* Includes alphanumeric characters and selected common symbols.
* Builds the set of characters that carry no meaning on their own.
* Digits and a few common symbols.
*
* @return an unmodifiable set of ignored characters
* @return an unmodifiable set of ignorable characters
*/
private static Set<Character> buildIgnoredChars() {
private static Set<Character> buildIgnorableChars() {
Set<Character> ignored = new HashSet<>(32);
// Add digits
for (char c = '0'; c <= '9'; c++) {
Expand Down Expand Up @@ -87,18 +87,30 @@ public static String formatQuery(String query) {
}

/**
* Checks whether all characters in the given string are within the allowed character set.
* Useful for validating usernames, identifiers, or safe input formats.
* Whether the given value consists entirely of characters that carry no meaning on their own,
* and therefore has nothing worth indexing or embedding. A bare id, a date or a run of
* punctuation is ignorable; anything containing a letter or a CJK character is not.
*
* @param str the string to validate
* @return true if all characters are allowed; false otherwise
* <p>Callers use this to skip values, so an empty or absent value is ignorable too: there is
* nothing in it to index.
*
* <p>This replaces {@code isAllAllowedChars}, whose loop returned on the first character
* <em>inside</em> the set rather than the first one outside it, making it the negation of both
* its own name and its own documentation. The practical effect was that ordinary prose was
* discarded while digit-only noise was kept, so an embedding store silently produced nothing
* for text that happened to contain no digit. The method is renamed rather than corrected in
* place, so that any caller depending on the previous meaning fails to compile instead of
* silently flipping behaviour.
*
* @param str the value to check
* @return true if every character is ignorable, or the value is null or empty
*/
public static boolean isAllAllowedChars(String str) {
public static boolean isAllIgnorableChars(String str) {

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.

First of all, thank you for your contribution. However, upon reviewing the repository, I noticed that there are still references to this function in some places. I have the following suggestions:

  1. Perform a full-code search before merging: Search the entire apache/geaflow repository for all call sites of isAllAllowedChars to confirm that all instances have been updated in the PR. You should run a search like this:
grep -r "isAllAllowedChars" --include="*.java" .
  1. Suggestion for a transition period regarding deprecated methods: Retain the old method but mark it as @Deprecated, and remove it only after 2–3 release cycles to avoid breaking downstream dependencies:
@Deprecated(since = "1.4.0", forRemoval = true)
public static boolean isAllAllowedChars(String str) {
return !isAllIgnorableChars(str);  // Delegate to the new method
}

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.

Thanks. I ran exactly the search you asked for, and there are no remaining call sites.

$ grep -r "isAllAllowedChars" --include="*.java" .
./geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SearchUtils.java:97:
     * <p>This replaces {@code isAllAllowedChars}, whose loop returned on the first character

One hit, and it is the Javadoc line in this file explaining what the method replaces. Widening the search to every file type gives two more, both generated build output rather than source: target/site/jacoco/.../SearchUtils.java.html and target/apidocs/.../SearchUtils.html. Those are rebuilt from this Javadoc. If the references you saw were in a target/ directory or in the doc text, that is what they were. The two real call sites, both in SubgraphSemanticPromptFunction, are updated in this PR and now read !isAllIgnorableChars(value).

On the deprecated shim, I would rather not add it, for three reasons.

The delegation is not behaviour preserving, which is the part that worries me most. isAllAllowedChars returned false as soon as it met one ignorable character, whereas isAllIgnorableChars returns true only when every character is ignorable. Those are not complements. Take "a1": the old method returned false, because 1 is in the set, while !isAllIgnorableChars("a1") is !false, that is true. So return !isAllIgnorableChars(str) would silently flip the answer for any mixed value, and mixed values are the common case in real data. LDBC entity values are exactly letters plus digits. A shim that quietly changes behaviour for the majority of inputs is the hazard the rename was meant to remove, not a mitigation of it.

It also does not compile here. since and forRemoval were added to @Deprecated in Java 9, and this project builds at <jdk.version>1.8</jdk.version> (root pom.xml). And 1.4.0 is not a GeaFlow version; the tree is at 0.8.0-SNAPSHOT.

There is no downstream contract to protect either. SearchUtils is not present in any release tag, v0.8.0-rc1 included:

$ for t in v0.8.0-rc1 v0.7.0 v0.7.0-rc3 v0.7.0-rc2 v0.7.0-rc1; do
    printf "%s: " $t; git cat-file -e "$t:geaflow-ai/.../operator/SearchUtils.java" 2>/dev/null \
      && echo present || echo absent; done
v0.8.0-rc1: absent
v0.7.0: absent
...

It arrived with #716 in January, after the rc, so no published artifact has ever exposed this method and nothing outside the repo can be calling it. A deprecation window over two or three release cycles would keep a method whose semantics are the defect this PR fixes, and it would protect no one.

If you would still like a transition period, the only shim I would be comfortable adding is one that preserves the original contract rather than approximating it, return containsAnyIgnorableChar(str) under a name that says what it does. Happy to add that if you want it, but given that the class is unreleased I think deleting the old name is the better trade. Your call.

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.

Okay, thank you very much for checking again. I've looked it over carefully, and indeed most of the references are documentation comments. You can delete the old names.

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.

Thanks for taking another look. Leaving the rename as it stands then, with the old name gone.

if (str == null || str.isEmpty()) {
return false; // Consider empty/null invalid; adjust based on use case
return true;
}
for (char c : str.toCharArray()) {
if (IGNORE_CHARS.contains(c)) {
if (!IGNORABLE_CHARS.contains(c)) {
return false;
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -72,12 +72,12 @@ public List<String> verbalize(GraphEntity entity) {
if (entity instanceof GraphVertex) {
GraphVertex graphVertex = (GraphVertex) entity;
return graphVertex.getVertex().getValues().stream()
.filter(str -> !SearchUtils.isAllAllowedChars(str))
.filter(str -> !SearchUtils.isAllIgnorableChars(str))
.map(SearchUtils::formatQuery).collect(Collectors.toList());
} else if (entity instanceof GraphEdge) {
GraphEdge graphEdge = (GraphEdge) entity;
return graphEdge.getEdge().getValues().stream()
.filter(str -> !SearchUtils.isAllAllowedChars(str))
.filter(str -> !SearchUtils.isAllIgnorableChars(str))
.map(SearchUtils::formatQuery).collect(Collectors.toList());
}
return new ArrayList<>();
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,225 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/

package org.apache.geaflow.ai.index;

import com.google.gson.Gson;
import com.sun.net.httpserver.HttpExchange;
import com.sun.net.httpserver.HttpServer;
import java.io.IOException;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CopyOnWriteArrayList;
import org.apache.geaflow.ai.common.model.EmbeddingResponse;
import org.apache.geaflow.ai.common.model.ModelConfig;
import org.apache.geaflow.ai.graph.GraphEntity;
import org.apache.geaflow.ai.graph.LocalMemoryGraphAccessor;
import org.apache.geaflow.ai.graph.io.Edge;
import org.apache.geaflow.ai.graph.io.EdgeGroup;
import org.apache.geaflow.ai.graph.io.EdgeSchema;
import org.apache.geaflow.ai.graph.io.EntityGroup;
import org.apache.geaflow.ai.graph.io.GraphSchema;
import org.apache.geaflow.ai.graph.io.MemoryGraph;
import org.apache.geaflow.ai.graph.io.Vertex;
import org.apache.geaflow.ai.graph.io.VertexGroup;
import org.apache.geaflow.ai.graph.io.VertexSchema;
import org.apache.geaflow.ai.verbalization.SubgraphSemanticPromptFunction;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;

/**
* The index over successive runs, against a local embeddings endpoint so the whole path is
* exercised without a model service: request and response serialisation, the on disk index file,
* and what a later run decides to embed.
*
* <p>Chiefly this pins down what happens to an entity that yields no embeddable text. Such an
* entity is registered with an empty vector list, which is also what makes it skippable within a
* run, since {@code scanEdge} returns an edge from both of its endpoints. Nothing is written to the
* index file for it, and {@code initStore} rebuilds its map from that file alone, so a later run
* treats it as unseen and embeds it once its value carries meaning. That last point is the one worth
* a test: the alternative, never revisiting it, would make the omission permanent.
*/
public class EmbeddingIndexLifecycleTest {

private static final String LABEL = "chunk";
private static final String EDGE_LABEL = "rel";
private static final String MEANINGFUL = "learning without thought is labour lost";
private static final String IGNORABLE = "2024-01-01";
private static final String EDGE_TEXT = "the master teaches the disciple";
private static final int DIMS = 4;

private HttpServer server;
private final List<String> requestBodies = new CopyOnWriteArrayList<>();

@BeforeEach
void startEndpoint() throws IOException {
server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
server.createContext("/v1/embeddings", this::handle);
server.setExecutor(null);
server.start();
}

@AfterEach
void stopEndpoint() {
if (server != null) {
server.stop(0);
}
}

private void handle(HttpExchange exchange) throws IOException {
byte[] requestBytes = readAll(exchange);
String body = new String(requestBytes, StandardCharsets.UTF_8);
requestBodies.add(body);

String[] inputs = new Gson().fromJson(body, Request.class).input;
EmbeddingResponse response = new EmbeddingResponse();
response.object = "list";
response.model = "test-local";
response.data = new ArrayList<>();
for (int i = 0; i < inputs.length; i++) {
EmbeddingResponse.EmbeddingVector vector = new EmbeddingResponse.EmbeddingVector();
vector.index = i;
vector.embedding = deterministicVector(inputs[i]);
response.data.add(vector);
}

byte[] out = new Gson().toJson(response).getBytes(StandardCharsets.UTF_8);
exchange.getResponseHeaders().add("Content-Type", "application/json; charset=utf-8");
exchange.sendResponseHeaders(200, out.length);
try (OutputStream os = exchange.getResponseBody()) {
os.write(out);
}
}

private static byte[] readAll(HttpExchange exchange) throws IOException {
java.io.ByteArrayOutputStream buffer = new java.io.ByteArrayOutputStream();
byte[] chunk = new byte[4096];
int read;
while ((read = exchange.getRequestBody().read(chunk)) > 0) {
buffer.write(chunk, 0, read);
}
return buffer.toByteArray();
}

/** Stable and non-zero, so a stored vector can be told apart from an absent one. */
private static double[] deterministicVector(String text) {
double[] vector = new double[DIMS];
for (int i = 0; i < DIMS; i++) {
vector[i] = ((text.hashCode() >> i) & 0xFF) / 255.0 + 0.1;
}
return vector;
}

private ModelConfig config() {
return new ModelConfig("test-local",
"http://127.0.0.1:" + server.getAddress().getPort(), "/v1/embeddings", "test-token");
}

@Test
public void testIndexLifecycleAcrossRuns(@TempDir Path tempDir) throws Exception {
String indexPath = tempDir.resolve("index.jsonl").toString();

// First run. One vertex and the edge carry meaning, the other vertex is only a date.
LocalMemoryGraphAccessor first = buildGraph(MEANINGFUL, IGNORABLE);
EmbeddingIndexStore store = new EmbeddingIndexStore();
store.initStore(first, new SubgraphSemanticPromptFunction(first), indexPath, config());

Assertions.assertEquals(1, requestBodies.size(),
"the two values that carry meaning belong to one batch, so one request");
Assertions.assertFalse(store.getEntityIndex(first.getVertex(LABEL, "v1")).isEmpty(),
"a value carrying meaning must hold a vector");
Assertions.assertTrue(store.getEntityIndex(first.getVertex(LABEL, "v2")).isEmpty(),
"a value that is only a date has nothing to embed");
Assertions.assertEquals(2, indexLines(indexPath),
"only the vertex and the edge that carry meaning are persisted");

// Second run over an unchanged graph. Everything with vectors comes back from the file, so
// the model is not consulted again.
requestBodies.clear();
LocalMemoryGraphAccessor second = buildGraph(MEANINGFUL, IGNORABLE);
EmbeddingIndexStore reloaded = new EmbeddingIndexStore();
reloaded.initStore(second, new SubgraphSemanticPromptFunction(second), indexPath, config());

Assertions.assertEquals(0, requestBodies.size(),
"an unchanged graph must not be re-embedded");
Assertions.assertFalse(reloaded.getEntityIndex(second.getVertex(LABEL, "v1")).isEmpty(),
"vectors must survive a rebuild from the index file");
Assertions.assertEquals(2, indexLines(indexPath), "no duplicate records appended");

// Third run, with the date replaced by text that carries meaning. Nothing was persisted for
// that entity, so this run must embed it, and must leave the rest alone.
requestBodies.clear();
String nowMeaningful = "the master replied in the temple";
LocalMemoryGraphAccessor third = buildGraph(MEANINGFUL, nowMeaningful);
EmbeddingIndexStore afterChange = new EmbeddingIndexStore();
afterChange.initStore(third, new SubgraphSemanticPromptFunction(third), indexPath, config());

Assertions.assertEquals(1, requestBodies.size(),
"exactly the entity that gained meaning is embedded");
Assertions.assertTrue(requestBodies.get(0).contains(nowMeaningful),
"the request must carry the newly meaningful text");
Assertions.assertFalse(requestBodies.get(0).contains(MEANINGFUL),
"an entity already in the index file must not be embedded again");
Assertions.assertFalse(afterChange.getEntityIndex(third.getVertex(LABEL, "v2")).isEmpty(),
"an entity that yielded nothing before must not be excluded for good");
Assertions.assertEquals(3, indexLines(indexPath), "the new record is appended");
}

private long indexLines(String indexPath) throws IOException {
return Files.readAllLines(java.nio.file.Paths.get(indexPath), StandardCharsets.UTF_8)
.stream().filter(line -> !line.trim().isEmpty()).count();
}

private LocalMemoryGraphAccessor buildGraph(String v1Text, String v2Text) {
GraphSchema schema = new GraphSchema();
schema.setName("lifecycle");
VertexSchema vs = new VertexSchema(LABEL, "id", Collections.singletonList("text"));
EdgeSchema es = new EdgeSchema(EDGE_LABEL, "srcId", "dstId",
Collections.singletonList("rel"));
schema.addVertex(vs);
schema.addEdge(es);

List<Vertex> vertices = new ArrayList<>(Arrays.asList(
new Vertex(LABEL, "v1", Collections.singletonList(v1Text)),
new Vertex(LABEL, "v2", Collections.singletonList(v2Text))));
List<Edge> edges = new ArrayList<>(Collections.singletonList(
new Edge(EDGE_LABEL, "v1", "v2", Collections.singletonList(EDGE_TEXT))));
Map<String, EntityGroup> entities = new HashMap<>();
entities.put(LABEL, new VertexGroup(vs, vertices));
entities.put(EDGE_LABEL, new EdgeGroup(es, edges));
return new LocalMemoryGraphAccessor(new MemoryGraph(schema, entities));
}

/** Mirrors the request the client sends, enough of it to read the inputs back. */
private static class Request {
private String[] input;
}
}
Loading
Loading