feat(ai): make the keyword search index resident and incremental

geaflow-ai/src/main/java/org/apache/geaflow/ai/GeaFlowMemoryServer.java

@@ -23,6 +23,8 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import org.apache.geaflow.ai.common.ErrorCode;
+import org.apache.geaflow.ai.common.model.ModelUtils;
 import org.apache.geaflow.ai.common.util.SeDeUtil;
 import org.apache.geaflow.ai.graph.*;
 import org.apache.geaflow.ai.graph.io.*;
@@ -123,6 +125,11 @@ public String addSchema(@Param("graphName") String graphName,
         } else {
             throw new RuntimeException("Cannot add schema: " + input);
         }
+        GraphMemoryServer schemaServer = CACHE.getServerByName(graphName);
+        if (schemaServer != null) {
+            // Verbalization is schema driven, so cached prompts and the index must be dropped.
+            schemaServer.onSchemaChanged();
+        }
         return "addSchema has been called, schemaName: " + schemaName;
     }
 
@@ -150,23 +157,41 @@ public String addEntity(@Param("graphName") String graphName,
         if (!(graph instanceof MemoryGraph)) {
             throw new RuntimeException("Graph cannot modify.");
         }
+        GraphMemoryServer insertServer = CACHE.getServerByName(graphName);
+        if (insertServer == null || insertServer.getGraphAccessors().isEmpty()) {
+            throw new RuntimeException("Server or graph accessor not available for graph: " + graphName);
+        }
         MemoryMutableGraph memoryMutableGraph = new MemoryMutableGraph((MemoryGraph) graph);
         List<GraphEntity> graphEntities = SeDeUtil.deserializeEntities(input);
 
+        // Opened before the writes and sealed right after them, so the resident index can verify
+        // that these entities really are every vertex level change it has not seen yet.
+        VertexVersionWindow window = insertServer.openVertexVersionWindow();
+        // Only entities the graph actually accepted may be reported. A rejected write (a duplicate
+        // id, for instance) leaves the graph holding the previous entity, so reporting the request
+        // object would index content the graph does not have.
+        List<GraphEntity> written = new ArrayList<>(graphEntities.size());
+        List<GraphEntity> rejected = new ArrayList<>();
         for (GraphEntity entity : graphEntities) {
-            if (entity instanceof GraphVertex) {
-                memoryMutableGraph.addVertex(((GraphVertex) entity).getVertex());
+            int code = entity instanceof GraphVertex
+                ? memoryMutableGraph.addVertex(((GraphVertex) entity).getVertex())
+                : memoryMutableGraph.addEdge(((GraphEdge) entity).getEdge());
+            if (code == ErrorCode.SUCCESS) {
+                written.add(entity);
             } else {
-                memoryMutableGraph.addEdge(((GraphEdge) entity).getEdge());
+                rejected.add(entity);
+                LOGGER.warn("Rejected entity {} on graph {}, code: {}",
+                    ModelUtils.getGraphEntityKey(entity), graphName, code);
             }
         }
-        GraphMemoryServer insertServer = CACHE.getServerByName(graphName);
-        if (insertServer == null || insertServer.getGraphAccessors().isEmpty()) {
-            throw new RuntimeException("Server or graph accessor not available for graph: " + graphName);
-        }
+        // Maintain the resident keyword index in place instead of rebuilding it on next query.
+        insertServer.onEntitiesUpserted(written, window.seal());
         CACHE.getConsolidateServer().executeConsolidateTask(
             insertServer.getGraphAccessors().get(0), memoryMutableGraph);
-        return "Success to add entities, num: " + graphEntities.size();
+        if (rejected.isEmpty()) {
+            return "Success to add entities, num: " + written.size();
+        }
+        return "Added entities, num: " + written.size() + ", rejected: " + rejected.size();
     }
 
     @Post
@@ -182,15 +207,34 @@ public String deleteEntity(@Param("graphName") String graphName,
         }
         MemoryMutableGraph memoryMutableGraph = new MemoryMutableGraph((MemoryGraph) graph);
         List<GraphEntity> graphEntities = SeDeUtil.deserializeEntities(input);
+        GraphMemoryServer deleteServer = CACHE.getServerByName(graphName);
+        VertexVersionWindow window = deleteServer == null
+            ? null : deleteServer.openVertexVersionWindow();
+        // Same reasoning as on the insert path, mirrored: a rejected delete leaves the entity in the
+        // graph, and reporting it would drop a document that is still supposed to be searchable.
+        List<GraphEntity> removed = new ArrayList<>(graphEntities.size());
+        List<GraphEntity> rejected = new ArrayList<>();
         for (GraphEntity entity : graphEntities) {
-            if (entity instanceof GraphVertex) {
-                memoryMutableGraph.removeVertex(entity.getLabel(),
-                    ((GraphVertex) entity).getVertex().getId());
+            int code = entity instanceof GraphVertex
+                ? memoryMutableGraph.removeVertex(entity.getLabel(),
+                    ((GraphVertex) entity).getVertex().getId())
+                : memoryMutableGraph.removeEdge(((GraphEdge) entity).getEdge());
+            if (code == ErrorCode.SUCCESS) {
+                removed.add(entity);
             } else {
-                memoryMutableGraph.removeEdge(((GraphEdge) entity).getEdge());
+                rejected.add(entity);
+                LOGGER.warn("Rejected removal of entity {} on graph {}, code: {}",
+                    ModelUtils.getGraphEntityKey(entity), graphName, code);
             }
         }
-        return "Success to remove entities, num: " + graphEntities.size();
+        if (deleteServer != null) {
+            // Deletes are applied to the index in place, no rebuild needed.
+            deleteServer.onEntitiesRemoved(removed, window.seal());
+        }
+        if (rejected.isEmpty()) {
+            return "Success to remove entities, num: " + removed.size();
+        }
+        return "Removed entities, num: " + removed.size() + ", rejected: " + rejected.size();
     }
 
     @Post

geaflow-ai/src/main/java/org/apache/geaflow/ai/GraphMemoryServer.java

@@ -20,16 +20,21 @@
 package org.apache.geaflow.ai;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashSet;
+import java.util.IdentityHashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 import org.apache.geaflow.ai.graph.GraphAccessor;
 import org.apache.geaflow.ai.graph.GraphEntity;
+import org.apache.geaflow.ai.graph.VertexVersionWindow;
 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.operator.EmbeddingOperator;
+import org.apache.geaflow.ai.operator.ResidentSearchIndex;
 import org.apache.geaflow.ai.operator.SearchOperator;
 import org.apache.geaflow.ai.operator.SessionOperator;
 import org.apache.geaflow.ai.search.VectorSearch;
@@ -44,6 +49,13 @@ public class GraphMemoryServer {
     private final List<GraphAccessor> graphAccessors = new ArrayList<>();
     private final List<IndexStore> indexStores = new ArrayList<>();
 
+    /**
+     * Keyword indexes kept alive across queries, one per keyword index store. Without this the
+     * global keyword index would be rebuilt from a full graph scan on every single query.
+     */
+    private final Map<IndexStore, ResidentSearchIndex> residentIndexes =
+        Collections.synchronizedMap(new IdentityHashMap<>());
+
     public void addGraphAccessor(GraphAccessor graph) {
         if (graph != null) {
             graphAccessors.add(graph);
@@ -57,6 +69,9 @@ public List<GraphAccessor> getGraphAccessors() {
     public void addIndexStore(IndexStore indexStore) {
         if (indexStore != null) {
             indexStores.add(indexStore);
+            if (indexStore instanceof EntityAttributeIndexStore) {
+                residentIndexes.put(indexStore, new ResidentSearchIndex());
+            }
         }
     }
 
@@ -86,7 +101,8 @@ public String search(VectorSearch search) {
         }
         for (IndexStore indexStore : indexStores) {
             if (indexStore instanceof EntityAttributeIndexStore) {
-                SessionOperator searchOperator = new SessionOperator(graphAccessors.get(0), indexStore);
+                SessionOperator searchOperator = new SessionOperator(graphAccessors.get(0),
+                    indexStore, residentIndexes.get(indexStore));
                 applySearch(sessionId, searchOperator, search);
             }
             if (indexStore instanceof EmbeddingIndexStore) {
@@ -121,6 +137,76 @@ public Context verbalize(String sessionId, VerbalizationFunction verbalizationFu
         return new Context(stringBuilder.toString());
     }
 
+    /**
+     * Captures the vertex version before a batch of graph writes. Pass the sealed window to
+     * {@link #onEntitiesUpserted} / {@link #onEntitiesRemoved} so the derived structures can tell
+     * whether the reported entities really are everything that changed.
+     */
+    public VertexVersionWindow openVertexVersionWindow() {
+        return VertexVersionWindow.open(graphAccessors.isEmpty() ? null : graphAccessors.get(0));
+    }
+
+    /**
+     * Applies written entities to the derived structures in place. Handles both new and rewritten
+     * entities, so callers do not need to distinguish them.
+     *
+     * <p>Memoized verbalizations need no explicit invalidation here: every entry carries the source
+     * version it was computed from, so the write itself makes the affected entries stale.
+     *
+     * @param window version range the batch covers, obtained from
+     *     {@link #openVertexVersionWindow()} and sealed after the writes
+     */
+    public void onEntitiesUpserted(List<GraphEntity> entities, VertexVersionWindow window) {
+        if (entities == null || entities.isEmpty() || graphAccessors.isEmpty()) {
+            return;
+        }
+        for (IndexStore indexStore : indexStores) {
+            if (!(indexStore instanceof EntityAttributeIndexStore)) {
+                continue;
+            }
+            ResidentSearchIndex residentIndex = residentIndexes.get(indexStore);
+            if (residentIndex != null) {
+                residentIndex.onEntitiesUpserted(graphAccessors.get(0), entities, indexStore,
+                    window);
+            }
+        }
+    }
+
+    /**
+     * Applies removed entities to the derived structures in place.
+     */
+    public void onEntitiesRemoved(List<GraphEntity> entities, VertexVersionWindow window) {
+        if (entities == null || entities.isEmpty() || graphAccessors.isEmpty()) {
+            return;
+        }
+        for (IndexStore indexStore : indexStores) {
+            if (!(indexStore instanceof EntityAttributeIndexStore)) {
+                continue;
+            }
+            ResidentSearchIndex residentIndex = residentIndexes.get(indexStore);
+            if (residentIndex != null) {
+                residentIndex.onEntitiesRemoved(graphAccessors.get(0), entities, indexStore, window);
+            }
+        }
+    }
+
+    /**
+     * Drops the derived structures wholesale. Used for changes that cannot be expressed per entity,
+     * such as a schema change altering how every entity is verbalized.
+     */
+    public void onSchemaChanged() {
+        for (IndexStore indexStore : indexStores) {
+            if (!(indexStore instanceof EntityAttributeIndexStore)) {
+                continue;
+            }
+            ((EntityAttributeIndexStore) indexStore).invalidateCache();
+            ResidentSearchIndex residentIndex = residentIndexes.get(indexStore);
+            if (residentIndex != null) {
+                residentIndex.invalidate();
+            }
+        }
+    }
+
     public List<GraphEntity> getSessionEntities(String sessionId) {
         List<SubGraph> subGraphList = sessionManagement.getSubGraph(sessionId);
         Set<GraphEntity> entitySet = new HashSet<>();

geaflow-ai/src/main/java/org/apache/geaflow/ai/common/config/Constants.java

@@ -47,6 +47,9 @@ public class Constants {
     public static int EMBEDDING_OPERATE_DEFAULT_TOPN = 50;
     public static int GRAPH_SEARCH_STORE_DEFAULT_TOPN = 30;
 
+    // Max number of memoized entity verbalizations kept by EntityAttributeIndexStore.
+    public static int ENTITY_ATTRIBUTE_INDEX_CACHE_MAX_SIZE = 200000;
+
     public static String CONSOLIDATE_KEYWORD_RELATION_LABEL = "consolidate_keyword_edge";
     public static String PREFIX_COMMON_KEYWORDS = "common_keywords";
 }

geaflow-ai/src/main/java/org/apache/geaflow/ai/common/model/ModelUtils.java

@@ -39,7 +39,55 @@ public static List<String> splitLongText(int maxChunkSize, String... textList) {
         return chunks;
     }
 
+    /**
+     * Separates the length of a key component from the component itself.
+     */
+    private static final char LENGTH_SEPARATOR = ':';
+
+    /**
+     * Stands in for a null component. Not a digit, so it cannot be confused with a length.
+     */
+    private static final String NULL_COMPONENT = "-:";
+
+    /**
+     * A key identifying an entity, unique across the graph.
+     *
+     * <p>Components are length prefixed rather than plainly concatenated, because the result is used
+     * as a primary key: callers delete and replace documents by this exact value, so two different
+     * entities mapping to the same key would make one silently destroy the other. Plain
+     * concatenation is not injective, {@code (id="a", label="bc")} and {@code (id="ab", label="c")}
+     * both yield {@code Vabc}. With lengths the encoding can be parsed back unambiguously, which is
+     * what makes it collision free.
+     */
     public static String getGraphEntityKey(GraphEntity entity) {
+        if (entity instanceof GraphVertex) {
+            return Constants.PREFIX_V
+                    + encodeKeyComponent(((GraphVertex) entity).getVertex().getId())
+                    + encodeKeyComponent(entity.getLabel());
+        } else if (entity instanceof GraphEdge) {
+            return Constants.PREFIX_E
+                    + encodeKeyComponent(((GraphEdge) entity).getEdge().getSrcId())
+                    + encodeKeyComponent(entity.getLabel())
+                    + encodeKeyComponent(((GraphEdge) entity).getEdge().getDstId());
+        }
+        return "";
+    }
+
+    private static String encodeKeyComponent(String component) {
+        if (component == null) {
+            return NULL_COMPONENT;
+        }
+        return component.length() + String.valueOf(LENGTH_SEPARATOR) + component;
+    }
+
+    /**
+     * The plainly concatenated key used before {@link #getGraphEntityKey} became collision free.
+     *
+     * <p>Only for reading persisted state written by an older version, such as an
+     * {@code EmbeddingIndexStore} index file. It is not injective, so a caller must handle the case
+     * where two entities produce the same value rather than trusting a lookup by it.
+     */
+    public static String getLegacyGraphEntityKey(GraphEntity entity) {
         if (entity instanceof GraphVertex) {
             return Constants.PREFIX_V + ((GraphVertex) entity).getVertex().getId() + entity.getLabel();
         } else if (entity instanceof GraphEdge) {

geaflow-ai/src/main/java/org/apache/geaflow/ai/graph/GraphAccessor.java

@@ -25,6 +25,34 @@
 
 public interface GraphAccessor {
 
+    /**
+     * Returned by {@link #getGraphVersion()} when the accessor cannot report content changes.
+     * Callers must then treat every read as potentially different and skip caching.
+     */
+    long VERSION_UNSUPPORTED = -1L;
+
+    /**
+     * A monotonically increasing counter bumped on every content or schema change of the underlying
+     * graph. Derived structures (verbalization caches, keyword indexes) compare it to decide whether
+     * they are still valid, so that direct mutations of the graph cannot silently go unnoticed.
+     *
+     * @return current graph version, or {@link #VERSION_UNSUPPORTED} if change tracking is not
+     *     available for this accessor
+     */
+    default long getGraphVersion() {
+        return VERSION_UNSUPPORTED;
+    }
+
+    /**
+     * Like {@link #getGraphVersion()} but only advanced by vertex and schema changes. Structures
+     * derived from vertices alone can watch this and survive edge writes.
+     *
+     * @return current vertex version, defaults to {@link #getGraphVersion()}
+     */
+    default long getVertexVersion() {
+        return getGraphVersion();
+    }
+
     GraphSchema getGraphSchema();
 
     GraphVertex getVertex(String label, String id);

geaflow-ai/src/main/java/org/apache/geaflow/ai/graph/LocalMemoryGraphAccessor.java

@@ -46,6 +46,16 @@ public LocalMemoryGraphAccessor(MemoryGraph memoryGraph) {
         this.graph = memoryGraph;
     }
 
+    @Override
+    public long getGraphVersion() {
+        return graph.getVersion();
+    }
+
+    @Override
+    public long getVertexVersion() {
+        return graph.getVertexVersion();
+    }
+
     @Override
     public GraphSchema getGraphSchema() {
         return graph.getGraphSchema();

geaflow-ai/src/main/java/org/apache/geaflow/ai/graph/MemoryMutableGraph.java

@@ -19,7 +19,6 @@
 
 package org.apache.geaflow.ai.graph;
 
-import java.util.ArrayList;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.geaflow.ai.common.ErrorCode;
 import org.apache.geaflow.ai.graph.io.*;
@@ -80,8 +79,7 @@ public int addVertexSchema(VertexSchema vertexSchema) {
         if (this.graph.entities.get(vertexSchema.getLabel()) != null) {
             return ErrorCode.GRAPH_ADD_VERTEX_SCHEMA_FAILED;
         }
-        this.graph.getGraphSchema().addVertex(vertexSchema);
-        this.graph.entities.put(vertexSchema.getLabel(), new VertexGroup(vertexSchema, new ArrayList<>()));
+        this.graph.registerVertexSchema(vertexSchema);
         return ErrorCode.SUCCESS;
     }
 
@@ -103,8 +101,7 @@ public int addEdgeSchema(EdgeSchema edgeSchema) {
         if (this.graph.entities.get(edgeSchema.getLabel()) != null) {
             return ErrorCode.GRAPH_ADD_EDGE_SCHEMA_FAILED;
         }
-        this.graph.getGraphSchema().addEdge(edgeSchema);
-        this.graph.entities.put(edgeSchema.getLabel(), new EdgeGroup(edgeSchema, new ArrayList<>()));
+        this.graph.registerEdgeSchema(edgeSchema);
         return ErrorCode.SUCCESS;
     }
 }