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

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

@@ -123,6 +123,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,20 +155,25 @@ 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();
         for (GraphEntity entity : graphEntities) {
             if (entity instanceof GraphVertex) {
                 memoryMutableGraph.addVertex(((GraphVertex) entity).getVertex());
             } else {
                 memoryMutableGraph.addEdge(((GraphEdge) entity).getEdge());
             }
         }
-        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(graphEntities, window.seal());
         CACHE.getConsolidateServer().executeConsolidateTask(
             insertServer.getGraphAccessors().get(0), memoryMutableGraph);
         return "Success to add entities, num: " + graphEntities.size();
@@ -182,6 +192,9 @@ 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();
         for (GraphEntity entity : graphEntities) {
             if (entity instanceof GraphVertex) {
                 memoryMutableGraph.removeVertex(entity.getLabel(),
@@ -190,6 +203,10 @@ public String deleteEntity(@Param("graphName") String graphName,
                 memoryMutableGraph.removeEdge(((GraphEdge) entity).getEdge());
             }
         }
+        if (deleteServer != null) {
+            // Deletes are applied to the index in place, no rebuild needed.
+            deleteServer.onEntitiesRemoved(graphEntities, window.seal());
+        }
         return "Success to remove entities, num: " + graphEntities.size();
     }
 

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, 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/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;
     }
 }

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

@@ -0,0 +1,108 @@
+/*
+ * 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.graph;
+
+/**
+ * The vertex version range that a reported batch of entity changes claims to cover.
+ *
+ * <p>A derived structure such as a resident keyword index can only apply a batch of changes in
+ * place if it can prove the batch is the <em>complete</em> set of vertex level changes since the
+ * structure was last known to be correct. Reading the current version after the fact is not a
+ * proof: any change made outside the reporting path would be silently accepted as already applied,
+ * and the structure would keep serving stale results forever.
+ *
+ * <p>So the writer states the range instead of the reader guessing it:
+ *
+ * <pre>
+ * VertexVersionWindow window = VertexVersionWindow.open(accessor);
+ * ... mutate the graph ...
+ * index.onEntitiesUpserted(entities, window.seal());
+ * </pre>
+ *
+ * <p>The consumer accepts the batch only when {@link #getFrom()} matches the version it last
+ * accepted <em>and</em> {@link #getTo()} still matches the graph, otherwise it falls back to a
+ * full rebuild. Changes that slip in between the writer's own mutations and {@link #seal()} cannot
+ * be detected this way; writers that mutate a graph concurrently must serialize themselves.
+ */
+public final class VertexVersionWindow {
+
+    /** Marks a window that has not been sealed yet, and can therefore not be trusted. */
+    private static final long UNSEALED = Long.MIN_VALUE;
+
+    private final GraphAccessor accessor;
+    private final long from;
+    private final long to;
+
+    private VertexVersionWindow(GraphAccessor accessor, long from, long to) {
+        this.accessor = accessor;
+        this.from = from;
+        this.to = to;
+    }
+
+    /**
+     * Captures the vertex version before a batch of writes.
+     *
+     * @param accessor graph the writes will be applied to, may be {@code null}
+     */
+    public static VertexVersionWindow open(GraphAccessor accessor) {
+        long from = accessor == null ? GraphAccessor.VERSION_UNSUPPORTED : accessor.getVertexVersion();
+        return new VertexVersionWindow(accessor, from, UNSEALED);
+    }
+
+    /**
+     * Captures the vertex version after the batch of writes. Call this as close to the last write
+     * as possible: everything between the write and this call is a blind spot.
+     */
+    public VertexVersionWindow seal() {
+        long to = accessor == null ? GraphAccessor.VERSION_UNSUPPORTED : accessor.getVertexVersion();
+        return new VertexVersionWindow(accessor, from, to);
+    }
+
+    public boolean isSealed() {
+        return to != UNSEALED;
+    }
+
+    public long getFrom() {
+        return from;
+    }
+
+    public long getTo() {
+        return to;
+    }
+
+    /**
+     * Whether this window can be trusted to describe every vertex level change between
+     * {@code acceptedVersion} and now.
+     *
+     * @param acceptedVersion version the consumer last accepted as fully applied
+     */
+    public boolean covers(long acceptedVersion) {
+        if (!isSealed() || from != acceptedVersion) {
+            return false;
+        }
+        // Anything that moved the version after the window was sealed is not described by it.
+        return accessor == null || accessor.getVertexVersion() == to;
+    }
+
+    @Override
+    public String toString() {
+        return "VertexVersionWindow{from=" + from + ", to=" + (isSealed() ? to : "unsealed") + '}';
+    }
+}